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 <mzb@tailscale.com>
This commit is contained in:
Michael Ben-Ami
2026-06-12 10:44:42 -04:00
committed by mzbenami
parent 2a0eafc20f
commit 6f281ccbcd
5 changed files with 234 additions and 32 deletions
+16 -6
View File
@@ -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()
+34 -14
View File
@@ -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
}
+3
View File
@@ -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")
+38 -12
View File
@@ -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
}
+143
View File
@@ -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)
})
}