From 6f281ccbcd01c05588d632aa348c88530f8aceba Mon Sep 17 00:00:00 2001 From: Michael Ben-Ami Date: Thu, 11 Jun 2026 12:44:02 -0400 Subject: [PATCH] feature/conn25: add on-remove hook for flows in FlowTable The hook fires when a flow is removed for any reason (LRU capacity eviction, tuple-collision displacement, or idle-time expiry). The hook is invoked exactly once per flow, after the flow table mutex is released, so callbacks may safely acquire other locks. We rename the IPMapper interface to Conn25Datapath, and add ClientFlowCreated/ClientFlowRemoved methods so *Conn25 can keep client-side address assignments alive while traffic is in flight. Those methods are currently stubbed for future work. Connector flows do not currently call these methods. Updates tailscale/corp#38630 Updates tailscale/corp#43180 Signed-off-by: Michael Ben-Ami --- feature/conn25/conn25.go | 22 +++-- feature/conn25/datapath.go | 48 ++++++++--- feature/conn25/datapath_test.go | 3 + feature/conn25/flowtable.go | 50 ++++++++--- feature/conn25/flowtable_test.go | 143 +++++++++++++++++++++++++++++++ 5 files changed, 234 insertions(+), 32 deletions(-) diff --git a/feature/conn25/conn25.go b/feature/conn25/conn25.go index 784b8a0c5..56b3affa7 100644 --- a/feature/conn25/conn25.go +++ b/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() diff --git a/feature/conn25/datapath.go b/feature/conn25/datapath.go index fac18667a..edfd2a048 100644 --- a/feature/conn25/datapath.go +++ b/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 } diff --git a/feature/conn25/datapath_test.go b/feature/conn25/datapath_test.go index b43894253..897d29f95 100644 --- a/feature/conn25/datapath_test.go +++ b/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") diff --git a/feature/conn25/flowtable.go b/feature/conn25/flowtable.go index 3e9415543..5cd883151 100644 --- a/feature/conn25/flowtable.go +++ b/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 } diff --git a/feature/conn25/flowtable_test.go b/feature/conn25/flowtable_test.go index 6ce289df0..7fbfc074f 100644 --- a/feature/conn25/flowtable_test.go +++ b/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) + }) +}