From 358152f758fdf68a03282aadfc7c70c94a7cabc0 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Mon, 6 Jul 2026 10:13:46 +0800 Subject: [PATCH] fix(go-document): add document and file access checks (#16592) ## Summary Adds ownership/access checks before updating or deleting documents, setting document metadata, and reading file contents from storage. Also adds tests for authorized and unauthorized access paths. --- internal/handler/document.go | 48 ++++- internal/handler/document_test.go | 266 ++++++++++++++++++++++++- internal/service/chat_pipeline.go | 24 ++- internal/service/chat_pipeline_test.go | 4 +- internal/service/chat_session.go | 8 +- internal/service/chat_session_test.go | 65 +++++- internal/service/file.go | 5 +- internal/service/file_test.go | 144 +++++++++++++ internal/service/openai_chat.go | 2 +- 9 files changed, 541 insertions(+), 25 deletions(-) diff --git a/internal/handler/document.go b/internal/handler/document.go index 13f04af7c9..e76c1d9736 100644 --- a/internal/handler/document.go +++ b/internal/handler/document.go @@ -330,7 +330,7 @@ func (h *DocumentHandler) GetDocumentPreview(c *gin.Context) { // @Success 200 {object} map[string]interface{} // @Router /api/v1/documents/{id} [put] func (h *DocumentHandler) UpdateDocument(c *gin.Context) { - _, errorCode, errorMessage := GetUser(c) + user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { jsonError(c, errorCode, errorMessage) return @@ -344,6 +344,19 @@ func (h *DocumentHandler) UpdateDocument(c *gin.Context) { return } + doc, err := h.documentService.GetDocumentByID(id) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "code": 1, + "message": "document not found", + }) + return + } + if !h.datasetService.Accessible(doc.KbID, user.ID) { + jsonError(c, common.CodeAuthenticationError, "No authorization.") + return + } + var req service.UpdateDocumentRequest if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusBadRequest, gin.H{ @@ -374,7 +387,7 @@ func (h *DocumentHandler) UpdateDocument(c *gin.Context) { // @Success 200 {object} map[string]interface{} // @Router /api/v1/documents/{id} [delete] func (h *DocumentHandler) DeleteDocument(c *gin.Context) { - _, errorCode, errorMessage := GetUser(c) + user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { jsonError(c, errorCode, errorMessage) return @@ -388,6 +401,19 @@ func (h *DocumentHandler) DeleteDocument(c *gin.Context) { return } + doc, err := h.documentService.GetDocumentByID(id) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "code": 1, + "message": "document not found", + }) + return + } + if !h.datasetService.Accessible(doc.KbID, user.ID) { + jsonError(c, common.CodeAuthenticationError, "No authorization.") + return + } + if err := h.documentService.DeleteDocument(id); err != nil { c.JSON(http.StatusInternalServerError, gin.H{ "error": err.Error(), @@ -1254,7 +1280,7 @@ type SetMetaRequest struct { // @Success 200 {object} map[string]interface{} // @Router /v1/document/set_meta [post] func (h *DocumentHandler) SetMeta(c *gin.Context) { - _, errorCode, errorMessage := GetUser(c) + user, errorCode, errorMessage := GetUser(c) if errorCode != common.CodeSuccess { jsonError(c, errorCode, errorMessage) return @@ -1321,7 +1347,21 @@ func (h *DocumentHandler) SetMeta(c *gin.Context) { } } - err := h.documentService.SetDocumentMetadata(req.DocID, meta) + // Authorization: user must be able to access the document's dataset. + doc, err := h.documentService.GetDocumentByID(req.DocID) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "code": 1, + "message": "document not found", + }) + return + } + if !h.datasetService.Accessible(doc.KbID, user.ID) { + jsonError(c, common.CodeAuthenticationError, "No authorization.") + return + } + + err = h.documentService.SetDocumentMetadata(req.DocID, meta) if err != nil { errMsg := err.Error() if strings.Contains(errMsg, "no such document") || strings.Contains(errMsg, "document not found") { diff --git a/internal/handler/document_test.go b/internal/handler/document_test.go index ef2d06d409..b06ab23c73 100644 --- a/internal/handler/document_test.go +++ b/internal/handler/document_test.go @@ -40,6 +40,12 @@ import ( type fakeDocumentService struct { deleted int err error + doc *service.DocumentResponse + docErr error + updateCalled bool + updatedID string + deleteCalled bool + deletedID string stopResult map[string]interface{} stopErr error thumbnails map[string]string @@ -50,6 +56,9 @@ type fakeDocumentService struct { metadataErr error metadataKBID string metadataDocIDs []string + setMetaCalled bool + setMetaDocID string + setMetaValue map[string]interface{} uploadLocalData []map[string]interface{} uploadLocalErrs []string uploadLocalKB *entity.Knowledgebase @@ -129,12 +138,22 @@ func (f *fakeDocumentService) CreateDocument(req *service.CreateDocumentRequest) return nil, nil } func (f *fakeDocumentService) GetDocumentByID(id string) (*service.DocumentResponse, error) { - return nil, nil + if f.docErr != nil { + return nil, f.docErr + } + if f.doc != nil { + return f.doc, nil + } + return nil, fmt.Errorf("document not found") } func (f *fakeDocumentService) UpdateDocument(id string, req *service.UpdateDocumentRequest) error { + f.updateCalled = true + f.updatedID = id return nil } func (f *fakeDocumentService) DeleteDocument(id string) error { + f.deleteCalled = true + f.deletedID = id return nil } func (f *fakeDocumentService) DeleteDocuments(ids []string, deleteAll bool, datasetID, userID string) (int, error) { @@ -193,6 +212,9 @@ func (f *fakeDocumentService) GetMetadataSummary(kbID string, docIDs []string) ( return f.metadataSummary, f.metadataErr } func (f *fakeDocumentService) SetDocumentMetadata(docID string, meta map[string]interface{}) error { + f.setMetaCalled = true + f.setMetaDocID = docID + f.setMetaValue = meta return nil } func (f *fakeDocumentService) DeleteDocumentMetadata(docID string, keys []string) error { @@ -242,6 +264,248 @@ func setupGinContextWithUser(method, path, body string) (*gin.Context, *httptest return c, w } +func setupDocumentPermissionDB(t *testing.T, accessible bool) { + t.Helper() + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ + TranslateError: true, + }) + if err != nil { + t.Fatalf("failed to open sqlite: %v", err) + } + if err := db.AutoMigrate( + &entity.Knowledgebase{}, + &entity.UserTenant{}, + ); err != nil { + t.Fatalf("failed to migrate permission tables: %v", err) + } + if err := db.Create(&entity.Knowledgebase{ + ID: "kb-owner", + TenantID: "tenant-owner", + Name: "owner-kb", + EmbdID: "embd-1", + CreatedBy: "owner-user", + Permission: string(entity.TenantPermissionTeam), + Status: sptr(string(entity.StatusValid)), + }).Error; err != nil { + t.Fatalf("insert knowledgebase: %v", err) + } + if accessible { + if err := db.Create(&entity.UserTenant{ + ID: "ut-user-1", + UserID: "user-1", + TenantID: "tenant-owner", + Role: "normal", + Status: sptr(string(entity.StatusValid)), + }).Error; err != nil { + t.Fatalf("insert user_tenant: %v", err) + } + } + + orig := dao.DB + dao.DB = db + t.Cleanup(func() { dao.DB = orig }) +} + +func TestSetMetaHandler_NotAccessible(t *testing.T) { + setupDocumentPermissionDB(t, false) + + fake := &fakeDocumentService{ + doc: &service.DocumentResponse{ID: "doc-1", KbID: "kb-owner"}, + } + h := &DocumentHandler{ + documentService: fake, + datasetService: service.NewDatasetService(), + } + + c, w := setupGinContextWithUser("POST", "/api/v1/document/set_meta", `{"doc_id":"doc-1","meta":"{\"poc\":\"blocked\"}"}`) + h.SetMeta(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp["code"] != float64(common.CodeAuthenticationError) { + t.Fatalf("expected auth error, got %v", resp) + } + if resp["message"] != "No authorization." { + t.Fatalf("unexpected message: %v", resp["message"]) + } + if fake.setMetaCalled { + t.Fatal("SetDocumentMetadata should not be called without dataset access") + } +} + +func TestSetMetaHandler_Accessible(t *testing.T) { + setupDocumentPermissionDB(t, true) + + fake := &fakeDocumentService{ + doc: &service.DocumentResponse{ID: "doc-1", KbID: "kb-owner"}, + } + h := &DocumentHandler{ + documentService: fake, + datasetService: service.NewDatasetService(), + } + + c, w := setupGinContextWithUser("POST", "/api/v1/document/set_meta", `{"doc_id":"doc-1","meta":"{\"category\":\"tech\",\"year\":2026}"}`) + h.SetMeta(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp["code"] != float64(common.CodeSuccess) || resp["data"] != true { + t.Fatalf("unexpected response: %v", resp) + } + if !fake.setMetaCalled { + t.Fatal("SetDocumentMetadata should be called with dataset access") + } + if fake.setMetaDocID != "doc-1" { + t.Fatalf("set meta doc id = %q, want doc-1", fake.setMetaDocID) + } + if fake.setMetaValue["category"] != "tech" || fake.setMetaValue["year"] != float64(2026) { + t.Fatalf("unexpected meta: %#v", fake.setMetaValue) + } +} + +func TestDeleteDocumentHandler_NotAccessible(t *testing.T) { + setupDocumentPermissionDB(t, false) + + fake := &fakeDocumentService{ + doc: &service.DocumentResponse{ID: "doc-1", KbID: "kb-owner"}, + } + h := &DocumentHandler{ + documentService: fake, + datasetService: service.NewDatasetService(), + } + + c, w := setupGinContextWithUser("DELETE", "/api/v1/documents/doc-1", "") + c.Params = gin.Params{{Key: "id", Value: "doc-1"}} + h.DeleteDocument(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp["code"] != float64(common.CodeAuthenticationError) { + t.Fatalf("expected auth error, got %v", resp) + } + if resp["message"] != "No authorization." { + t.Fatalf("unexpected message: %v", resp["message"]) + } + if fake.deleteCalled { + t.Fatal("DeleteDocument should not be called without dataset access") + } +} + +func TestDeleteDocumentHandler_Accessible(t *testing.T) { + setupDocumentPermissionDB(t, true) + + fake := &fakeDocumentService{ + doc: &service.DocumentResponse{ID: "doc-1", KbID: "kb-owner"}, + } + h := &DocumentHandler{ + documentService: fake, + datasetService: service.NewDatasetService(), + } + + c, w := setupGinContextWithUser("DELETE", "/api/v1/documents/doc-1", "") + c.Params = gin.Params{{Key: "id", Value: "doc-1"}} + h.DeleteDocument(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if !fake.deleteCalled { + t.Fatal("DeleteDocument should be called with dataset access") + } + if fake.deletedID != "doc-1" { + t.Fatalf("deleted id = %q, want doc-1", fake.deletedID) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp["message"] != "deleted successfully" { + t.Fatalf("unexpected response: %v", resp) + } +} + +func TestUpdateDocumentHandler_NotAccessible(t *testing.T) { + setupDocumentPermissionDB(t, false) + + fake := &fakeDocumentService{ + doc: &service.DocumentResponse{ID: "doc-1", KbID: "kb-owner"}, + } + h := &DocumentHandler{ + documentService: fake, + datasetService: service.NewDatasetService(), + } + + c, w := setupGinContextWithUser("PUT", "/api/v1/documents/doc-1", `{"name":"blocked"}`) + c.Params = gin.Params{{Key: "id", Value: "doc-1"}} + h.UpdateDocument(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp["code"] != float64(common.CodeAuthenticationError) { + t.Fatalf("expected auth error, got %v", resp) + } + if resp["message"] != "No authorization." { + t.Fatalf("unexpected message: %v", resp["message"]) + } + if fake.updateCalled { + t.Fatal("UpdateDocument should not be called without dataset access") + } +} + +func TestUpdateDocumentHandler_Accessible(t *testing.T) { + setupDocumentPermissionDB(t, true) + + fake := &fakeDocumentService{ + doc: &service.DocumentResponse{ID: "doc-1", KbID: "kb-owner"}, + } + h := &DocumentHandler{ + documentService: fake, + datasetService: service.NewDatasetService(), + } + + c, w := setupGinContextWithUser("PUT", "/api/v1/documents/doc-1", `{"name":"allowed"}`) + c.Params = gin.Params{{Key: "id", Value: "doc-1"}} + h.UpdateDocument(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + if !fake.updateCalled { + t.Fatal("UpdateDocument should be called with dataset access") + } + if fake.updatedID != "doc-1" { + t.Fatalf("updated id = %q, want doc-1", fake.updatedID) + } + var resp map[string]interface{} + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp["message"] != "updated successfully" { + t.Fatalf("unexpected response: %v", resp) + } +} + func setupUploadHandlerDB(t *testing.T, role string) *gorm.DB { t.Helper() diff --git a/internal/service/chat_pipeline.go b/internal/service/chat_pipeline.go index d535e2696a..4f9d539ce2 100644 --- a/internal/service/chat_pipeline.go +++ b/internal/service/chat_pipeline.go @@ -151,6 +151,7 @@ type AsyncChatResult struct { // - kwargs: extra parameters (doc_ids, knowledge, quote, etc.). func (s *ChatPipelineService) AsyncChat( ctx context.Context, + userID string, chat *entity.Chat, messages []map[string]interface{}, stream bool, @@ -188,7 +189,7 @@ func (s *ChatPipelineService) AsyncChat( } if !hasKBs && !useWebSearch { - return s.AsyncChatSolo(ctx, chat, messages, stream) + return s.AsyncChatSolo(ctx, userID, chat, messages, stream) } // Spawn goroutine for the async pipeline. All remaining phases run inside. @@ -324,9 +325,9 @@ func (s *ChatPipelineService) AsyncChat( } } if modelType == "chat" { - textAttachmentsList, imageAttachments = splitFileAttachments(files, false) + textAttachmentsList, imageAttachments = splitFileAttachments(userID, files, false) } else { - textAttachmentsList, imageFiles = splitFileAttachments(files, true) + textAttachmentsList, imageFiles = splitFileAttachments(userID, files, true) } attachments = strings.Join(textAttachmentsList, "\n\n") common.Debug("Resolved attachments", @@ -1278,6 +1279,7 @@ func (s *ChatPipelineService) AsyncChat( // Equivalent to Python's async_chat_solo() in dialog_service.py:289-337. func (s *ChatPipelineService) AsyncChatSolo( ctx context.Context, + userID string, chat *entity.Chat, messages []map[string]interface{}, stream bool, @@ -1321,11 +1323,11 @@ func (s *ChatPipelineService) AsyncChatSolo( isImage2Text := modelType == "image2text" if len(messages) > 0 { if files, hasFiles := messages[len(messages)-1]["files"]; hasFiles { - attachmentsStr = s.processFileAttachments(files) + attachmentsStr = s.processFileAttachments(userID, files) if isImage2Text { imageFiles = s.extractRawImageURLs(files) } else { - imageFiles = s.extractImageFiles(files) + imageFiles = s.extractImageFiles(userID, files) } } } @@ -1581,12 +1583,12 @@ func (s *ChatPipelineService) AsyncChatSolo( // extractImageFiles extracts data-URI image attachments from the files list. // Mirrors Python split_file_attachments raw mode. -func (s *ChatPipelineService) extractImageFiles(files interface{}) []string { +func (s *ChatPipelineService) extractImageFiles(userID string, files interface{}) []string { // ── File-dict mode ── if fileDicts, ok := parseFileDicts(files); ok { fileSvc := NewFileService() // Use raw=false to get base64 data URIs for images. - _, images, err := fileSvc.GetFileContents(fileDicts, false) + _, images, err := fileSvc.GetFileContents(userID, fileDicts, false) if err != nil { common.Warn("GetFileContents failed in extractImageFiles", zap.Error(err)) @@ -2045,11 +2047,11 @@ func lastUserQuestion(messages []map[string]interface{}) string { // // When files are file dicts (Python-compatible format), calls // FileService.GetFileContents to fetch actual blobs from storage. -func (s *ChatPipelineService) processFileAttachments(files interface{}) string { +func (s *ChatPipelineService) processFileAttachments(userID string, files interface{}) string { // ── File-dict mode ── if fileDicts, ok := parseFileDicts(files); ok { fileSvc := NewFileService() - texts, _, err := fileSvc.GetFileContents(fileDicts, false) + texts, _, err := fileSvc.GetFileContents(userID, fileDicts, false) if err != nil { common.Warn("GetFileContents failed in processFileAttachments", zap.Error(err)) @@ -2100,11 +2102,11 @@ func (s *ChatPipelineService) processFileAttachments(files interface{}) string { // URIs → image files. // - raw=true: all items go to textAttachments (Python's FileService.get_files // with raw=True pre-separates images, so non-image content arrives here). -func splitFileAttachments(files interface{}, raw bool) (textAttachments []string, imageAttachments []string) { +func splitFileAttachments(userID string, files interface{}, raw bool) (textAttachments []string, imageAttachments []string) { // ── Mode 1: file dicts (Python-compatible) ── if fileDicts, ok := parseFileDicts(files); ok { fileSvc := NewFileService() - texts, images, err := fileSvc.GetFileContents(fileDicts, raw) + texts, images, err := fileSvc.GetFileContents(userID, fileDicts, raw) if err != nil { common.Warn("GetFileContents failed, falling back to string splitting", zap.Error(err)) diff --git a/internal/service/chat_pipeline_test.go b/internal/service/chat_pipeline_test.go index 0bb2025173..8338eb9993 100644 --- a/internal/service/chat_pipeline_test.go +++ b/internal/service/chat_pipeline_test.go @@ -80,7 +80,7 @@ func TestAsyncChat_RejectsNonUserLastMessage(t *testing.T) { {"role": "user", "content": "first"}, {"role": "assistant", "content": "last message must not be assistant"}, } - _, err := s.AsyncChat(context.Background(), dialForTest(""), messages, false, nil) + _, err := s.AsyncChat(context.Background(), "user-1", dialForTest(""), messages, false, nil) if err == nil { t.Fatal("expected error for non-user last message, got nil") } @@ -93,7 +93,7 @@ func TestAsyncChat_RejectsNonUserLastMessage(t *testing.T) { // service should return an error before spawning the goroutine. func TestAsyncChat_EmptyMessages(t *testing.T) { s := &ChatPipelineService{} - _, err := s.AsyncChat(context.Background(), dialForTest(""), nil, false, nil) + _, err := s.AsyncChat(context.Background(), "user-1", dialForTest(""), nil, false, nil) if err == nil { t.Fatal("expected error for empty messages, got nil") } diff --git a/internal/service/chat_session.go b/internal/service/chat_session.go index 786296fa50..7e14b3fe75 100644 --- a/internal/service/chat_session.go +++ b/internal/service/chat_session.go @@ -56,7 +56,7 @@ type userTenantStore interface { } type chatPipelineRunner interface { - AsyncChat(ctx context.Context, chat *entity.Chat, messages []map[string]interface{}, stream bool, kwargs map[string]interface{}) (<-chan AsyncChatResult, error) + AsyncChat(ctx context.Context, userID string, chat *entity.Chat, messages []map[string]interface{}, stream bool, kwargs map[string]interface{}) (<-chan AsyncChatResult, error) } // chunkFeedbackApplier is the dispatch seam for chunk-level feedback @@ -1297,7 +1297,7 @@ func (s *ChatSessionService) Completion(userID string, conversationID string, me if kwargs == nil { kwargs = map[string]interface{}{} } - resultChan, err := s.pipeline.AsyncChat(context.Background(), dialog, messages, false, kwargs) + resultChan, err := s.pipeline.AsyncChat(context.Background(), userID, dialog, messages, false, kwargs) if err != nil { return nil, err } @@ -1383,7 +1383,7 @@ func (s *ChatSessionService) CompletionStream(ctx context.Context, userID string if kwargs == nil { kwargs = map[string]interface{}{} } - resultChan, err := s.pipeline.AsyncChat(ctx, dialog, messages, true, kwargs) + resultChan, err := s.pipeline.AsyncChat(ctx, userID, dialog, messages, true, kwargs) if err != nil { streamChan <- fmt.Sprintf("data: %s\n\n", fmt.Sprintf(`{"code": 500, "message": "%s", "data": {"answer": "**ERROR**: %s", "reference": []}}`, err.Error(), err.Error())) return err @@ -1556,7 +1556,7 @@ func (s *ChatSessionService) ChatCompletions( } // --- 6. Run pipeline --- - resultChan, err := s.pipeline.AsyncChat(ctx, dialog, requestMsg, stream, kwargs) + resultChan, err := s.pipeline.AsyncChat(ctx, userID, dialog, requestMsg, stream, kwargs) if err != nil { return fail(err) } diff --git a/internal/service/chat_session_test.go b/internal/service/chat_session_test.go index 969a63ed1e..1fad7a3e6c 100644 --- a/internal/service/chat_session_test.go +++ b/internal/service/chat_session_test.go @@ -158,9 +158,11 @@ func (f *fakeTenantStore) GetTenantIDsByUserID(userID string) ([]string, error) type fakePipeline struct { resultChan <-chan AsyncChatResult err error + userID string } -func (f *fakePipeline) AsyncChat(ctx context.Context, chat *entity.Chat, messages []map[string]interface{}, stream bool, kwargs map[string]interface{}) (<-chan AsyncChatResult, error) { +func (f *fakePipeline) AsyncChat(ctx context.Context, userID string, chat *entity.Chat, messages []map[string]interface{}, stream bool, kwargs map[string]interface{}) (<-chan AsyncChatResult, error) { + f.userID = userID return f.resultChan, f.err } @@ -824,6 +826,9 @@ func TestCompletion_Success(t *testing.T) { if ans != "Hello world" { t.Fatalf("expected answer 'Hello world', got %q", ans) } + if pipeline.userID != "user-1" { + t.Fatalf("pipeline userID = %q, want user-1", pipeline.userID) + } got := parseMessages(store.sessions["session-1"].Message) if len(got) != 3 { @@ -840,6 +845,64 @@ func TestCompletion_Success(t *testing.T) { } } +func TestChatCompletionsPassesRequestUserIDToPipeline(t *testing.T) { + store := newFakeSessionStore() + store.sessions["session-1"] = &entity.ChatSession{ + ID: "session-1", + DialogID: "dialog-1", + Message: json.RawMessage(`[{"role":"assistant","content":"Welcome!"}]`), + Reference: json.RawMessage(`[]`), + } + store.dialogs["dialog-1"] = &entity.Chat{ + ID: "dialog-1", + TenantID: "tenant-owner", + LLMID: "chat@factory", + LLMSetting: entity.JSONMap{}, + PromptConfig: entity.JSONMap{ + "parameters": []interface{}{}, + }, + } + store.dialogExists["tenant-owner|dialog-1"] = true + + pipeline := &fakePipeline{ + resultChan: makeResultChan( + AsyncChatResult{Answer: "ok", Final: true, Reference: map[string]interface{}{"chunks": []interface{}{}}}, + ), + } + + svc := &ChatSessionService{ + chatSessionDAO: store, + userTenantDAO: &fakeTenantStore{tenantIDs: []string{"tenant-owner"}}, + pipeline: pipeline, + } + + _, err := svc.ChatCompletions( + context.Background(), + "user-1", + "dialog-1", + "session-1", + []map[string]interface{}{{"role": "user", "content": "hi"}}, + "", + nil, + "", + nil, + nil, + false, + false, + false, + nil, + ) + if err != nil { + t.Fatalf("ChatCompletions failed: %v", err) + } + if pipeline.userID != "user-1" { + t.Fatalf("pipeline userID = %q, want request user user-1", pipeline.userID) + } + if pipeline.userID == store.dialogs["dialog-1"].TenantID { + t.Fatalf("pipeline used dialog tenant %q instead of request user", pipeline.userID) + } +} + func TestCompletion_EmptyMessages(t *testing.T) { svc := &ChatSessionService{ chatSessionDAO: &fakeSessionStore{}, diff --git a/internal/service/file.go b/internal/service/file.go index 34d683cbaa..24dc5f70be 100644 --- a/internal/service/file.go +++ b/internal/service/file.go @@ -1090,7 +1090,7 @@ func (s *FileService) DownloadAgentFile(tenantID, location string) ([]byte, erro // for the given file dicts. // - raw=false: images returned as base64 data URIs in images; non-images parsed and returned as text. // - raw=true: images returned as raw bytes in images; non-images parsed and returned as text. -func (s *FileService) GetFileContents(fileDicts []map[string]interface{}, raw bool) (texts []string, images []string, err error) { +func (s *FileService) GetFileContents(uid string, fileDicts []map[string]interface{}, raw bool) (texts []string, images []string, err error) { storageImpl := storage.GetStorageFactory().GetStorage() if storageImpl == nil { return nil, nil, fmt.Errorf("storage not initialized") @@ -1105,6 +1105,9 @@ func (s *FileService) GetFileContents(fileDicts []map[string]interface{}, raw bo if ferr != nil || file == nil || file.Location == nil || *file.Location == "" { continue } + if !s.checkFileTeamPermission(file, uid) { + return nil, nil, fmt.Errorf("No authorization.") + } data, derr := storageImpl.Get(file.ParentID, *file.Location) if derr != nil || len(data) == 0 { continue diff --git a/internal/service/file_test.go b/internal/service/file_test.go index cacfd126c6..2e9664e41b 100644 --- a/internal/service/file_test.go +++ b/internal/service/file_test.go @@ -11,6 +11,9 @@ import ( "testing" "time" + "github.com/glebarez/sqlite" + "gorm.io/gorm" + "ragflow/internal/dao" "ragflow/internal/entity" "ragflow/internal/storage" @@ -23,6 +26,7 @@ type fakeStorage struct { blob []byte err error exists bool + getCalls int } func testFileService() *FileService { @@ -46,6 +50,7 @@ func (f *fakeStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) } func (f *fakeStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) { + f.getCalls++ f.lastBucket = bucket f.lastFnm = fnm return f.blob, f.err @@ -79,6 +84,145 @@ func (f *fakeStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool panic("not implemented in fakeStorage") } +func setupFileContentPermissionDB(t *testing.T, accessible bool) { + t.Helper() + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ + TranslateError: true, + }) + if err != nil { + t.Fatalf("failed to open sqlite: %v", err) + } + if err := db.AutoMigrate( + &entity.File{}, + &entity.File2Document{}, + &entity.Document{}, + &entity.Knowledgebase{}, + &entity.Tenant{}, + &entity.UserTenant{}, + ); err != nil { + t.Fatalf("failed to migrate file content tables: %v", err) + } + + location := "doc.txt" + if err := db.Create(&entity.File{ + ID: "file-1", + ParentID: "bucket-1", + TenantID: "owner-user", + Name: "doc.txt", + Location: &location, + Type: "file", + }).Error; err != nil { + t.Fatalf("insert file: %v", err) + } + if err := db.Create(&entity.Document{ + ID: "doc-1", + KbID: "kb-owner", + ParserID: "naive", + ParserConfig: entity.JSONMap{}, + Status: sptr(string(entity.StatusValid)), + }).Error; err != nil { + t.Fatalf("insert document: %v", err) + } + fileID := "file-1" + docID := "doc-1" + if err := db.Create(&entity.File2Document{ + ID: "f2d-1", + FileID: &fileID, + DocumentID: &docID, + }).Error; err != nil { + t.Fatalf("insert file2document: %v", err) + } + if err := db.Create(&entity.Knowledgebase{ + ID: "kb-owner", + TenantID: "tenant-owner", + Name: "owner-kb", + EmbdID: "embd-1", + CreatedBy: "owner-user", + Permission: string(entity.TenantPermissionTeam), + Status: sptr(string(entity.StatusValid)), + }).Error; err != nil { + t.Fatalf("insert knowledgebase: %v", err) + } + if err := db.Create(&entity.Tenant{ + ID: "tenant-owner", + LLMID: "llm-1", + EmbdID: "embd-1", + ASRID: "asr-1", + Status: sptr(string(entity.StatusValid)), + }).Error; err != nil { + t.Fatalf("insert tenant: %v", err) + } + if accessible { + if err := db.Create(&entity.UserTenant{ + ID: "ut-user-1", + UserID: "user-1", + TenantID: "tenant-owner", + Role: "normal", + Status: sptr(string(entity.StatusValid)), + }).Error; err != nil { + t.Fatalf("insert user_tenant: %v", err) + } + } + + orig := dao.DB + dao.DB = db + t.Cleanup(func() { dao.DB = orig }) +} + +func TestFileService_GetFileContents_NotAccessible(t *testing.T) { + setupFileContentPermissionDB(t, false) + + mockStorage := &fakeStorage{blob: []byte("secret")} + factory := storage.GetStorageFactory() + originalStorage := factory.GetStorage() + factory.SetStorage(mockStorage) + t.Cleanup(func() { factory.SetStorage(originalStorage) }) + + svc := testFileService() + texts, images, err := svc.GetFileContents("user-1", []map[string]interface{}{{"id": "file-1"}}, false) + if err == nil { + t.Fatal("expected authorization error") + } + if err.Error() != "No authorization." { + t.Fatalf("unexpected error: %v", err) + } + if len(texts) != 0 || len(images) != 0 { + t.Fatalf("expected no content, got texts=%v images=%v", texts, images) + } + if mockStorage.getCalls != 0 { + t.Fatalf("storage should not be read without permission, got %d calls", mockStorage.getCalls) + } +} + +func TestFileService_GetFileContents_Accessible(t *testing.T) { + setupFileContentPermissionDB(t, true) + + mockStorage := &fakeStorage{blob: []byte("allowed content")} + factory := storage.GetStorageFactory() + originalStorage := factory.GetStorage() + factory.SetStorage(mockStorage) + t.Cleanup(func() { factory.SetStorage(originalStorage) }) + + svc := testFileService() + texts, images, err := svc.GetFileContents("user-1", []map[string]interface{}{{"id": "file-1"}}, false) + if err != nil { + t.Fatalf("GetFileContents failed: %v", err) + } + if len(images) != 0 { + t.Fatalf("expected no images, got %v", images) + } + if len(texts) != 1 || !strings.Contains(texts[0], "allowed content") { + t.Fatalf("unexpected texts: %v", texts) + } + if mockStorage.getCalls != 1 { + t.Fatalf("storage get calls = %d, want 1", mockStorage.getCalls) + } + if mockStorage.lastBucket != "bucket-1" || mockStorage.lastFnm != "doc.txt" { + t.Fatalf("storage read %s/%s, want bucket-1/doc.txt", mockStorage.lastBucket, mockStorage.lastFnm) + } +} + func TestFileService_DownloadAgentFile_Success(t *testing.T) { // Setup mock storage expectedBlob := []byte("fake file content") diff --git a/internal/service/openai_chat.go b/internal/service/openai_chat.go index f298277924..0964bf72a4 100644 --- a/internal/service/openai_chat.go +++ b/internal/service/openai_chat.go @@ -319,7 +319,7 @@ func (s *OpenAIChatService) OpenAIChatCompletions(c *gin.Context, userID, chatID chatKwargs["doc_ids"] = docIDsStr } - asyncResults, asyncErr := s.pipeline.AsyncChat(ctx, dialog, filteredMessages, openaiReq.Stream, chatKwargs) + asyncResults, asyncErr := s.pipeline.AsyncChat(ctx, userID, dialog, filteredMessages, openaiReq.Stream, chatKwargs) if asyncErr != nil { s.writeDataError(c, asyncErr.Error()) return