mirror of
https://github.com/Awuqing/BackupX.git
synced 2026-05-13 09:00:15 +08:00
功能: 新增 InstallTokenService 含输入校验、限流、GC
This commit is contained in:
Awuqing
187
server/internal/service/install_token_service.go
Normal file
187
server/internal/service/install_token_service.go
Normal file
@@ -0,0 +1,187 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"backupx/server/internal/apperror"
|
||||
"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
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// 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 只读查询（不消费），供 compose 端点预检 Mode。
|
||||
func (s *InstallTokenService) Peek(ctx context.Context, token string) (*model.AgentInstallToken, error) {
|
||||
if strings.TrimSpace(token) == "" {
|
||||
return nil, nil
|
||||
}
|
||||
return s.repo.FindByToken(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 strings.TrimSpace(in.AgentVersion) == "" {
|
||||
return apperror.BadRequest("INSTALL_TOKEN_INVALID", "agentVersion 必填", nil)
|
||||
}
|
||||
if in.TTLSeconds < InstallTokenMinTTL || in.TTLSeconds > InstallTokenMaxTTL {
|
||||
return apperror.BadRequest("INSTALL_TOKEN_INVALID",
|
||||
fmt.Sprintf("ttlSeconds 需在 %d-%d", InstallTokenMinTTL, InstallTokenMaxTTL), 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
|
||||
}
|
||||
156
server/internal/service/install_token_service_test.go
Normal file
156
server/internal/service/install_token_service_test.go
Normal file
@@ -0,0 +1,156 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"backupx/server/internal/model"
|
||||
"backupx/server/internal/repository"
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/gorm"
|
||||
gormlogger "gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
func openInstallTokenTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "it.db")),
|
||||
&gorm.Config{Logger: gormlogger.Default.LogMode(gormlogger.Silent)})
|
||||
if err != nil {
|
||||
t.Fatalf("open: %v", err)
|
||||
}
|
||||
if err := db.AutoMigrate(&model.AgentInstallToken{}, &model.Node{}); err != nil {
|
||||
t.Fatalf("migrate: %v", err)
|
||||
}
|
||||
return db
|
||||
}
|
||||
|
||||
func TestInstallTokenServiceCreateAndConsume(t *testing.T) {
|
||||
db := openInstallTokenTestDB(t)
|
||||
repo := repository.NewAgentInstallTokenRepository(db)
|
||||
nodeRepo := repository.NewNodeRepository(db)
|
||||
|
||||
node := &model.Node{Name: "n1", Token: "agent-token"}
|
||||
if err := nodeRepo.Create(context.Background(), node); err != nil {
|
||||
t.Fatalf("create node: %v", err)
|
||||
}
|
||||
|
||||
svc := NewInstallTokenService(repo, nodeRepo)
|
||||
created, err := svc.Create(context.Background(), InstallTokenInput{
|
||||
NodeID: node.ID,
|
||||
Mode: model.InstallModeSystemd,
|
||||
Arch: model.InstallArchAuto,
|
||||
AgentVersion: "v1.7.0",
|
||||
DownloadSrc: model.InstallSourceGitHub,
|
||||
TTLSeconds: 900,
|
||||
CreatedByID: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
if created.Token == "" || created.ExpiresAt.Before(time.Now().UTC()) {
|
||||
t.Fatalf("invalid token: %+v", created)
|
||||
}
|
||||
|
||||
consumed, err := svc.Consume(context.Background(), created.Token)
|
||||
if err != nil {
|
||||
t.Fatalf("consume: %v", err)
|
||||
}
|
||||
if consumed == nil || consumed.Node.ID != node.ID {
|
||||
t.Fatalf("expected consumed token for node, got %+v", consumed)
|
||||
}
|
||||
|
||||
again, err := svc.Consume(context.Background(), created.Token)
|
||||
if err != nil {
|
||||
t.Fatalf("second consume err: %v", err)
|
||||
}
|
||||
if again != nil {
|
||||
t.Fatalf("expected nil on second consume")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstallTokenServicePeekDoesNotConsume(t *testing.T) {
|
||||
db := openInstallTokenTestDB(t)
|
||||
repo := repository.NewAgentInstallTokenRepository(db)
|
||||
nodeRepo := repository.NewNodeRepository(db)
|
||||
node := &model.Node{Name: "n2", Token: "tok2"}
|
||||
_ = nodeRepo.Create(context.Background(), node)
|
||||
|
||||
svc := NewInstallTokenService(repo, nodeRepo)
|
||||
out, err := svc.Create(context.Background(), InstallTokenInput{
|
||||
NodeID: node.ID, Mode: "docker", Arch: "auto",
|
||||
AgentVersion: "v1", DownloadSrc: "github", TTLSeconds: 300, CreatedByID: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
|
||||
// Peek 两次都应成功（不消费）
|
||||
for i := 0; i < 2; i++ {
|
||||
rec, err := svc.Peek(context.Background(), out.Token)
|
||||
if err != nil {
|
||||
t.Fatalf("peek %d: %v", i, err)
|
||||
}
|
||||
if rec == nil || rec.Mode != "docker" {
|
||||
t.Fatalf("peek %d bad: %+v", i, rec)
|
||||
}
|
||||
}
|
||||
|
||||
// 之后仍可消费
|
||||
consumed, _ := svc.Consume(context.Background(), out.Token)
|
||||
if consumed == nil {
|
||||
t.Fatalf("consume after peek failed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstallTokenServiceValidatesInput(t *testing.T) {
|
||||
db := openInstallTokenTestDB(t)
|
||||
nodeRepo := repository.NewNodeRepository(db)
|
||||
node := &model.Node{Name: "valid", Token: "t"}
|
||||
_ = nodeRepo.Create(context.Background(), node)
|
||||
|
||||
svc := NewInstallTokenService(repository.NewAgentInstallTokenRepository(db), nodeRepo)
|
||||
cases := []struct {
|
||||
name string
|
||||
in InstallTokenInput
|
||||
}{
|
||||
{"bad mode", InstallTokenInput{NodeID: node.ID, Mode: "xxx", Arch: "auto", AgentVersion: "v1", DownloadSrc: "github", TTLSeconds: 300, CreatedByID: 1}},
|
||||
{"bad arch", InstallTokenInput{NodeID: node.ID, Mode: "systemd", Arch: "risc", AgentVersion: "v1", DownloadSrc: "github", TTLSeconds: 300, CreatedByID: 1}},
|
||||
{"bad source", InstallTokenInput{NodeID: node.ID, Mode: "systemd", Arch: "auto", AgentVersion: "v1", DownloadSrc: "bogus", TTLSeconds: 300, CreatedByID: 1}},
|
||||
{"bad ttl low", InstallTokenInput{NodeID: node.ID, Mode: "systemd", Arch: "auto", AgentVersion: "v1", DownloadSrc: "github", TTLSeconds: 10, CreatedByID: 1}},
|
||||
{"bad ttl high", InstallTokenInput{NodeID: node.ID, Mode: "systemd", Arch: "auto", AgentVersion: "v1", DownloadSrc: "github", TTLSeconds: 999999, CreatedByID: 1}},
|
||||
{"missing version", InstallTokenInput{NodeID: node.ID, Mode: "systemd", Arch: "auto", AgentVersion: "", DownloadSrc: "github", TTLSeconds: 300, CreatedByID: 1}},
|
||||
{"missing node id", InstallTokenInput{NodeID: 0, Mode: "systemd", Arch: "auto", AgentVersion: "v1", DownloadSrc: "github", TTLSeconds: 300, CreatedByID: 1}},
|
||||
{"node not exists", InstallTokenInput{NodeID: 999, Mode: "systemd", Arch: "auto", AgentVersion: "v1", DownloadSrc: "github", TTLSeconds: 300, CreatedByID: 1}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
if _, err := svc.Create(context.Background(), tc.in); err == nil {
|
||||
t.Errorf("%s: expected validation error", tc.name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstallTokenServiceRateLimit(t *testing.T) {
|
||||
db := openInstallTokenTestDB(t)
|
||||
nodeRepo := repository.NewNodeRepository(db)
|
||||
node := &model.Node{Name: "rl", Token: "rl"}
|
||||
_ = nodeRepo.Create(context.Background(), node)
|
||||
|
||||
svc := NewInstallTokenService(repository.NewAgentInstallTokenRepository(db), nodeRepo)
|
||||
base := InstallTokenInput{
|
||||
NodeID: node.ID, Mode: "systemd", Arch: "auto",
|
||||
AgentVersion: "v1", DownloadSrc: "github", TTLSeconds: 300, CreatedByID: 1,
|
||||
}
|
||||
// 前 5 次成功
|
||||
for i := 0; i < 5; i++ {
|
||||
if _, err := svc.Create(context.Background(), base); err != nil {
|
||||
t.Fatalf("iter %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
// 第 6 次应被限流
|
||||
_, err := svc.Create(context.Background(), base)
|
||||
if err == nil {
|
||||
t.Fatalf("expected rate limit error")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user