feat: reuseport activator

activator/activator.go

@@ -15,6 +15,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"slices"
 	"strconv"
 	"strings"
 	"sync"
@@ -46,10 +47,11 @@ type Server struct {
 	forwardToTarget bool
 	targetAddr      string
 	kubeletAddr     *netip.Addr
+	lns             Listeners
 }
 
 type ConnHook func(net.Conn) (conn net.Conn, cont bool, err error)
-type RestoreHook func() error
+type RestoreHook func() (int, error)
 
 type Option func(s *Server)
 
@@ -59,13 +61,15 @@ func SetTargetAddr(addr string) Option {
 	}
 }
 
-func NewServer(ctx context.Context, nn ns.NetNS, opts ...Option) (*Server, error) {
+func NewServer(ctx context.Context, nn ns.NetNS, connHook ConnHook, restoreHook RestoreHook, opts ...Option) (*Server, error) {
 	s := &Server{
 		quit:           make(chan any),
 		connectTimeout: time.Second * 5,
 		proxyTimeout:   time.Second * 5,
 		ns:             nn,
 		sandboxPid:     parsePidFromNetNS(nn),
+		connHook:       connHook,
+		restoreHook:    restoreHook,
 	}
 	return s, nil
 }
@@ -94,10 +98,8 @@ var (
 	DefaultIfaces  = []string{IfaceLoopback, IfaceETH0}
 )
 
-func (s *Server) Start(ctx context.Context, connHook ConnHook, restoreHook RestoreHook, ports ...uint16) error {
-	s.connHook = connHook
-	s.restoreHook = restoreHook
-	s.ports = ports
+func (s *Server) Start(ctx context.Context, _ int, listeners Listeners, skipStart bool) error {
+	s.ports = listeners.Ports()
 
 	if err := s.loadPinnedMaps(); err != nil {
 		return err
@@ -121,10 +123,35 @@ func (s *Server) Start(ctx context.Context, connHook ConnHook, restoreHook Resto
 		}
 	}
 
+	if skipStart {
+		if err := s.Reset(); err != nil {
+			return err
+		}
+	}
+
 	s.started = true
 	return nil
 }
 
+func (s *Server) GetListeners(ctx context.Context, pid int) []Listener {
+	if s.lns != nil {
+		return s.lns
+	}
+	lns, err := GetListenersOfPID(ctx, pid)
+	if err != nil {
+		return s.lns
+	}
+	listeners := []Listener{}
+	for _, ln := range lns {
+		if !slices.Contains(s.ports, ln.Port) {
+			continue
+		}
+		listeners = append(listeners, ln)
+	}
+	s.lns = listeners
+	return s.lns
+}
+
 const AttachActivatorFlag = "-zeropod-attach-activator"
 
 // AttachExec attaches the activator using exec on itself.
@@ -175,12 +202,13 @@ func (s *Server) SetPeekBufferSize(size int) {
 
 // ForwardToTarget instructs the activator to forward any incoming traffic to
 // the specified address. The connHook and restoreHook will both be disabled.
-func (s *Server) ForwardToTarget(addr string) {
+func (s *Server) ForwardToTarget(_ context.Context, addr string) error {
 	// disable hooks
 	s.connHook = func(c net.Conn) (net.Conn, bool, error) { return c, true, nil }
-	s.restoreHook = func() error { return nil }
+	s.restoreHook = func() (int, error) { return 0, nil }
 	s.targetAddr = addr
 	s.forwardToTarget = true
+	return nil
 }
 
 func (s *Server) listen(ctx context.Context, port uint16) (int, error) {
@@ -294,7 +322,7 @@ func (s *Server) handleConnection(ctx context.Context, netConn net.Conn, port ui
 		}
 	}()
 
-	if err := s.restoreHook(); err != nil {
+	if _, err := s.restoreHook(); err != nil {
 		log.G(ctx).Errorf("restoreHook: %s", err)
 		return
 	}
@@ -509,7 +537,7 @@ func (s *Server) LastActivity(port uint16) (time.Time, error) {
 		return time.Time{}, NoActivityRecordedErr{}
 	}
 
-	return convertBPFTime(val)
+	return ConvertBPFTime(val)
 }
 
 func (s *Server) initActivityTracker() error {
@@ -526,9 +554,9 @@ func netNSPath(pid int) string {
 	return fmt.Sprintf("/proc/%d/ns/net", pid)
 }
 
-// convertBPFTime takes the value of bpf_ktime_get_ns and converts it to a
+// ConvertBPFTime takes the value of bpf_ktime_get_ns and converts it to a
 // time.Time.
-func convertBPFTime(t uint64) (time.Time, error) {
+func ConvertBPFTime(t uint64) (time.Time, error) {
 	b, err := getBootTimeNS()
 	if err != nil {
 		return time.Time{}, err

activator/activator_test.go

@@ -145,7 +145,12 @@ func TestActivator(t *testing.T) {
 	for name, tc := range tests {
 		t.Run(name, func(t *testing.T) {
 			ctx, cancel := context.WithCancel(context.Background())
-			s, err := NewServer(ctx, nn)
+			if tc.connHook == nil {
+				tc.connHook = func(c net.Conn) (net.Conn, bool, error) {
+					return c, true, nil
+				}
+			}
+			s, err := NewServer(ctx, nn, tc.connHook, func() (int, error) { return 0, nil })
 			require.NoError(t, err)
 
 			port, err := freePort()
@@ -232,56 +237,52 @@ func startServer(t *testing.T, ctx context.Context, s *Server, port uint16, tc *
 	ts := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		fmt.Fprint(w, response)
 	}))
-	if tc.connHook == nil {
-		tc.connHook = func(c net.Conn) (net.Conn, bool, error) {
-			return c, true, nil
-		}
-	}
 
 	once := sync.Once{}
 	loopIterations := 0
-	err := s.Start(
-		ctx,
-		tc.connHook,
-		func() error {
-			if tc.loopConnection {
-				loopIterations += 1
-				if loopIterations > 10 {
-					t.Error("loop detection failed")
-					return fmt.Errorf("loop detection failed")
-				}
-				// return nil
+	s.restoreHook = func() (int, error) {
+		if tc.loopConnection {
+			loopIterations += 1
+			if loopIterations > 10 {
+				t.Error("loop detection failed")
+				return 0, fmt.Errorf("loop detection failed")
 			}
-			once.Do(func() {
-				// simulate a delay until our server is started
-				time.Sleep(time.Millisecond * 200)
-				network := "tcp4"
-				if tc.ipv6 {
-					network = "tcp6"
-				}
-				l, err := net.Listen(network, fmt.Sprintf(":%d", port))
-				require.NoError(t, err)
+			// return nil
+		}
+		once.Do(func() {
+			// simulate a delay until our server is started
+			time.Sleep(time.Millisecond * 200)
+			network := "tcp4"
+			if tc.ipv6 {
+				network = "tcp6"
+			}
+			l, err := net.Listen(network, fmt.Sprintf(":%d", port))
+			require.NoError(t, err)
 
-				if !tc.loopConnection {
-					if err := s.DisableRedirects(); err != nil {
-						t.Errorf("could not disable redirects: %s", err)
-					}
+			if !tc.loopConnection {
+				if err := s.DisableRedirects(); err != nil {
+					t.Errorf("could not disable redirects: %s", err)
 				}
+			}
 
-				// replace listener of server
-				ts.Listener.Close()
-				ts.Listener = l
-				ts.Start()
-				t.Logf("listening on %s", l.Addr().String())
+			// replace listener of server
+			ts.Listener.Close()
+			ts.Listener = l
+			ts.Start()
+			t.Logf("listening on %s", l.Addr().String())
 
-				t.Cleanup(func() {
-					l.Close()
-					ts.Close()
-				})
+			t.Cleanup(func() {
+				l.Close()
+				ts.Close()
 			})
-			return nil
-		},
-		port,
+		})
+		return 0, nil
+	}
+	err := s.Start(
+		ctx,
+		os.Getpid(),
+		Listeners{{Port: port}},
+		false,
 	)
 	require.NoError(t, err)
 	s.enableRedirect(port)

activator/interface.go

@@ -0,0 +1,20 @@
+package activator
+
+import (
+	"context"
+	"time"
+)
+
+type Activator interface {
+	Start(ctx context.Context, pid int, listeners Listeners, skipStart bool) error
+	Started() bool
+	Reset() error
+	DisableRedirects() error
+	AttachExec() error
+	SetProxyTimeout(d time.Duration)
+	SetConnectTimeout(d time.Duration)
+	LastActivity(port uint16) (time.Time, error)
+	Stop(ctx context.Context)
+	GetListeners(ctx context.Context, pid int) []Listener
+	ForwardToTarget(ctx context.Context, addr string) error
+}