123

cmii-uav-watchdog/services/auth_service.go

@@ -12,6 +12,13 @@ import (
 	"time"
 )
 
+// 单例相关变量
+var (
+	authServiceInstance *AuthService
+	authServiceOnce     sync.Once
+	authServiceMutex    sync.Mutex
+)
+
 // AuthService 授权服务
 type AuthService struct {
 	mu                sync.RWMutex
@@ -21,36 +28,51 @@ type AuthService struct {
 	initialized       bool
 }
 
-// NewAuthService 创建授权服务
+// NewAuthService 创建授权服务（单例模式）
 func NewAuthService() *AuthService {
-	service := &AuthService{
-		hostInfoSet: make(map[string]models2.HostInfo),
-		totpService: NewTOTPService(),
-		initialized: false,
-	}
+	// 使用sync.Once确保初始化逻辑只执行一次
+	authServiceOnce.Do(func() {
+		authServiceMutex.Lock()
+		defer authServiceMutex.Unlock()
 
-	// 尝试从本地加载授权信息
-	service.loadAuthorizationInfo()
-
-	// 判断 项目级别的 TOTP密钥是否为空
-	// 若为空 则生成一个 二级TOTP密钥 然后持久化写入到授权文件中
-	if service.authorizationInfo.SecondTOTPSecret == "" {
-		secondTOTPSecret, err := service.totpService.GenerateTOTPSecret()
-		if err != nil {
-			log.Printf("生成二级TOTP密钥失败: %v", err)
-			return nil
+		// 如果实例已存在，直接返回
+		if authServiceInstance != nil {
+			return
 		}
-		service.authorizationInfo.SecondTOTPSecret = secondTOTPSecret
 
-		// 持久化写入到授权文件中
-		err = service.saveAuthorizationInfo()
-		if err != nil {
-			log.Printf("持久化写入授权文件失败: %v", err)
-			return nil
+		// 创建新实例
+		service := &AuthService{
+			hostInfoSet: make(map[string]models2.HostInfo),
+			totpService: NewTOTPService(),
+			initialized: false,
 		}
-	}
 
-	return service
+		// 尝试从本地加载授权信息
+		service.loadAuthorizationInfo()
+
+		// 判断 项目级别的 TOTP密钥是否为空
+		// 若为空 则生成一个 二级TOTP密钥 然后持久化写入到授权文件中
+		if service.authorizationInfo.SecondTOTPSecret == "" {
+			secondTOTPSecret, err := service.totpService.GenerateTierTwoTOTPSecret()
+			if err != nil {
+				log.Printf("生成二级TOTP密钥失败: %v", err)
+				return
+			}
+			service.authorizationInfo.SecondTOTPSecret = secondTOTPSecret
+
+			// 持久化写入到授权文件中
+			err = service.saveAuthorizationInfo()
+			if err != nil {
+				log.Printf("持久化写入授权文件失败: %v", err)
+				return
+			}
+		}
+
+		// 设置全局实例
+		authServiceInstance = service
+	})
+
+	return authServiceInstance
 }
 
 // AddHostInfo 添加主机信息
@@ -73,7 +95,7 @@ func (as *AuthService) GenerateAuthorizationFile() (*models2.AuthorizationFile,
 	}
 
 	// 生成TOTP验证码
-	totpCode, err := as.totpService.GenerateTOTP()
+	totpCode, err := as.totpService.GenerateTierTwoTOTPCode(as.authorizationInfo.SecondTOTPSecret)
 	if err != nil {
 		return nil, err
 	}
@@ -118,7 +140,7 @@ func (as *AuthService) ProcessAuthorizationCode(code models2.AuthorizationCode)
 	defer as.mu.Unlock()
 
 	// 验证TOTP
-	if err := as.totpService.VerifyTOTP(code.TOTPCode); err != nil {
+	if !as.totpService.VerifyTierTwoTOTPCode(code.TOTPCode, as.authorizationInfo.SecondTOTPSecret) {
 		return errors.New("无效的授权码: TOTP验证失败")
 	}