ai_site/platform/internal/tenantperm/store.go

package tenantperm

import (
	"context"
	"database/sql"
	"sync"

	"aijianzhan/platform/internal/authx"
)

// Store 公司权限额度：超管授予，公司内分配不得越界。
// configured=false 表示从未配置 → 读时视为全量默认；configured=true 时允许空额度。
type Store interface {
	Get(ctx context.Context, tenantID int64) ([]string, error)
	Set(ctx context.Context, tenantID int64, perms []string) error
	Allows(ctx context.Context, tenantID int64, perm string) (bool, error)
	EnsureDefault(ctx context.Context, tenantID int64) error
	Configured(ctx context.Context, tenantID int64) (bool, error)
}

type MemoryStore struct {
	mu   sync.Mutex
	by   map[int64][]string // key 存在即已配置
}

func NewMemoryStore() *MemoryStore {
	return &MemoryStore{by: map[int64][]string{}}
}

func (s *MemoryStore) Configured(_ context.Context, tenantID int64) (bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	_, ok := s.by[tenantID]
	return ok, nil
}

func (s *MemoryStore) Get(_ context.Context, tenantID int64) ([]string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if ps, ok := s.by[tenantID]; ok {
		return append([]string{}, ps...), nil
	}
	return authx.CompanyPermCatalog(), nil
}

func (s *MemoryStore) Set(_ context.Context, tenantID int64, perms []string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.by[tenantID] = intersectCatalog(perms) // 可为空切片
	return nil
}

func (s *MemoryStore) Allows(ctx context.Context, tenantID int64, perm string) (bool, error) {
	ps, err := s.Get(ctx, tenantID)
	if err != nil {
		return false, err
	}
	want := authx.NormalizePerm(perm)
	for _, p := range ps {
		if p == want {
			return true, nil
		}
	}
	return false, nil
}

func (s *MemoryStore) EnsureDefault(ctx context.Context, tenantID int64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.by[tenantID]; ok {
		return nil
	}
	s.by[tenantID] = authx.CompanyPermCatalog()
	return nil
}

type PostgresStore struct{ DB *sql.DB }

func NewPostgresStore(db *sql.DB) *PostgresStore { return &PostgresStore{DB: db} }

func (s *PostgresStore) Configured(ctx context.Context, tenantID int64) (bool, error) {
	var n int
	err := s.DB.QueryRowContext(ctx, `
SELECT COUNT(*) FROM platform_meta.tenant_entitlement_state WHERE tenant_id=$1`, tenantID).Scan(&n)
	return n > 0, err
}

func (s *PostgresStore) Get(ctx context.Context, tenantID int64) ([]string, error) {
	ok, err := s.Configured(ctx, tenantID)
	if err != nil {
		return nil, err
	}
	if !ok {
		return authx.CompanyPermCatalog(), nil
	}
	rows, err := s.DB.QueryContext(ctx, `
SELECT perm FROM platform_meta.tenant_permissions WHERE tenant_id=$1 ORDER BY perm`, tenantID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := make([]string, 0)
	for rows.Next() {
		var p string
		if err := rows.Scan(&p); err != nil {
			return nil, err
		}
		out = append(out, authx.NormalizePerm(p))
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return authx.NormalizePerms(out), nil
}

func (s *PostgresStore) Set(ctx context.Context, tenantID int64, perms []string) error {
	allowed := intersectCatalog(perms)
	tx, err := s.DB.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()
	if _, err := tx.ExecContext(ctx, `DELETE FROM platform_meta.tenant_permissions WHERE tenant_id=$1`, tenantID); err != nil {
		return err
	}
	for _, p := range allowed {
		if _, err := tx.ExecContext(ctx, `
INSERT INTO platform_meta.tenant_permissions(tenant_id, perm) VALUES($1,$2)`, tenantID, p); err != nil {
			return err
		}
	}
	if _, err := tx.ExecContext(ctx, `
INSERT INTO platform_meta.tenant_entitlement_state(tenant_id, updated_at)
VALUES($1, now())
ON CONFLICT (tenant_id) DO UPDATE SET updated_at=now()`, tenantID); err != nil {
		return err
	}
	return tx.Commit()
}

func (s *PostgresStore) Allows(ctx context.Context, tenantID int64, perm string) (bool, error) {
	ps, err := s.Get(ctx, tenantID)
	if err != nil {
		return false, err
	}
	want := authx.NormalizePerm(perm)
	for _, p := range ps {
		if p == want {
			return true, nil
		}
	}
	return false, nil
}

func (s *PostgresStore) EnsureDefault(ctx context.Context, tenantID int64) error {
	ok, err := s.Configured(ctx, tenantID)
	if err != nil {
		return err
	}
	if ok {
		return nil
	}
	return s.Set(ctx, tenantID, authx.CompanyPermCatalog())
}

func intersectCatalog(perms []string) []string {
	cat := map[string]struct{}{}
	for _, p := range authx.CompanyPermCatalog() {
		cat[p] = struct{}{}
	}
	out := make([]string, 0, len(perms))
	seen := map[string]struct{}{}
	for _, p := range authx.NormalizePerms(perms) {
		if _, ok := cat[p]; !ok {
			continue
		}
		if _, ok := seen[p]; ok {
			continue
		}
		seen[p] = struct{}{}
		out = append(out, p)
	}
	return out
}

func MustAllow(ctx context.Context, st Store, tenantID int64, want []string) error {
	if st == nil || tenantID <= 0 {
		return nil
	}
	allowance, err := st.Get(ctx, tenantID)
	if err != nil {
		return err
	}
	return authx.AssertWithinEntitlement(want, allowance)
}