feat: add agent start code

2025-12-28 11:42:58 +08:00
parent 6d7997c92a
commit 78c3707cce
19 changed files with 1491 additions and 501 deletions
@@ -12,6 +12,8 @@ import (
 	"path/filepath"
 	"strings"
 	"time"
+
+	"gorm.io/gorm"
 )

 // AgentService Agent 服务
@@ -29,75 +31,171 @@ func generateToken() string {
 	return hex.EncodeToString(bytes)
 }

-// Register Agent 注册（进入待审核状态）
-func (s *AgentService) Register(req *models.AgentRegisterRequest, ip string) (*models.Agent, error) {
-	// 检查是否已存在同名待审核的 Agent
-	var existing models.Agent
-	if err := database.DB.Where("name = ? AND status = ?", req.Name, "pending").First(&existing).Error; err == nil {
-		// 更新现有记录
-		now := models.LocalTime(time.Now())
-		database.DB.Model(&existing).Updates(map[string]interface{}{
-			"hostname":  req.Hostname,
-			"version":   req.Version,
-			"ip":        ip,
-			"last_seen": now,
-		})
-		return &existing, nil
+// generateRegCode 生成令牌（64位，与认证 Token 相同）
+func generateRegCode() string {
+	bytes := make([]byte, 32)
+	rand.Read(bytes)
+	return hex.EncodeToString(bytes)
+}
+
+// ========== 注册码管理 ==========
+
+// CreateRegCode 创建令牌（同时创建 Agent 记录）
+func (s *AgentService) CreateRegCode(remark string, maxUses int, expiresAt *time.Time) (*models.AgentRegCode, error) {
+	var expires *models.LocalTime
+	if expiresAt != nil {
+		t := models.LocalTime(*expiresAt)
+		expires = &t
 	}

-	// 创建新的待审核 Agent
+	token := generateRegCode()
+
+	regCode := &models.AgentRegCode{
+		Code:      token,
+		Remark:    remark,
+		MaxUses:   maxUses,
+		ExpiresAt: expires,
+		Enabled:   true,
+	}
+
+	if err := database.DB.Create(regCode).Error; err != nil {
+		return nil, err
+	}
+
+	logger.Infof("[Agent] 创建令牌: %s (max_uses=%d)", token[:8]+"...", maxUses)
+	return regCode, nil
+}
+
+// ListRegCodes 获取注册码列表
+func (s *AgentService) ListRegCodes() []models.AgentRegCode {
+	var codes []models.AgentRegCode
+	database.DB.Order("id DESC").Find(&codes)
+	return codes
+}
+
+// DeleteRegCode 删除注册码
+func (s *AgentService) DeleteRegCode(id uint) error {
+	return database.DB.Delete(&models.AgentRegCode{}, id).Error
+}
+
+// ValidateRegCode 验证注册码
+func (s *AgentService) ValidateRegCode(code string) (*models.AgentRegCode, error) {
+	var regCode models.AgentRegCode
+	if err := database.DB.Where("code = ?", code).First(&regCode).Error; err != nil {
+		return nil, &ServiceError{Message: "无效的注册码"}
+	}
+
+	if !regCode.Enabled {
+		return nil, &ServiceError{Message: "注册码已禁用"}
+	}
+
+	// 检查使用次数
+	if regCode.MaxUses > 0 && regCode.UsedCount >= regCode.MaxUses {
+		return nil, &ServiceError{Message: "注册码已达到使用上限"}
+	}
+
+	// 检查过期时间
+	if regCode.ExpiresAt != nil && time.Time(*regCode.ExpiresAt).Before(time.Now()) {
+		return nil, &ServiceError{Message: "注册码已过期"}
+	}
+
+	return &regCode, nil
+}
+
+// UseRegCode 使用注册码（增加使用计数）
+func (s *AgentService) UseRegCode(id uint) {
+	database.DB.Model(&models.AgentRegCode{}).Where("id = ?", id).UpdateColumn("used_count", gorm.Expr("used_count + 1"))
+}
+
+// ========== Agent 注册 ==========
+
+// RegisterByToken 通过令牌注册 Agent（首次 WebSocket 连接时调用）
+// 返回: agent, isNewAgent, error
+func (s *AgentService) RegisterByToken(token string, machineID string, ip string) (*models.Agent, bool, error) {
+	// 验证令牌
+	regCode, err := s.ValidateRegCode(token)
+	if err != nil {
+		return nil, false, err
+	}
+
+	// 如果提供了 machine_id，先检查是否已存在
+	if machineID != "" {
+		var existing models.Agent
+		if err := database.DB.Where("machine_id = ?", machineID).First(&existing).Error; err == nil {
+			// 已存在，更新 token 和状态，复用已有 Agent
+			now := models.LocalTime(time.Now())
+			database.DB.Model(&existing).Updates(map[string]interface{}{
+				"token":    token,
+				"ip":       ip,
+				"status":   "online",
+				"last_seen": now,
+			})
+			s.UseRegCode(regCode.ID)
+			logger.Infof("[Agent] Agent #%d 通过 machine_id 复用 (%s)", existing.ID, machineID[:8]+"...")
+			return &existing, false, nil
+		}
+	}
+
+	// 创建 Agent，使用令牌作为认证 Token
 	now := models.LocalTime(time.Now())
 	agent := &models.Agent{
-		Name:     req.Name,
-		Hostname: req.Hostname,
-		Version:  req.Version,
-		IP:       ip,
-		Status:   "pending",
-		LastSeen: &now,
-		Enabled:  true,
+		Name:      fmt.Sprintf("agent-%d", time.Now().Unix()),
+		Token:     token,
+		MachineID: machineID,
+		IP:        ip,
+		Status:    "online",
+		LastSeen:  &now,
+		Enabled:   true,
 	}

 	if err := database.DB.Create(agent).Error; err != nil {
-		return nil, err
+		return nil, false, err
 	}

-	logger.Infof("[Agent] 新 Agent 注册: %s (%s)", req.Name, ip)
-	return agent, nil
+	s.UseRegCode(regCode.ID)
+	logger.Infof("[Agent] Agent 通过令牌注册: #%d (%s)", agent.ID, ip)
+	return agent, true, nil
 }

-// Approve 审核通过 Agent，生成 Token
-func (s *AgentService) Approve(id uint) (*models.Agent, error) {
-	agent := s.GetByID(id)
-	if agent == nil {
-		return nil, &ServiceError{Message: "Agent 不存在"}
+// Register Agent 注册（必须使用令牌）- 保留兼容旧版本
+func (s *AgentService) Register(req *models.AgentRegisterRequest, ip string) (*models.Agent, string, error) {
+	// 必须提供令牌
+	if req.Token == "" {
+		return nil, "", &ServiceError{Message: "缺少注册令牌"}
 	}

-	if agent.Status != "pending" {
-		return nil, &ServiceError{Message: "Agent 状态不是待审核"}
+	regCode, err := s.ValidateRegCode(req.Token)
+	if err != nil {
+		return nil, "", err
 	}

-	token := generateToken()
+	// 检查是否已存在同名 Agent
+	var existing models.Agent
+	if err := database.DB.Where("name = ?", req.Name).First(&existing).Error; err == nil {
+		return nil, "", &ServiceError{Message: "Agent 名称已存在"}
+	}
+
+	// 创建新 Agent，使用令牌作为认证 Token
 	now := models.LocalTime(time.Now())
-
-	if err := database.DB.Model(agent).Updates(map[string]interface{}{
-		"token":     token,
-		"status":    "online",
-		"last_seen": now,
-	}).Error; err != nil {
-		return nil, err
+	agent := &models.Agent{
+		Name:      req.Name,
+		Token:     req.Token,
+		Hostname:  req.Hostname,
+		Version:   req.Version,
+		BuildTime: req.BuildTime,
+		IP:        ip,
+		Status:    "online",
+		LastSeen:  &now,
+		Enabled:   true,
 	}

-	agent.Token = token
-	agent.Status = "online"
-	agent.LastSeen = &now
+	if err := database.DB.Create(agent).Error; err != nil {
+		return nil, "", err
+	}

-	logger.Infof("[Agent] Agent 已审核通过: %s (#%d)", agent.Name, agent.ID)
-	return agent, nil
-}
-
-// Reject 拒绝 Agent
-func (s *AgentService) Reject(id uint) error {
-	return database.DB.Delete(&models.Agent{}, id).Error
+	s.UseRegCode(regCode.ID)
+	logger.Infof("[Agent] Agent 注册成功: %s (%s)", req.Name, ip)
+	return agent, req.Token, nil
 }

 // Update 更新 Agent
@@ -109,7 +207,7 @@ func (s *AgentService) Update(id uint, name, description string, enabled bool) e
 	}).Error
 }

-// Delete 删除 Agent
+// Delete 删除 Agent（物理删除）
 func (s *AgentService) Delete(id uint) error {
 	// 检查是否有关联任务
 	var count int64
@@ -118,7 +216,7 @@ func (s *AgentService) Delete(id uint) error {
 		return &ServiceError{Message: "该 Agent 下还有关联任务，无法删除"}
 	}

-	return database.DB.Delete(&models.Agent{}, id).Error
+	return database.DB.Unscoped().Delete(&models.Agent{}, id).Error
 }

 // GetByID 根据 ID 获取 Agent
@@ -139,27 +237,16 @@ func (s *AgentService) GetByToken(token string) *models.Agent {
 	return &agent
 }

-// List 获取已审核的 Agent 列表
+// List 获取 Agent 列表
 func (s *AgentService) List() []models.Agent {
 	var agents []models.Agent
-	database.DB.Where("status != ?", "pending").Order("id DESC").Find(&agents)
+	database.DB.Order("id DESC").Find(&agents)
 	return agents
 }

-// ListPending 获取待审核的 Agent 列表
-func (s *AgentService) ListPending() []models.Agent {
-	var agents []models.Agent
-	database.DB.Where("status = ?", "pending").Order("id DESC").Find(&agents)
-	return agents
-}
-
-// RegenerateToken 重新生成 Token
+// RegenerateToken 重新生成 Token - 已废弃，保留空实现避免路由错误
 func (s *AgentService) RegenerateToken(id uint) (string, error) {
-	newToken := generateToken()
-	if err := database.DB.Model(&models.Agent{}).Where("id = ?", id).Update("token", newToken).Error; err != nil {
-		return "", err
-	}
-	return newToken, nil
+	return "", &ServiceError{Message: "此功能已禁用"}
 }

 // Heartbeat Agent 心跳
@@ -209,20 +296,6 @@ func (s *AgentService) Heartbeat(token, ip, version, buildTime, hostname, osType
 	return agent, nil
 }

-// CheckPendingAgent 检查待审核 Agent 的状态（用于 Agent 轮询）
-func (s *AgentService) CheckPendingAgent(name, ip string) (*models.Agent, error) {
-	var agent models.Agent
-	if err := database.DB.Where("name = ? AND ip = ?", name, ip).First(&agent).Error; err != nil {
-		return nil, &ServiceError{Message: "Agent 未注册"}
-	}
-
-	// 更新最后心跳时间
-	now := models.LocalTime(time.Now())
-	database.DB.Model(&agent).Update("last_seen", now)
-
-	return &agent, nil
-}
-
 // GetTasks 获取 Agent 的任务列表
 func (s *AgentService) GetTasks(agentID uint) []models.AgentTask {
 	var tasks []models.Task
@@ -0,0 +1,293 @@
+package services
+
+import (
+	"baihu/internal/database"
+	"baihu/internal/logger"
+	"baihu/internal/models"
+	"encoding/json"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// AgentWSManager WebSocket 连接管理器
+type AgentWSManager struct {
+	connections   map[uint]*AgentConnection // agentID -> connection
+	ipConnections map[string]int            // IP -> 连接数
+	ipLastAttempt map[string]time.Time      // IP -> 最后连接尝试时间
+	ipFailCount   map[string]int            // IP -> 连续失败次数
+	mu            sync.RWMutex
+}
+
+// 限流配置
+const (
+	maxConnectionsPerIP = 10              // 每个 IP 最大连接数
+	minConnectInterval  = 5 * time.Second // 同一 IP 最小连接间隔
+	maxFailCount        = 5               // 最大连续失败次数
+	failBlockDuration   = 5 * time.Minute // 失败后封禁时长
+)
+
+// AgentConnection Agent WebSocket 连接
+type AgentConnection struct {
+	AgentID   uint
+	IP        string
+	Conn      *websocket.Conn
+	Send      chan []byte
+	LastPing  time.Time
+	mu        sync.Mutex
+}
+
+// WSMessage WebSocket 消息结构
+type WSMessage struct {
+	Type string          `json:"type"`
+	Data json.RawMessage `json:"data,omitempty"`
+}
+
+// 消息类型常量
+const (
+	WSTypeHeartbeat    = "heartbeat"
+	WSTypeHeartbeatAck = "heartbeat_ack"
+	WSTypeTasks        = "tasks"
+	WSTypeTaskResult   = "task_result"
+	WSTypeUpdate       = "update"
+	WSTypeDisconnect   = "disconnect"
+	WSTypeConnected    = "connected"   // 连接成功，包含注册状态
+	WSTypeDisabled     = "disabled"    // Agent 被禁用
+	WSTypeEnabled      = "enabled"     // Agent 被启用
+	WSTypeFetchTasks   = "fetch_tasks" // Agent 请求任务列表
+)
+
+var agentWSManager *AgentWSManager
+var agentWSOnce sync.Once
+
+// GetAgentWSManager 获取单例
+func GetAgentWSManager() *AgentWSManager {
+	agentWSOnce.Do(func() {
+		agentWSManager = &AgentWSManager{
+			connections:   make(map[uint]*AgentConnection),
+			ipConnections: make(map[string]int),
+			ipLastAttempt: make(map[string]time.Time),
+			ipFailCount:   make(map[string]int),
+		}
+		go agentWSManager.cleanupLoop()
+	})
+	return agentWSManager
+}
+
+// CheckRateLimit 检查 IP 限流，返回是否允许连接
+func (m *AgentWSManager) CheckRateLimit(ip string) (bool, string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	now := time.Now()
+
+	// 检查是否被封禁（连续失败过多）
+	if failCount, exists := m.ipFailCount[ip]; exists && failCount >= maxFailCount {
+		if lastAttempt, ok := m.ipLastAttempt[ip]; ok {
+			if now.Sub(lastAttempt) < failBlockDuration {
+				remaining := failBlockDuration - now.Sub(lastAttempt)
+				return false, "连接失败次数过多，请 " + remaining.Round(time.Second).String() + " 后重试"
+			}
+			// 封禁时间已过，重置计数
+			delete(m.ipFailCount, ip)
+		}
+	}
+
+	// 检查连接频率
+	if lastAttempt, exists := m.ipLastAttempt[ip]; exists {
+		if now.Sub(lastAttempt) < minConnectInterval {
+			return false, "连接过于频繁，请稍后重试"
+		}
+	}
+
+	// 检查 IP 连接数
+	if count, exists := m.ipConnections[ip]; exists && count >= maxConnectionsPerIP {
+		return false, "该 IP 连接数已达上限"
+	}
+
+	m.ipLastAttempt[ip] = now
+	return true, ""
+}
+
+// RecordConnectFail 记录连接失败
+func (m *AgentWSManager) RecordConnectFail(ip string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.ipFailCount[ip]++
+	m.ipLastAttempt[ip] = time.Now()
+	if m.ipFailCount[ip] >= maxFailCount {
+		logger.Warnf("[AgentWS] IP %s 连续失败 %d 次，已封禁 %v", ip, m.ipFailCount[ip], failBlockDuration)
+	}
+}
+
+// RecordConnectSuccess 记录连接成功，重置失败计数
+func (m *AgentWSManager) RecordConnectSuccess(ip string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	delete(m.ipFailCount, ip)
+}
+
+// Register 注册连接
+func (m *AgentWSManager) Register(agentID uint, conn *websocket.Conn, ip string) *AgentConnection {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	// 关闭旧连接
+	if old, exists := m.connections[agentID]; exists {
+		// 减少旧 IP 的连接计数
+		if old.IP != "" {
+			if count, ok := m.ipConnections[old.IP]; ok && count > 0 {
+				m.ipConnections[old.IP] = count - 1
+			}
+		}
+		old.Close()
+	}
+
+	ac := &AgentConnection{
+		AgentID:  agentID,
+		IP:       ip,
+		Conn:     conn,
+		Send:     make(chan []byte, 256),
+		LastPing: time.Now(),
+	}
+	m.connections[agentID] = ac
+
+	// 增加 IP 连接计数
+	m.ipConnections[ip]++
+
+	logger.Infof("[AgentWS] Agent #%d 已连接 (%s)", agentID, ip)
+	return ac
+}
+
+// Unregister 注销连接
+func (m *AgentWSManager) Unregister(agentID uint) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if conn, exists := m.connections[agentID]; exists {
+		// 减少 IP 连接计数
+		if conn.IP != "" {
+			if count, ok := m.ipConnections[conn.IP]; ok && count > 0 {
+				m.ipConnections[conn.IP] = count - 1
+			}
+		}
+		conn.Close()
+		delete(m.connections, agentID)
+		logger.Infof("[AgentWS] Agent #%d 已断开", agentID)
+	}
+}
+
+// GetConnection 获取连接
+func (m *AgentWSManager) GetConnection(agentID uint) *AgentConnection {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.connections[agentID]
+}
+
+// SendToAgent 发送消息给指定 Agent
+func (m *AgentWSManager) SendToAgent(agentID uint, msgType string, data interface{}) error {
+	conn := m.GetConnection(agentID)
+	if conn == nil {
+		return nil // Agent 不在线
+	}
+
+	dataBytes, _ := json.Marshal(data)
+	msg := WSMessage{Type: msgType, Data: dataBytes}
+	msgBytes, _ := json.Marshal(msg)
+
+	select {
+	case conn.Send <- msgBytes:
+		return nil
+	default:
+		return nil // 缓冲区满，丢弃
+	}
+}
+
+// BroadcastTasks 广播任务更新给指定 Agent
+func (m *AgentWSManager) BroadcastTasks(agentID uint) {
+	agentService := NewAgentService()
+	tasks := agentService.GetTasks(agentID)
+	m.SendToAgent(agentID, WSTypeTasks, map[string]interface{}{
+		"tasks": tasks,
+	})
+}
+
+// OnlineCount 在线 Agent 数量
+func (m *AgentWSManager) OnlineCount() int {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return len(m.connections)
+}
+
+// cleanupLoop 清理超时连接
+func (m *AgentWSManager) cleanupLoop() {
+	ticker := time.NewTicker(30 * time.Second)
+	defer ticker.Stop()
+
+	for range ticker.C {
+		m.mu.Lock()
+		now := time.Now()
+
+		// 清理超时连接
+		for agentID, conn := range m.connections {
+			if now.Sub(conn.LastPing) > 2*time.Minute {
+				// 减少 IP 连接计数
+				if conn.IP != "" {
+					if count, ok := m.ipConnections[conn.IP]; ok && count > 0 {
+						m.ipConnections[conn.IP] = count - 1
+					}
+				}
+				conn.Close()
+				delete(m.connections, agentID)
+				// 更新数据库状态
+				database.DB.Model(&models.Agent{}).Where("id = ?", agentID).Update("status", "offline")
+				logger.Infof("[AgentWS] Agent #%d 心跳超时，已断开", agentID)
+			}
+		}
+
+		// 清理过期的限流记录（超过 10 分钟未活动）
+		for ip, lastAttempt := range m.ipLastAttempt {
+			if now.Sub(lastAttempt) > 10*time.Minute {
+				delete(m.ipLastAttempt, ip)
+				delete(m.ipFailCount, ip)
+				// 只清理没有活跃连接的 IP 计数
+				if m.ipConnections[ip] == 0 {
+					delete(m.ipConnections, ip)
+				}
+			}
+		}
+
+		m.mu.Unlock()
+	}
+}
+
+// Close 关闭连接
+func (c *AgentConnection) Close() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.Conn != nil {
+		c.Conn.Close()
+		c.Conn = nil
+	}
+	if c.Send != nil {
+		close(c.Send)
+		c.Send = nil
+	}
+}
+
+// WriteMessage 写入消息
+func (c *AgentConnection) WriteMessage(data []byte) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.Conn == nil {
+		return nil
+	}
+	c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+	return c.Conn.WriteMessage(websocket.TextMessage, data)
+}
+
+// UpdatePing 更新心跳时间
+func (c *AgentConnection) UpdatePing() {
+	c.LastPing = time.Now()
+}
@@ -11,7 +11,7 @@ func NewTaskService() *TaskService {
 	return &TaskService{}
 }

-func (ts *TaskService) CreateTask(name, command, schedule string, timeout int, workDir, cleanConfig, envs, taskType, config string) *models.Task {
+func (ts *TaskService) CreateTask(name, command, schedule string, timeout int, workDir, cleanConfig, envs, taskType, config string, agentID *uint) *models.Task {
 	if taskType == "" {
 		taskType = "task"
 	}
@@ -25,6 +25,7 @@ func (ts *TaskService) CreateTask(name, command, schedule string, timeout int, w
 		WorkDir:     workDir,
 		CleanConfig: cleanConfig,
 		Envs:        envs,
+		AgentID:     agentID,
 		Enabled:     true,
 	}
 	database.DB.Create(task)
@@ -61,7 +62,7 @@ func (ts *TaskService) GetTaskByID(id int) *models.Task {
 	return &task
 }

-func (ts *TaskService) UpdateTask(id int, name, command, schedule string, timeout int, workDir, cleanConfig, envs string, enabled bool, taskType, config string) *models.Task {
+func (ts *TaskService) UpdateTask(id int, name, command, schedule string, timeout int, workDir, cleanConfig, envs string, enabled bool, taskType, config string, agentID *uint) *models.Task {
 	var task models.Task
 	if err := database.DB.First(&task, id).Error; err != nil {
 		return nil
@@ -74,6 +75,7 @@ func (ts *TaskService) UpdateTask(id int, name, command, schedule string, timeou
 	task.CleanConfig = cleanConfig
 	task.Envs = envs
 	task.Enabled = enabled
+	task.AgentID = agentID
 	if taskType != "" {
 		task.Type = taskType
 	}