tstest/natlab: refactor PacketHandler into a larger interface.

The new interface lets implementors more precisely distinguish
local traffic from forwarded traffic, and applies different
forwarding logic within Machines for each type. This allows
Machines to be packet forwarders, which didn't quite work
with the implementation of Inject.

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

tstest/natlab/natlab.go

@@ -12,6 +12,7 @@
 package natlab
 
 import (
+	"bytes"
 	"context"
 	"crypto/sha256"
 	"encoding/base64"
@@ -40,6 +41,12 @@ type Packet struct {
 	locator string
 }
 
+// Equivalent returns true if Src, Dst and Payload are the same in p
+// and p2.
+func (p *Packet) Equivalent(p2 *Packet) bool {
+	return p.Src == p2.Src && p.Dst == p2.Dst && bytes.Equal(p.Payload, p2.Payload)
+}
+
 // Clone returns a copy of p that shares nothing with p.
 func (p *Packet) Clone() *Packet {
 	return &Packet{
@@ -266,8 +273,41 @@ func (v PacketVerdict) String() string {
 	}
 }
 
-// A PacketHandler is a function that can process packets.
-type PacketHandler func(p *Packet, inIf *Interface) PacketVerdict
+// A PacketHandler can look at packets arriving at, departing, and
+// transiting a Machine, and filter or mutate them.
+//
+// Each method is invoked with a Packet that natlab would like to keep
+// processing. Handlers can return that same Packet to allow
+// processing to continue; nil to drop the Packet; or a different
+// Packet that should be processed instead of the original.
+//
+// Packets passed to handlers share no state with anything else, and
+// are therefore safe to mutate. It's safe to return the original
+// packet mutated in-place, or a brand new packet initialized from
+// scratch.
+//
+// Packets mutated by a PacketHandler are processed anew by the
+// associated Machine, as if the packet had always been the mutated
+// one. For example, if HandleForward is invoked with a Packet, and
+// the handler changes the destination IP address to one of the
+// Machine's own IPs, the Machine restarts delivery, but this time
+// going to a local PacketConn (which in turn will invoke HandleIn,
+// since the packet is now destined for local delivery).
+type PacketHandler interface {
+	// HandleIn processes a packet arriving on iif, whose destination
+	// is an IP address owned by the attached Machine. If p is
+	// returned unmodified, the Machine will go on to deliver the
+	// Packet to the appropriate listening PacketConn, if one exists.
+	HandleIn(p *Packet, iif *Interface) *Packet
+	// HandleOut processes a packet about to depart on oif from a
+	// local PacketConn. If p is returned unmodified, the Machine will
+	// transmit the Packet on oif.
+	HandleOut(p *Packet, oif *Interface) *Packet
+	// HandleForward is called when the Machine wants to forward a
+	// packet from iif to oif. If p is returned unmodified, the
+	// Machine will transmit the packet on oif.
+	HandleForward(p *Packet, iif, oif *Interface) *Packet
+}
 
 // A Machine is a representation of an operating system's network
 // stack. It has a network routing table and can have multiple
@@ -278,19 +318,14 @@ type Machine struct {
 	// not be globally unique.
 	Name string
 
-	// HandlePacket, if not nil, is a function that gets invoked for
-	// every packet this Machine receives, and every packet sent by a
-	// local PacketConn. Returns a verdict for how the packet should
-	// continue to be handled (or not).
+	// PacketHandler, if not nil, is a PacketHandler implementation
+	// that inspects all packets arriving, departing, or transiting
+	// the Machine. See the definition of the PacketHandler interface
+	// for semantics.
 	//
-	// HandlePacket's interface parameter is the interface on which
-	// the packet was received, or nil for a packet sent by a local
-	// PacketConn or Inject call.
-	//
-	// The packet provided to HandlePacket can safely be mutated and
-	// Inject()ed if desired. This can be used to implement things
-	// like stateful firewalls and NAT boxes.
-	HandlePacket PacketHandler
+	// If PacketHandler is nil, the machine allows all inbound
+	// traffic, all outbound traffic, and drops forwarded packets.
+	PacketHandler PacketHandler
 
 	mu         sync.Mutex
 	interfaces []*Interface
@@ -300,26 +335,42 @@ 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 *Packet) error {
-	p = p.Clone()
-	p.setLocator("mach=%s", m.Name)
-	p.Trace("Machine.Inject")
-	_, err := m.writePacket(p)
-	return err
+func (m *Machine) isLocalIP(ip netaddr.IP) bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	for _, intf := range m.interfaces {
+		for _, iip := range intf.ips {
+			if ip == iip {
+				return true
+			}
+		}
+	}
+	return false
 }
 
 func (m *Machine) deliverIncomingPacket(p *Packet, iface *Interface) {
 	p.setLocator("mach=%s if=%s", m.Name, iface.name)
+
+	if m.isLocalIP(p.Dst.IP) {
+		m.deliverLocalPacket(p, iface)
+	} else {
+		m.forwardPacket(p, iface)
+	}
+}
+
+func (m *Machine) deliverLocalPacket(p *Packet, iface *Interface) {
 	// 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 {
-		p.Trace("Machine.HandlePacket")
-		verdict := m.HandlePacket(p.Clone(), iface)
-		p.Trace("Machine.HandlePacket verdict=%s", verdict)
-		if verdict == Drop {
-			// Custom packet handler ate the packet, we're done.
+	if m.PacketHandler != nil {
+		p2 := m.PacketHandler.HandleIn(p.Clone(), iface)
+		if p2 == nil {
+			// Packet dropped, nothing left to do.
+			return
+		}
+		if !p.Equivalent(p2) {
+			// Restart delivery, this packet might be a forward packet
+			// now.
+			m.deliverIncomingPacket(p2, iface)
 			return
 		}
 	}
@@ -353,6 +404,35 @@ func (m *Machine) deliverIncomingPacket(p *Packet, iface *Interface) {
 	p.Trace("dropped, no listening conn")
 }
 
+func (m *Machine) forwardPacket(p *Packet, iif *Interface) {
+	oif, err := m.interfaceForIP(p.Dst.IP)
+	if err != nil {
+		p.Trace("%v", err)
+		return
+	}
+
+	if m.PacketHandler == nil {
+		// Forwarding not allowed by default
+		p.Trace("drop, forwarding not allowed")
+		return
+	}
+	p2 := m.PacketHandler.HandleForward(p.Clone(), iif, oif)
+	if p2 == nil {
+		p.Trace("drop")
+		// Packet dropped, done.
+		return
+	}
+	if !p.Equivalent(p2) {
+		// Packet changed, restart delivery.
+		p2.Trace("PacketHandler mutated packet")
+		m.deliverIncomingPacket(p2, iif)
+		return
+	}
+
+	p.Trace("-> net=%s oif=%s", oif.net.Name, oif)
+	oif.net.write(p)
+}
+
 func unspecOf(ip netaddr.IP) netaddr.IP {
 	if ip.Is4() {
 		return v4unspec
@@ -455,13 +535,17 @@ func (m *Machine) writePacket(p *Packet) (n int, err error) {
 		return 0, err
 	}
 
-	if m.HandlePacket != nil {
-		p.Trace("Machine.HandlePacket")
-		verdict := m.HandlePacket(p.Clone(), nil)
-		p.Trace("Machine.HandlePacket verdict=%s", verdict)
-		if verdict == Drop {
+	if m.PacketHandler != nil {
+		p2 := m.PacketHandler.HandleOut(p.Clone(), iface)
+		if p2 == nil {
+			// Packet dropped, done.
 			return len(p.Payload), nil
 		}
+		if !p.Equivalent(p2) {
+			// Restart transmission, src may have changed weirdly
+			m.writePacket(p2)
+			return
+		}
 	}
 
 	p.Trace("-> net=%s if=%s", iface.net.Name, iface)