TUN-6696: Refactor flow into funnel and close idle funnels

A funnel is an abstraction for 1 source to many destinations.
As part of this refactoring, shared logic between Darwin and Linux are moved into icmp_posix

packet/flow.go

@@ -1,94 +0,0 @@
-package packet
-
-import (
-	"errors"
-	"fmt"
-	"net/netip"
-	"sync"
-)
-
-var (
-	ErrFlowNotFound = errors.New("flow not found")
-)
-
-// FlowID represents a key type that can be used by FlowTracker
-type FlowID interface {
-	// Type returns the name of the type that implements the FlowID
-	Type() string
-	fmt.Stringer
-}
-
-type Flow struct {
-	Src       netip.Addr
-	Dst       netip.Addr
-	Responder FlowResponder
-}
-
-func isSameFlow(f1, f2 *Flow) bool {
-	if f1 == nil || f2 == nil {
-		return false
-	}
-	return *f1 == *f2
-}
-
-// FlowResponder sends response packets to the flow
-type FlowResponder interface {
-	// SendPacket returns a packet to the flow. It must not modify the packet,
-	// and after return it must not read the packet
-	SendPacket(pk RawPacket) error
-}
-
-// FlowTracker tracks flow from the perspective of eyeball to origin
-type FlowTracker struct {
-	lock  sync.RWMutex
-	flows map[FlowID]*Flow
-}
-
-func NewFlowTracker() *FlowTracker {
-	return &FlowTracker{
-		flows: make(map[FlowID]*Flow),
-	}
-}
-
-func (sft *FlowTracker) Get(id FlowID) (*Flow, bool) {
-	sft.lock.RLock()
-	defer sft.lock.RUnlock()
-	flow, ok := sft.flows[id]
-	return flow, ok
-}
-
-// Registers a flow. If shouldReplace = true, replace the current flow
-func (sft *FlowTracker) Register(id FlowID, flow *Flow, shouldReplace bool) (replaced bool) {
-	sft.lock.Lock()
-	defer sft.lock.Unlock()
-	currentFlow, ok := sft.flows[id]
-	if !ok {
-		sft.flows[id] = flow
-		return false
-	}
-
-	if shouldReplace && !isSameFlow(currentFlow, flow) {
-		sft.flows[id] = flow
-		return true
-	}
-	return false
-}
-
-// Unregisters a flow. If force = true, delete it even if it maps to a different flow
-func (sft *FlowTracker) Unregister(id FlowID, flow *Flow, force bool) (forceDeleted bool) {
-	sft.lock.Lock()
-	defer sft.lock.Unlock()
-	currentFlow, ok := sft.flows[id]
-	if !ok {
-		return false
-	}
-	if isSameFlow(currentFlow, flow) {
-		delete(sft.flows, id)
-		return false
-	}
-	if force {
-		delete(sft.flows, id)
-		return true
-	}
-	return false
-}

packet/funnel.go

@@ -0,0 +1,191 @@
+package packet
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/netip"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+var (
+	ErrFunnelNotFound = errors.New("funnel not found")
+)
+
+// Funnel is an abstraction to pipe from 1 src to 1 or more destinations
+type Funnel interface {
+	// SendToDst sends a raw packet to a destination
+	SendToDst(dst netip.Addr, pk RawPacket) error
+	// ReturnToSrc returns a raw packet to the source
+	ReturnToSrc(pk RawPacket) error
+	// LastActive returns the last time SendToDst or ReturnToSrc is called
+	LastActive() time.Time
+	// Close closes the funnel. Further call to SendToDst or ReturnToSrc should return an error
+	Close() error
+	// Equal compares if 2 funnels are equivalent
+	Equal(other Funnel) bool
+}
+
+// FunnelUniPipe is a unidirectional pipe for sending raw packets
+type FunnelUniPipe interface {
+	// SendPacket sends a packet to/from the funnel. It must not modify the packet,
+	// and after return it must not read the packet
+	SendPacket(dst netip.Addr, pk RawPacket) error
+	Close() error
+}
+
+// RawPacketFunnel is an implementation of Funnel that sends raw packets. It can be embedded in other structs to
+// satisfy the Funnel interface.
+type RawPacketFunnel struct {
+	Src netip.Addr
+	// last active unix time. Unit is seconds
+	lastActive int64
+	sendPipe   FunnelUniPipe
+	returnPipe FunnelUniPipe
+}
+
+func NewRawPacketFunnel(src netip.Addr, sendPipe, returnPipe FunnelUniPipe) *RawPacketFunnel {
+	return &RawPacketFunnel{
+		Src:        src,
+		lastActive: time.Now().Unix(),
+		sendPipe:   sendPipe,
+		returnPipe: returnPipe,
+	}
+}
+
+func (rpf *RawPacketFunnel) SendToDst(dst netip.Addr, pk RawPacket) error {
+	rpf.updateLastActive()
+	return rpf.sendPipe.SendPacket(dst, pk)
+}
+
+func (rpf *RawPacketFunnel) ReturnToSrc(pk RawPacket) error {
+	rpf.updateLastActive()
+	return rpf.returnPipe.SendPacket(rpf.Src, pk)
+}
+
+func (rpf *RawPacketFunnel) updateLastActive() {
+	atomic.StoreInt64(&rpf.lastActive, time.Now().Unix())
+}
+
+func (rpf *RawPacketFunnel) LastActive() time.Time {
+	lastActive := atomic.LoadInt64(&rpf.lastActive)
+	return time.Unix(lastActive, 0)
+}
+
+func (rpf *RawPacketFunnel) Close() error {
+	sendPipeErr := rpf.sendPipe.Close()
+	returnPipeErr := rpf.returnPipe.Close()
+	if sendPipeErr != nil {
+		return sendPipeErr
+	}
+	if returnPipeErr != nil {
+		return returnPipeErr
+	}
+	return nil
+}
+
+func (rpf *RawPacketFunnel) Equal(other Funnel) bool {
+	otherRawFunnel, ok := other.(*RawPacketFunnel)
+	if !ok {
+		return false
+	}
+	if rpf.Src != otherRawFunnel.Src {
+		return false
+	}
+	if rpf.sendPipe != otherRawFunnel.sendPipe {
+		return false
+	}
+	if rpf.returnPipe != otherRawFunnel.returnPipe {
+		return false
+	}
+	return true
+}
+
+// FunnelID represents a key type that can be used by FunnelTracker
+type FunnelID interface {
+	// Type returns the name of the type that implements the FunnelID
+	Type() string
+	fmt.Stringer
+}
+
+// FunnelTracker tracks funnel from the perspective of eyeball to origin
+type FunnelTracker struct {
+	lock    sync.RWMutex
+	funnels map[FunnelID]Funnel
+}
+
+func NewFunnelTracker() *FunnelTracker {
+	return &FunnelTracker{
+		funnels: make(map[FunnelID]Funnel),
+	}
+}
+
+func (ft *FunnelTracker) ScheduleCleanup(ctx context.Context, idleTimeout time.Duration) {
+	checkIdleTicker := time.NewTicker(idleTimeout)
+	defer checkIdleTicker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-checkIdleTicker.C:
+			ft.cleanup(idleTimeout)
+		}
+	}
+}
+
+func (ft *FunnelTracker) cleanup(idleTimeout time.Duration) {
+	ft.lock.Lock()
+	defer ft.lock.Unlock()
+
+	now := time.Now()
+	for id, funnel := range ft.funnels {
+		lastActive := funnel.LastActive()
+		if now.After(lastActive.Add(idleTimeout)) {
+			funnel.Close()
+			delete(ft.funnels, id)
+		}
+	}
+}
+
+func (ft *FunnelTracker) Get(id FunnelID) (Funnel, bool) {
+	ft.lock.RLock()
+	defer ft.lock.RUnlock()
+	funnel, ok := ft.funnels[id]
+	return funnel, ok
+}
+
+// Registers a funnel. It replaces the current funnel.
+func (ft *FunnelTracker) Register(id FunnelID, funnel Funnel) (replaced bool) {
+	ft.lock.Lock()
+	defer ft.lock.Unlock()
+	currentFunnel, exists := ft.funnels[id]
+	if !exists {
+		ft.funnels[id] = funnel
+		return false
+	}
+	replaced = !currentFunnel.Equal(funnel)
+	if replaced {
+		currentFunnel.Close()
+	}
+	ft.funnels[id] = funnel
+	return replaced
+}
+
+// Unregisters a funnel if the funnel equals to the current funnel
+func (ft *FunnelTracker) Unregister(id FunnelID, funnel Funnel) (deleted bool) {
+	ft.lock.Lock()
+	defer ft.lock.Unlock()
+	currentFunnel, exists := ft.funnels[id]
+	if !exists {
+		return true
+	}
+	if currentFunnel.Equal(funnel) {
+		currentFunnel.Close()
+		delete(ft.funnels, id)
+		return true
+	}
+	return false
+}