From 36b742a7e9a81ed782363dd19efd77b7030fbd0f Mon Sep 17 00:00:00 2001 From: naiba Date: Thu, 1 Oct 2026 04:28:34 +0000 Subject: [PATCH] fix: harden dashboard security boundaries --- .../controller/credential_redaction_test.go | 44 ++++++---- cmd/dashboard/controller/ddns.go | 10 ++- cmd/dashboard/controller/fm.go | 14 ++++ cmd/dashboard/controller/mcp_test.go | 7 +- cmd/dashboard/controller/mcp_transfer.go | 45 +++++++++-- .../controller/mcp_transfer_host_test.go | 80 +++++++++++++++++++ cmd/dashboard/controller/server.go | 25 +++--- .../controller/server_config_oracle_test.go | 54 +++++++++++++ .../controller/terminal_input_limit_test.go | 19 +++++ cmd/dashboard/controller/transfer.go | 9 ++- service/singleton/server_transfer.go | 37 ++++++++- service/singleton/server_transfer_test.go | 59 +++++++++++++- 12 files changed, 357 insertions(+), 46 deletions(-) create mode 100644 cmd/dashboard/controller/mcp_transfer_host_test.go create mode 100644 cmd/dashboard/controller/server_config_oracle_test.go diff --git a/cmd/dashboard/controller/credential_redaction_test.go b/cmd/dashboard/controller/credential_redaction_test.go index 14b289da..0daaa6fd 100644 --- a/cmd/dashboard/controller/credential_redaction_test.go +++ b/cmd/dashboard/controller/credential_redaction_test.go @@ -109,12 +109,14 @@ func TestListDDNS_RedactsCredentials(t *testing.T) { defer setupTenancyTest(t)() p := model.DDNSProfile{ - Common: model.Common{UserID: 10}, - Name: "cf", - Provider: "cloudflare", - AccessID: "id", - AccessSecret: "super-secret-token", - WebhookHeaders: `{"Authorization":"Bearer xxx"}`, + Common: model.Common{UserID: 10}, + Name: "cf", + Provider: "cloudflare", + AccessID: "id", + AccessSecret: "super-secret-token", + WebhookURL: "https://ddns.example/update?token=embedded-secret", + WebhookRequestBody: `{"token":"embedded-secret"}`, + WebhookHeaders: `{"Authorization":"Bearer xxx"}`, } require.NoError(t, singleton.DB.Create(&p).Error) singleton.DDNSShared.InsertForTest(&p) @@ -124,12 +126,18 @@ func TestListDDNS_RedactsCredentials(t *testing.T) { require.NoError(t, err) require.Len(t, out, 1) require.Empty(t, out[0].AccessSecret, "access_secret must be redacted in list response") + require.Empty(t, out[0].WebhookURL, "webhook_url must be redacted in list response") + require.Empty(t, out[0].WebhookRequestBody, "webhook_request_body must be redacted in list response") require.Empty(t, out[0].WebhookHeaders, "webhook_headers must be redacted in list response") require.Equal(t, "id", out[0].AccessID, "non-secret fields must be preserved") var stored model.DDNSProfile require.NoError(t, singleton.DB.First(&stored, p.ID).Error) require.Equal(t, "super-secret-token", stored.AccessSecret, "redaction must not mutate stored data") + require.Equal(t, "https://ddns.example/update?token=embedded-secret", stored.WebhookURL, + "redaction must not mutate stored webhook URL") + require.Equal(t, `{"token":"embedded-secret"}`, stored.WebhookRequestBody, + "redaction must not mutate stored webhook body") } func TestListNotification_RedactsCredentials(t *testing.T) { @@ -164,14 +172,15 @@ func TestUpdateDDNS_EmptySecretPreservesStored(t *testing.T) { defer setupTenancyTest(t)() existing := model.DDNSProfile{ - Common: model.Common{UserID: 10}, - Name: "cf", - Provider: "webhook", - AccessID: "id", - AccessSecret: "keep-me", - WebhookURL: "http://127.0.0.1/", - WebhookMethod: 1, - WebhookHeaders: `{"X-Token":"keep-header"}`, + Common: model.Common{UserID: 10}, + Name: "cf", + Provider: "webhook", + AccessID: "id", + AccessSecret: "keep-me", + WebhookURL: "http://127.0.0.1/", + WebhookRequestBody: `{"X-Token":"keep-body"}`, + WebhookMethod: 1, + WebhookHeaders: `{"X-Token":"keep-header"}`, } require.NoError(t, singleton.DB.Create(&existing).Error) singleton.DDNSShared.InsertForTest(&existing) @@ -181,10 +190,11 @@ func TestUpdateDDNS_EmptySecretPreservesStored(t *testing.T) { "provider": "webhook", "access_id": "id", "access_secret": "", - "webhook_url": "http://127.0.0.1/", + "webhook_url": "", "webhook_method": 1, "webhook_request_type": 1, "webhook_headers": "", + "webhook_request_body": "", "max_retries": 3, }) c.Params = gin.Params{{Key: "id", Value: itoa(existing.ID)}} @@ -197,6 +207,10 @@ func TestUpdateDDNS_EmptySecretPreservesStored(t *testing.T) { require.Equal(t, "keep-me", after.AccessSecret, "empty submitted secret must preserve stored value") require.Equal(t, `{"X-Token":"keep-header"}`, after.WebhookHeaders, "empty submitted headers must preserve stored value") + require.Equal(t, "http://127.0.0.1/", after.WebhookURL, + "empty submitted webhook URL must preserve stored value") + require.Equal(t, `{"X-Token":"keep-body"}`, after.WebhookRequestBody, + "empty submitted webhook body must preserve stored value") } func TestUpdateDDNS_NonEmptySecretOverwrites(t *testing.T) { diff --git a/cmd/dashboard/controller/ddns.go b/cmd/dashboard/controller/ddns.go index a2a2c20e..b03ba5ca 100644 --- a/cmd/dashboard/controller/ddns.go +++ b/cmd/dashboard/controller/ddns.go @@ -34,6 +34,8 @@ func listDDNS(c *gin.Context) ([]*model.DDNSProfile, error) { // 不影响 singleton 内原始数据。 for _, p := range ddnsProfiles { p.AccessSecret = "" + p.WebhookURL = "" + p.WebhookRequestBody = "" p.WebhookHeaders = "" } @@ -144,10 +146,8 @@ func updateDDNS(c *gin.Context) (any, error) { p.Provider = df.Provider p.Domains = df.Domains p.AccessID = df.AccessID - p.WebhookURL = df.WebhookURL p.WebhookMethod = df.WebhookMethod p.WebhookRequestType = df.WebhookRequestType - p.WebhookRequestBody = df.WebhookRequestBody // 凭据在列表接口已脱敏，前端无法回填；空值视为"不修改"，保留旧值避免误清空。 if df.AccessSecret != "" { @@ -156,6 +156,12 @@ func updateDDNS(c *gin.Context) (any, error) { if df.WebhookHeaders != "" { p.WebhookHeaders = df.WebhookHeaders } + if df.WebhookURL != "" { + p.WebhookURL = df.WebhookURL + } + if df.WebhookRequestBody != "" { + p.WebhookRequestBody = df.WebhookRequestBody + } for n, domain := range p.Domains { // IDN to ASCII diff --git a/cmd/dashboard/controller/fm.go b/cmd/dashboard/controller/fm.go index dfdb2728..51d5dd58 100644 --- a/cmd/dashboard/controller/fm.go +++ b/cmd/dashboard/controller/fm.go @@ -15,6 +15,19 @@ import ( "github.com/nezhahq/nezha/service/singleton" ) +// The official file-manager client uploads in 1 MiB WebSocket messages. Cap +// each complete message at that protocol boundary so gorilla/websocket cannot +// buffer an attacker-sized frame before the IO stream relay sees it. +const fileManagerWebSocketInputLimit int64 = 1024 * 1024 + +type websocketReadLimiter interface { + SetReadLimit(limit int64) +} + +func limitFileManagerWebSocketInput(conn websocketReadLimiter) { + conn.SetReadLimit(fileManagerWebSocketInputLimit) +} + // Create FM session // @Summary Create FM session // @Description Create an "attached" FM. It is advised to only call this within a terminal session. @@ -99,6 +112,7 @@ func fmStream(c *gin.Context) (any, error) { if err != nil { return nil, newWsError("%v", err) } + limitFileManagerWebSocketInput(wsConn) conn := websocketx.NewConn(wsConn) pingTransport := newWebsocketPingTransport(conn, wsConn.Close) stopPing := startWebsocketPingTicker(c.Request.Context(), time.Second*10, pingTransport) diff --git a/cmd/dashboard/controller/mcp_test.go b/cmd/dashboard/controller/mcp_test.go index 1b78f0d9..f90f4e23 100644 --- a/cmd/dashboard/controller/mcp_test.go +++ b/cmd/dashboard/controller/mcp_test.go @@ -39,7 +39,12 @@ func setupMCPTest(t *testing.T) (func(), uint64) { require.NoError(t, err) require.NoError(t, db.AutoMigrate(&model.User{}, &model.APIToken{}, &model.MCPAuditLog{}, &model.Server{}, &model.WAF{})) singleton.DB = db - singleton.Conf = &singleton.ConfigClass{Config: &model.Config{JWTTimeout: 1}} + singleton.Conf = &singleton.ConfigClass{Config: &model.Config{ + JWTTimeout: 1, + ConfigDashboard: model.ConfigDashboard{ + DashboardHost: "example.com", + }, + }} singleton.Conf.SetMCPEnabled(true) user := model.User{Common: model.Common{ID: 100}, Username: "alice", Role: model.RoleMember} diff --git a/cmd/dashboard/controller/mcp_transfer.go b/cmd/dashboard/controller/mcp_transfer.go index 6f204811..8aa09670 100644 --- a/cmd/dashboard/controller/mcp_transfer.go +++ b/cmd/dashboard/controller/mcp_transfer.go @@ -13,6 +13,7 @@ import ( "io" "math" "net/http" + "net/url" "strconv" "strings" "sync" @@ -91,6 +92,37 @@ var ( transferSecretVal string ) +var errMCPTransferHostNotConfigured = errors.New("MCP transfer URL host is not an operator-declared dashboard host") + +// mcpTransferURLBase never reflects an untrusted request Host into a bearer +// transfer URL. DashboardHost is the canonical public authority; deployments +// that intentionally serve multiple names must declare them as reserved hosts. +func mcpTransferURLBase(c *gin.Context) (string, error) { + if c == nil || c.Request == nil || singleton.Conf == nil { + return "", errMCPTransferHostNotConfigured + } + + requestHost := strings.TrimSpace(c.Request.Host) + host := strings.TrimSpace(singleton.Conf.DashboardHost) + if host == "" { + if !singleton.IsReservedDashboardHost(requestHost) { + return "", errMCPTransferHostNotConfigured + } + host = requestHost + } + + parsed, err := url.Parse("//" + host) + if err != nil || parsed.Host == "" || parsed.Host != host || parsed.User != nil || parsed.Path != "" || parsed.RawQuery != "" || parsed.Fragment != "" { + return "", errMCPTransferHostNotConfigured + } + + scheme := "https" + if c.Request.TLS == nil && c.Request.Header.Get("X-Forwarded-Proto") != "https" { + scheme = "http" + } + return scheme + "://" + host, nil +} + // transferHMACSecret 返回进程内随机生成的 HMAC key。 // 这是有意设计：transferEntries 本身也只活在内存 sync.Map 里，dashboard // 重启等价于全部 token 失效；让 secret 也随进程随机，可以避免“secret 来自 @@ -315,6 +347,10 @@ func mintTransferTool(c *gin.Context, serverID uint64, path string, ttlSeconds i if err := validateTransferPath(path); err != nil { return nil, err } + baseURL, err := mcpTransferURLBase(c) + if err != nil { + return nil, err + } ttl := time.Duration(ttlSeconds) * time.Second if ttl <= 0 { ttl = transferTokenTTLDefault @@ -348,14 +384,9 @@ func mintTransferTool(c *gin.Context, serverID uint64, path string, ttlSeconds i if err != nil { return nil, err } - scheme := "https" - if c.Request.TLS == nil && c.Request.Header.Get("X-Forwarded-Proto") != "https" { - scheme = "http" - } - host := c.Request.Host - url := fmt.Sprintf("%s://%s/mcp/%s/%s", scheme, host, dir, t) + transferURL := fmt.Sprintf("%s/mcp/%s/%s", baseURL, dir, t) return map[string]any{ - "url": url, + "url": transferURL, "method": map[transferDirection]string{transferDirDownload: "GET", transferDirUpload: "POST"}[dir], "expires_at": entry.ExpiresAt, }, nil diff --git a/cmd/dashboard/controller/mcp_transfer_host_test.go b/cmd/dashboard/controller/mcp_transfer_host_test.go new file mode 100644 index 00000000..58340fa0 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_host_test.go @@ -0,0 +1,80 @@ +package controller + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestMCPTransferURLPinsUntrustedRequestHostToDashboardHost(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + PurgeTransferEntries() + t.Cleanup(func() { PurgeTransferEntries() }) + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + + c, _ := mcpCallCtx(t, tok, uid, nil) + c.Request.Host = "attacker.example" + c.Request.Header.Set("X-Forwarded-Proto", "https") + + out, err := mintTransferTool(c, 7, "/srv/file", 60, transferDirDownload, transferEntry{}) + require.NoError(t, err) + result := out.(map[string]any) + require.Contains(t, result["url"], "https://example.com/mcp/download/") + require.NotContains(t, result["url"], "attacker.example") +} + +func TestMCPTransferURLRejectsUntrustedHostWithoutCanonicalHost(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + PurgeTransferEntries() + t.Cleanup(func() { PurgeTransferEntries() }) + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + singleton.Conf.DashboardHost = "" + singleton.Conf.InstallHost = "" + singleton.Conf.ListenHost = "" + singleton.Conf.ReservedHosts = "" + + c, _ := mcpCallCtx(t, tok, uid, nil) + c.Request = c.Request.Clone(c.Request.Context()) + c.Request.Method = http.MethodPost + c.Request.Host = "attacker.example" + + _, err := mintTransferTool(c, 7, "/srv/file", 60, transferDirDownload, transferEntry{}) + require.ErrorIs(t, err, errMCPTransferHostNotConfigured) + entries := 0 + transferEntries.Range(func(_, _ any) bool { + entries++ + return true + }) + require.Zero(t, entries, "host validation must happen before a bearer transfer token is minted") +} + +func TestMCPTransferURLAllowsOperatorReservedHost(t *testing.T) { + cleanup, _ := setupMCPTest(t) + defer cleanup() + singleton.Conf.DashboardHost = "" + singleton.Conf.ReservedHosts = "mcp.example.com" + + c, _ := mcpCallCtx(t, nil, 0, nil) + c.Request.Host = "mcp.example.com" + c.Request.Header.Set("X-Forwarded-Proto", "https") + + base, err := mcpTransferURLBase(c) + require.NoError(t, err) + require.Equal(t, "https://mcp.example.com", base) +} + +func TestMCPTransferURLRejectsMalformedConfiguredHost(t *testing.T) { + cleanup, _ := setupMCPTest(t) + defer cleanup() + singleton.Conf.DashboardHost = "https://example.com/path" + + c, _ := mcpCallCtx(t, nil, 0, nil) + _, err := mcpTransferURLBase(c) + require.ErrorIs(t, err, errMCPTransferHostNotConfigured) +} diff --git a/cmd/dashboard/controller/server.go b/cmd/dashboard/controller/server.go index c5c6a4ac..ccbebf5b 100644 --- a/cmd/dashboard/controller/server.go +++ b/cmd/dashboard/controller/server.go @@ -236,10 +236,9 @@ func getServerConfig(c *gin.Context) (string, error) { } s, ok := singleton.ServerShared.Get(id) - if !ok { - return "", nil - } - if !s.HasPermission(c) { + // Foreign and unknown IDs deliberately share one response so sequential + // server IDs cannot be enumerated across tenants. + if !ok || !s.HasPermission(c) || !patAllowsServer(c, id) { return "", singleton.Localizer.ErrorT("permission denied") } if s.GetTaskStream() == nil { @@ -293,16 +292,16 @@ func setServerConfig(c *gin.Context) (*model.ServerTaskResponse, error) { slist := singleton.ServerShared.GetList() servers := make([]*model.Server, 0, len(configForm.Servers)) for _, sid := range configForm.Servers { - if s, ok := slist[sid]; ok { - if !s.HasPermission(c) { - return nil, singleton.Localizer.ErrorT("permission denied") - } - if s.GetTaskStream() == nil { - resp.Offline = append(resp.Offline, s.ID) - continue - } - servers = append(servers, s) + s, ok := slist[sid] + // Match getServerConfig: do not reveal whether a denied ID exists. + if !ok || !s.HasPermission(c) || !patAllowsServer(c, sid) { + return nil, singleton.Localizer.ErrorT("permission denied") } + if s.GetTaskStream() == nil { + resp.Offline = append(resp.Offline, s.ID) + continue + } + servers = append(servers, s) } var wg sync.WaitGroup diff --git a/cmd/dashboard/controller/server_config_oracle_test.go b/cmd/dashboard/controller/server_config_oracle_test.go new file mode 100644 index 00000000..6ffe0542 --- /dev/null +++ b/cmd/dashboard/controller/server_config_oracle_test.go @@ -0,0 +1,54 @@ +package controller + +import ( + "net/http" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestGetServerConfigForeignAndUnknownIDsAreIndistinguishable(t *testing.T) { + cleanup, _ := setupMCPTest(t) + defer cleanup() + require.NoError(t, singleton.DB.Create(&model.User{Common: model.Common{ID: 200}, Username: "bob", Role: model.RoleMember}).Error) + tok, _ := mkToken(t, 200, []string{model.ScopeServerRead}, nil) + + probe := func(id string) error { + c, _ := patRequestCtx(t, tok, 200, http.MethodGet, "/api/v1/server/config/"+id, nil) + c.Params = gin.Params{{Key: "id", Value: id}} + _, err := getServerConfig(c) + return err + } + + foreignErr := probe("7") + unknownErr := probe("999999") + require.Error(t, foreignErr) + require.Error(t, unknownErr) + require.Equal(t, foreignErr.Error(), unknownErr.Error()) +} + +func TestSetServerConfigForeignAndUnknownIDsAreIndistinguishable(t *testing.T) { + cleanup, _ := setupMCPTest(t) + defer cleanup() + require.NoError(t, singleton.DB.Create(&model.User{Common: model.Common{ID: 200}, Username: "bob", Role: model.RoleMember}).Error) + tok, _ := mkToken(t, 200, []string{model.ScopeServerWrite}, nil) + + probe := func(id uint64) error { + c, _ := patRequestCtx(t, tok, 200, http.MethodPost, "/api/v1/server/config", model.ServerConfigForm{ + Servers: []uint64{id}, + Config: "{}", + }) + _, err := setServerConfig(c) + return err + } + + foreignErr := probe(7) + unknownErr := probe(999999) + require.Error(t, foreignErr) + require.Error(t, unknownErr) + require.Equal(t, foreignErr.Error(), unknownErr.Error()) +} diff --git a/cmd/dashboard/controller/terminal_input_limit_test.go b/cmd/dashboard/controller/terminal_input_limit_test.go index 7ed3dcc9..8fb4cbab 100644 --- a/cmd/dashboard/controller/terminal_input_limit_test.go +++ b/cmd/dashboard/controller/terminal_input_limit_test.go @@ -2,8 +2,27 @@ package controller import "testing" +type readLimitRecorder struct { + limit int64 +} + +func (r *readLimitRecorder) SetReadLimit(limit int64) { + r.limit = limit +} + func TestTerminalWebSocketInputLimitAllowsBoundedPaste(t *testing.T) { if terminalWebSocketInputLimit != 512*1024+64 { t.Fatalf("terminal WebSocket input limit = %d, want 512 KiB plus control-byte allowance", terminalWebSocketInputLimit) } } + +func TestFileManagerWebSocketInputLimitMatchesUploadChunk(t *testing.T) { + if fileManagerWebSocketInputLimit != 1024*1024 { + t.Fatalf("file-manager WebSocket input limit = %d, want official client's 1 MiB upload chunk", fileManagerWebSocketInputLimit) + } + recorder := new(readLimitRecorder) + limitFileManagerWebSocketInput(recorder) + if recorder.limit != fileManagerWebSocketInputLimit { + t.Fatalf("file-manager WebSocket configured limit = %d, want %d", recorder.limit, fileManagerWebSocketInputLimit) + } +} diff --git a/cmd/dashboard/controller/transfer.go b/cmd/dashboard/controller/transfer.go index 0e8f3ca7..f3c76b59 100644 --- a/cmd/dashboard/controller/transfer.go +++ b/cmd/dashboard/controller/transfer.go @@ -171,6 +171,12 @@ const transferStreamPingInterval = 30 * time.Second // @Success 200 {object} model.ServerTransfer // @Router /ws/transfer [get] func transferStream(c *gin.Context) (any, error) { + subID, ch, err := singleton.ServerTransferShared.Subscribe(getUid(c)) + if err != nil { + return nil, err + } + defer singleton.ServerTransferShared.Unsubscribe(subID) + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) if err != nil { return nil, newWsError("%v", err) @@ -180,9 +186,6 @@ func transferStream(c *gin.Context) (any, error) { deregisterPAT := registerPATConnection(c, func() { _ = conn.Close() }) defer deregisterPAT() - subID, ch := singleton.ServerTransferShared.Subscribe() - defer singleton.ServerTransferShared.Unsubscribe(subID) - // Pings keep the socket warm even when the broker is quiet. Without this // a long idle period followed by a transfer event would race against // upstream proxy idle-timeouts that may have already closed the conn. diff --git a/service/singleton/server_transfer.go b/service/singleton/server_transfer.go index cef41476..ff3f93ef 100644 --- a/service/singleton/server_transfer.go +++ b/service/singleton/server_transfer.go @@ -152,6 +152,8 @@ type ServerTransferClass struct { subMu sync.Mutex subs map[uint64]chan *model.ServerTransfer + subOwners map[uint64]uint64 + subCounts map[uint64]int nextSubID uint64 timeout time.Duration @@ -174,6 +176,18 @@ var ErrServerAlreadyTransferring = errors.New("server already has an in-flight t // deferred to OnAgentReconnect / PushIfOnline. var ErrAgentTooOldForTransfer = fmt.Errorf("agent build older than %s does not support server transfer (TaskTypeServerTransferApply)", MinServerTransferAgentVersion) +const ( + // A normal dashboard needs one subscription per open transfer page. Keep a + // small per-user allowance for tabs/devices plus a process-wide hard stop. + maxServerTransferSubscribersPerUser = 8 + maxServerTransferSubscribersGlobal = 1024 +) + +var ( + errTooManyServerTransferSubscribersForUser = errors.New("too many concurrent server transfer subscribers for this user") + errTooManyServerTransferSubscribersGlobal = errors.New("too many concurrent server transfer subscribers") +) + // agentSupportsTransfer reports whether s has reported a build version >= // MinServerTransferAgentVersion. Returns true when version is unknown (agent // never reported) so callers can defer the decision; PushIfOnline re-checks @@ -210,6 +224,8 @@ func NewServerTransferClass() *ServerTransferClass { verifiedHandshakes: make(map[uint64]string), initiating: make(map[uint64]bool), subs: make(map[uint64]chan *model.ServerTransfer), + subOwners: make(map[uint64]uint64), + subCounts: make(map[uint64]int), timeout: defaultServerTransferTimeout, stopCh: make(chan struct{}), } @@ -1410,20 +1426,37 @@ func (c *ServerTransferClass) sweepTimeouts() { // Subscribe registers a channel that will receive every transfer transition // event from this point forward. The caller MUST Unsubscribe when done or // the broker will block forever if the channel is unbuffered or full. -func (c *ServerTransferClass) Subscribe() (uint64, <-chan *model.ServerTransfer) { +func (c *ServerTransferClass) Subscribe(userID uint64) (uint64, <-chan *model.ServerTransfer, error) { c.subMu.Lock() defer c.subMu.Unlock() + if len(c.subs) >= maxServerTransferSubscribersGlobal { + return 0, nil, errTooManyServerTransferSubscribersGlobal + } + if c.subCounts[userID] >= maxServerTransferSubscribersPerUser { + return 0, nil, errTooManyServerTransferSubscribersForUser + } id := atomic.AddUint64(&c.nextSubID, 1) ch := make(chan *model.ServerTransfer, 16) c.subs[id] = ch - return id, ch + c.subOwners[id] = userID + c.subCounts[userID]++ + return id, ch, nil } func (c *ServerTransferClass) Unsubscribe(id uint64) { c.subMu.Lock() ch, ok := c.subs[id] + ownerID, ownerOK := c.subOwners[id] delete(c.subs, id) + delete(c.subOwners, id) + if ownerOK { + if c.subCounts[ownerID] <= 1 { + delete(c.subCounts, ownerID) + } else { + c.subCounts[ownerID]-- + } + } c.subMu.Unlock() if ok { close(ch) diff --git a/service/singleton/server_transfer_test.go b/service/singleton/server_transfer_test.go index 2be8eb05..fa6fecf2 100644 --- a/service/singleton/server_transfer_test.go +++ b/service/singleton/server_transfer_test.go @@ -455,7 +455,8 @@ func TestServerTransferBroadcastReachesSubscribers(t *testing.T) { defer cleanup() seedServerForTransfer(t, 1, 100) - id, ch := c.Subscribe() + id, ch, err := c.Subscribe(100) + require.NoError(t, err) defer c.Unsubscribe(id) tr := initiateAndRegister(t, c, 1, 100, 200, 1) @@ -469,7 +470,7 @@ func TestServerTransferBroadcastReachesSubscribers(t *testing.T) { t.Fatal("expected Pending broadcast within 1s") } - _, _, err := markPendingVerified(t, c, 1) + _, _, err = markPendingVerified(t, c, 1) require.NoError(t, err) select { case ev := <-ch: @@ -479,6 +480,57 @@ func TestServerTransferBroadcastReachesSubscribers(t *testing.T) { } } +func TestServerTransferSubscribeEnforcesPerUserLimitAndReleasesQuota(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + + ids := make([]uint64, 0, maxServerTransferSubscribersPerUser) + for range maxServerTransferSubscribersPerUser { + id, ch, err := c.Subscribe(100) + require.NoError(t, err) + require.NotNil(t, ch) + ids = append(ids, id) + } + + _, _, err := c.Subscribe(100) + require.ErrorIs(t, err, errTooManyServerTransferSubscribersForUser) + + otherID, _, err := c.Subscribe(200) + require.NoError(t, err, "one user's quota must not block another user") + c.Unsubscribe(otherID) + + c.Unsubscribe(ids[0]) + replacementID, _, err := c.Subscribe(100) + require.NoError(t, err, "disconnecting must release the user's quota") + c.Unsubscribe(replacementID) + for _, id := range ids[1:] { + c.Unsubscribe(id) + } +} + +func TestServerTransferSubscribeEnforcesGlobalLimit(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + + ids := make([]uint64, 0, maxServerTransferSubscribersGlobal) + for userID := uint64(1); len(ids) < maxServerTransferSubscribersGlobal; userID++ { + for range maxServerTransferSubscribersPerUser { + if len(ids) == maxServerTransferSubscribersGlobal { + break + } + id, _, err := c.Subscribe(userID) + require.NoError(t, err) + ids = append(ids, id) + } + } + + _, _, err := c.Subscribe(999999) + require.ErrorIs(t, err, errTooManyServerTransferSubscribersGlobal) + for _, id := range ids { + c.Unsubscribe(id) + } +} + // The "one active transfer per server" invariant must hold under concurrent // callers (two operators batch-moving the same server, or batch-move racing // retry). The old flow had a TOCTOU between HasPending() and Initiate() that @@ -1626,7 +1678,8 @@ func TestOnServersDeletedTerminatesPendingTransfersAndClearsIndexes(t *testing.T tr := initiateAndRegister(t, c, 1, 100, 200, 1) require.True(t, c.HasPending(1)) - subID, ch := c.Subscribe() + subID, ch, err := c.Subscribe(100) + require.NoError(t, err) defer c.Unsubscribe(subID) require.NoError(t, DB.Unscoped().Delete(&model.Server{}, tr.ServerID).Error)