[tunnel] Prepare for multi-conn routing and egress NAT

* Refactor Router interface in prep for adding ECMP routing. * Split implementation into client and server flavors - each side will use its own implementation. * Some minor naming refactor.

Author: dilyevsky

api/core/v1alpha/tunnel_types.go

@@ -25,7 +25,21 @@ type TunnelNode struct {
 Status TunnelNodeStatus `json:"status,omitempty"`
 }
 
+type EgressGatewaySpec struct {
+// Whether the egress gateway is enabled. Default is false.
+// When enabled, the egress gateway will be used to route traffic from the tunnel
+// node to the internet. Traffic will be SNAT'ed.
+// +optional
+Enabled bool `json:"enabled,omitempty"`
+}
+
 type TunnelNodeSpec struct {
+// Configures Egress Gateway mode on the Tunnel Node. In this mode, the Tunnel
+// Node acts as a gateway for outbound connections originating from the
+// Agent side in addition to its default mode (where the connections arrive in the
+// direction of the Agent).
+// +optional
+EgressGateway *EgressGatewaySpec `json:"egressGateway,omitempty"`
 }
 
 type AgentStatus struct {

pkg/tunnel/client.go

@@ -278,10 +278,10 @@ func (c *TunnelClient) Start(ctx context.Context) error {
 }
 
 if c.options.mode == TunnelClientModeKernel {
-c.router, err = router.NewNetlinkRouter(routerOpts...)
-if err != nil {
-return fmt.Errorf("failed to create kernel router: %w", err)
-}
+//c.router, err = router.NewNetlinkRouter(routerOpts...)
+//if err != nil {
+//return fmt.Errorf("failed to create kernel router: %w", err)
+//}
 } else if c.options.mode == TunnelClientModeUser {
 c.router, err = router.NewNetstackRouter(routerOpts...)
 if err != nil {
@@ -298,9 +298,8 @@ func (c *TunnelClient) Start(ctx context.Context) error {
 for _, prefix := range route.Prefixes() {
 slog.Info("Adding route", slog.String("prefix", prefix.String()))
 
-_, _, err := c.router.AddPeer(prefix, c.conn)
-if err != nil {
-return fmt.Errorf("failed to add peer route %s: %w", prefix.String(), err)
+if err := c.router.Add(prefix, c.conn); err != nil {
+return fmt.Errorf("failed to add route %s: %w", prefix.String(), err)
 }
 }
 }

pkg/tunnel/connection/conn.go

@@ -0,0 +1,12 @@
+package connection
+
+import "io"
+
+// Connection is a simple interface implemented by connect-ip-go and custom
+// connection types.
+type Connection interface {
+io.Closer
+
+ReadPacket([]byte) (int, error)
+WritePacket([]byte) ([]byte, error)
+}

pkg/tunnel/connection/muxed_connection.go renamed to pkg/tunnel/connection/muxed_conn.go

@@ -3,7 +3,6 @@ package connection
 import (
 "errors"
 "fmt"
-"io"
 "log/slog"
 "net"
 "net/netip"
@@ -15,19 +14,8 @@ import (
 "github.com/apoxy-dev/apoxy/pkg/netstack"
 )
 
-// Connection is a simple interface implemented by connect-ip-go and custom
-// connection types.
-type Connection interface {
-io.Closer
-ReadPacket([]byte) (int, error)
-WritePacket([]byte) ([]byte, error)
-}
-
-var _ Connection = (*MuxedConnection)(nil)
-
-// MuxedConnection is a connection that multiplexes multiple downstream
-// connections over a single virtual connection.
-type MuxedConnection struct {
+// MuxedConn multiplexes multiple connection.Conn objects.
+type MuxedConn struct {
 // Maps tunnel destination address to CONNECT-IP connection.
 conns *triemap.TrieMap[Connection]
 incomingPackets chan *[]byte
@@ -37,8 +25,9 @@ type MuxedConnection struct {
 closed atomic.Bool
 }
 
-func NewMuxedConnection() *MuxedConnection {
-return &MuxedConnection{
+// NewMuxedConn creates a new muxedConn.
+func NewMuxedConn() *MuxedConn {
+return &MuxedConn{
 conns: triemap.New[Connection](),
 incomingPackets: make(chan *[]byte, 100),
 packetBufferPool: sync.Pool{
@@ -50,7 +39,7 @@ func NewMuxedConnection() *MuxedConnection {
 }
 }
 
-func (m *MuxedConnection) AddConnection(prefix netip.Prefix, conn Connection) {
+func (m *MuxedConn) AddConnection(prefix netip.Prefix, conn Connection) {
 if prefix.IsValid() && prefix.Addr().Is6() {
 m.conns.Insert(prefix, conn)
 go m.readPackets(conn)
@@ -59,7 +48,7 @@ func (m *MuxedConnection) AddConnection(prefix netip.Prefix, conn Connection) {
 }
 }
 
-func (m *MuxedConnection) RemoveConnection(prefix netip.Prefix) error {
+func (m *MuxedConn) RemoveConnection(prefix netip.Prefix) error {
 // Has the connection already been closed?
 if m.closed.Load() {
 // Then this becomes a no-op.
@@ -85,7 +74,7 @@ func (m *MuxedConnection) RemoveConnection(prefix netip.Prefix) error {
 return nil
 }
 
-func (m *MuxedConnection) Prefixes() []netip.Prefix {
+func (m *MuxedConn) Prefixes() []netip.Prefix {
 var prefixes []netip.Prefix
 m.conns.ForEach(func(prefix netip.Prefix, value Connection) bool {
 prefixes = append(prefixes, prefix)
@@ -94,7 +83,7 @@ func (m *MuxedConnection) Prefixes() []netip.Prefix {
 return prefixes
 }
 
-func (m *MuxedConnection) Close() error {
+func (m *MuxedConn) Close() error {
 var firstErr error
 m.closeOnce.Do(func() {
 // Close all connections in the map.
@@ -122,7 +111,7 @@ func (m *MuxedConnection) Close() error {
 return firstErr
 }
 
-func (m *MuxedConnection) ReadPacket(pkt []byte) (int, error) {
+func (m *MuxedConn) ReadPacket(pkt []byte) (int, error) {
 p, ok := <-m.incomingPackets
 if !ok {
 return 0, net.ErrClosed
@@ -138,7 +127,7 @@ func (m *MuxedConnection) ReadPacket(pkt []byte) (int, error) {
 return n, nil
 }
 
-func (m *MuxedConnection) WritePacket(pkt []byte) ([]byte, error) {
+func (m *MuxedConn) WritePacket(pkt []byte) ([]byte, error) {
 slog.Debug("Write packet to connection", slog.Int("bytes", len(pkt)))
 
 var dstIP netip.Addr
@@ -170,7 +159,7 @@ func (m *MuxedConnection) WritePacket(pkt []byte) ([]byte, error) {
 return conn.WritePacket(pkt)
 }
 
-func (m *MuxedConnection) readPackets(conn Connection) {
+func (m *MuxedConn) readPackets(conn Connection) {
 for {
 pkt := m.packetBufferPool.Get().(*[]byte)
 

pkg/tunnel/connection/muxed_connection_test.go renamed to pkg/tunnel/connection/muxed_conn_test.go

@@ -40,7 +40,7 @@ func (m *MockConnection) Close() error {
 
 func TestMuxedConnection(t *testing.T) {
 t.Run("Add and Remove Connection", func(t *testing.T) {
-mux := connection.NewMuxedConnection()
+mux := connection.NewMuxedConn()
 mockConn := new(MockConnection)
 mockConn.On("ReadPacket", mock.Anything).Return(0, []byte{}, nil).Maybe()
 mockConn.On("Close").Return(nil).Once()
@@ -56,14 +56,14 @@ func TestMuxedConnection(t *testing.T) {
 })
 
 t.Run("Remove Connection - Invalid Prefix", func(t *testing.T) {
-mux := connection.NewMuxedConnection()
+mux := connection.NewMuxedConn()
 prefix := netip.MustParsePrefix("192.0.2.0/24")
 err := mux.RemoveConnection(prefix)
 assert.Error(t, err)
 })
 
 t.Run("WritePacket - Success", func(t *testing.T) {
-mux := connection.NewMuxedConnection()
+mux := connection.NewMuxedConn()
 mockConn := new(MockConnection)
 mockConn.On("ReadPacket", mock.Anything).Return(0, []byte{}, nil).Maybe()
 
@@ -83,7 +83,7 @@ func TestMuxedConnection(t *testing.T) {
 })
 
 t.Run("WritePacket - No Connection Found", func(t *testing.T) {
-mux := connection.NewMuxedConnection()
+mux := connection.NewMuxedConn()
 
 pkt := make([]byte, 40)
 pkt[0] = 0x60
@@ -95,7 +95,7 @@ func TestMuxedConnection(t *testing.T) {
 })
 
 t.Run("ReadPacket - Success", func(t *testing.T) {
-mux := connection.NewMuxedConnection()
+mux := connection.NewMuxedConn()
 mockConn := new(MockConnection)
 
 expected := []byte("hello")
@@ -115,7 +115,7 @@ func TestMuxedConnection(t *testing.T) {
 })
 
 t.Run("ReadPacket - Closed Channel", func(t *testing.T) {
-mux := connection.NewMuxedConnection()
+mux := connection.NewMuxedConn()
 _ = mux.Close()
 
 buf := make([]byte, 1500)
@@ -125,7 +125,7 @@ func TestMuxedConnection(t *testing.T) {
 })
 
 t.Run("Close - All Connections", func(t *testing.T) {
-mux := connection.NewMuxedConnection()
+mux := connection.NewMuxedConn()
 mockConn := new(MockConnection)
 mockConn.On("ReadPacket", mock.Anything).Return(0, []byte{}, nil).Maybe()
 mockConn.On("Close").Return(nil).Once()

pkg/tunnel/router/netstack.go

@@ -22,9 +22,10 @@ var (
 )
 
 // NetstackRouter implements Router using a userspace network stack.
+// This router can be used for both client and server sides.
 type NetstackRouter struct {
 tunDev *netstack.TunDevice
-mux *connection.MuxedConnection
+mux *connection.MuxedConn
 proxy *socksproxy.ProxyServer
 localAddresses []netip.Prefix
 resolveConf *network.ResolveConfig
@@ -48,7 +49,7 @@ func NewNetstackRouter(opts ...Option) (*NetstackRouter, error) {
 
 return &NetstackRouter{
 tunDev: tunDev,
-mux: connection.NewMuxedConnection(),
+mux: connection.NewMuxedConn(),
 proxy: proxy,
 localAddresses: options.localAddresses,
 resolveConf: options.resolveConf,
@@ -105,23 +106,42 @@ func (r *NetstackRouter) Start(ctx context.Context) error {
 return g.Wait()
 }
 
-// AddPeer adds a peer route to the tunnel.
-func (r *NetstackRouter) AddPeer(peer netip.Prefix, conn connection.Connection) (netip.Addr, []netip.Prefix, error) {
-r.mux.AddConnection(peer, conn)
-return peer.Addr(), r.localAddresses, nil
+// Add adds a dst route to the tunnel.
+func (r *NetstackRouter) Add(dst netip.Prefix, conn connection.Connection) error {
+r.mux.AddConnection(dst, conn)
+return nil
+}
+
+// Del removes a dst route from the tunnel.
+func (r *NetstackRouter) Del(dst netip.Prefix, _ string) error {
+return r.DelAll(dst) // TODO: implement multi-conn routes.
 }
 
-// RemovePeer removes a peer route from the tunnel.
-func (r *NetstackRouter) RemovePeer(peer netip.Prefix) error {
-if err := r.mux.RemoveConnection(peer); err != nil {
+// DelAll removes all routes for the dst.
+func (r *NetstackRouter) DelAll(dst netip.Prefix) error {
+if err := r.mux.RemoveConnection(dst); err != nil {
 slog.Error("failed to remove connection", slog.Any("error", err))
 }
 
 return nil
 }
 
+// ListRoutes returns a list of all routes in the tunnel.
+func (r *NetstackRouter) ListRoutes() ([]TunnelRoute, error) {
+ps := r.mux.Prefixes()
+rts := make([]TunnelRoute, 0, len(ps))
+for _, p := range ps {
+rts = append(rts, TunnelRoute{
+Dst: p,
+// TODO: Add connID,
+State: TunnelRouteStateActive,
+})
+}
+return rts, nil
+}
+
 // GetMuxedConnection returns the muxed connection for adding/removing connections.
-func (r *NetstackRouter) GetMuxedConnection() *connection.MuxedConnection {
+func (r *NetstackRouter) GetMuxedConnection() *connection.MuxedConn {
 return r.mux
 }
 

pkg/tunnel/router/netstack_test.go

@@ -36,12 +36,11 @@ func TestNetstackRouter(t *testing.T) {
 
 // Test AddPeer
 prefix := netip.MustParsePrefix("fd00::1/128")
-conn := connection.NewMuxedConnection()
-_, _, err = r.AddPeer(prefix, conn)
-require.NoError(t, err)
+conn := connection.NewMuxedConn()
+require.NoError(t, r.Add(prefix, conn))
 
 // Test RemovePeer
-err = r.RemovePeer(prefix)
+err = r.DelAll(prefix)
 require.NoError(t, err)
 
 // Test Close

pkg/tunnel/router/router.go

@@ -8,6 +8,19 @@ import (
 "github.com/apoxy-dev/apoxy/pkg/tunnel/connection"
 )
 
+type TunnelRouteState int
+
+const (
+TunnelRouteStateActive TunnelRouteState = iota
+TunnelRouteStateDraining
+)
+
+type TunnelRoute struct {
+Dst netip.Prefix
+TunID string
+State TunnelRouteState
+}
+
 // Router is an interface for managing tunnel routing.
 type Router interface {
 io.Closer
@@ -16,12 +29,24 @@ type Router interface {
 // It's a blocking call that should be run in a separate goroutine.
 Start(ctx context.Context) error
 
-// AddPeer adds a peer route to the tunnel. Returns the list of IP prefixes
-// to be advertised to the peer (if none returned, no prefixes should be advertised).
-AddPeer(peer netip.Prefix, conn connection.Connection) (privAddr netip.Addr, routes []netip.Prefix, err error)
+// AddRoute adds a dst prefix to be routed through the given tunnel connection.
+// If multiple tunnels are provided, the router will distribute traffic across them
+// using ECMP (Equal Cost Multi-Path Routing) hashing.
+Add(dst netip.Prefix, tun connection.Connection) error
+
+// Del removes a routing associations for a given destination prefix and Connection name.
+// New matching flows will stop being routed through the tunnel immediately while
+// existing flows may continue to use the tunnel for some draining period before
+// getting re-routed via a different tunnel or dropped (if no tunnel is available for
+// the given dst).
+Del(dst netip.Prefix, name string) error
+
+// DelAll removes all routing associations for a given destination prefix.
+// Connections will be drained same way as Del.
+DelAll(dst netip.Prefix) error
 
-// RemovePeer removes a peer route from the tunnel identified by the given prefix.
-RemovePeer(peer netip.Prefix) error
+// ListRoutes returns a list of all routes currently managed by the router.
+ListRoutes() ([]TunnelRoute, error)
 
 // LocalAddresses returns the list of local addresses that are assigned to the router.
 LocalAddresses() ([]netip.Prefix, error)