-- 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 open_ports : Set<Port> = { :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;
    };

-- Port-forward map: incoming proto+port -> internal addr+port
portforward wan_forwards
    on  wan
    via Map<(Protocol, Port), (IPv4, Port)> = {
        (tcp, :8080) -> (10.17.1.10, :80),
        (tcp, :2222) -> (10.17.1.11, :22)
    };

-- Masquerade outbound traffic from RFC1918 sources
masquerade wan_snat
    on  wan
    src rfc1918;

-- Inbound to router
policy input : Frame hook Input = {
    | Frame(_, IPv4(_, TCP(tcp, _)))
        if tcp.dport in open_ports                              -> Allow;
    | Frame(_, IPv4(_, UDP(udp, _)))
        if udp.dport == :51944                                  -> Allow;
    | _                                                         -> Drop;
};

-- Forwarded traffic
policy forward : Frame hook Forward = {
    | frame if iif in lan_zone && oif == wan                    -> blockOutboundWG(frame);
    | Frame(iif in lan_zone -> wan, _)                          -> Allow;
    | Frame(iif in lan_zone -> lan_zone, _)                     -> Allow;
    | _                                                         -> Drop;
};