Go and Python: fix IDOR issue of search completions (#16864)

### Summary

In Go and python implementation, the dataset / KB id isn't validated if
it is accessible by this user.

---------

Signed-off-by: Jin Hai <haijin.chn@gmail.com>

api/apps/restful_apis/search_api.py

@@ -25,6 +25,7 @@ from api.apps import current_user, login_required
 from api.constants import DATASET_NAME_LIMIT
 from api.db.db_models import DB
 from api.db.services import duplicate_name
+from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.services.search_service import SearchService
 from api.db.services.user_service import TenantService, UserTenantService
 from common.misc_utils import get_uuid
@@ -220,6 +221,11 @@ async def completion(search_id):
     if not kb_ids:
         return get_data_error_result(message="`kb_ids` is required.")
 
+    # check if the kb_ids is accessible for this user
+    for kb_id in kb_ids:
+        if not KnowledgebaseService.accessible(kb_id=kb_id, user_id=uid):
+            return get_data_error_result(message=f"You don't own the dataset {kb_id}")
+
     async def stream():
         nonlocal req, uid, kb_ids, search_config
         try:

internal/dao/kb.go

@@ -259,9 +259,9 @@ func (dao *KnowledgebaseDAO) GetDetail(kbID string) (*entity.KnowledgebaseDetail
 // 2. If user is the owner tenant, return true
 // 3. If permission is "me", only owner tenant can access
 // 4. If permission is "team", user must be a member of the tenant
-func (dao *KnowledgebaseDAO) Accessible(kbID, userID string) bool {
+func (dao *KnowledgebaseDAO) Accessible(datasetID, userID string) bool {
 	var kb entity.Knowledgebase
-	err := DB.Where("id = ? AND status = ?", kbID, string(entity.StatusValid)).First(&kb).Error
+	err := DB.Where("id = ? AND status = ?", datasetID, string(entity.StatusValid)).First(&kb).Error
 	if err != nil {
 		return false
 	}

internal/dao/search.go

@@ -197,8 +197,8 @@ func (dao *SearchDAO) QueryByTenantIDAndID(tenantID string, searchID string) ([]
 
 // DeleteByID deletes a search by ID (soft delete by setting status to "0")
 // Reference: Python common_service.py::delete_by_id
-func (dao *SearchDAO) DeleteByID(id string) error {
-	return DB.Model(&entity.Search{}).Where("id = ?", id).Update("status", "0").Error
+func (dao *SearchDAO) DeleteByID(tenantID, id string) error {
+	return DB.Model(&entity.Search{}).Where("tenant_id = ? AND id = ?", tenantID, id).Update("status", "0").Error
 }
 
 // Accessible4Deletion checks if a search can be deleted by a specific user

internal/handler/search.go

@@ -427,7 +427,7 @@ func (h *SearchHandler) Completion(c *gin.Context) {
 	adapter := &service.TenantStreamAdapter{LLM: h.streamLLM, TenantID: plan.UserID, ModelID: plan.ModelID}
 
 	hadError := false
-	for delta := range h.askService.StreamWithOptions(c.Request.Context(), adapter, plan.UserID, plan.Question, plan.KBIDs, plan.Options) {
+	for delta := range h.askService.StreamWithOptions(c.Request.Context(), adapter, plan.UserID, plan.Question, plan.DatasetIDs, plan.Options) {
 		switch delta.Kind {
 		case service.AskDeltaAnswer:
 			writer.Write(c, sseAnswer(delta.Value, nil, false))

internal/service/ask_service.go

@@ -113,16 +113,16 @@ func (s *AskService) Stream(ctx context.Context, llm StreamingLLM, userID, quest
 
 // StreamWithOptions runs Stream while allowing callers such as saved Search
 // apps to pass search_config retrieval options through to RetrievalTest.
-func (s *AskService) StreamWithOptions(ctx context.Context, llm StreamingLLM, userID, question string, kbIDs []string, opts AskStreamOptions) <-chan AskDelta {
+func (s *AskService) StreamWithOptions(ctx context.Context, llm StreamingLLM, userID, question string, datasetIDs []string, opts AskStreamOptions) <-chan AskDelta {
 	out := make(chan AskDelta, 32)
 	go func() {
 		defer close(out)
-		s.run(ctx, llm, userID, question, kbIDs, opts, out)
+		s.run(ctx, llm, userID, question, datasetIDs, opts, out)
 	}()
 	return out
 }
 
-func (s *AskService) run(ctx context.Context, llm StreamingLLM, userID, question string, kbIDs []string, opts AskStreamOptions, out chan<- AskDelta) {
+func (s *AskService) run(ctx context.Context, llm StreamingLLM, userID, question string, datasetIDs []string, opts AskStreamOptions, out chan<- AskDelta) {
 	// Phase 1: Retrieval.
 	topK := DefaultAskTopK
 	if opts.TopK != nil {
@@ -138,7 +138,7 @@ func (s *AskService) run(ctx context.Context, llm StreamingLLM, userID, question
 	}
 
 	req := &RetrievalTestRequest{
-		Datasets:               common.StringSlice(kbIDs),
+		Datasets:               common.StringSlice(datasetIDs),
 		Question:               question,
 		DocIDs:                 opts.DocIDs,
 		UseKG:                  opts.UseKG,

internal/service/search.go

@@ -32,6 +32,7 @@ import (
 type SearchService struct {
 	searchDAO     *dao.SearchDAO
 	userTenantDAO *dao.UserTenantDAO
+	datasetDAO    *dao.KnowledgebaseDAO
 	tenantService *TenantService
 }
 
@@ -40,6 +41,7 @@ func NewSearchService() *SearchService {
 	return &SearchService{
 		searchDAO:     dao.NewSearchDAO(),
 		userTenantDAO: dao.NewUserTenantDAO(),
+		datasetDAO:    dao.NewKnowledgebaseDAO(),
 		tenantService: NewTenantService(),
 	}
 }
@@ -326,7 +328,7 @@ func (s *SearchService) DeleteSearch(userID string, searchID string) error {
 
 	// Step 2: Execute delete (same as Python SearchService.delete_by_id)
 	// Python: cls.model.delete().where(cls.model.id == pid).execute()
-	if err = s.searchDAO.DeleteByID(searchID); err != nil {
+	if err = s.searchDAO.DeleteByID(userID, searchID); err != nil {
 		return fmt.Errorf("failed to delete search App %s: %w", searchID, err)
 	}
 
@@ -346,12 +348,12 @@ func (s *SearchService) AccessibleForCompletion(userID string, searchID string)
 }
 
 type SearchCompletionPlan struct {
-	UserID   string
-	SearchID string
-	Question string
-	KBIDs    []string
-	ModelID  string
-	Options  AskStreamOptions
+	UserID     string
+	SearchID   string
+	Question   string
+	DatasetIDs []string
+	ModelID    string
+	Options    AskStreamOptions
 }
 
 func (s *SearchService) PrepareCompletion(userID, searchID string, req *SearchCompletionsRequest) (*SearchCompletionPlan, common.ErrorCode, error) {
@@ -376,21 +378,28 @@ func (s *SearchService) PrepareCompletion(userID, searchID string, req *SearchCo
 		return nil, common.CodeServerError, err
 	}
 	if !accessible {
-		return nil, common.CodeAuthenticationError, fmt.Errorf("No authorization.")
+		return nil, common.CodeAuthenticationError, fmt.Errorf("no authorization")
 	}
 
 	searchDetail, err := s.GetDetail(searchID)
 	if err != nil || searchDetail == nil {
-		return nil, common.CodeDataError, fmt.Errorf("Cannot find search %s", searchID)
+		return nil, common.CodeDataError, fmt.Errorf("cannot find search %s", searchID)
 	}
 	searchConfig := searchConfigMapFromValue(searchDetail["search_config"])
 
-	kbIDs := stringSliceFromSearchConfig(searchConfig["kb_ids"])
-	if len(kbIDs) == 0 {
-		kbIDs = stringSliceFromSearchConfig(req.KBIDs)
+	datasetIDs := stringSliceFromSearchConfig(searchConfig["kb_ids"])
+	if len(datasetIDs) == 0 {
+		datasetIDs = stringSliceFromSearchConfig(req.KBIDs)
 	}
-	if len(kbIDs) == 0 {
-		return nil, common.CodeDataError, fmt.Errorf("`kb_ids` is required.")
+	if len(datasetIDs) == 0 {
+		return nil, common.CodeDataError, fmt.Errorf("`kb_ids` is required")
+	}
+
+	for _, datasetID := range datasetIDs {
+		accessible = s.datasetDAO.Accessible(datasetID, userID)
+		if !accessible {
+			return nil, common.CodeAuthenticationError, fmt.Errorf("no authorization for dataset %s", datasetID)
+		}
 	}
 
 	modelID, _ := stringFromSearchConfig(searchConfig["chat_id"])
@@ -399,19 +408,20 @@ func (s *SearchService) PrepareCompletion(userID, searchID string, req *SearchCo
 		if tenantSvc == nil {
 			tenantSvc = NewTenantService()
 		}
-		defaultModel, err := tenantSvc.GetDefaultModelName(userID, entity.ModelTypeChat)
+		var defaultModelName string
+		defaultModelName, err = tenantSvc.GetDefaultModelName(userID, entity.ModelTypeChat)
 		if err == nil {
-			modelID = strings.TrimSpace(defaultModel)
+			modelID = strings.TrimSpace(defaultModelName)
 		}
 	}
 
 	return &SearchCompletionPlan{
-		UserID:   userID,
-		SearchID: searchID,
-		Question: question,
-		KBIDs:    kbIDs,
-		ModelID:  modelID,
-		Options:  askOptionsFromSearchConfig(searchID, searchConfig),
+		UserID:     userID,
+		SearchID:   searchID,
+		Question:   question,
+		DatasetIDs: datasetIDs,
+		ModelID:    modelID,
+		Options:    askOptionsFromSearchConfig(searchID, searchConfig),
 	}, common.CodeSuccess, nil
 }
 

test/testcases/restful_api/conftest.py

@@ -247,3 +247,12 @@ def ensure_parsed_document(rest_client, create_document):
         return dataset_id, document_id
 
     return _ensure
+
+
+@pytest.hookimpl(tryfirst=True, hookwrapper=True)
+def pytest_runtest_protocol(item, nextitem):
+    import time
+    start = time.perf_counter()
+    yield
+    duration = time.perf_counter() - start
+    print(f"  [{duration:.3f}s]")

test/testcases/restful_api/test_openai_compatible.py

@@ -117,7 +117,7 @@ def test_openai_compatible_invalid_chat(rest_client):
     assert expected_message in payload["message"], payload
 
 
-@pytest.mark.p2
+@pytest.mark.p3
 def test_openai_compatible_nonstream_shape(rest_client, create_chat):
     chat_id = create_chat("restful_openai_nonstream_chat")
     res = rest_client.post(
@@ -127,7 +127,7 @@ def test_openai_compatible_nonstream_shape(rest_client, create_chat):
             "messages": [{"role": "user", "content": "hello"}],
             "stream": False,
         },
-        timeout=60,
+        timeout=120,
     )
     assert res.status_code == 200
     payload = res.json()
@@ -157,7 +157,7 @@ def test_openai_compatible_defaults_to_nonstream_when_stream_is_missing(rest_cli
             "model": "model",
             "messages": [{"role": "user", "content": "hello"}],
         },
-        timeout=60,
+        timeout=120,
     )
     assert res.status_code == 200
     assert "application/json" in res.headers.get("Content-Type", ""), res.headers.get("Content-Type", "")
@@ -182,7 +182,7 @@ def test_openai_compatible_nonstream_with_reference_output_shape(rest_client, cr
                 "reference_metadata": {"include": True, "fields": ["author"]},
             },
         },
-        timeout=60,
+        timeout=120,
     )
     assert res.status_code == 200
     payload = res.json()

test/testcases/restful_api/test_searches.py

@@ -148,11 +148,7 @@ def test_search_completion_sse_shape_when_kb_ids_provided(rest_client, search_re
         timeout=60,
     )
     assert res.status_code == 200
-    content_type = res.headers.get("Content-Type", "")
-    assert "text/event-stream" in content_type, content_type
+    payload = res.json()
+    assert payload["code"] == 102, payload
+    assert "You don't own the dataset nonexistent_dataset" in payload["message"], payload
 
-    events = _sse_events(res.text)
-    assert events, res.text
-    parsed = [json.loads(evt) for evt in events]
-    assert isinstance(parsed[0], dict), parsed
-    assert parsed[-1].get("data") is True, parsed[-1]