/*-
 * Copyright (c) 2016 Yandex LLC
 * Copyright (c) 2016 Andrey V. Elsukov <ae@FreeBSD.org>
 * Copyright (c) 2025-2026 Gleb Smirnoff <glebius@FreeBSD.org>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include <sys/param.h>
#include <sys/systm.h>
#include <sys/mbuf.h>
#include <sys/kernel.h>
#include <sys/lock.h>
#include <sys/sx.h>
#include <sys/socket.h>
#include <sys/tree.h>
#include <net/ethernet.h>
#include <net/if.h>
#include <net/if_pflog.h>
#include <net/vnet.h>
#include <net/bpf.h>

#include <netinet/ip.h>
#include <netinet/ip_fw.h>
#include <netinet/ip_var.h>
#include <netpfil/ipfw/ip_fw_private.h>

static bool
bpf_ipfw_chkdir(void *arg __unused, const struct mbuf *m, int dir)
{
        return ((dir == BPF_D_IN && m_rcvif(m) == NULL) ||
            (dir == BPF_D_OUT && m_rcvif(m) != NULL));
}

static const struct bif_methods bpf_ipfw_methods = {
        .bif_chkdir = bpf_ipfw_chkdir,
};

struct ipfw_tap {
        RB_ENTRY(ipfw_tap)      entry;
        uint32_t                rule;
        u_int                   refs;
        struct bpf_if           *bpf;
        char                    name[sizeof("ipfw4294967295")];
};

static inline int
tap_compare(const struct ipfw_tap *a, const struct ipfw_tap *b)
{
        return (a->rule != b->rule ? (a->rule < b->rule ? -1 : 1) : 0);
}
RB_GENERATE_STATIC(tap_tree, ipfw_tap, entry, tap_compare);
VNET_DEFINE_STATIC(struct ipfw_tap, default_tap) = { .name = "ipfw0" };
#define V_default_tap   VNET(default_tap)

void
ipfw_tap_alloc(struct ip_fw_chain *ch, uint32_t rule)
{
        struct ipfw_tap *tap, key = { .rule = rule };
        int n __diagused;

        MPASS(rule > 0 && rule < IPFW_DEFAULT_RULE);
        IPFW_UH_WLOCK_ASSERT(ch);

        tap = RB_FIND(tap_tree, &ch->taps, &key);
        if (tap != NULL) {
                MPASS(tap->rule == rule);
                tap->refs++;
                return;
        }
        tap = malloc(sizeof(*tap), M_IPFW, M_WAITOK);
        tap->rule = rule;
        tap->refs = 1;
        n = snprintf(tap->name, sizeof(tap->name), "ipfw%u", rule);
        MPASS(n > 4 && n < sizeof("ipfw4294967295"));
        tap->bpf = bpf_attach(tap->name, DLT_EN10MB, PFLOG_HDRLEN,
            &bpf_ipfw_methods, NULL);
        tap = RB_INSERT(tap_tree, &ch->taps, tap);
        MPASS(tap == NULL);
}

void
ipfw_tap_free(struct ip_fw_chain *ch, uint32_t rule)
{
        struct ipfw_tap *tap, key = { .rule = rule };

        MPASS(rule > 0 && rule < IPFW_DEFAULT_RULE);
        IPFW_UH_WLOCK_ASSERT(ch);

        tap = RB_FIND(tap_tree, &ch->taps, &key);
        MPASS(tap != NULL);
        if (--tap->refs == 0) {
                bpf_detach(tap->bpf);
                RB_REMOVE(tap_tree, &ch->taps, tap);
                free(tap, M_IPFW);
        }
}

void
ipfw_bpf_tap(struct ip_fw_chain *ch, struct ip_fw_args *args,
    struct ip *ip, uint32_t rulenum)
{
        struct ipfw_tap *tap;

        if (rulenum == IPFW_DEFAULT_RULE) {
                tap = &V_default_tap;
        } else {
                struct ipfw_tap key = { .rule = rulenum };

                tap = RB_FIND(tap_tree, &ch->taps, &key);
                MPASS(tap != NULL);
                /*
                 * Compatibility: if user is not using per-rule taps, fallback
                 * to the default tap.
                 */
                if (!bpf_peers_present(tap->bpf))
                        tap = &V_default_tap;
        }
        if (args->flags & IPFW_ARGS_LENMASK) {
                bpf_tap(tap->bpf, args->mem, IPFW_ARGS_LENGTH(args->flags));
        } else if (args->flags & IPFW_ARGS_ETHER) {
                /* layer2, use orig hdr */
                bpf_mtap(tap->bpf, args->m);
        } else {
                char *fakehdr;

                /* Add fake header. Later we will store
                 * more info in the header.
                 */
                if (ip->ip_v == 4)
                        fakehdr = "DDDDDDSSSSSS\x08\x00";
                else if (ip->ip_v == 6)
                        fakehdr = "DDDDDDSSSSSS\x86\xdd";
                else
                        /* Obviously bogus EtherType. */
                        fakehdr = "DDDDDDSSSSSS\xff\xff";

                bpf_mtap2(tap->bpf, fakehdr, ETHER_HDR_LEN, args->m);
        }
}

VNET_DEFINE_STATIC(struct bpf_if *, bpf_pflog);
#define V_bpf_pflog     VNET(bpf_pflog)
void
ipfw_pflog_tap(void *data, struct mbuf *m)
{
        bpf_mtap2(V_bpf_pflog, data, PFLOG_HDRLEN, m);
}

void
ipfw_bpf_init(int first __unused)
{
        V_default_tap.bpf = bpf_attach(V_default_tap.name, DLT_EN10MB,
            PFLOG_HDRLEN, &bpf_ipfw_methods, NULL);
        V_bpf_pflog = bpf_attach("ipfwlog0", DLT_PFLOG, PFLOG_HDRLEN,
            &bpf_ipfw_methods, NULL);
}

void
ipfw_bpf_uninit(int last __unused)
{
        bpf_detach(V_default_tap.bpf);
        bpf_detach(V_bpf_pflog);
}