BackupX/server/internal/service/install_token_service.go

package service

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"strings"
	"time"

	"backupx/server/internal/apperror"
	"backupx/server/internal/installscript"
	"backupx/server/internal/model"
	"backupx/server/internal/repository"
)

// InstallTokenService 负责一次性安装令牌的创建/消费/校验。
type InstallTokenService struct {
	repo     repository.AgentInstallTokenRepository
	nodeRepo repository.NodeRepository
}

func NewInstallTokenService(repo repository.AgentInstallTokenRepository, nodeRepo repository.NodeRepository) *InstallTokenService {
	return &InstallTokenService{repo: repo, nodeRepo: nodeRepo}
}

// InstallTokenInput 生成一次性安装令牌的输入。
type InstallTokenInput struct {
	NodeID       uint
	Mode         string
	Arch         string
	AgentVersion string
	DownloadSrc  string
	TTLSeconds   int
	CreatedByID  uint
}

// InstallTokenOutput 生成结果。
type InstallTokenOutput struct {
	Token     string
	ExpiresAt time.Time
	Node      *model.Node
	Record    *model.AgentInstallToken
}

// InstallCommandInput 生成可展示安装命令所需的完整业务输入。
type InstallCommandInput struct {
	InstallTokenInput
	MasterURL string
}

// InstallCommandOutput 是 UI 生成安装命令所需的完整业务输出。
type InstallCommandOutput struct {
	Token              string
	ExpiresAt          time.Time
	Node               *model.Node
	Record             *model.AgentInstallToken
	URL                string
	FallbackURL        string
	ComposeURL         string
	FallbackComposeURL string
	ScriptBase64       string
}

// ConsumedInstallToken 消费成功后返回给 handler 的组合体。
type ConsumedInstallToken struct {
	Record *model.AgentInstallToken
	Node   *model.Node
}

// 校验与限流常量。
const (
	InstallTokenMinTTL     = 300   // 5 分钟
	InstallTokenMaxTTL     = 86400 // 24 小时
	InstallTokenRateWindow = 60 * time.Second
	InstallTokenRatePerWin = 5
)

var (
	validInstallModes   = map[string]bool{model.InstallModeSystemd: true, model.InstallModeDocker: true, model.InstallModeForeground: true}
	validInstallArches  = map[string]bool{model.InstallArchAmd64: true, model.InstallArchArm64: true, model.InstallArchAuto: true}
	validInstallSources = map[string]bool{model.InstallSourceGitHub: true, model.InstallSourceGhproxy: true}
)

// Create 生成一次性安装令牌。
func (s *InstallTokenService) Create(ctx context.Context, in InstallTokenInput) (*InstallTokenOutput, error) {
	if err := s.validate(in); err != nil {
		return nil, err
	}
	node, err := s.nodeRepo.FindByID(ctx, in.NodeID)
	if err != nil {
		return nil, err
	}
	if node == nil {
		return nil, apperror.New(404, "NODE_NOT_FOUND", "节点不存在", nil)
	}

	since := time.Now().UTC().Add(-InstallTokenRateWindow)
	count, err := s.repo.CountCreatedSince(ctx, in.NodeID, since)
	if err != nil {
		return nil, err
	}
	if count >= InstallTokenRatePerWin {
		return nil, apperror.TooManyRequests("INSTALL_TOKEN_RATE_LIMITED",
			fmt.Sprintf("每 %d 秒最多生成 %d 次", int(InstallTokenRateWindow.Seconds()), InstallTokenRatePerWin), nil)
	}

	token, err := generateInstallToken()
	if err != nil {
		return nil, fmt.Errorf("generate token: %w", err)
	}
	expiresAt := time.Now().UTC().Add(time.Duration(in.TTLSeconds) * time.Second)
	record := &model.AgentInstallToken{
		Token:       token,
		NodeID:      in.NodeID,
		Mode:        in.Mode,
		Arch:        in.Arch,
		AgentVer:    in.AgentVersion,
		DownloadSrc: in.DownloadSrc,
		ExpiresAt:   expiresAt,
		CreatedByID: in.CreatedByID,
	}
	if err := s.repo.Create(ctx, record); err != nil {
		return nil, err
	}
	return &InstallTokenOutput{Token: token, ExpiresAt: expiresAt, Node: node, Record: record}, nil
}

// CreateCommand 创建 install token，并返回 UI 展示安装命令所需的 URL 与嵌入式脚本。
func (s *InstallTokenService) CreateCommand(ctx context.Context, in InstallCommandInput) (*InstallCommandOutput, error) {
	masterURL := strings.TrimRight(strings.TrimSpace(in.MasterURL), "/")
	if masterURL == "" {
		return nil, apperror.BadRequest("INSTALL_TOKEN_INVALID", "masterURL 必填", nil)
	}
	if err := s.validate(in.InstallTokenInput); err != nil {
		return nil, err
	}
	node, err := s.nodeRepo.FindByID(ctx, in.NodeID)
	if err != nil {
		return nil, err
	}
	if node == nil {
		return nil, apperror.New(404, "NODE_NOT_FOUND", "节点不存在", nil)
	}
	if _, err := renderInstallCommandScript(masterURL, node, &model.AgentInstallToken{
		Mode:        in.Mode,
		Arch:        in.Arch,
		AgentVer:    in.AgentVersion,
		DownloadSrc: in.DownloadSrc,
	}); err != nil {
		return nil, err
	}
	out, err := s.Create(ctx, in.InstallTokenInput)
	if err != nil {
		return nil, err
	}
	script, err := renderInstallCommandScript(masterURL, out.Node, out.Record)
	if err != nil {
		return nil, err
	}
	result := &InstallCommandOutput{
		Token:        out.Token,
		ExpiresAt:    out.ExpiresAt,
		Node:         out.Node,
		Record:       out.Record,
		URL:          masterURL + "/api/install/" + out.Token,
		FallbackURL:  masterURL + "/install/" + out.Token,
		ScriptBase64: base64.StdEncoding.EncodeToString([]byte(script)),
	}
	if out.Record.Mode == model.InstallModeDocker {
		result.ComposeURL = masterURL + "/api/install/" + out.Token + "/compose.yml"
		result.FallbackComposeURL = masterURL + "/install/" + out.Token + "/compose.yml"
	}
	return result, nil
}

func renderInstallCommandScript(masterURL string, node *model.Node, record *model.AgentInstallToken) (string, error) {
	return installscript.RenderScript(installscript.Context{
		MasterURL:     masterURL,
		AgentToken:    node.Token,
		AgentVersion:  record.AgentVer,
		Mode:          record.Mode,
		Arch:          record.Arch,
		DownloadBase:  installscript.DownloadBaseFor(record.DownloadSrc),
		InstallPrefix: "/opt/backupx-agent",
		NodeID:        node.ID,
	})
}

// Consume 原子消费令牌。未命中/已过期/已消费均返回 (nil, nil)。
func (s *InstallTokenService) Consume(ctx context.Context, token string) (*ConsumedInstallToken, error) {
	if strings.TrimSpace(token) == "" {
		return nil, nil
	}
	record, err := s.repo.ConsumeByToken(ctx, token)
	if err != nil {
		return nil, err
	}
	if record == nil {
		return nil, nil
	}
	node, err := s.nodeRepo.FindByID(ctx, record.NodeID)
	if err != nil {
		return nil, err
	}
	if node == nil {
		return nil, apperror.New(404, "NODE_NOT_FOUND", "节点已被删除", nil)
	}
	return &ConsumedInstallToken{Record: record, Node: node}, nil
}

// Peek 只读查询（不消费）且仅返回有效 token（未消费、未过期），供 compose 端点预检 Mode。
// 对已过期/已消费的 token 返回 (nil, nil)，与 Consume 语义保持一致，
// 避免 compose handler 误放行"僵尸 token"造成后续 Consume 必然失败的迷惑链路。
func (s *InstallTokenService) Peek(ctx context.Context, token string) (*model.AgentInstallToken, error) {
	if strings.TrimSpace(token) == "" {
		return nil, nil
	}
	return s.repo.FindValidByToken(ctx, token)
}

// StartGC 启动后台 GC，按 interval 扫描并删 ExpiresAt < now-7d 的记录。
func (s *InstallTokenService) StartGC(ctx context.Context, interval time.Duration) {
	if interval <= 0 {
		interval = time.Hour
	}
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				_, _ = s.repo.DeleteExpiredBefore(ctx, time.Now().UTC().Add(-7*24*time.Hour))
			}
		}
	}()
}

func (s *InstallTokenService) validate(in InstallTokenInput) error {
	if in.NodeID == 0 {
		return apperror.BadRequest("INSTALL_TOKEN_INVALID", "nodeId 必填", nil)
	}
	if !validInstallModes[in.Mode] {
		return apperror.BadRequest("INSTALL_TOKEN_INVALID", "mode 非法", nil)
	}
	if !validInstallArches[in.Arch] {
		return apperror.BadRequest("INSTALL_TOKEN_INVALID", "arch 非法", nil)
	}
	if !validInstallSources[in.DownloadSrc] {
		return apperror.BadRequest("INSTALL_TOKEN_INVALID", "downloadSrc 非法", nil)
	}
	if err := validateInstallAgentVersion(in.AgentVersion); err != nil {
		return err
	}
	if in.TTLSeconds < InstallTokenMinTTL || in.TTLSeconds > InstallTokenMaxTTL {
		return apperror.BadRequest("INSTALL_TOKEN_INVALID",
			fmt.Sprintf("ttlSeconds 需在 %d-%d", InstallTokenMinTTL, InstallTokenMaxTTL), nil)
	}
	return nil
}

func validateInstallAgentVersion(v string) error {
	v = strings.TrimSpace(v)
	if v == "" {
		return apperror.BadRequest("INSTALL_TOKEN_INVALID", "agentVersion 必填", nil)
	}
	if len(v) > 64 {
		return apperror.BadRequest("INSTALL_TOKEN_INVALID", "agentVersion 不能超过 64 字符", nil)
	}
	for _, c := range v {
		switch {
		case c >= '0' && c <= '9':
		case c >= 'a' && c <= 'z':
		case c >= 'A' && c <= 'Z':
		case c == '.' || c == '-' || c == '_' || c == '+':
		default:
			return apperror.BadRequest("INSTALL_TOKEN_INVALID", "agentVersion 包含非法字符", nil)
		}
	}
	return nil
}

func generateInstallToken() (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}