tstest/natlab: add PacketHandler and Inject.

Together, they can be used to plug custom packet processors into
Machines.

Signed-off-by: David Anderson <danderson@tailscale.com>

tstest/natlab/natlab.go

@@ -67,10 +67,17 @@ type Network struct {
 	Prefix4 netaddr.IPPrefix
 	Prefix6 netaddr.IPPrefix
 
-	mu      sync.Mutex
+	mu        sync.Mutex
-	machine map[netaddr.IP]*Machine
+	machine   map[netaddr.IP]*Machine
-	lastV4  netaddr.IP
+	defaultGW *Machine // optional
-	lastV6  netaddr.IP
+	lastV4    netaddr.IP
+	lastV6    netaddr.IP
+}
+
+func (n *Network) SetDefaultGateway(gw *Machine) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	n.defaultGW = gw
 }
 
 func (n *Network) addMachineLocked(ip netaddr.IP, m *Machine) {
@@ -135,8 +142,11 @@ func (n *Network) write(p []byte, dst, src netaddr.IPPort) (num int, err error)
 	defer n.mu.Unlock()
 	m, ok := n.machine[dst.IP]
 	if !ok {
-		trace(p, "net=%s dropped, no route to %v", n.Name, dst.IP)
+		if n.defaultGW == nil {
-		return len(p), nil
+			trace(p, "net=%s dropped, no route to %v", n.Name, dst.IP)
+			return len(p), nil
+		}
+		m = n.defaultGW
 	}
 
 	// Pretend it went across the network. Make a copy so nobody
@@ -191,13 +201,48 @@ type routeEntry struct {
 	iface  *Interface
 }
 
-// A Machine is a representation of an operating system's network stack.
+// A PacketVerdict is a decision of what to do with a packet.
-// It has a network routing table and can have multiple attached networks.
+type PacketVerdict int
+
+const (
+	// Continue means the packet should be processed by the "local
+	// sockets" logic of the Machine.
+	Continue PacketVerdict = iota
+	// Drop means the packet should not be handled further.
+	Drop
+)
+
+func (v PacketVerdict) String() string {
+	switch v {
+	case Continue:
+		return "Continue"
+	case Drop:
+		return "Drop"
+	default:
+		return fmt.Sprintf("<unknown verdict %d>", v)
+	}
+}
+
+// A PacketHandler is a function that can process packets.
+type PacketHandler func(p []byte, dst, src netaddr.IPPort) PacketVerdict
+
+// A Machine is a representation of an operating system's network
+// stack. It has a network routing table and can have multiple
+// attached networks. The zero value is valid, but lacks any
+// networking capability until Attach is called.
 type Machine struct {
 	// Name is a pretty name for debugging and packet tracing. It need
 	// not be globally unique.
 	Name string
 
+	// HandlePacket, if not nil, is a function that gets invoked for
+	// every packet this Machine receives. Returns a verdict for how
+	// the packet should continue to be handled (or not).
+	//
+	// This can be used to implement things like stateful firewalls
+	// and NAT boxes.
+	HandlePacket PacketHandler
+
 	mu         sync.Mutex
 	interfaces []*Interface
 	routes     []routeEntry // sorted by longest prefix to shortest
@@ -206,7 +251,26 @@ type Machine struct {
 	conns6 map[netaddr.IPPort]*conn // conns that want IPv6 packets
 }
 
+// Inject transmits p from src to dst, without the need for a local socket.
+// It's useful for implementing e.g. NAT boxes that need to mangle IPs.
+func (m *Machine) Inject(p []byte, dst, src netaddr.IPPort) error {
+	trace(p, "mach=%s src=%s dst=%s packet injected", m.Name, src, dst)
+	_, err := m.writePacket(p, dst, src)
+	return err
+}
+
 func (m *Machine) deliverIncomingPacket(p []byte, dst, src netaddr.IPPort) {
+	// TODO: can't hold lock while handling packet. This is safe as
+	// long as you set HandlePacket before traffic starts flowing.
+	if m.HandlePacket != nil {
+		verdict := m.HandlePacket(p, dst, src)
+		trace(p, "mach=%s src=%v dst=%v packethandler verdict=%s", m.Name, src, dst, verdict)
+		if verdict == Drop {
+			// Custom packet handler ate the packet, we're done.
+			return
+		}
+	}
+
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
@@ -555,6 +619,7 @@ func (c *conn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 	select {
 	case pkt := <-c.in:
 		n = copy(p, pkt.p)
+		trace(pkt.p, "mach=%s src=%s PacketConn.ReadFrom", c.m.Name, pkt.src)
 		return n, pkt.src.UDPAddr(), nil
 	case <-ctx.Done():
 		return 0, nil, context.DeadlineExceeded

tstest/natlab/natlab_test.go

@@ -142,3 +142,75 @@ func TestMultiNetwork(t *testing.T) {
 		t.Errorf("addr = %q; want %q", addr, natLANAddr)
 	}
 }
+
+func TestPacketHandler(t *testing.T) {
+	lan := &Network{
+		Name:    "lan",
+		Prefix4: mustPrefix("192.168.0.0/24"),
+		Prefix6: mustPrefix("fd00:916::/64"),
+	}
+	internet := NewInternet()
+
+	client := &Machine{Name: "client"}
+	nat := &Machine{Name: "nat"}
+	lan.SetDefaultGateway(nat)
+	server := &Machine{Name: "server"}
+
+	ifClient := client.Attach("eth0", lan)
+	ifNATWAN := nat.Attach("wan", internet)
+	_ = nat.Attach("lan", lan)
+	ifServer := server.Attach("server", internet)
+
+	// This HandlePacket implements a basic (some might say "broken")
+	// 1:1 NAT, where client's IP gets replaced with the NAT's WAN IP,
+	// and vice versa.
+	//
+	// This NAT is not suitable for actual use, since it doesn't do
+	// port remappings or any other things that NATs usually to. But
+	// it works as a demonstrator for a single client behind the NAT,
+	// where the NAT box itself doesn't also make PacketConns.
+	nat.HandlePacket = func(p []byte, dst, src netaddr.IPPort) PacketVerdict {
+		switch {
+		case dst.IP.Is6():
+			return Continue // no NAT for ipv6
+		case src.IP == ifClient.V4():
+			nat.Inject(p, dst, netaddr.IPPort{IP: ifNATWAN.V4(), Port: src.Port})
+			return Drop
+		case dst.IP == ifNATWAN.V4():
+			nat.Inject(p, netaddr.IPPort{IP: ifClient.V4(), Port: dst.Port}, src)
+			return Drop
+		default:
+			return Continue
+		}
+	}
+
+	clientPC, err := client.ListenPacket("udp4", ":123")
+	if err != nil {
+		t.Fatal(err)
+	}
+	serverPC, err := server.ListenPacket("udp4", ":456")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	const msg = "some message"
+	serverAddr := netaddr.IPPort{IP: ifServer.V4(), Port: 456}
+	if _, err := clientPC.WriteTo([]byte(msg), serverAddr.UDPAddr()); err != nil {
+		t.Fatal(err)
+	}
+
+	buf := make([]byte, 1500) // TODO: care about MTUs in the natlab package somewhere
+	n, addr, err := serverPC.ReadFrom(buf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	buf = buf[:n]
+	if string(buf) != msg {
+		t.Errorf("read %q; want %q", buf, msg)
+	}
+	mappedAddr := netaddr.IPPort{IP: ifNATWAN.V4(), Port: 123}
+	if addr.String() != mappedAddr.String() {
+		t.Errorf("addr = %q; want %q", addr, mappedAddr)
+	}
+
+}