server: validate control pool counts (#5459)

server/control.go

@@ -17,6 +17,7 @@ package server
 import (
 	"context"
 	"fmt"
+	"math"
 	"net"
 	"runtime/debug"
 	"sync"
@@ -45,6 +46,8 @@ type ControlID uint64
 
 var nextControlID atomic.Uint64
 
+const workConnPoolCapacityOffset = 10
+
 type controlEntry struct {
 	ctl *Control
 	id  ControlID
@@ -431,10 +434,27 @@ type Control struct {
 }
 
 func NewControl(ctx context.Context, sessionCtx *SessionContext) (*Control, error) {
-	poolCount := min(sessionCtx.LoginMsg.PoolCount, int(sessionCtx.ServerCfg.Transport.MaxPoolCount))
+	if sessionCtx.LoginMsg.PoolCount < 0 {
+		return nil, fmt.Errorf("invalid pool count %d, must be non-negative", sessionCtx.LoginMsg.PoolCount)
+	}
+	if sessionCtx.ServerCfg.Transport.MaxPoolCount < 0 {
+		return nil, fmt.Errorf(
+			"invalid max pool count %d, must be non-negative",
+			sessionCtx.ServerCfg.Transport.MaxPoolCount,
+		)
+	}
+	effectivePoolCount := min(int64(sessionCtx.LoginMsg.PoolCount), sessionCtx.ServerCfg.Transport.MaxPoolCount)
+	maxPoolCountForChannel := int64(math.MaxInt) - int64(workConnPoolCapacityOffset)
+	if effectivePoolCount > maxPoolCountForChannel {
+		return nil, fmt.Errorf(
+			"invalid effective pool count %d, cannot safely add %d for work connection pool capacity",
+			effectivePoolCount, workConnPoolCapacityOffset,
+		)
+	}
+	poolCount := int(effectivePoolCount)
 	ctl := &Control{
 		sessionCtx:    sessionCtx,
-		workConnCh:    make(chan *proxy.WorkConn, poolCount+10),
+		workConnCh:    make(chan *proxy.WorkConn, poolCount+workConnPoolCapacityOffset),
 		proxies:       make(map[string]proxy.Proxy),
 		poolCount:     poolCount,
 		portsUsedNum:  0,

server/control_test.go

@@ -17,6 +17,7 @@ package server
 import (
 	"context"
 	"errors"
+	"math"
 	"net"
 	"os"
 	"sync"
@@ -52,6 +53,57 @@ func TestControlPendingReplacementFinishesWithoutStarting(t *testing.T) {
 	require.Equal(t, int64(0), metrics.closedClients())
 }
 
+func TestNewControlPoolCountBoundaries(t *testing.T) {
+	for _, tc := range []struct {
+		name          string
+		poolCount     int
+		maxPoolCount  int64
+		wantErr       string
+		wantPoolCount int
+		wantCapacity  int
+	}{
+		{name: "negative pool count below offset", poolCount: -11, maxPoolCount: 5, wantErr: "invalid pool count"},
+		{name: "negative pool count at offset", poolCount: -10, maxPoolCount: 5, wantErr: "invalid pool count"},
+		{name: "negative pool count", poolCount: -1, maxPoolCount: 5, wantErr: "invalid pool count"},
+		{name: "zero pool count", poolCount: 0, maxPoolCount: 5, wantPoolCount: 0, wantCapacity: 10},
+		{name: "pool count capped", poolCount: 10, maxPoolCount: 5, wantPoolCount: 5, wantCapacity: 15},
+		{name: "maximum int pool count capped", poolCount: math.MaxInt, maxPoolCount: 5, wantPoolCount: 5, wantCapacity: 15},
+		{name: "negative maximum", poolCount: 1, maxPoolCount: -1, wantErr: "invalid max pool count"},
+		{name: "maximum int64 with small client pool", poolCount: 1, maxPoolCount: math.MaxInt64, wantPoolCount: 1, wantCapacity: 11},
+		{name: "maximum int client and server overflow", poolCount: math.MaxInt, maxPoolCount: math.MaxInt64, wantErr: "cannot safely add"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			conn := newDeadlineReadConn()
+			msgConn := msg.NewConn(conn, msg.NewV1ReadWriter(conn))
+			cfg := &v1.ServerConfig{}
+			cfg.Transport.MaxPoolCount = tc.maxPoolCount
+
+			ctl, err := NewControl(context.Background(), &SessionContext{
+				RC:            &controller.ResourceController{},
+				PxyManager:    proxy.NewManager(),
+				PluginManager: plugin.NewManager(),
+				AuthVerifier:  auth.AlwaysPassVerifier,
+				Conn:          msgConn,
+				LoginMsg: &msg.Login{
+					RunID:     "pool-count-run",
+					PoolCount: tc.poolCount,
+				},
+				ServerCfg: cfg,
+			})
+			if tc.wantErr != "" {
+				require.Nil(t, ctl)
+				require.ErrorContains(t, err, tc.wantErr)
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, tc.wantPoolCount, ctl.poolCount)
+			require.Equal(t, tc.wantCapacity, cap(ctl.workConnCh))
+			require.NoError(t, ctl.Close())
+		})
+	}
+}
+
 func TestControlRunningReplacementFinishesInWorker(t *testing.T) {
 	clientRegistry := registry.NewClientRegistry()
 	manager := NewControlManager(clientRegistry)

server/service_test.go

@@ -17,6 +17,7 @@ package server
 import (
 	"context"
 	"errors"
+	"math"
 	"net"
 	"net/http"
 	"runtime"
@@ -631,6 +632,48 @@ func TestServiceRegisterControlRejectsInvalidCodecSelection(t *testing.T) {
 	}
 }
 
+func TestServiceRegisterControlPoolCountBoundaries(t *testing.T) {
+	for _, tc := range []struct {
+		name      string
+		poolCount int
+		wantErr   bool
+		wantPool  int
+	}{
+		{name: "less than channel offset", poolCount: -11, wantErr: true},
+		{name: "channel offset", poolCount: -10, wantErr: true},
+		{name: "negative", poolCount: -1, wantErr: true},
+		{name: "zero", poolCount: 0, wantPool: 0},
+		{name: "capped by server maximum", poolCount: 10, wantPool: 5},
+		{name: "maximum int capped", poolCount: math.MaxInt, wantPool: 5},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			svr := newControlTestService(t)
+			svr.cfg.Transport.MaxPoolCount = 5
+			conn := newDeadlineReadConn()
+			msgConn := msg.NewConn(conn, msg.NewV1ReadWriter(conn))
+			const timestamp = int64(1)
+
+			ctl, err := svr.RegisterControl(msgConn, &msg.Login{
+				RunID:        "pool-count-run",
+				ClientID:     "client",
+				Timestamp:    timestamp,
+				PrivilegeKey: util.GetAuthKey("", timestamp),
+				PoolCount:    tc.poolCount,
+			}, false, wire.ProtocolV1, "")
+			if tc.wantErr {
+				require.Nil(t, ctl)
+				require.ErrorContains(t, err, "unexpected error when creating new controller")
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, tc.wantPool, ctl.poolCount)
+			require.NoError(t, ctl.Close())
+			waitForControlDone(t, ctl)
+		})
+	}
+}
+
 func TestServiceWorkConnRoutingRejectsLostGeneration(t *testing.T) {
 	for _, action := range []string{"replace", "close"} {
 		t.Run(action, func(t *testing.T) {