Agent-System/internal/auth/bootstrap.go

package auth

import (
	"context"
	"errors"
	"fmt"
	"os"
	"strings"

	"chat-agent-system/internal/config"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
)

func isUniqueViolation(err error) bool {
	var pgErr *pgconn.PgError
	if errors.As(err, &pgErr) && pgErr.Code == "23505" {
		return true
	}
	msg := err.Error()
	return strings.Contains(msg, "23505") || strings.Contains(msg, "duplicate key")
}

// BootstrapAdminFromConfig 确保配置的管理员账户在 PG 中存在且具备 role-admin。
func BootstrapAdminFromConfig(ctx context.Context, store AuthStore, cfg *config.Config) error {
	if store == nil || cfg == nil {
		return nil
	}
	adminUser := cfg.AdminUsername
	adminPass := cfg.AdminPassword
	if adminUser == "" {
		adminUser = "admin"
	}
	if adminPass == "" {
		return nil
	}

	existing, err := store.GetUserByUsername(ctx, adminUser)
	if err == nil && existing != nil {
		if existing.RoleID != "role-admin" {
			if _, roleErr := store.UpdateUserRole(ctx, existing.ID, "role-admin", "系统管理员"); roleErr != nil {
				return fmt.Errorf("bootstrap admin UpdateUserRole: %w", roleErr)
			}
		}
		if existing.PasswordHash == "" {
			hash, hashErr := HashPassword(adminPass)
			if hashErr != nil {
				return hashErr
			}
			_, err = store.UpdateUserPassword(ctx, existing.ID, hash, false)
			return err
		}
		return nil
	}

	hash, err := HashPassword(adminPass)
	if err != nil {
		return fmt.Errorf("bootstrap admin hash: %w", err)
	}
	_, err = store.CreateUser(ctx, &UserAccount{
		Username:           adminUser,
		Name:               "系统管理员",
		PasswordHash:       hash,
		Status:             StatusActive,
		MustChangePassword: false,
		RoleID:             "role-admin",
		Position:           "系统管理员",
		Department:         "系统运维部",
		EmployeeID:         "SYS-000",
	})
	return err
}

// BootstrapDevSeedUsers 在 dev/test 环境补齐演示账号（与 MemoryStore 种子对齐）。
func BootstrapDevSeedUsers(ctx context.Context, store AuthStore, cfg *config.Config) error {
	if store == nil || cfg == nil {
		return nil
	}
	if os.Getenv("SEED_DEMO_USERS") == "false" || os.Getenv("DISABLE_DEV_SEED_USERS") == "true" {
		return nil
	}
	if cfg.Env != "dev" && cfg.Env != "test" && cfg.Env != "" {
		return nil
	}
	seeds := []struct {
		username string
		name     string
		roleID   string
		position string
		dept     string
		empID    string
	}{
		{"zhanggong", "张工", "role-tech-lead", "技术主管", "骨科耗材研发组", "RD-001"},
		{"lizhuguan", "李主管", "role-tech-lead", "技术主管", "骨科耗材研发组", "RD-002"},
		{"wangzong", "王总监", "role-director", "研发总监", "研发中心", "RD-003"},
		{"guest", "访客", "role-guest", "待指派", "外协中心", "EXT-999"},
	}
	passHash, err := HashPassword("123456")
	if err != nil {
		return err
	}
	for _, su := range seeds {
		if _, err := store.GetUserByUsername(ctx, su.username); err == nil {
			continue
		}
		_, err := store.CreateUser(ctx, &UserAccount{
			Username:     su.username,
			Name:         su.name,
			PasswordHash: passHash,
			Status:       StatusActive,
			RoleID:       su.roleID,
			Position:     su.position,
			Department:   su.dept,
			EmployeeID:   su.empID,
		})
		if err != nil {
			// Concurrent NewServer bootstrap and sequence drift after fixed-id
			// goose seeds can race on users_pkey/username. If the demo user
			// exists afterwards, treat the seed as successful.
			if _, getErr := store.GetUserByUsername(ctx, su.username); getErr == nil {
				continue
			}
			if isUniqueViolation(err) {
				if _, retryErr := store.CreateUser(ctx, &UserAccount{
					Username:     su.username,
					Name:         su.name,
					PasswordHash: passHash,
					Status:       StatusActive,
					RoleID:       su.roleID,
					Position:     su.position,
					Department:   su.dept,
					EmployeeID:   su.empID,
				}); retryErr == nil {
					continue
				} else if _, getErr := store.GetUserByUsername(ctx, su.username); getErr == nil {
					continue
				}
			}
			return fmt.Errorf("seed user %s: %w", su.username, err)
		}
	}
	return nil
}

// ReactivateDevSeedUsers 将 dev/test 演示账号恢复为 active（集成测试可能临时停用）。
func ReactivateDevSeedUsers(ctx context.Context, store AuthStore, cfg *config.Config) error {
	if store == nil || cfg == nil {
		return nil
	}
	if os.Getenv("SEED_DEMO_USERS") == "false" || os.Getenv("DISABLE_DEV_SEED_USERS") == "true" {
		return nil
	}
	if cfg.Env != "dev" && cfg.Env != "test" && cfg.Env != "" {
		return nil
	}
	for _, username := range []string{"zhanggong", "lizhuguan", "wangzong", "guest"} {
		user, err := store.GetUserByUsername(ctx, username)
		if err != nil || user == nil {
			continue
		}
		if user.Status != StatusActive {
			if _, err := store.UpdateUserStatus(ctx, user.ID, StatusActive); err != nil {
				return fmt.Errorf("reactivate %s: %w", username, err)
			}
		}
	}
	return nil
}

// BootstrapIAM 在 PG IAM 就绪后执行管理员与 dev 种子初始化。
func BootstrapIAM(ctx context.Context, pool *pgxpool.Pool, store AuthStore, cfg *config.Config) error {
	if !IAMTablesReady(ctx, pool) {
		return nil
	}
	if err := BootstrapAdminFromConfig(ctx, store, cfg); err != nil {
		return err
	}
	if err := BootstrapDevSeedUsers(ctx, store, cfg); err != nil {
		return err
	}
	return ReactivateDevSeedUsers(ctx, store, cfg)
}