package store import ( "context" "reflect" "gorm.io/gorm" "gorm.io/gorm/clause" ) // tenantScopedMarker:模型实现此(空)接口即"受租户隔离"——gorm tenant 插件据此对其 // 查询/更新/删除自动加 `tenant_id = `、创建时自动填 tenant_id。 // 未实现的模型(User/Tenant/TenantMember/Prompt/Model/Pricing…)不受影响。 // 未实现方法用小写→只有本包模型能标记(KB/Agent 已标;Doc/DocLink/Task 待增量铺开)。 type tenantScopedMarker interface{ isTenantScoped() } // ---- 请求上下文携带 tenant_id(中间件注入 → 传到 store → 插件读取)---- type ctxKeyTenant struct{} // WithTenant 把 tenant_id 放进 context(空则原样返回,避免污染系统/回填查询)。 func WithTenant(ctx context.Context, tenantID string) context.Context { if tenantID == "" { return ctx } return context.WithValue(ctx, ctxKeyTenant{}, tenantID) } func tenantFromCtx(ctx context.Context) string { if ctx == nil { return "" } s, _ := ctx.Value(ctxKeyTenant{}).(string) return s } // registerTenantScope 挂 gorm 回调,实现"统一强制、别靠人肉"的租户隔离。 func registerTenantScope(db *gorm.DB) { q := func(d *gorm.DB) { addTenantWhere(d) } _ = db.Callback().Query().Before("gorm:query").Register("tenant:query", q) _ = db.Callback().Update().Before("gorm:update").Register("tenant:update", q) _ = db.Callback().Delete().Before("gorm:delete").Register("tenant:delete", q) _ = db.Callback().Create().Before("gorm:create").Register("tenant:create", setTenantOnCreate) } func isTenantScopedStmt(db *gorm.DB) bool { if db.Statement.Schema == nil { return false } _, ok := reflect.New(db.Statement.Schema.ModelType).Interface().(tenantScopedMarker) return ok } // addTenantWhere:受租户模型 + ctx 有 tenant → 追加 tenant_id 过滤。 // ctx 无 tenant(系统/回填/未登录)→ 不过滤(这些路径本就需跨租户;用户面由中间件保证有 tenant)。 func addTenantWhere(db *gorm.DB) { if !isTenantScopedStmt(db) { return } if tid := tenantFromCtx(db.Statement.Context); tid != "" { db.Statement.AddClause(clause.Where{Exprs: []clause.Expression{ clause.Eq{Column: clause.Column{Table: db.Statement.Table, Name: "tenant_id"}, Value: tid}, }}) } } // setTenantOnCreate:受租户模型 + ctx 有 tenant → 强制把 tenant_id 设为 ctx 租户(防越权写他租)。 func setTenantOnCreate(db *gorm.DB) { if !isTenantScopedStmt(db) { return } if tid := tenantFromCtx(db.Statement.Context); tid != "" { db.Statement.SetColumn("TenantID", tid) } } // BackfillRowTenants 给受租户表的存量行回填 tenant_id = owner 的默认租户(幂等:只填空的)。 // 系统操作、无 ctx 租户 → 不受插件过滤影响。当前覆盖 KB / Agent(同步路径已接插件)。 func (p *Postgres) BackfillRowTenants(ctx context.Context) error { if p.db == nil { return errStoreDisabled } // 子查询:owner(user) → 其默认(最早 active)租户。 sub := "SELECT tenant_id FROM sundynix_tenant_member m WHERE m.user_id = t.owner AND m.status = 'active' ORDER BY m.created_at ASC LIMIT 1" for _, tbl := range []string{"sundynix_kb", "sundynix_agent"} { q := "UPDATE " + tbl + " t SET tenant_id = (" + sub + ") WHERE (t.tenant_id IS NULL OR t.tenant_id = '')" if err := p.db.WithContext(ctx).Exec(q).Error; err != nil { return err } } return nil }