123

cmii-uav-watchdog-otp/totp/totp.go

@@ -77,6 +77,14 @@ type ValidateOpts struct {
 	Encoder otp.Encoder
 }
 
+func (opts *ValidateOpts) ConvertToValidateOpts(generateOpts GenerateOpts) {
+	opts.Period = generateOpts.Period
+	opts.Skew = 1
+	opts.Digits = generateOpts.Digits
+	opts.Algorithm = generateOpts.Algorithm
+	opts.Encoder = otp.EncoderDefault
+}
+
 // GenerateCodeCustom takes a timepoint and produces a passcode using a
 // secret and the provided opts. (Under the hood, this is making an adapted
 // call to hcmii-uav-watchdog-otp.GenerateCodeCustom)

cmii-uav-watchdog/main.go

@@ -17,7 +17,7 @@ func main() {
 		return
 	}
 
-	// 初始化授权服务
+	// 初始化授权服务（使用单例模式）
 	authService := services.NewAuthService()
 	if authService == nil {
 		log.Fatalf("初始化授权服务失败")

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,8 +28,19 @@ type AuthService struct {
 	initialized       bool
 }
 
-// NewAuthService 创建授权服务
+// NewAuthService 创建授权服务（单例模式）
 func NewAuthService() *AuthService {
+	// 使用sync.Once确保初始化逻辑只执行一次
+	authServiceOnce.Do(func() {
+		authServiceMutex.Lock()
+		defer authServiceMutex.Unlock()
+
+		// 如果实例已存在，直接返回
+		if authServiceInstance != nil {
+			return
+		}
+
+		// 创建新实例
 		service := &AuthService{
 			hostInfoSet: make(map[string]models2.HostInfo),
 			totpService: NewTOTPService(),
@@ -35,10 +53,10 @@ func NewAuthService() *AuthService {
 		// 判断 项目级别的 TOTP密钥是否为空
 		// 若为空 则生成一个 二级TOTP密钥 然后持久化写入到授权文件中
 		if service.authorizationInfo.SecondTOTPSecret == "" {
-		secondTOTPSecret, err := service.totpService.GenerateTOTPSecret()
+			secondTOTPSecret, err := service.totpService.GenerateTierTwoTOTPSecret()
 			if err != nil {
 				log.Printf("生成二级TOTP密钥失败: %v", err)
-			return nil
+				return
 			}
 			service.authorizationInfo.SecondTOTPSecret = secondTOTPSecret
 
@@ -46,11 +64,15 @@ func NewAuthService() *AuthService {
 			err = service.saveAuthorizationInfo()
 			if err != nil {
 				log.Printf("持久化写入授权文件失败: %v", err)
-			return nil
+				return
 			}
 		}
 
-	return service
+		// 设置全局实例
+		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验证失败")
 	}
 

cmii-uav-watchdog/services/heartbeat_service.go

@@ -3,6 +3,7 @@ package services
 import (
 	"cmii-uav-watchdog-common/models"
 	"errors"
+	"log"
 	"time"
 )
 
@@ -27,6 +28,11 @@ func (hs *HeartbeatService) ProcessHeartbeat(req models.HeartbeatRequest) (*mode
 		return nil, errors.New("无效的时间戳")
 	}
 
+	secondTOTPSecret := hs.authService.authorizationInfo.SecondTOTPSecret
+	if secondTOTPSecret == "" {
+		return nil, errors.New("二级TOTP密钥为空")
+	}
+
 	// 添加主机信息到集合
 	hs.authService.AddHostInfo(req.HostInfo)
 
@@ -36,17 +42,26 @@ func (hs *HeartbeatService) ProcessHeartbeat(req models.HeartbeatRequest) (*mode
 			Authorized:       false,
 			TOTPCode:         "",
 			Timestamp:        time.Now().Unix(),
-			SecondTOTPSecret: "",
+			SecondTOTPSecret: secondTOTPSecret,
 		}, nil
 	}
 
-	// 检查totp密码是都有效
+	// 检查totp验证码是否有效
+	if !hs.totpService.VerifyTierTwoTOTPCode(req.TOTPCode, secondTOTPSecret) {
+		// 解析认证主机的相关信息
+
+		// 计算 请求时间与当前时间的时间差
+		diff := time.Now().Unix() - req.Timestamp
+		log.Printf("心跳请求时间与当前时间的时间差: %d", diff)
+
+		return nil, errors.New("无效的TOTP验证码，请检查系统时间是否正确！")
+	}
 
 	// 检查主机是否已授权
 	authorized := hs.authService.IsHostAuthorized(req.HostInfo)
 
 	// 生成TOTP验证码
-	totpCode, err := hs.totpService.GenerateTOTP()
+	totpCode, err := hs.totpService.GenerateTierTwoTOTPCode(secondTOTPSecret)
 	if err != nil {
 		return nil, err
 	}

cmii-uav-watchdog/services/totp_service.go

@@ -2,7 +2,6 @@ package services
 
 import (
 	"cmii-uav-watchdog/config"
-	"errors"
 	"log"
 	"time"
 
@@ -10,6 +9,17 @@ import (
 	"cmii-uav-watchdog-otp/totp"
 )
 
+var tierTwoTOTPSecretOpts = totp.GenerateOpts{
+	SecretSize:  32,
+	Issuer:      "cmii-uav-watchdog",
+	AccountName: "cmii-uav-watchdog",
+	Period:      30,
+	Secret:      []byte{},
+	Digits:      otp.DigitsSix,
+	Algorithm:   otp.AlgorithmSHA1,
+	Rand:        nil,
+}
+
 // TOTPService TOTP服务
 type TOTPService struct {
 	secret string
@@ -22,8 +32,8 @@ func NewTOTPService() *TOTPService {
 	}
 }
 
-// GenerateTOTP 生成TOTP验证码
+// GenerateTierOneTOTP 生成一级TOTP验证码
-func (ts *TOTPService) GenerateTOTP() (string, error) {
+func (ts *TOTPService) GenerateTierOneTOTP() (string, error) {
 	// 使用当前时间生成TOTP
 	code, err := totp.GenerateCode(ts.secret, time.Now())
 	if err != nil {
@@ -33,29 +43,20 @@ func (ts *TOTPService) GenerateTOTP() (string, error) {
 	return code, nil
 }
 
-// VerifyTOTP 验证TOTP验证码
+// VerifyTierOneTOTP 验证一级TOTP验证码
-func (ts *TOTPService) VerifyTOTP(code string) error {
+func (ts *TOTPService) VerifyTierOneTOTP(code string) bool {
 	// 验证TOTP
 	valid := totp.Validate(code, ts.secret)
 	if !valid {
-		return errors.New("无效的TOTP验证码")
+		return false
 	}
 
-	return nil
+	return true
 }
 
-// GenerateTOTPSecret 生成TOTP密钥
+// GenerateTierTwoTOTPSecret 生成二级TOTP密钥
-func (ts *TOTPService) GenerateTOTPSecret() (string, error) {
+func (ts *TOTPService) GenerateTierTwoTOTPSecret() (string, error) {
-	secret, err := totp.Generate(totp.GenerateOpts{
+	secret, err := totp.Generate(tierTwoTOTPSecretOpts)
-		SecretSize:  32,
-		Issuer:      "cmii-uav-watchdog",
-		AccountName: "cmii-uav-watchdog",
-		Period:      30,
-		Secret:      []byte{},
-		Digits:      otp.DigitsSix,
-		Algorithm:   otp.AlgorithmSHA1,
-		Rand:        nil,
-	})
 	if err != nil {
 		log.Printf("生成TOTP密钥失败: %v", err)
 		return "", err
@@ -63,3 +64,23 @@ func (ts *TOTPService) GenerateTOTPSecret() (string, error) {
 
 	return secret.Secret(), nil
 }
+
+// GenerateTierTwoTOTPCode 生成二级TOTP验证码
+func (ts *TOTPService) GenerateTierTwoTOTPCode(secret string) (string, error) {
+	code, err := totp.GenerateCode(secret, time.Now())
+	if err != nil {
+		return "", err
+	}
+	return code, nil
+}
+
+// VerifyTierTwoTOTPCode 验证二级TOTP验证码
+func (ts *TOTPService) VerifyTierTwoTOTPCode(code string, secret string) bool {
+	validateOpts := totp.ValidateOpts{}
+	validateOpts.ConvertToValidateOpts(tierTwoTOTPSecretOpts)
+	valid, err := totp.ValidateCustom(code, secret, time.Now(), validateOpts)
+	if err != nil {
+		return false
+	}
+	return valid
+}

go.mod

@@ -1,3 +1,3 @@
-module cmii-uav-watchdog
+module cmii-uav-watchdog-project
 
-go 1.23
+go 1.24