feature/conn25/conn25.go

@@ -257,7 +257,7 @@ func (e *extension) installHooks(dph *datapathHandler) error {
 	return nil
 }
 
-// ClientTransitIPForMagicIP implements [IPMapper].
+// ClientTransitIPForMagicIP implements [Conn25Datapath].
 func (c *Conn25) ClientTransitIPForMagicIP(m netip.Addr) (netip.Addr, error) {
 	if addr, ok := c.client.transitIPForMagicIP(m); ok {
 		return addr, nil
@@ -272,7 +272,17 @@ func (c *Conn25) ClientTransitIPForMagicIP(m netip.Addr) (netip.Addr, error) {
 	return netip.Addr{}, ErrUnmappedMagicIP
 }
 
-// ConnectorRealIPForTransitIPConnection implements [IPMapper].
+// ClientFlowCreated implements [Conn25Datapath].
+func (c *Conn25) ClientFlowCreated(transitIP netip.Addr) {
+	// TODO(tailscale/corp#43180): manage state for address assignment expiry
+}
+
+// ClientFlowRemoved implements [Conn25Datapath].
+func (c *Conn25) ClientFlowRemoved(transitIP netip.Addr) {
+	// TODO(tailscale/corp#43180): manage state for address assignment expiry
+}
+
+// ConnectorRealIPForTransitIPConnection implements [Conn25Datapath].
 func (c *Conn25) ConnectorRealIPForTransitIPConnection(src, transit netip.Addr) (netip.Addr, error) {
 	if addr, ok := c.connector.realIPForTransitIPConnection(src, transit); ok {
 		return addr, nil
@@ -704,8 +714,8 @@ type client struct {
 	byConnKey       map[key.NodePublic]set.Set[netip.Prefix]
 }
 
-// transitIPForMagicIP is part of the implementation of the IPMapper interface for dataflows lookups.
-// See also [IPMapper.ClientTransitIPForMagicIP].
+// transitIPForMagicIP is part of the implementation of the [Conn25Datapath] interface for dataflow lookups.
+// See also [Conn25Datapath.ClientTransitIPForMagicIP].
 func (c *client) transitIPForMagicIP(magicIP netip.Addr) (netip.Addr, bool) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
@@ -1262,8 +1272,8 @@ type connector struct {
 	transitIPs map[netip.Addr]map[netip.Addr]appAddr
 }
 
-// realIPForTransitIPConnection is part of the implementation of the IPMapper interface for dataflows lookups.
-// See also [IPMapper.ConnectorRealIPForTransitIPConnection].
+// realIPForTransitIPConnection is part of the implementation of the [Conn25Datapath] interface for dataflow lookups.
+// See also [Conn25Datapath.ConnectorRealIPForTransitIPConnection].
 func (c *connector) realIPForTransitIPConnection(srcIP netip.Addr, transitIP netip.Addr) (netip.Addr, bool) {
 	c.mu.Lock()
 	defer c.mu.Unlock()

feature/conn25/datapath.go

@@ -23,9 +23,14 @@ var (
 	ErrUnmappedSrcAndTransitIP = errors.New("unmapped src and transit IP")
 )
 
-// IPMapper provides methods for mapping special app connector IPs to each other
-// in aid of performing DNAT and SNAT on app connector packets.
-type IPMapper interface {
+// Conn25Datapath is the interface for the surface of [*Conn25] that the datapath
+// handler needs. It provides methods for address mapping to help the datapath handler
+// implement DNAT/SNAT, and flow lifecycle handlers so that *Conn25 can keep address
+// assignments active for active flows.
+//
+// [*Conn25] is the only production implementation; the interface exists to let
+// datapath tests substitute a lightweight fake.
+type Conn25Datapath interface {
 	// ClientTransitIPForMagicIP returns a Transit IP for the given magicIP on a client.
 	// If the magicIP is within a configured Magic IP range for an app on the client,
 	// but not mapped to an active Transit IP, implementations should return [ErrUnmappedMagicIP].
@@ -42,6 +47,14 @@ type IPMapper interface {
 	// a nil error, and a zero-value [netip.Addr] to indicate this is potentially valid,
 	// non-app-connector traffic.
 	ConnectorRealIPForTransitIPConnection(srcIP netip.Addr, transitIP netip.Addr) (netip.Addr, error)
+
+	// ClientFlowCreated is called after a client-side flow for transitIP has
+	// been installed in the client flow table.
+	ClientFlowCreated(transitIP netip.Addr)
+	// ClientFlowRemoved is called after such a flow is removed. For each
+	// flow installed in the client flow table, ClientFlowCreated is called
+	// before any ClientFlowRemoved that fires for it.
+	ClientFlowRemoved(transitIP netip.Addr)
 }
 
 // datapathHandler handles packets from the datapath,
@@ -65,16 +78,16 @@ type IPMapper interface {
 // There are two exposed methods, one for handling packets from the tun device,
 // and one for handling packets from WireGuard, but through the use of flow tables,
 // we can handle four cases: client outbound, client return, connector outbound,
-// connector return. The first packet goes through IPMapper, which is where Connectors
-// 2025 authoritative state is stored. For valid packets relevant to connectors,
+// connector return. The first packet goes through [Conn25Datapath], which is where
+// Connectors 2025 authoritative state is stored. For valid packets relevant to connectors,
 // a bidirectional flow entry is installed, so that subsequent packets (and all return traffic)
 // hit that cache. Only outbound (towards internet) packets create new flows; return (from internet)
 // packets either match a cached entry or pass through.
 //
-// We check the cache before IPMapper both for performance, and so that existing flows stay alive
-// even if address mappings change mid-flow.
+// We check the cache before [Conn25Datapath] both for performance, and so that existing flows
+// stay alive even if address mappings change mid-flow.
 type datapathHandler struct {
-	ipMapper IPMapper
+	conn25 Conn25Datapath
 
 	// Flow caches. One for the client, and one for the connector.
 	clientFlowTable    *FlowTable
@@ -89,9 +102,9 @@ const (
 	maxConnectorFlows = 100_000
 )
 
-func newDatapathHandler(ipMapper IPMapper, logf logger.Logf) *datapathHandler {
+func newDatapathHandler(conn25 Conn25Datapath, logf logger.Logf) *datapathHandler {
 	return &datapathHandler{
-		ipMapper:           ipMapper,
+		conn25:             conn25,
 		clientFlowTable:    NewFlowTable(maxClientFlows),
 		connectorFlowTable: NewFlowTable(maxConnectorFlows),
 		logf:               logf,
@@ -139,7 +152,7 @@ func (dh *datapathHandler) HandlePacketFromWireGuard(p *packet.Parsed, tun *tstu
 	// other (non-app-connector) traffic, or broken app-connector traffic
 	// that needs to be re-established by a new outbound packet.
 	transitIP := p.Dst.Addr()
-	realIP, err := dh.ipMapper.ConnectorRealIPForTransitIPConnection(p.Src.Addr(), transitIP)
+	realIP, err := dh.conn25.ConnectorRealIPForTransitIPConnection(p.Src.Addr(), transitIP)
 	if err != nil {
 		if errors.Is(err, ErrUnmappedSrcAndTransitIP) {
 			rj := packet.TailscaleRejectedHeader{
@@ -179,7 +192,6 @@ func (dh *datapathHandler) HandlePacketFromWireGuard(p *packet.Parsed, tun *tstu
 		FromTun: incoming,
 		FromWG:  outgoing,
 	})
-
 	outgoing.Action(p)
 	return filter.Accept
 }
@@ -216,7 +228,7 @@ func (dh *datapathHandler) HandlePacketFromTunDevice(p *packet.Parsed) filter.Re
 	// or broken return app-connector traffic on a connector, which needs to be re-established
 	// with a new outbound packet.
 	magicIP := p.Dst.Addr()
-	transitIP, err := dh.ipMapper.ClientTransitIPForMagicIP(magicIP)
+	transitIP, err := dh.conn25.ClientTransitIPForMagicIP(magicIP)
 	if err != nil {
 		if errors.Is(err, ErrUnmappedMagicIP) {
 			// TODO(tailscale/corp#34257): This path should deliver an ICMP error to the client.
@@ -242,11 +254,19 @@ func (dh *datapathHandler) HandlePacketFromTunDevice(p *packet.Parsed) filter.Re
 		Tuple:  flowtrack.MakeTuple(p.IPProto, netip.AddrPortFrom(transitIP, p.Dst.Port()), p.Src),
 		Action: dh.snatAction(magicIP),
 	}
+
+	// Notify Conn25 that a flow for transitIP is being established before
+	// installing it in the flow table. This guarantees that ClientFlowCreated
+	// for this flow precedes any ClientFlowRemoved that fires for it.
+	dh.conn25.ClientFlowCreated(transitIP)
+
 	dh.clientFlowTable.NewFlow(FlowData{
 		FromTun: outgoing,
 		FromWG:  incoming,
+		OnRemove: func() {
+			dh.conn25.ClientFlowRemoved(transitIP)
+		},
 	})
-
 	outgoing.Action(p)
 	return filter.Accept
 }

feature/conn25/datapath_test.go

@@ -33,6 +33,9 @@ func (tc *testConn25) ConnectorRealIPForTransitIPConnection(srcIP netip.Addr, tr
 	return tc.connectorRealIPForTransitIPConnectionFn(srcIP, transitIP)
 }
 
+func (tc *testConn25) ClientFlowCreated(transitIP netip.Addr) {}
+func (tc *testConn25) ClientFlowRemoved(transitIP netip.Addr) {}
+
 func TestHandlePacketFromTunDevice(t *testing.T) {
 	clientSrcIP := netip.MustParseAddr("100.70.0.1")
 	magicIP := netip.MustParseAddr("10.64.0.1")

feature/conn25/flowtable.go

@@ -32,6 +32,12 @@ type TupleAndAction struct {
 type FlowData struct {
 	FromTun TupleAndAction
 	FromWG  TupleAndAction
+
+	// OnRemove, if non-nil, is invoked when the flow is removed from the
+	// table for any reason (idle expiration, tuple-collision displacement
+	// in [FlowTable.NewFlow], or capacity eviction). It is called once,
+	// outside the table's mutex, so it may safely acquire other locks.
+	OnRemove func()
 }
 
 // Origin is used to track the direction of a flow.
@@ -52,7 +58,6 @@ type cachedFlow struct {
 	data FlowData // user-defined tuples and actions for both directions
 
 	lastSeen mono.Time // tracks when the flow was last hit for expiration management
-	// onRemove func()    // fires on removal/expiration (e.g. update watchers, send RST to client)
 }
 
 // FlowTable stores and retrieves [FlowData] that can be looked up
@@ -196,27 +201,36 @@ func (t *FlowTable) lookup(k flowtrack.Tuple, dir Origin) (PacketAction, bool) {
 // would cause the table to exceed its maximum size, the least recently used
 // (looked-up or created) flow is evicted. data is not validated, the caller must
 // supply non-nil packet actions.
+//
+// Any [FlowData.OnRemove] callbacks belonging to displaced or evicted flows are
+// invoked after the table's mutex is released, before NewFlow returns.
 func (t *FlowTable) NewFlow(data FlowData) {
-	t.mu.Lock()
-	defer t.mu.Unlock()
+	var onRemoves []func()
 
+	t.mu.Lock()
 	// If either tuple leads to anything existing, remove it.
-	t.removeFlowLocked(t.fromTunCache[data.FromTun.Tuple])
-	t.removeFlowLocked(t.fromWGCache[data.FromWG.Tuple])
+	onRemoves = append(onRemoves, t.removeFlowLocked(t.fromTunCache[data.FromTun.Tuple]))
+	onRemoves = append(onRemoves, t.removeFlowLocked(t.fromWGCache[data.FromWG.Tuple]))
 
 	flow := &cachedFlow{
 		data:     data,
 		lastSeen: mono.Now(),
-		// Populate onRemove()
 	}
 
 	ele := t.lru.PushFront(flow)
 	if t.maxEntries > 0 && t.lru.Len() > t.maxEntries {
-		t.removeFlowLocked(t.lru.Back())
+		onRemoves = append(onRemoves, t.removeFlowLocked(t.lru.Back()))
 	}
 
 	t.fromTunCache[data.FromTun.Tuple] = ele
 	t.fromWGCache[data.FromWG.Tuple] = ele
+	t.mu.Unlock()
+
+	for _, onRemove := range onRemoves {
+		if onRemove != nil {
+			onRemove()
+		}
+	}
 }
 
 // StartExpiredSweeper starts a sweeper that removes idle flows that have
@@ -246,8 +260,8 @@ func (t *FlowTable) removeIdle(now mono.Time) int {
 		return 0
 	}
 
+	var onRemoves []func()
 	t.mu.Lock()
-	defer t.mu.Unlock()
 
 	removed := 0
 	for ele := t.lru.Back(); ele != nil; ele = t.lru.Back() {
@@ -258,20 +272,32 @@ func (t *FlowTable) removeIdle(now mono.Time) int {
 		if now.Sub(flow.lastSeen) <= t.idleTimeout {
 			break
 		}
-		t.removeFlowLocked(ele)
+		onRemoves = append(onRemoves, t.removeFlowLocked(ele))
 		removed++
 	}
+	t.mu.Unlock()
+
+	for _, onRemove := range onRemoves {
+		if onRemove != nil {
+			onRemove()
+		}
+	}
+
 	return removed
 }
 
-func (t *FlowTable) removeFlowLocked(ele *list.Element) {
+// removeFlowLocked detaches the flow at ele from t, and returns the flow's
+// [FlowData.OnRemove] callback, which may be nil. The caller must hold the
+// mutex while calling removeFlowLocked, and release it before invoking the
+// callback.
+func (t *FlowTable) removeFlowLocked(ele *list.Element) func() {
 	if ele == nil {
-		return
+		return nil
 	}
 
 	flow := t.lru.Remove(ele).(*cachedFlow)
 	delete(t.fromTunCache, flow.data.FromTun.Tuple)
 	delete(t.fromWGCache, flow.data.FromWG.Tuple)
 
-	// TODO(mzb): run flow.onRemove()
+	return flow.data.OnRemove
 }

feature/conn25/flowtable_test.go

@@ -5,6 +5,7 @@ package conn25
 
 import (
 	"fmt"
+	"maps"
 	"net/netip"
 	"testing"
 	"testing/synctest"
@@ -372,3 +373,145 @@ func TestFlowTable_removeIdle(t *testing.T) {
 		assertFlowMiss(t, ft, FromTun, flows[1].FromTun.Tuple)
 	})
 }
+
+// recordOnRemove returns an OnRemove that increments fired[name] when invoked.
+// Used to verify OnRemove fires for the right flows the right number of times.
+func recordOnRemove(fired map[string]int, name string) func() {
+	return func() { fired[name]++ }
+}
+
+func TestFlowTable_OnRemove(t *testing.T) {
+	tun1 := mkTuple("1.1.1.1:1000", "2.2.2.2:80")
+	wg1 := mkTuple("2.2.2.2:80", "1.1.1.1:1000")
+	tun2 := mkTuple("3.3.3.3:1000", "4.4.4.4:80")
+	wg2 := mkTuple("4.4.4.4:80", "3.3.3.3:1000")
+
+	t.Run("displacement", func(t *testing.T) {
+		// fd2 collides with fd1 on the FromTun tuple; fd1 is displaced and
+		// its OnRemove fires. fd2 stays installed and its OnRemove does not.
+		ft := NewFlowTable(0)
+		fired := map[string]int{}
+
+		fd1 := mkFlow(tun1, wg1)
+		fd1.OnRemove = recordOnRemove(fired, "fd1")
+		ft.NewFlow(fd1)
+
+		fd2 := mkFlow(tun1, wg2)
+		fd2.OnRemove = recordOnRemove(fired, "fd2")
+		ft.NewFlow(fd2)
+
+		if want := (map[string]int{"fd1": 1}); !maps.Equal(fired, want) {
+			t.Errorf("fired = %v, want %v", fired, want)
+		}
+	})
+
+	t.Run("one-replaces-two", func(t *testing.T) {
+		// fd3's FromTun matches fd1's, and fd3's FromWG matches fd2's. Both
+		// fd1 and fd2 should be displaced and both OnRemoves should fire.
+		ft := NewFlowTable(0)
+		fired := map[string]int{}
+
+		fd1 := mkFlow(tun1, wg1)
+		fd1.OnRemove = recordOnRemove(fired, "fd1")
+		ft.NewFlow(fd1)
+
+		fd2 := mkFlow(tun2, wg2)
+		fd2.OnRemove = recordOnRemove(fired, "fd2")
+		ft.NewFlow(fd2)
+
+		fd3 := mkFlow(tun1, wg2)
+		fd3.OnRemove = recordOnRemove(fired, "fd3")
+		ft.NewFlow(fd3)
+
+		if want := (map[string]int{"fd1": 1, "fd2": 1}); !maps.Equal(fired, want) {
+			t.Errorf("fired = %v, want %v", fired, want)
+		}
+	})
+
+	t.Run("reinstall-same-tuples-fires-once", func(t *testing.T) {
+		// Re-installing a flow with identical tuples from both directions
+		// causes removeFlowLocked to be called twice (once from each direction).
+		// Only one OnRemove should be called for the single flow.
+		ft := NewFlowTable(0)
+		fired := map[string]int{}
+
+		fd1 := mkFlow(tun1, wg1)
+		fd1.OnRemove = recordOnRemove(fired, "fd1")
+		ft.NewFlow(fd1)
+
+		fd2 := mkFlow(tun1, wg1) // identical tuples
+		fd2.OnRemove = recordOnRemove(fired, "fd2")
+		ft.NewFlow(fd2)
+
+		if want := (map[string]int{"fd1": 1}); !maps.Equal(fired, want) {
+			t.Errorf("fired = %v, want %v", fired, want)
+		}
+	})
+
+	t.Run("capacity-eviction", func(t *testing.T) {
+		// With capacity 1, installing fd2 evicts fd1 from the back of the
+		// LRU; fd1's OnRemove fires.
+		ft := NewFlowTable(1)
+		fired := map[string]int{}
+
+		fd1 := mkFlow(tun1, wg1)
+		fd1.OnRemove = recordOnRemove(fired, "fd1")
+		ft.NewFlow(fd1)
+
+		fd2 := mkFlow(tun2, wg2)
+		fd2.OnRemove = recordOnRemove(fired, "fd2")
+		ft.NewFlow(fd2)
+
+		if want := (map[string]int{"fd1": 1}); !maps.Equal(fired, want) {
+			t.Errorf("fired = %v, want %v", fired, want)
+		}
+	})
+
+	syncSubtest(t, "remove-idle", func(t *testing.T) {
+		ft := NewFlowTable(0, WithFlowIdleTimeout(time.Minute))
+		fired := map[string]int{}
+
+		fd1 := mkFlow(tun1, wg1)
+		fd1.OnRemove = recordOnRemove(fired, "fd1")
+		ft.NewFlow(fd1)
+
+		time.Sleep(2 * time.Minute) // advance synthetic clock past idleTimeout
+		if got, want := ft.removeIdle(mono.Now()), 1; got != want {
+			t.Errorf("removeIdle returned %d, want %d", got, want)
+		}
+
+		if want := (map[string]int{"fd1": 1}); !maps.Equal(fired, want) {
+			t.Errorf("fired = %v, want %v", fired, want)
+		}
+	})
+
+	t.Run("nil-onremove-no-panic", func(t *testing.T) {
+		ft := NewFlowTable(0)
+
+		ft.NewFlow(mkFlow(tun1, wg1)) // OnRemove unset
+		ft.NewFlow(mkFlow(tun1, wg2)) // displaces the first flow
+	})
+
+	t.Run("runs-outside-table-lock", func(t *testing.T) {
+		ft := NewFlowTable(0)
+
+		var onRemoveRan bool
+		fd1 := mkFlow(tun1, wg1)
+		fd1.OnRemove = func() {
+			// NewFlow is used here because we know it acquires the mutex,
+			// So this will prove OnRemove() is called with the mutex released.
+			ft.NewFlow(mkFlow(tun2, wg2))
+			onRemoveRan = true
+		}
+		ft.NewFlow(fd1)
+
+		// This should cause displacement of the first flow, and OnRemove to fire.
+		ft.NewFlow(mkFlow(tun1, mkTuple("9.9.9.9:99", "8.8.8.8:88")))
+
+		if !onRemoveRan {
+			t.Errorf("OnRemove did not run")
+		}
+		// The new install should be visible.
+		assertFlowHit(t, ft, FromTun, tun2)
+	})
+}