From 2e3a896cab8a7a4ccf56a654c0b8cbd62c98d043 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Wed, 7 Aug 2024 09:18:55 -0700
Subject: [PATCH] add stateful firewall

Change-Id: I4a963f144f24481746c50a2aa97671b7bfc1f267
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 tstest/natlab/vnet/nat.go | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/tstest/natlab/vnet/nat.go b/tstest/natlab/vnet/nat.go
index 91b01adf2..179feb733 100644
--- a/tstest/natlab/vnet/nat.go
+++ b/tstest/natlab/vnet/nat.go
@@ -5,6 +5,7 @@ package vnet
 
 import (
 	"errors"
+	"log"
 	"math/rand/v2"
 	"net/netip"
 	"time"
@@ -111,7 +112,7 @@ func (n *oneToOneNAT) PickIncomingDst(src, dst netip.AddrPort, at time.Time) (la
 	return netip.AddrPortFrom(n.lanIP, dst.Port())
 }
 
-type hardKeyOut struct {
+type srcDstTuple struct {
 	src netip.AddrPort
 	dst netip.AddrPort
 }
@@ -137,7 +138,7 @@ type lanAddrAndTime struct {
 type hardNAT struct {
 	wanIP netip.Addr
 
-	out map[hardKeyOut]portMappingAndTime
+	out map[srcDstTuple]portMappingAndTime
 	in  map[hardKeyIn]lanAddrAndTime
 }
 
@@ -148,7 +149,7 @@ func init() {
 }
 
 func (n *hardNAT) PickOutgoingSrc(src, dst netip.AddrPort, at time.Time) (wanSrc netip.AddrPort) {
-	ko := hardKeyOut{src, dst}
+	ko := srcDstTuple{src, dst}
 	if pm, ok := n.out[ko]; ok {
 		// Existing flow.
 		// TODO: bump timestamp
@@ -196,9 +197,10 @@ func (n *hardNAT) PickIncomingDst(src, dst netip.AddrPort, at time.Time) (lanDst
 // Unlike Linux, this implementation is capped at 32k entries and doesn't resort
 // to other allocation strategies when all 32k WAN ports are taken.
 type easyNAT struct {
-	wanIP netip.Addr
-	out   map[netip.AddrPort]portMappingAndTime
-	in    map[uint16]lanAddrAndTime
+	wanIP   netip.Addr
+	out     map[netip.AddrPort]portMappingAndTime
+	in      map[uint16]lanAddrAndTime
+	lastOut map[srcDstTuple]time.Time // (lan:port, wan:port) => last packet out time
 }
 
 func init() {
@@ -208,6 +210,7 @@ func init() {
 }
 
 func (n *easyNAT) PickOutgoingSrc(src, dst netip.AddrPort, at time.Time) (wanSrc netip.AddrPort) {
+	mak.Set(&n.lastOut, srcDstTuple{src, dst}, at)
 	if pm, ok := n.out[src]; ok {
 		// Existing flow.
 		// TODO: bump timestamp
@@ -235,5 +238,14 @@ func (n *easyNAT) PickIncomingDst(src, dst netip.AddrPort, at time.Time) (lanDst
 	if dst.Addr() != n.wanIP {
 		return netip.AddrPort{} // drop; not for us. shouldn't happen if natlabd routing isn't broken.
 	}
-	return n.in[dst.Port()].lanAddr
+	lanDst = n.in[dst.Port()].lanAddr
+
+	// Stateful firewall: drop incoming packets that don't have traffic out.
+	// TODO(bradfitz): verify Linux does this in the router code, not in the NAT code.
+	if t, ok := n.lastOut[srcDstTuple{lanDst, src}]; !ok || at.Sub(t) > 300*time.Second {
+		log.Printf("Drop incoming packet from %v to %v; no recent outgoing packet", src, dst)
+		return netip.AddrPort{}
+	}
+
+	return lanDst
 }