feat(govai): 0617 优化首批 — 安全/私有化/深度研究/服务层/可观测性

借鉴 odysseus 的能力设计，全程净室实现、零 AGPL 代码、不引入 AGPL 依赖。 T1 提示注入防护: pkg/promptguard 包裹外部/知识库内容为不可信数据，buildMessages 移出 system 指令区。 T2 安全 CI: .github/workflows(ci+security: govulncheck/gitleaks/actionlint/hadolint/trivy)+dependabot+.hadolint.yaml；go.mod 加 toolchain go1.25.11 修复 20 个 stdlib CVE。 T3 管理员 2FA: 迁移 000016 + RFC6238 TOTP/备份码(pkg/auth, 零依赖) + 登录流程集成(后端)。 T4 本地模型: LLM/embedding 支持本地 vLLM/Ollama(OpenAI 兼容, 鉴权头条件发送, NoAuth) + docs/local-deploy.md。 T6 深度研究: 迁移 000017 + Python research-worker(净室多步流水线, 检索避开 SearXNG) + Go research 服务/handler/路由。 T7 service 层: 新增 internal/service/{research,twofa}, 2FA 业务逻辑从胖 handler 下沉, 接口注入可单测。 T10 缓存/可观测性: internal/cache(Redis+内存, 优雅降级) 接入 store 热点列表; Prometheus 指标+/metrics; docs/openapi.yaml。验证: go build/vet/test ./... 全绿(8 包); research-worker 12 单测过; 真实 PG 应用迁移并烟测。
2026-06-17 17:52:47 +08:00
parent 97feb42afb
commit c949204662
55 changed files with 4341 additions and 25 deletions
@@ -0,0 +1,139 @@
+// Package twofa 提供两步验证（2FA / TOTP）的业务编排（service 层）。
+//
+// 业务规则（生成密钥/备份码、校验、启用/关闭判定）集中在此，DB 操作通过 Store 接口
+// 注入，便于单测。TOTP/备份码算法复用 pkg/auth。HTTP handler 仅做参数解析与响应。
+package twofa
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/enterprise-ai-platform/server/pkg/auth"
+)
+
+// Issuer 显示在认证器 App 中的发行方名称。
+const Issuer = "政智通 GovAI"
+
+// Store 2FA 持久化接口。
+type Store interface {
+	// GetStatus 返回是否启用与剩余可用备份码数量。
+	GetStatus(ctx context.Context, userID string) (enabled bool, remaining int, err error)
+	// GetSecret 返回用户的 TOTP 密钥（可能为空）。
+	GetSecret(ctx context.Context, userID string) (secret string, err error)
+	// GetSecretAndEnabled 返回密钥与启用状态。
+	GetSecretAndEnabled(ctx context.Context, userID string) (secret string, enabled bool, err error)
+	// SaveEnrollment 原子写入新密钥并重置备份码（未启用）。
+	SaveEnrollment(ctx context.Context, userID, secret string, codeHashes []string) error
+	// Enable 置 totp_enabled=true。
+	Enable(ctx context.Context, userID string) error
+	// Disable 关闭 2FA：清除密钥并删除所有备份码。
+	Disable(ctx context.Context, userID string) error
+	// ConsumeBackupCode 校验并一次性消费备份码，命中返回 true。
+	ConsumeBackupCode(ctx context.Context, userID, code string) (bool, error)
+}
+
+var (
+	ErrAlreadyEnabled = errors.New("两步验证已启用")
+	ErrNotSetup       = errors.New("尚未开始两步验证设置")
+	ErrBadCode        = errors.New("验证码或备份码错误")
+)
+
+// EnrollResult 是开始设置 2FA 的返回。
+type EnrollResult struct {
+	Secret      string
+	OtpauthURI  string
+	BackupCodes []string // 明文，仅返回一次
+}
+
+type Service struct {
+	store       Store
+	backupCount int
+	nowUnix     func() int64
+}
+
+func NewService(store Store) *Service {
+	return &Service{store: store, backupCount: 8, nowUnix: func() int64 { return time.Now().Unix() }}
+}
+
+// Status 返回当前 2FA 状态。
+func (s *Service) Status(ctx context.Context, userID string) (enabled bool, remaining int, err error) {
+	return s.store.GetStatus(ctx, userID)
+}
+
+// Enroll 生成新密钥与备份码并落库（未启用，需 Verify 确认）。
+func (s *Service) Enroll(ctx context.Context, userID, email string) (*EnrollResult, error) {
+	_, enabled, err := s.store.GetSecretAndEnabled(ctx, userID)
+	if err != nil {
+		return nil, err
+	}
+	if enabled {
+		return nil, ErrAlreadyEnabled
+	}
+
+	secret, err := auth.GenerateTOTPSecret()
+	if err != nil {
+		return nil, err
+	}
+	plain, hashes, err := auth.GenerateBackupCodes(s.backupCount)
+	if err != nil {
+		return nil, err
+	}
+	if err := s.store.SaveEnrollment(ctx, userID, secret, hashes); err != nil {
+		return nil, err
+	}
+	return &EnrollResult{
+		Secret:      secret,
+		OtpauthURI:  auth.TOTPProvisioningURI(secret, email, Issuer),
+		BackupCodes: plain,
+	}, nil
+}
+
+// EnableAfterVerify 校验首个验证码并启用 2FA。
+func (s *Service) EnableAfterVerify(ctx context.Context, userID, code string) error {
+	secret, err := s.store.GetSecret(ctx, userID)
+	if err != nil {
+		return err
+	}
+	if secret == "" {
+		return ErrNotSetup
+	}
+	if !auth.ValidateTOTP(secret, code, s.nowUnix()) {
+		return ErrBadCode
+	}
+	return s.store.Enable(ctx, userID)
+}
+
+// Disable 校验 TOTP 或备份码后关闭 2FA。未启用时视为成功（幂等）。
+func (s *Service) Disable(ctx context.Context, userID, code, backupCode string) error {
+	secret, enabled, err := s.store.GetSecretAndEnabled(ctx, userID)
+	if err != nil {
+		return err
+	}
+	if !enabled {
+		return nil
+	}
+	if !s.verify(ctx, userID, secret, code, backupCode) {
+		return ErrBadCode
+	}
+	return s.store.Disable(ctx, userID)
+}
+
+// VerifyLogin 在登录流程中校验 2FA：先试 TOTP，再试备份码（一次性消费）。
+// secret 由调用方在登录查询时一并取出，避免重复查库。
+func (s *Service) VerifyLogin(ctx context.Context, userID, secret, totpCode, backupCode string) bool {
+	return s.verify(ctx, userID, secret, totpCode, backupCode)
+}
+
+func (s *Service) verify(ctx context.Context, userID, secret, totpCode, backupCode string) bool {
+	if totpCode != "" && secret != "" && auth.ValidateTOTP(secret, totpCode, s.nowUnix()) {
+		return true
+	}
+	if backupCode != "" {
+		ok, err := s.store.ConsumeBackupCode(ctx, userID, backupCode)
+		if err == nil && ok {
+			return true
+		}
+	}
+	return false
+}
@@ -0,0 +1,139 @@
+package twofa
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/enterprise-ai-platform/server/pkg/auth"
+)
+
+type fakeStore struct {
+	enabled    bool
+	remaining  int
+	secret     string
+	saved      bool
+	enabled2   bool // Enable 被调用
+	disabled   bool // Disable 被调用
+	backupOK   bool
+	getErr     error
+}
+
+func (f *fakeStore) GetStatus(ctx context.Context, userID string) (bool, int, error) {
+	return f.enabled, f.remaining, f.getErr
+}
+func (f *fakeStore) GetSecret(ctx context.Context, userID string) (string, error) {
+	return f.secret, f.getErr
+}
+func (f *fakeStore) GetSecretAndEnabled(ctx context.Context, userID string) (string, bool, error) {
+	return f.secret, f.enabled, f.getErr
+}
+func (f *fakeStore) SaveEnrollment(ctx context.Context, userID, secret string, codeHashes []string) error {
+	f.saved = true
+	f.secret = secret
+	return nil
+}
+func (f *fakeStore) Enable(ctx context.Context, userID string) error  { f.enabled2 = true; return nil }
+func (f *fakeStore) Disable(ctx context.Context, userID string) error { f.disabled = true; return nil }
+func (f *fakeStore) ConsumeBackupCode(ctx context.Context, userID, code string) (bool, error) {
+	return f.backupOK, nil
+}
+
+const fixedNow int64 = 1_700_000_000
+
+func newSvc(store Store) *Service {
+	s := NewService(store)
+	s.nowUnix = func() int64 { return fixedNow }
+	return s
+}
+
+func TestEnroll_RejectsWhenAlreadyEnabled(t *testing.T) {
+	svc := newSvc(&fakeStore{enabled: true})
+	if _, err := svc.Enroll(context.Background(), "u1", "a@b.c"); !errors.Is(err, ErrAlreadyEnabled) {
+		t.Fatalf("已启用应返回 ErrAlreadyEnabled，实际: %v", err)
+	}
+}
+
+func TestEnroll_GeneratesAndSaves(t *testing.T) {
+	store := &fakeStore{}
+	svc := newSvc(store)
+	res, err := svc.Enroll(context.Background(), "u1", "admin@govai.gov.cn")
+	if err != nil {
+		t.Fatalf("Enroll 出错: %v", err)
+	}
+	if res.Secret == "" || len(res.BackupCodes) != 8 {
+		t.Fatalf("应返回密钥与 8 个备份码，实际 codes=%d", len(res.BackupCodes))
+	}
+	if !store.saved {
+		t.Fatal("应调用 SaveEnrollment")
+	}
+	if res.OtpauthURI == "" {
+		t.Fatal("应返回 otpauth URI")
+	}
+}
+
+func TestEnableAfterVerify(t *testing.T) {
+	secret, _ := auth.GenerateTOTPSecret()
+	code, _ := auth.TOTPCodeAt(secret, fixedNow)
+
+	// 未设置密钥
+	if err := newSvc(&fakeStore{secret: ""}).EnableAfterVerify(context.Background(), "u1", code); !errors.Is(err, ErrNotSetup) {
+		t.Fatalf("无密钥应返回 ErrNotSetup，实际: %v", err)
+	}
+	// 错误验证码
+	if err := newSvc(&fakeStore{secret: secret}).EnableAfterVerify(context.Background(), "u1", "000000"); !errors.Is(err, ErrBadCode) {
+		t.Fatalf("错误码应返回 ErrBadCode，实际: %v", err)
+	}
+	// 正确验证码
+	store := &fakeStore{secret: secret}
+	if err := newSvc(store).EnableAfterVerify(context.Background(), "u1", code); err != nil {
+		t.Fatalf("正确码应成功，实际: %v", err)
+	}
+	if !store.enabled2 {
+		t.Fatal("应调用 Enable")
+	}
+}
+
+func TestDisable(t *testing.T) {
+	secret, _ := auth.GenerateTOTPSecret()
+	code, _ := auth.TOTPCodeAt(secret, fixedNow)
+
+	// 未启用 → 幂等成功，不调用 Disable
+	store0 := &fakeStore{enabled: false}
+	if err := newSvc(store0).Disable(context.Background(), "u1", "", ""); err != nil || store0.disabled {
+		t.Fatalf("未启用应幂等返回 nil 且不调用 Disable，err=%v disabled=%v", err, store0.disabled)
+	}
+	// 启用 + 正确 TOTP
+	store1 := &fakeStore{enabled: true, secret: secret}
+	if err := newSvc(store1).Disable(context.Background(), "u1", code, ""); err != nil {
+		t.Fatalf("正确 TOTP 应成功: %v", err)
+	}
+	if !store1.disabled {
+		t.Fatal("应调用 Disable")
+	}
+	// 启用 + 备份码
+	store2 := &fakeStore{enabled: true, secret: secret, backupOK: true}
+	if err := newSvc(store2).Disable(context.Background(), "u1", "", "backup-xxxx"); err != nil || !store2.disabled {
+		t.Fatalf("备份码应可关闭，err=%v disabled=%v", err, store2.disabled)
+	}
+	// 启用 + 错误码
+	store3 := &fakeStore{enabled: true, secret: secret, backupOK: false}
+	if err := newSvc(store3).Disable(context.Background(), "u1", "000000", "bad"); !errors.Is(err, ErrBadCode) {
+		t.Fatalf("错误码应返回 ErrBadCode，实际: %v", err)
+	}
+}
+
+func TestVerifyLogin(t *testing.T) {
+	secret, _ := auth.GenerateTOTPSecret()
+	code, _ := auth.TOTPCodeAt(secret, fixedNow)
+
+	if !newSvc(&fakeStore{}).VerifyLogin(context.Background(), "u1", secret, code, "") {
+		t.Fatal("正确 TOTP 应通过")
+	}
+	if !newSvc(&fakeStore{backupOK: true}).VerifyLogin(context.Background(), "u1", secret, "", "backup") {
+		t.Fatal("有效备份码应通过")
+	}
+	if newSvc(&fakeStore{backupOK: false}).VerifyLogin(context.Background(), "u1", secret, "000000", "bad") {
+		t.Fatal("错误码应不通过")
+	}
+}
@@ -0,0 +1,121 @@
+package twofa
+
+import (
+	"context"
+
+	"github.com/enterprise-ai-platform/server/pkg/auth"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+// pgxStore 基于 pgx 连接池的 Store 实现。
+type pgxStore struct {
+	pool *pgxpool.Pool
+}
+
+func NewPgxStore(pool *pgxpool.Pool) Store {
+	return &pgxStore{pool: pool}
+}
+
+func (s *pgxStore) GetStatus(ctx context.Context, userID string) (bool, int, error) {
+	var enabled bool
+	var remaining int
+	err := s.pool.QueryRow(ctx,
+		`SELECT u.totp_enabled,
+		        (SELECT COUNT(*) FROM user_backup_codes b WHERE b.user_id = u.id AND b.used_at IS NULL)
+		 FROM users u WHERE u.id = $1`, userID).Scan(&enabled, &remaining)
+	return enabled, remaining, err
+}
+
+func (s *pgxStore) GetSecret(ctx context.Context, userID string) (string, error) {
+	var secret *string
+	if err := s.pool.QueryRow(ctx,
+		`SELECT totp_secret FROM users WHERE id = $1`, userID).Scan(&secret); err != nil {
+		return "", err
+	}
+	if secret == nil {
+		return "", nil
+	}
+	return *secret, nil
+}
+
+func (s *pgxStore) GetSecretAndEnabled(ctx context.Context, userID string) (string, bool, error) {
+	var secret *string
+	var enabled bool
+	if err := s.pool.QueryRow(ctx,
+		`SELECT totp_secret, totp_enabled FROM users WHERE id = $1`, userID).Scan(&secret, &enabled); err != nil {
+		return "", false, err
+	}
+	if secret == nil {
+		return "", enabled, nil
+	}
+	return *secret, enabled, nil
+}
+
+func (s *pgxStore) SaveEnrollment(ctx context.Context, userID, secret string, codeHashes []string) error {
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback(ctx)
+
+	if _, err = tx.Exec(ctx,
+		`UPDATE users SET totp_secret = $2, totp_enabled = false WHERE id = $1`, userID, secret); err != nil {
+		return err
+	}
+	if _, err = tx.Exec(ctx, `DELETE FROM user_backup_codes WHERE user_id = $1`, userID); err != nil {
+		return err
+	}
+	for _, h := range codeHashes {
+		if _, err = tx.Exec(ctx,
+			`INSERT INTO user_backup_codes (user_id, code_hash) VALUES ($1, $2)`, userID, h); err != nil {
+			return err
+		}
+	}
+	return tx.Commit(ctx)
+}
+
+func (s *pgxStore) Enable(ctx context.Context, userID string) error {
+	_, err := s.pool.Exec(ctx, `UPDATE users SET totp_enabled = true WHERE id = $1`, userID)
+	return err
+}
+
+func (s *pgxStore) Disable(ctx context.Context, userID string) error {
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback(ctx)
+	if _, err = tx.Exec(ctx,
+		`UPDATE users SET totp_enabled = false, totp_secret = NULL WHERE id = $1`, userID); err != nil {
+		return err
+	}
+	if _, err = tx.Exec(ctx, `DELETE FROM user_backup_codes WHERE user_id = $1`, userID); err != nil {
+		return err
+	}
+	return tx.Commit(ctx)
+}
+
+func (s *pgxStore) ConsumeBackupCode(ctx context.Context, userID, code string) (bool, error) {
+	rows, err := s.pool.Query(ctx,
+		`SELECT id, code_hash FROM user_backup_codes WHERE user_id = $1 AND used_at IS NULL`, userID)
+	if err != nil {
+		return false, err
+	}
+	type bc struct{ id, hash string }
+	var list []bc
+	for rows.Next() {
+		var x bc
+		if rows.Scan(&x.id, &x.hash) == nil {
+			list = append(list, x)
+		}
+	}
+	rows.Close() // 先释放连接再执行更新
+
+	for _, x := range list {
+		if auth.CheckBackupCode(code, x.hash) {
+			_, _ = s.pool.Exec(ctx, `UPDATE user_backup_codes SET used_at = NOW() WHERE id = $1`, x.id)
+			return true, nil
+		}
+	}
+	return false, nil
+}