fix(envd): log net.Connections error in port scanner and add unit tests

Author: chill-czar

packages/envd/internal/port/forward_test.go

@@ -17,7 +17,7 @@ func TestStartForwarding_StopsOnClosedMessages(t *testing.T) {
 	t.Parallel()
 
 	l := zerolog.Nop()
-	scanner := NewScanner(time.Hour)
+	scanner := NewScanner(&l, time.Hour)
 	f := &Forwarder{
 		logger:            &l,
 		ports:             make(map[string]*PortToForward),

packages/envd/internal/port/scan.go

@@ -3,12 +3,14 @@ package port
 import (
 	"time"
 
+	"github.com/rs/zerolog"
 	"github.com/shirou/gopsutil/v4/net"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
 )
 
 type Scanner struct {
+	logger   *zerolog.Logger
 	scanExit chan struct{}
 	subs     *smap.Map[*ScannerSubscriber]
 	period   time.Duration
@@ -18,8 +20,9 @@ func (s *Scanner) Destroy() {
 	close(s.scanExit)
 }
 
-func NewScanner(period time.Duration) *Scanner {
+func NewScanner(logger *zerolog.Logger, period time.Duration) *Scanner {
 	return &Scanner{
+		logger:   logger,
 		period:   period,
 		subs:     smap.New[*ScannerSubscriber](),
 		scanExit: make(chan struct{}),
@@ -40,11 +43,17 @@ func (s *Scanner) ScanAndBroadcast() {
 
 	for {
 		// tcp monitors both ipv4 and ipv6 connections.
-		processes, _ := net.Connections("tcp")
-		for _, sub := range s.subs.Items() {
-			// Pass scanExit so a subscriber that stopped consuming can't park this
-			// loop on a send and keep it from reaching the select below.
-			sub.Signal(processes, s.scanExit)
+		processes, err := net.Connections("tcp")
+		if err != nil {
+			if s.logger != nil {
+				s.logger.Error().Err(err).Msg("Failed to scan open TCP connections")
+			}
+		} else {
+			for _, sub := range s.subs.Items() {
+				// Pass scanExit so a subscriber that stopped consuming can't park this
+				// loop on a send and keep it from reaching the select below.
+				sub.Signal(processes, s.scanExit)
+			}
 		}
 		select {
 		case <-s.scanExit:

packages/envd/internal/port/scan_test.go

@@ -0,0 +1,114 @@
+package port
+
+import (
+	"bytes"
+	"testing"
+	"time"
+
+	"github.com/rs/zerolog"
+	"github.com/shirou/gopsutil/v4/net"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewScanner(t *testing.T) {
+	buf := &bytes.Buffer{}
+	logger := zerolog.New(buf)
+
+	s := NewScanner(&logger, 50*time.Millisecond)
+	require.NotNil(t, s)
+	assert.NotNil(t, s.subs)
+	assert.Equal(t, 50*time.Millisecond, s.period)
+
+	sub := s.AddSubscriber("test-sub", &ScannerFilter{
+		IPs:   []string{"127.0.0.1"},
+		State: "LISTEN",
+	})
+	require.NotNil(t, sub)
+	assert.Equal(t, "test-sub", sub.id)
+}
+
+func TestScanAndBroadcastLifecycle(t *testing.T) {
+	buf := &bytes.Buffer{}
+	logger := zerolog.New(buf)
+
+	s := NewScanner(&logger, 10*time.Millisecond)
+	sub := s.AddSubscriber("sub-1", nil)
+
+	received := make(chan struct{}, 1)
+	stopDrain := make(chan struct{})
+	go func() {
+		for {
+			select {
+			case <-stopDrain:
+				return
+			case <-sub.Messages:
+				select {
+				case received <- struct{}{}:
+				default:
+				}
+			}
+		}
+	}()
+
+	done := make(chan struct{})
+	go func() {
+		s.ScanAndBroadcast()
+		close(done)
+	}()
+
+	select {
+	case <-received:
+		// Received at least one scan message
+	case <-time.After(1 * time.Second):
+		t.Fatal("timed out waiting for scan message")
+	}
+
+	s.Destroy()
+
+	select {
+	case <-done:
+		// Scanner goroutine exited promptly
+	case <-time.After(1 * time.Second):
+		t.Fatal("ScanAndBroadcast did not exit after Destroy()")
+	}
+
+	close(stopDrain)
+}
+
+func TestSubscriberFilter(t *testing.T) {
+	sub := NewScannerSubscriber("sub-filter", &ScannerFilter{
+		IPs:   []string{"127.0.0.1", "localhost"},
+		State: "LISTEN",
+	})
+
+	conns := []net.ConnectionStat{
+		{
+			Status: "LISTEN",
+			Laddr:  net.Addr{IP: "127.0.0.1", Port: 8080},
+		},
+		{
+			Status: "ESTABLISHED",
+			Laddr:  net.Addr{IP: "127.0.0.1", Port: 8080},
+		},
+		{
+			Status: "LISTEN",
+			Laddr:  net.Addr{IP: "192.168.1.1", Port: 8080},
+		},
+	}
+
+	exit := make(chan struct{})
+	defer close(exit)
+
+	go sub.Signal(conns, exit)
+
+	select {
+	case filtered := <-sub.Messages:
+		require.Len(t, filtered, 1)
+		assert.Equal(t, uint32(8080), filtered[0].Laddr.Port)
+		assert.Equal(t, "127.0.0.1", filtered[0].Laddr.IP)
+		assert.Equal(t, "LISTEN", filtered[0].Status)
+	case <-time.After(1 * time.Second):
+		t.Fatal("timed out waiting for filtered messages")
+	}
+}

packages/envd/internal/port/scanner_test.go

@@ -73,7 +73,7 @@ func TestScannerSubscriber_Signal_DeliversToSlowReceiver(t *testing.T) {
 func TestScanner_Destroy_ExitsWithStalledSubscriber(t *testing.T) {
 	t.Parallel()
 
-	scanner := NewScanner(10 * time.Millisecond)
+	scanner := NewScanner(nil, 10*time.Millisecond)
 	scanner.AddSubscriber("stalled-sub", nil)
 
 	scanDone := make(chan struct{})

packages/envd/main.go

@@ -303,11 +303,12 @@ func run() error {
 		IdleTimeout:  idleTimeout,
 	}
 
+	portLogger := l.With().Str("logger", "port-forwarder").Logger()
+
 	// Bind all open ports on 127.0.0.1 and localhost to the eth0 interface
-	portScanner := publicport.NewScanner(portScannerInterval)
+	portScanner := publicport.NewScanner(&portLogger, portScannerInterval)
 	defer portScanner.Destroy()
 
-	portLogger := l.With().Str("logger", "port-forwarder").Logger()
 	portForwarder := publicport.NewForwarder(&portLogger, portScanner, cgroupManager)
 	if resumeHandover {
 		// Re-adopt the socats carried across the upgrade before the forwarder's