From 92816b46c2e09b2daddc5d42e31e6d1b7c2b8159 Mon Sep 17 00:00:00 2001 From: buua436 Date: Fri, 6 Nov 2026 14:55:00 +0800 Subject: [PATCH] fix: scope upload quotas and lock WeCom listener binding (#22076) --- internal/channels/wecom.go | 28 +++--- internal/channels/wecom_test.go | 73 +++++++++++++-- internal/dao/document.go | 17 +++- internal/dao/document_tenant_count_test.go | 90 +++++++++++++++++++ internal/service/file/file_upload.go | 2 +- .../service/file/file_upload_quota_test.go | 54 +++++++++++ 6 files changed, 241 insertions(+), 23 deletions(-) create mode 100644 internal/dao/document_tenant_count_test.go create mode 100644 internal/service/file/file_upload_quota_test.go diff --git a/internal/channels/wecom.go b/internal/channels/wecom.go index 2f04db1f33..88c1e95c09 100644 --- a/internal/channels/wecom.go +++ b/internal/channels/wecom.go @@ -154,7 +154,6 @@ func newWeComChannelFromConfig(accountID string, cfg map[string]any) (*wecomChan AccountID: accountID, ConnectionType: connectionType, Secret: firstString(cfg, "secret"), - WebhookHost: firstString(cfg, "webhook_host"), APIBaseURL: firstString(cfg, "api_base_url"), WSURL: firstString(cfg, "ws_url"), } @@ -189,29 +188,36 @@ func newWeComChannelFromConfig(accountID string, cfg map[string]any) (*wecomChan if len(account.AESKey) != 43 { return nil, fmt.Errorf("wecom account %q aes_key (EncodingAESKey) must be 43 characters, got %d", accountID, len(account.AESKey)) } + if rawHost, ok := cfg["webhook_host"]; ok { + host, ok := rawHost.(string) + if !ok { + return nil, fmt.Errorf("wecom account %q webhook_host must be %q", accountID, defaultWeComWebhookHost) + } + host = strings.TrimSpace(host) + if host != "" && host != defaultWeComWebhookHost { + return nil, fmt.Errorf("wecom account %q webhook_host must use the shared address %q", accountID, defaultWeComWebhookHost) + } + } agentID, err := parseWeComInt(cfg["agent_id"], "agent_id") if err != nil { return nil, fmt.Errorf("wecom account %q %w", accountID, err) } account.AgentID = agentID - var configuredWebhookPort *int if raw, ok := cfg["webhook_port"]; ok { port, err := parseWeComInt(raw, "webhook_port") if err != nil { return nil, fmt.Errorf("wecom account %q %w", accountID, err) } + if port <= 0 || port > 65535 { + return nil, fmt.Errorf("wecom account %q webhook_port must be between 1 and 65535", accountID) + } + if port != defaultWeComWebhookPort { + return nil, fmt.Errorf("wecom account %q webhook_port must use the shared port %d", accountID, defaultWeComWebhookPort) + } account.WebhookPort = int(port) - configuredWebhookPort = &account.WebhookPort } - channel, err := newWeComChannel(account) - if err != nil { - return nil, err - } - if configuredWebhookPort != nil { - channel.account.WebhookPort = *configuredWebhookPort - } - return channel, nil + return newWeComChannel(account) } // ChannelID returns the platform identifier used by the chat-channel runtime. diff --git a/internal/channels/wecom_test.go b/internal/channels/wecom_test.go index dd0686ecc1..f0071f419d 100644 --- a/internal/channels/wecom_test.go +++ b/internal/channels/wecom_test.go @@ -90,15 +90,14 @@ func TestNewWeComChannelFromConfigTreatsUnknownConnectionTypeAsWebhook(t *testin } } -func TestNewWeComChannelFromConfigAcceptsIntegerAgentID(t *testing.T) { +func TestNewWeComChannelFromConfigAcceptsIntegerAgentIDAndDefaultsListener(t *testing.T) { encodingAESKey, _ := testWeComAESKey() channel, err := newWeComChannelFromConfig("account-1", map[string]any{ - "corp_id": "corp-id", - "agent_id": 0, - "secret": "secret", - "token": "token", - "aes_key": encodingAESKey, - "webhook_port": 0, + "corp_id": "corp-id", + "agent_id": 0, + "secret": "secret", + "token": "token", + "aes_key": encodingAESKey, }) if err != nil { t.Fatalf("agent_id=0 error = %v", err) @@ -106,8 +105,64 @@ func TestNewWeComChannelFromConfigAcceptsIntegerAgentID(t *testing.T) { if channel.account.AgentID != 0 { t.Fatalf("agent_id = %d, want 0", channel.account.AgentID) } - if channel.account.WebhookPort != 0 { - t.Fatalf("webhook_port = %d, want 0", channel.account.WebhookPort) + if channel.account.WebhookPort != defaultWeComWebhookPort { + t.Fatalf("webhook_port = %d, want %d", channel.account.WebhookPort, defaultWeComWebhookPort) + } + if channel.account.WebhookHost != defaultWeComWebhookHost { + t.Fatalf("webhook_host = %q, want %q", channel.account.WebhookHost, defaultWeComWebhookHost) + } +} + +func TestNewWeComChannelFromConfigRejectsTenantControlledListener(t *testing.T) { + encodingAESKey, _ := testWeComAESKey() + baseConfig := map[string]any{ + "corp_id": "corp-id", + "agent_id": "1000001", + "secret": "secret", + "token": "token", + "aes_key": encodingAESKey, + } + for _, tc := range []struct { + name string + field string + value any + }{ + {name: "ephemeral port", field: "webhook_port", value: 0}, + {name: "out of range port", field: "webhook_port", value: 65536}, + {name: "tenant selected port", field: "webhook_port", value: 3003}, + {name: "loopback host", field: "webhook_host", value: "127.0.0.1"}, + {name: "hostname", field: "webhook_host", value: "localhost"}, + {name: "non-string host", field: "webhook_host", value: []string{"0.0.0.0"}}, + } { + t.Run(tc.name, func(t *testing.T) { + config := make(map[string]any, len(baseConfig)+1) + for key, value := range baseConfig { + config[key] = value + } + config[tc.field] = tc.value + if _, err := newWeComChannelFromConfig("account-1", config); err == nil { + t.Fatalf("accepted tenant controlled %s=%v", tc.field, tc.value) + } + }) + } +} + +func TestNewWeComChannelFromConfigAllowsSharedListenerConfig(t *testing.T) { + encodingAESKey, _ := testWeComAESKey() + channel, err := newWeComChannelFromConfig("account-1", map[string]any{ + "corp_id": "corp-id", + "agent_id": "1000001", + "secret": "secret", + "token": "token", + "aes_key": encodingAESKey, + "webhook_host": defaultWeComWebhookHost, + "webhook_port": defaultWeComWebhookPort, + }) + if err != nil { + t.Fatalf("shared listener config error = %v", err) + } + if channel.account.WebhookHost != defaultWeComWebhookHost || channel.account.WebhookPort != defaultWeComWebhookPort { + t.Fatalf("listener address = %s:%d, want %s:%d", channel.account.WebhookHost, channel.account.WebhookPort, defaultWeComWebhookHost, defaultWeComWebhookPort) } } diff --git a/internal/dao/document.go b/internal/dao/document.go index 32b649542d..ac17eaaba4 100644 --- a/internal/dao/document.go +++ b/internal/dao/document.go @@ -512,10 +512,23 @@ func (dao *DocumentDAO) GetByDocumentIDAndDatasetID(ctx context.Context, db *gor return &document, err } -// CountByTenantID counts documents by tenant ID +// CountByTenantID counts documents belonging to datasets owned by a tenant. func (dao *DocumentDAO) CountByTenantID(ctx context.Context, db *gorm.DB, tenantID string) (int64, error) { var count int64 - err := db.WithContext(ctx).Model(&entity.Document{}).Where("created_by = ?", tenantID).Count(&count).Error + err := db.WithContext(ctx).Model(&entity.Document{}). + Joins("JOIN knowledgebase ON knowledgebase.id = document.kb_id"). + Where("knowledgebase.tenant_id = ?", tenantID). + Count(&count).Error + return count, err +} + +// CountByUserID counts documents attributed to a user. created_by may contain +// either a user ID or a tenant ID, so it is not a reliable user scope. +func (dao *DocumentDAO) CountByUserID(ctx context.Context, db *gorm.DB, userID string) (int64, error) { + var count int64 + err := db.WithContext(ctx).Model(&entity.Document{}). + Where("user_id = ?", userID). + Count(&count).Error return count, err } diff --git a/internal/dao/document_tenant_count_test.go b/internal/dao/document_tenant_count_test.go new file mode 100644 index 0000000000..21a2cb68e9 --- /dev/null +++ b/internal/dao/document_tenant_count_test.go @@ -0,0 +1,90 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package dao + +import ( + "testing" + + "github.com/glebarez/sqlite" + "gorm.io/gorm" +) + +func documentCountTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatal("open sqlite:", err) + } + if err := db.Exec("CREATE TABLE knowledgebase (id TEXT PRIMARY KEY, tenant_id TEXT NOT NULL)").Error; err != nil { + t.Fatal("create knowledgebase table:", err) + } + if err := db.Exec("CREATE TABLE document (id TEXT PRIMARY KEY, kb_id TEXT NOT NULL, created_by TEXT NOT NULL, user_id TEXT NOT NULL)").Error; err != nil { + t.Fatal("create document table:", err) + } + if err := db.Exec("INSERT INTO knowledgebase (id, tenant_id) VALUES (?, ?), (?, ?)", "kb-1", "tenant-1", "kb-2", "tenant-2").Error; err != nil { + t.Fatal("insert knowledgebases:", err) + } + return db +} + +func TestDocumentDAOCountByTenantIDUsesKnowledgebaseTenant(t *testing.T) { + db := documentCountTestDB(t) + if err := db.Exec("INSERT INTO document (id, kb_id, created_by, user_id) VALUES (?, ?, ?, ?), (?, ?, ?, ?), (?, ?, ?, ?)", + "doc-1", "kb-1", "user-1", "user-1", + "doc-2", "kb-1", "user-2", "user-2", + "doc-3", "kb-2", "user-3", "user-3", + ).Error; err != nil { + t.Fatal("insert documents:", err) + } + + count, err := NewDocumentDAO().CountByTenantID(t.Context(), db, "tenant-1") + if err != nil { + t.Fatalf("CountByTenantID: %v", err) + } + if count != 2 { + t.Fatalf("CountByTenantID = %d, want 2 documents owned by tenant-1", count) + } +} + +func TestDocumentDAOCountByUserIDUsesUserIDWhenCreatedByIsTenant(t *testing.T) { + db := documentCountTestDB(t) + if err := db.Exec("INSERT INTO document (id, kb_id, created_by, user_id) VALUES (?, ?, ?, ?), (?, ?, ?, ?), (?, ?, ?, ?), (?, ?, ?, ?), (?, ?, ?, ?), (?, ?, ?, ?)", + "doc-1", "kb-1", "tenant-1", "user-1", + "doc-2", "kb-1", "user-2", "user-1", + "doc-3", "kb-2", "user-1", "user-3", + "doc-4", "kb-2", "user-4", "user-4", + "doc-5", "kb-1", "tenant-1", "user-2", + "doc-6", "kb-1", "tenant-1", "tenant-1", + ).Error; err != nil { + t.Fatalf("insert documents: %v", err) + } + + count, err := NewDocumentDAO().CountByUserID(t.Context(), db, "user-1") + if err != nil { + t.Fatalf("CountByUserID: %v", err) + } + if count != 2 { + t.Fatalf("CountByUserID = %d, want 2 documents attributed to user-1", count) + } + + tenantOwnerCount, err := NewDocumentDAO().CountByUserID(t.Context(), db, "tenant-1") + if err != nil { + t.Fatalf("CountByUserID for tenant owner: %v", err) + } + if tenantOwnerCount != 1 { + t.Fatalf("CountByUserID for tenant owner = %d, want 1 owner document", tenantOwnerCount) + } +} diff --git a/internal/service/file/file_upload.go b/internal/service/file/file_upload.go index 93971117e6..2e858106ea 100644 --- a/internal/service/file/file_upload.go +++ b/internal/service/file/file_upload.go @@ -289,7 +289,7 @@ func (s *FileService) checkUploadInfoHealth(ctx context.Context, filename string var maxNum int64 if _, err := fmt.Sscanf(maxFileNumPerUser, "%d", &maxNum); err == nil && maxNum > 0 { var docCount int64 - docCount, err = s.GetDocCount(ctx, userID) + docCount, err = dao.NewDocumentDAO().CountByUserID(ctx, dao.DB, userID) if err != nil { return fmt.Errorf("failed to get document count: %w", err) } diff --git a/internal/service/file/file_upload_quota_test.go b/internal/service/file/file_upload_quota_test.go new file mode 100644 index 0000000000..510d4fffe4 --- /dev/null +++ b/internal/service/file/file_upload_quota_test.go @@ -0,0 +1,54 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package file + +import ( + "strings" + "testing" + + "github.com/glebarez/sqlite" + "gorm.io/gorm" + + "ragflow/internal/common" + "ragflow/internal/dao" + "ragflow/internal/permission" +) + +func TestCheckUploadInfoHealthEnforcesUserQuotaForTenantCreatedDocuments(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.Exec("CREATE TABLE document (id TEXT PRIMARY KEY, kb_id TEXT NOT NULL, created_by TEXT NOT NULL, user_id TEXT NOT NULL)").Error; err != nil { + t.Fatalf("create document table: %v", err) + } + if err := db.Exec("INSERT INTO document (id, kb_id, created_by, user_id) VALUES (?, ?, ?, ?), (?, ?, ?, ?)", + "doc-1", "kb-1", "tenant-1", "user-1", + "doc-2", "kb-1", "tenant-1", "user-1", + ).Error; err != nil { + t.Fatalf("insert documents: %v", err) + } + oldDB := dao.DB + dao.DB = db + t.Cleanup(func() { dao.DB = oldDB }) + t.Setenv(common.EnvMaxFileNumPerUser, "2") + + ctx := permission.WithIdentity(t.Context(), "user-1", "tenant-1") + err = testFileService().checkUploadInfoHealth(ctx, "upload.txt") + if err == nil || !strings.Contains(err.Error(), "exceed the maximum file number") { + t.Fatalf("checkUploadInfoHealth() error = %v, want quota rejection", err) + } +}