-- Example: home router firewall in FWL
-- Compile with: fwlc compile examples/router.fwl

interface wan  : WAN { dynamic; };
interface lan  : LAN { cidr4 = { 10.17.1.0/24 }; };
interface wg0  : WireGuard {};

zone lan_zone = { lan, wg0 };

let rfc1918 : Set<IPv4> = { 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16 };

let forwards : Map<(Protocol, Port), (IP, Port)> = {
  (tcp, :8080) -> (10.17.1.10, :80),
  (tcp, :2222) -> (10.17.1.11, :22)
};

-- WireGuard handshake detection (compiles to ct mark state machine)
pattern WGInitiation : (UDPHeader, Bytes) =
  (udp { length = 156 }, [0x01 _*]);

pattern WGResponse : (UDPHeader, Bytes) =
  (udp { length = 100 }, [0x02 _*]);

flow WireGuardHandshake : FlowPattern =
  WGInitiation . WGResponse within 5s;

-- Block LAN clients from tunnelling out via WireGuard
rule blockOutboundWG : Frame -> <FlowMatch, Log> Action =
  \frame ->
    case frame of {
      | Frame(iif in lan_zone -> wan, IPv4(ip, UDP(udp, payload)))
          if matches(WGInitiation, (udp, payload)) ->
            case perform FlowMatch.check(flowOf(ip, wg), WireGuardHandshake) of {
              | Matched -> do {
                  perform Log.emit(Warn, "WG blocked");
                  Drop
                };
              | _ -> Continue;
            };
      | _ -> Continue;
    };

-- Inbound to router
policy input : Frame
  on { hook = Input, table = Filter, priority = Filter }
  = {
    | _ if ct.state in { Established, Related } -> Allow;
    | Frame(lo, _)                              -> Allow;
    | Frame(_, IPv6(ip6, ICMPv6(_, _)))
        if ip6.src in fe80::/10                 -> Allow;
    | Frame(_, IPv4(_, TCP(tcp, _)))
        if tcp.dport == :22                     -> Allow;
    | Frame(_, IPv4(_, UDP(udp, _)))
        if udp.dport == :51944                  -> Allow;
    | _                                         -> Drop;
  };

-- Forwarded traffic
policy forward : Frame
  on { hook = Forward, table = Filter, priority = Filter }
  = {
    | _ if ct.state in { Established, Related }       -> Allow;
    | frame if iif in lan_zone && oif == wan          -> blockOutboundWG(frame);
    | _ if ct.status == DNAT                          -> Allow;
    | Frame(iif in lan_zone -> wan, _)                -> Allow;
    | Frame(iif in lan_zone -> lan_zone, _)           -> Allow;
    | Frame(wan -> lan_zone, IPv4(ip, TCP(th, _) | UDP(th, _)))
        if (ip.protocol, th.dport) in forwards        -> Allow;
    | _                                               -> Drop;
  };

-- Outbound from router
policy output : Frame
  on { hook = Output, table = Filter, priority = Filter }
  = {
    | _ -> Allow;
  };

-- NAT
policy nat_prerouting : Frame
  on { hook = Prerouting, table = NAT, priority = DstNat }
  = {
    | Frame(_, IPv4(ip, TCP(th, _) | UDP(th, _))) ->
        if perform FIB.daddrLocal(ip.dst)
        then DNATMap((ip.protocol, th.dport), forwards)
        else Allow;
    | _ -> Allow;
  };

policy nat_postrouting : Frame
  on { hook = Postrouting, table = NAT, priority = SrcNat }
  = {
    | Frame(_ -> wan, IPv4(ip, _)) if ip.src in rfc1918 -> Masquerade;
    | _ -> Allow;
  };