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>

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
 }