backend/internal/service/signature_service.go

package service

import (
	"context"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/x509"
	"encoding/base64"
	"encoding/binary"
	"encoding/pem"
	"fmt"
	"strconv"
	"sync"
	"time"

	"carrotskin/internal/model"
	"carrotskin/internal/repository"
	"carrotskin/pkg/redis"

	"go.uber.org/zap"
)

// 常量定义
const (
	KeySize               = 4096
	ExpirationDays        = 90
	RefreshDays           = 60
	PublicKeyRedisKey     = "yggdrasil:public_key"
	PrivateKeyRedisKey    = "yggdrasil:private_key"
	KeyExpirationRedisKey = "yggdrasil:key_expiration"
	RedisTTL              = 0 // 永不过期，由应用程序管理过期时间
)

// SignatureService 签名服务（导出以便依赖注入）
type SignatureService struct {
	profileRepo repository.ProfileRepository
	redis       *redis.Client
	logger      *zap.Logger

	// 缓存已解析的根密钥对，避免每次签名都重复 PEM 解析（PKCS1 解析开销较大）
	rootKeyMu       sync.RWMutex
	rootPrivateKey  *rsa.PrivateKey
	rootPublicKeyPEM string
}

// NewSignatureService 创建SignatureService实例
func NewSignatureService(
	profileRepo repository.ProfileRepository,
	redisClient *redis.Client,
	logger *zap.Logger,
) *SignatureService {
	return &SignatureService{
		profileRepo: profileRepo,
		redis:       redisClient,
		logger:      logger,
	}
}

// NewKeyPair 生成新的RSA密钥对
func (s *SignatureService) NewKeyPair(ctx context.Context) (*model.KeyPair, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, KeySize)
	if err != nil {
		return nil, fmt.Errorf("生成RSA密钥对失败: %w", err)
	}

	// 获取公钥
	publicKey := &privateKey.PublicKey

	// PEM编码私钥
	privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	privateKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: privateKeyBytes,
	})

	// PEM编码公钥
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return nil, fmt.Errorf("编码公钥失败: %w", err)
	}
	publicKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: publicKeyBytes,
	})

	// 计算时间
	now := time.Now().UTC()
	expiration := now.AddDate(0, 0, ExpirationDays)
	refresh := now.AddDate(0, 0, RefreshDays)

	// 获取Yggdrasil根密钥并签名公钥
	yggPublicKey, yggPrivateKey, err := s.GetOrCreateYggdrasilKeyPair(ctx)
	if err != nil {
		return nil, fmt.Errorf("获取Yggdrasil根密钥失败: %w", err)
	}

	// 构造签名消息（PEM 公钥 + 过期时间戳）
	expiresAtMillis := expiration.UnixMilli()
	message := make([]byte, 0, len(publicKeyPEM)+20)
	message = append(message, publicKeyPEM...)
	message = strconv.AppendInt(message, expiresAtMillis, 10)

	// 使用SHA1withRSA签名
	publicKeySignature, err := signWithRootKey(yggPrivateKey, message)
	if err != nil {
		return nil, fmt.Errorf("签名失败: %w", err)
	}

	// V2签名：timestamp (8 bytes, big endian) + publicKey (DER)
	messageV2 := make([]byte, 8+len(publicKeyBytes))
	binary.BigEndian.PutUint64(messageV2[0:8], uint64(expiresAtMillis))
	copy(messageV2[8:], publicKeyBytes)

	publicKeySignatureV2, err := signWithRootKey(yggPrivateKey, messageV2)
	if err != nil {
		return nil, fmt.Errorf("V2签名失败: %w", err)
	}

	return &model.KeyPair{
		PrivateKey:           string(privateKeyPEM),
		PublicKey:            string(publicKeyPEM),
		PublicKeySignature:   base64.StdEncoding.EncodeToString(publicKeySignature),
		PublicKeySignatureV2: base64.StdEncoding.EncodeToString(publicKeySignatureV2),
		YggdrasilPublicKey:   yggPublicKey,
		Expiration:           expiration,
		Refresh:              refresh,
	}, nil
}

// signWithRootKey 使用根密钥对消息做 SHA1withRSA 签名
func signWithRootKey(privateKey *rsa.PrivateKey, message []byte) ([]byte, error) {
	hashed := sha1.Sum(message)
	return rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA1, hashed[:])
}

// GetOrCreateYggdrasilKeyPair 获取或创建Yggdrasil根密钥对
// 优先使用进程内缓存的已解析私钥；缓存未命中时从Redis加载并解析，再缓存。
func (s *SignatureService) GetOrCreateYggdrasilKeyPair(ctx context.Context) (string, *rsa.PrivateKey, error) {
	// 1. 读缓存（快路径）
	s.rootKeyMu.RLock()
	if s.rootPrivateKey != nil && s.isRootKeyValid() {
		pub := s.rootPublicKeyPEM
		priv := s.rootPrivateKey
		s.rootKeyMu.RUnlock()
		return pub, priv, nil
	}
	s.rootKeyMu.RUnlock()

	// 2. 缓存未命中，加锁准备从 Redis 加载/生成
	s.rootKeyMu.Lock()
	defer s.rootKeyMu.Unlock()

	// 双重检查：可能在等待锁期间已被其他 goroutine 填充
	if s.rootPrivateKey != nil && s.isRootKeyValidLocked() {
		return s.rootPublicKeyPEM, s.rootPrivateKey, nil
	}

	// 尝试从Redis获取密钥（MGet 一次往返拿三个字段，减少网络往返）
	pub, priv, err := s.loadRootKeyFromRedis(ctx)
	if err == nil && priv != nil {
		s.rootPrivateKey = priv
		s.rootPublicKeyPEM = pub
		s.logger.Info("从Redis加载Yggdrasil根密钥")
		return pub, priv, nil
	}

	// Redis 没有，生成新的根密钥对
	s.logger.Info("生成新的Yggdrasil根密钥对")
	privateKey, err := rsa.GenerateKey(rand.Reader, KeySize)
	if err != nil {
		return "", nil, fmt.Errorf("生成RSA密钥失败: %w", err)
	}

	// PEM编码私钥
	privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	privateKeyPEM := string(pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: privateKeyBytes,
	}))

	// PEM编码公钥
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	if err != nil {
		return "", nil, fmt.Errorf("编码公钥失败: %w", err)
	}
	publicKeyPEM := string(pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: publicKeyBytes,
	}))

	// 计算过期时间（90天）
	expiration := time.Now().AddDate(0, 0, ExpirationDays)

	// 保存到Redis
	if err := s.redis.Set(ctx, PublicKeyRedisKey, publicKeyPEM, RedisTTL); err != nil {
		s.logger.Warn("保存公钥到Redis失败", zap.Error(err))
	}
	if err := s.redis.Set(ctx, PrivateKeyRedisKey, privateKeyPEM, RedisTTL); err != nil {
		s.logger.Warn("保存私钥到Redis失败", zap.Error(err))
	}
	if err := s.redis.Set(ctx, KeyExpirationRedisKey, expiration.Format(time.RFC3339), RedisTTL); err != nil {
		s.logger.Warn("保存密钥过期时间到Redis失败", zap.Error(err))
	}

	// 缓存已解析的私钥与公钥PEM
	s.rootPrivateKey = privateKey
	s.rootPublicKeyPEM = publicKeyPEM
	return publicKeyPEM, privateKey, nil
}

// loadRootKeyFromRedis 从Redis加载根密钥，使用 MGet 单次往返获取三字段
func (s *SignatureService) loadRootKeyFromRedis(ctx context.Context) (string, *rsa.PrivateKey, error) {
	results, err := s.redis.MGet(ctx, PublicKeyRedisKey, PrivateKeyRedisKey, KeyExpirationRedisKey).Result()
	if err != nil {
		return "", nil, err
	}
	if len(results) < 3 {
		return "", nil, fmt.Errorf("MGet 返回字段数不足")
	}

	publicKeyPEM, _ := results[0].(string)
	privateKeyPEM, _ := results[1].(string)
	expStr, _ := results[2].(string)

	if publicKeyPEM == "" || privateKeyPEM == "" || expStr == "" {
		return "", nil, fmt.Errorf("Redis中密钥字段不完整")
	}

	expTime, err := time.Parse(time.RFC3339, expStr)
	if err != nil || !time.Now().Before(expTime) {
		return "", nil, fmt.Errorf("密钥已过期或过期时间解析失败")
	}

	block, _ := pem.Decode([]byte(privateKeyPEM))
	if block == nil {
		return "", nil, fmt.Errorf("解析PEM私钥失败")
	}
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return "", nil, fmt.Errorf("解析RSA私钥失败: %w", err)
	}
	return publicKeyPEM, privateKey, nil
}

// isRootKeyValid 在持有读锁时判断缓存密钥是否过期。
// 注意：缓存的私钥本身不带过期时间，过期信息存于 Redis（由 loadRootKeyFromRedis 校验）。
// 一旦加载成功即认为在进程生命周期内有效，由 Redis 过期触发重新生成。
func (s *SignatureService) isRootKeyValid() bool {
	return s.isRootKeyValidLocked()
}

func (s *SignatureService) isRootKeyValidLocked() bool {
	return s.rootPrivateKey != nil && s.rootPublicKeyPEM != ""
}

// GetPublicKeyFromRedis 从Redis获取公钥
func (s *SignatureService) GetPublicKeyFromRedis(ctx context.Context) (string, error) {
	publicKey, err := s.redis.Get(ctx, PublicKeyRedisKey)
	if err != nil {
		return "", fmt.Errorf("从Redis获取公钥失败: %w", err)
	}
	if publicKey == "" {
		// 如果Redis中没有，创建新的密钥对
		publicKey, _, err = s.GetOrCreateYggdrasilKeyPair(ctx)
		if err != nil {
			return "", fmt.Errorf("创建新密钥对失败: %w", err)
		}
	}
	return publicKey, nil
}

// SignStringWithSHA1withRSA 使用SHA1withRSA签名字符串
// 使用缓存的已解析根私钥；若缓存为空则触发加载/生成。
func (s *SignatureService) SignStringWithSHA1withRSA(ctx context.Context, data string) (string, error) {
	_, privateKey, err := s.GetOrCreateYggdrasilKeyPair(ctx)
	if err != nil {
		return "", fmt.Errorf("获取签名私钥失败: %w", err)
	}

	signature, err := signWithRootKey(privateKey, []byte(data))
	if err != nil {
		return "", fmt.Errorf("签名失败: %w", err)
	}
	return base64.StdEncoding.EncodeToString(signature), nil
}