Merge pull request #1265 from trheyi/main

Add Yao custom fields and enhance permission checks
This commit is contained in:
Max 2025-11-05 11:32:47 +08:00 committed by GitHub
commit 79cffc5cd4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 168 additions and 20 deletions

View file

@ -412,6 +412,11 @@ func makeJob(data []byte) (*Job, error) {
job.CreatedBy = "system"
}
// If YaoCreatedBy is set, use it as the created by
if job.YaoCreatedBy != "" {
job.CreatedBy = job.YaoCreatedBy
}
// Set default enabled to true if not specified
// Note: Go's zero value for bool is false, so we need to explicitly check if it was set
// Since we can't distinguish between explicitly set false and zero value,

View file

@ -135,6 +135,12 @@ type Job struct {
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
// Yao custom fields
YaoCreatedBy string `json:"__yao_created_by,omitempty"` // nullable: true
YaoUpdatedBy string `json:"__yao_updated_by,omitempty"` // nullable: true
YaoTeamID string `json:"__yao_team_id,omitempty"` // nullable: true
YaoTenantID string `json:"__yao_tenant_id,omitempty"`
// Relationships
Category *Category `json:"category,omitempty"`
Executions []Execution `json:"executions,omitempty"`

View file

@ -14,12 +14,14 @@ import (
"github.com/yaoapp/yao/attachment"
"github.com/yaoapp/yao/job"
"github.com/yaoapp/yao/kb"
"github.com/yaoapp/yao/openapi/oauth/authorized"
oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/openapi/response"
)
// CreateDocumentRecord creates a document record in the database immediately
// This is called synchronously when the API request comes in
func CreateDocumentRecord(ctx context.Context, req *AddFileRequest, jobID string) error {
func CreateDocumentRecord(ctx context.Context, authInfo *oauthtypes.AuthorizedInfo, req *AddFileRequest, jobID string) error {
// Check if kb.Instance is available
if kb.Instance == nil {
return fmt.Errorf("knowledge base not initialized")
@ -70,6 +72,11 @@ func CreateDocumentRecord(ctx context.Context, req *AddFileRequest, jobID string
"job_id": jobID,
}
// With create scope
if authInfo != nil {
documentData = authInfo.WithCreateScope(documentData)
}
// Add base request fields
req.BaseUpsertRequest.AddBaseFields(documentData)
@ -153,7 +160,7 @@ func HandleFileContent(ctx context.Context, req *AddFileRequest) error {
// AddFileHandler processes a file addition request with business logic only
// This function combines both document creation and content processing for sync operations
func AddFileHandler(ctx context.Context, req *AddFileRequest, jobID ...string) error {
func AddFileHandler(ctx context.Context, authInfo *oauthtypes.AuthorizedInfo, req *AddFileRequest, jobID ...string) error {
// Validate request
if err := req.Validate(); err != nil {
return err
@ -171,7 +178,7 @@ func AddFileHandler(ctx context.Context, req *AddFileRequest, jobID ...string) e
}
// Create document record
if err := CreateDocumentRecord(ctx, req, jid); err != nil {
if err := CreateDocumentRecord(ctx, authInfo, req, jid); err != nil {
return err
}
@ -181,8 +188,31 @@ func AddFileHandler(ctx context.Context, req *AddFileRequest, jobID ...string) e
// addFileWithRequest processes a file addition with pre-parsed request using Gin context
func addFileWithRequest(c *gin.Context, req *AddFileRequest) {
// Check collection permission
authInfo := authorized.GetInfo(c)
hasPermission, err := checkCollectionPermission(authInfo, req.CollectionID)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// 403 Forbidden
if !hasPermission {
errorResp := &response.ErrorResponse{
Code: response.ErrAccessDenied.Code,
ErrorDescription: "Forbidden: No permission to update collection",
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// Use the business logic function
err := AddFileHandler(c.Request.Context(), req)
err = AddFileHandler(c.Request.Context(), authInfo, req)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
@ -304,6 +334,28 @@ func AddFileAsync(c *gin.Context) {
log.Info("AddFileAsync: Generated doc_id: %s", req.DocID)
// Check collection permission
authInfo := authorized.GetInfo(c)
hasPermission, err := checkCollectionPermission(authInfo, req.CollectionID)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// 403 Forbidden
if !hasPermission {
errorResp := &response.ErrorResponse{
Code: response.ErrAccessDenied.Code,
ErrorDescription: "Forbidden: No permission to update collection",
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// Step 1: Get job options with defaults
jobName, jobDescription, jobIcon, jobCategory := req.GetJobOptions(
"Knowledge Base File Processing", // default name
@ -322,6 +374,11 @@ func AddFileAsync(c *gin.Context) {
jobCreateData["icon"] = jobIcon
}
// With create scope
if authInfo != nil {
jobCreateData = authInfo.WithCreateScope(jobCreateData)
}
// Create and save Job in one step to get JobID
j, err := job.OnceAndSave(job.GOROUTINE, jobCreateData)
if err != nil {
@ -337,7 +394,7 @@ func AddFileAsync(c *gin.Context) {
log.Info("AddFileAsync: Job created and saved with ID: %s", j.JobID)
// Step 2: Create document record immediately with job_id
err = CreateDocumentRecord(c.Request.Context(), &req, j.JobID)
err = CreateDocumentRecord(c.Request.Context(), authInfo, &req, j.JobID)
if err != nil {
log.Error("AddFileAsync: Failed to create document record: %v", err)
errorResp := &response.ErrorResponse{
@ -446,8 +503,7 @@ func ProcessAddFile(process *process.Process) interface{} {
if jid, ok := reqMap["job_id"].(string); ok {
jobID = jid
}
err = CreateDocumentRecord(ctx, req, jobID)
err = CreateDocumentRecord(ctx, authorized.ProcessAuthInfo(process), req, jobID)
if err != nil {
exception.New("failed to create document record: %s", 500, err.Error()).Throw()
}

View file

@ -15,6 +15,7 @@ import (
"github.com/yaoapp/yao/openapi/oauth/authorized"
oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/openapi/response"
"github.com/yaoapp/yao/openapi/utils"
)
// Collection Management Handlers
@ -752,7 +753,7 @@ func getProviderSettings(providerID, optionValue, locale string) (*ProviderSetti
}
// checkCollectionPermission checks if the user has permission to access the collection
func checkCollectionPermission(authInfo *oauthtypes.AuthorizedInfo, collectionID string) (bool, error) {
func checkCollectionPermission(authInfo *oauthtypes.AuthorizedInfo, collectionID string, readable ...bool) (bool, error) {
// Team Permission validation)
if authInfo == nil {
@ -771,7 +772,7 @@ func checkCollectionPermission(authInfo *oauthtypes.AuthorizedInfo, collectionID
}
collection, err := config.FindCollection(collectionID, model.QueryParam{
Select: []interface{}{"collection_id", "__yao_created_by", "__yao_updated_by", "__yao_team_id"},
Select: []interface{}{"collection_id", "__yao_created_by", "__yao_updated_by", "__yao_team_id", "public", "share"},
Wheres: []model.QueryWhere{
{Column: "collection_id", Value: collectionID},
},
@ -786,6 +787,18 @@ func checkCollectionPermission(authInfo *oauthtypes.AuthorizedInfo, collectionID
return false, fmt.Errorf("failed to find collection: %v", err)
}
// if readable is true, check if the collection is readable
if len(readable) > 0 && readable[0] {
if utils.ToBool(collection["public"]) {
return true, nil
}
// Team only permission validation
if collection["share"] == "team" && authInfo.Constraints.TeamOnly {
return true, nil
}
}
// Combined Team and Owner permission validation
if authInfo.Constraints.TeamOnly && authInfo.Constraints.OwnerOnly {
if collection["__yao_created_by"] == authInfo.UserID && collection["__yao_team_id"] == authInfo.TeamID {

View file

@ -7,6 +7,7 @@ import (
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/graphrag/types"
kbutils "github.com/yaoapp/gou/graphrag/utils"
"github.com/yaoapp/gou/model"
"github.com/yaoapp/yao/attachment"
"github.com/yaoapp/yao/kb"
@ -103,15 +104,9 @@ func ListDocuments(c *gin.Context) {
Select: selectFields,
}
// Get authorized information
authInfo := authorized.GetInfo(c)
// Add filters
var wheres []model.QueryWhere
// Apply permission-based filtering
wheres = append(wheres, AuthFilter(c, authInfo)...)
// Filter by keywords (search in name and description)
if keywords := strings.TrimSpace(c.Query("keywords")); keywords != "" {
wheres = append(wheres, model.QueryWhere{
@ -137,12 +132,41 @@ func ListDocuments(c *gin.Context) {
})
}
// Get authorized information
authInfo := authorized.GetInfo(c)
// Filter by collection_id
if collectionID := strings.TrimSpace(c.Query("collection_id")); collectionID != "" {
wheres = append(wheres, model.QueryWhere{
Column: "collection_id",
Value: collectionID,
})
// If collection_id is provided, validate collection permission
// If not provided, filter by authorization constraints (TeamOnly or OwnerOnly)
collectionID := strings.TrimSpace(c.Query("collection_id"))
if collectionID != "" {
// Validate collection permission
hasPermission, err := checkCollectionPermission(authInfo, collectionID, true)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// 403 Forbidden
if !hasPermission {
errorResp := &response.ErrorResponse{
Code: response.ErrAccessDenied.Code,
ErrorDescription: "Forbidden: No permission to update collection",
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
wheres = append(wheres, model.QueryWhere{Column: "collection_id", Value: collectionID})
} else {
// Filter by authorization constraints
wheres = append(wheres, AuthFilter(c, authInfo)...)
}
// Filter by status (support multiple values separated by comma)
@ -392,6 +416,41 @@ func RemoveDocs(c *gin.Context) {
return
}
// Validate document permissions
collectionIDs := []string{}
authInfo := authorized.GetInfo(c)
for _, docID := range validDocIDs {
collectionID, _ := kbutils.ExtractCollectionIDFromDocID(docID)
if collectionID == "" {
collectionID = "default"
}
collectionIDs = append(collectionIDs, collectionID)
}
for _, collectionID := range collectionIDs {
// Check update permission
hasPermission, err := checkCollectionPermission(authInfo, collectionID)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// 403 Forbidden
if !hasPermission {
errorResp := &response.ErrorResponse{
Code: response.ErrAccessDenied.Code,
ErrorDescription: "Forbidden: No permission to update collection",
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
}
// Remove documents using GraphRAG
deletedCount, err := kb.Instance.RemoveDocs(c.Request.Context(), validDocIDs)
if err != nil {

View file

@ -2,9 +2,18 @@ package authorized
import (
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/yao/openapi/oauth/types"
)
// ProcessAuthInfo extracts authorized information from the process
func ProcessAuthInfo(p *process.Process) *types.AuthorizedInfo {
// TODO: Implement this function
// Get authorized information from the process context
info := &types.AuthorizedInfo{}
return info
}
// GetInfo extracts authorized information from the gin context
// This function reads authorization data that was set by the OAuth guard middleware
func GetInfo(c *gin.Context) *types.AuthorizedInfo {