sundynix-agentix/sundynix-gateway/internal/store/tenant_scope.go

package store

import (
	"context"
	"reflect"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// tenantScopedMarker：模型实现此(空)接口即"受租户隔离"——gorm tenant 插件据此对其
// 查询/更新/删除自动加 `tenant_id = <ctx租户>`、创建时自动填 tenant_id。
// 未实现的模型（User/Tenant/TenantMember/Prompt/Model/Pricing…）不受影响。
// 未实现方法用小写→只有本包模型能标记（KB/Agent 已标；Doc/DocLink/Task 待增量铺开）。
type tenantScopedMarker interface{ isTenantScoped() }

// ---- 请求上下文携带 tenant_id（中间件注入 → 传到 store → 插件读取）----
type ctxKeyTenant struct{}
type ctxKeySystem struct{}

// WithTenant 把 tenant_id 放进 context（空则原样返回，避免污染系统/回填查询）。
func WithTenant(ctx context.Context, tenantID string) context.Context {
	if tenantID == "" {
		return ctx
	}
	return context.WithValue(ctx, ctxKeyTenant{}, tenantID)
}

// WithoutTenant 标记本次操作为「系统/跨租户」——即使 ctx 里带着某租户（如 admin 自己的），
// 插件也**不**加 tenant 过滤。用于 admin 系统级聚合、跨租户后台任务等需全平台可见的路径。
func WithoutTenant(ctx context.Context) context.Context {
	return context.WithValue(ctx, ctxKeySystem{}, true)
}

func isSystemCtx(ctx context.Context) bool {
	if ctx == nil {
		return false
	}
	b, _ := ctx.Value(ctxKeySystem{}).(bool)
	return b
}

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) || isSystemCtx(db.Statement.Context) {
		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) || isSystemCtx(db.Statement.Context) {
		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
}