#!/usr/sbin/nft -f
# Compiled from examples/router.fwl
# Single inet table: fwl

flush ruleset

table inet fwl {

    # ── Data: let rfc1918 ────────────────────────────────────────────────────
    set rfc1918 {
        type ipv4_addr
        flags interval
        elements = {
            10.0.0.0/8,
            172.16.0.0/12,
            192.168.0.0/16
        }
    }

    # ── Data: let forwards ──────────────────────────────────────────────────
    map forwards {
        type inet_proto . inet_service : ipv4_addr . inet_service
        elements = {
            tcp . 8080 : 10.17.1.10 . 80,
            tcp . 2222 : 10.17.1.11 . 22
        }
    }

    # ── WireGuard ct mark state machine ─────────────────────────────────────
    # Compiles: flow WireGuardHandshake = WGInitiation . WGResponse within 5s
    # State: ct mark 0 = Idle, 1 = SawInitiation, 2 = Confirmed
    #
    # WGInitiation: UDP, udp length == 156 (8 hdr + 148 payload), payload[0] == 0x01
    # WGResponse:   UDP, udp length == 100 (8 hdr + 92 payload),  payload[0] == 0x02
    # @th,64,8 = first byte of UDP payload (offset 64 bits past transport header start)

    chain wg_flow {
        # Packet 1: Idle → SawInitiation
        ct state new ct mark 0 \
            meta l4proto udp udp length 156 \
            @th,64,8 0x01 \
            ct mark set 1 \
            return

        # Packet 2: SawInitiation → Confirmed
        ct mark 1 \
            meta l4proto udp udp length 100 \
            @th,64,8 0x02 \
            ct mark set 2 \
            return
    }

    # ── rule blockOutboundWG ─────────────────────────────────────────────────
    # Compiles: rule blockOutboundWG : Frame -> <FlowMatch, Log> Action
    # Called via jump from forward. Drops confirmed WG handshakes, returns otherwise.

    chain blockOutboundWG {
        # Feed matching UDP into the WG state machine
        meta nfproto ipv4 meta l4proto udp \
            udp length 156 \
            @th,64,8 0x01 \
            jump wg_flow

        # If handshake is now Confirmed (ct mark 2): log + drop
        ct mark 2 \
            log prefix "WG blocked: " level warn \
            drop

        # Continue: return to forward chain (no verdict)
        return
    }

    # ── policy input ─────────────────────────────────────────────────────────
    # hook = Input, table = Filter, priority = filter (0), default = drop

    chain input {
        type filter hook input priority filter; policy drop;

        # | _ if ct.state in { Established, Related } -> Allow
        ct state { established, related } accept

        # | Frame(lo, _) -> Allow
        iifname "lo" accept

        # | Frame(_, IPv6(ip6, ICMPv6(_, _))) if ip6.src in fe80::/10 -> Allow
        meta nfproto ipv6 ip6 nexthdr ipv6-icmp ip6 saddr fe80::/10 accept

        # | Frame(_, IPv4(_, TCP(tcp, _))) if tcp.dport == :22 -> Allow
        meta nfproto ipv4 meta l4proto tcp tcp dport 22 accept

        # | Frame(_, IPv4(_, UDP(udp, _))) if udp.dport == :51944 -> Allow
        meta nfproto ipv4 meta l4proto udp udp dport 51944 accept

        # | _ -> Drop  (chain policy)
    }

    # ── policy forward ───────────────────────────────────────────────────────
    # hook = Forward, table = Filter, priority = filter (0), default = drop

    chain forward {
        type filter hook forward priority filter; policy drop;

        # | _ if ct.state in { Established, Related } -> Allow
        ct state { established, related } accept

        # | frame if iif in lan_zone && oif == wan -> blockOutboundWG(frame)
        meta iifname { "lan", "wg0" } meta oifname "wan" jump blockOutboundWG

        # | _ if ct.status == DNAT -> Allow
        ct status dnat accept

        # | Frame(iif in lan_zone -> wan, _) -> Allow
        meta iifname { "lan", "wg0" } meta oifname "wan" accept

        # | Frame(iif in lan_zone -> lan_zone, _) -> Allow
        meta iifname { "lan", "wg0" } meta oifname { "lan", "wg0" } accept

        # | Frame(wan -> lan_zone, IPv4(ip, TCP|UDP)) if (proto, dport) in forwards -> Allow
        # Membership test only — the actual DNAT is done in nat_prerouting.
        meta iifname "wan" meta oifname { "lan", "wg0" } \
            meta nfproto ipv4 \
            meta l4proto { tcp, udp } \
            meta l4proto . th dport @forwards \
            accept

        # | _ -> Drop  (chain policy)
    }

    # ── policy output ────────────────────────────────────────────────────────
    # hook = Output, table = Filter, priority = filter (0), default = accept

    chain output {
        type filter hook output priority filter; policy accept;
        # | _ -> Allow  (chain policy)
    }

    # ── policy nat_prerouting ────────────────────────────────────────────────
    # hook = Prerouting, table = NAT, priority = dstnat (-100), default = accept

    chain nat_prerouting {
        type nat hook prerouting priority dstnat; policy accept;

        # | Frame(_, IPv4(ip, TCP|UDP)) ->
        #     if FIB.daddrLocal(ip.dst) then DNATMap((proto, dport), forwards) else Allow
        meta nfproto ipv4 meta l4proto { tcp, udp } \
            fib daddr type local \
            dnat ip to meta l4proto . th dport map @forwards

        # | _ -> Allow  (chain policy)
    }

    # ── policy nat_postrouting ───────────────────────────────────────────────
    # hook = Postrouting, table = NAT, priority = srcnat (100), default = accept

    chain nat_postrouting {
        type nat hook postrouting priority srcnat; policy accept;

        # | Frame(_ -> wan, IPv4(ip, _)) if ip.src in rfc1918 -> Masquerade
        meta oifname "wan" meta nfproto ipv4 ip saddr @rfc1918 masquerade

        # | _ -> Allow  (chain policy)
    }
}