[proxy] Keep custom TCP listeners alive after mapping batches (#6415)

proxy/server_test.go

@@ -3,14 +3,20 @@ package proxy
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io"
+	"net"
 	"testing"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/metric/noop"
+	"google.golang.org/grpc"
 
+	proxymetrics "github.com/netbirdio/netbird/proxy/internal/metrics"
+	"github.com/netbirdio/netbird/proxy/internal/types"
 	"github.com/netbirdio/netbird/shared/management/proto"
 )
 
@@ -202,3 +208,117 @@ func TestRedactMappingForLog_HandlesEmptyOrNilFields(t *testing.T) {
 	assert.Nil(t, redacted.Auth, "nil Auth must remain nil")
 	assert.Empty(t, redacted.Path, "empty Path must remain empty")
 }
+
+type statusUpdateOnlyClient struct {
+	proto.ProxyServiceClient
+}
+
+func (statusUpdateOnlyClient) SendStatusUpdate(context.Context, *proto.SendStatusUpdateRequest, ...grpc.CallOption) (*proto.SendStatusUpdateResponse, error) {
+	return &proto.SendStatusUpdateResponse{}, nil
+}
+
+func TestSetupTCPMappingBindsCustomListenPort(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	port := uint16(ln.Addr().(*net.TCPAddr).Port) //nolint:gosec // test port allocated by the OS
+	require.NoError(t, ln.Close())
+
+	meter, err := proxymetrics.New(context.Background(), noop.Meter{})
+	require.NoError(t, err)
+
+	srv := &Server{
+		Logger:      quietLifecycleLogger(),
+		mgmtClient:  statusUpdateOnlyClient{},
+		meter:       meter,
+		mainPort:    8443,
+		portRouters: make(map[uint16]*portRouter),
+		svcPorts:    make(map[types.ServiceID][]uint16),
+	}
+	t.Cleanup(func() {
+		srv.portMu.Lock()
+		for p, pr := range srv.portRouters {
+			pr.cancel()
+			require.NoError(t, pr.listener.Close())
+			delete(srv.portRouters, p)
+		}
+		srv.portMu.Unlock()
+		srv.portRouterWg.Wait()
+	})
+
+	mapping := &proto.ProxyMapping{
+		Type:       proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED,
+		Id:         "svc-tcp",
+		AccountId:  "acct-1",
+		Domain:     "ssh.example.com",
+		Mode:       "tcp",
+		ListenPort: int32(port),
+		Path: []*proto.PathMapping{
+			{Target: "10.0.0.5:22"},
+		},
+	}
+
+	require.NoError(t, srv.setupTCPMapping(context.Background(), mapping))
+
+	srv.portMu.RLock()
+	pr := srv.portRouters[port]
+	ports := append([]uint16(nil), srv.svcPorts[types.ServiceID("svc-tcp")]...)
+	srv.portMu.RUnlock()
+
+	require.NotNil(t, pr, "custom TCP mapping must create a per-port router")
+	assert.Equal(t, []uint16{port}, ports, "service must track the custom listen port for cleanup")
+
+	second, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
+	if err == nil {
+		_ = second.Close()
+	}
+	require.Error(t, err, "custom TCP listen port must be bound after setup")
+}
+
+func TestCustomTCPPortRouterOutlivesMappingBatchContext(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	port := uint16(ln.Addr().(*net.TCPAddr).Port) //nolint:gosec // test port allocated by the OS
+	require.NoError(t, ln.Close())
+
+	meter, err := proxymetrics.New(context.Background(), noop.Meter{})
+	require.NoError(t, err)
+
+	srvCtx, srvCancel := context.WithCancel(context.Background())
+	t.Cleanup(srvCancel)
+
+	srv := &Server{
+		ctx:         srvCtx,
+		Logger:      quietLifecycleLogger(),
+		meter:       meter,
+		mainPort:    8443,
+		portRouters: make(map[uint16]*portRouter),
+		svcPorts:    make(map[types.ServiceID][]uint16),
+	}
+	t.Cleanup(func() {
+		srv.portMu.Lock()
+		for p, pr := range srv.portRouters {
+			pr.cancel()
+			if err := pr.listener.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
+				require.NoError(t, err)
+			}
+			delete(srv.portRouters, p)
+		}
+		srv.portMu.Unlock()
+		srv.portRouterWg.Wait()
+	})
+
+	batchCtx, cancelBatch := context.WithCancel(context.Background())
+	_, err = srv.getOrCreatePortRouter(batchCtx, port)
+	require.NoError(t, err)
+
+	cancelBatch()
+
+	assert.Never(t, func() bool {
+		second, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
+		if err == nil {
+			_ = second.Close()
+			return true
+		}
+		return false
+	}, 200*time.Millisecond, 10*time.Millisecond, "custom TCP listener must outlive mapping-batch context cancellation")
+}