admin/TaskPool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add agent start code

Browse Source
This commit is contained in:
engigu
2025-12-28 11:42:58 +08:00
parent 6d7997c92a
commit 78c3707cce
19 changed files with 1491 additions and 501 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/controllers/agent_controller.go
+300 -74
View File
@@ -1,72 +1,46 @@
package controllers
import (
"baihu/internal/logger"
"baihu/internal/models"
"baihu/internal/services"
"baihu/internal/utils"
"encoding/json"
"net/http"
"strconv"
"strings"
"time"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
)
var agentUpgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
// AgentController Agent 控制器
type AgentController struct {
agentService *services.AgentService
wsManager *services.AgentWSManager
}
// NewAgentController 创建 Agent 控制器
func NewAgentController() *AgentController {
return &AgentController{
agentService: services.NewAgentService(),
wsManager: services.GetAgentWSManager(),
}
}
// List 获取已审核的 Agent 列表
// List 获取 Agent 列表
func (c *AgentController) List(ctx *gin.Context) {
agents := c.agentService.List()
utils.Success(ctx, agents)
}
// ListPending 获取待审核的 Agent 列表
func (c *AgentController) ListPending(ctx *gin.Context) {
agents := c.agentService.ListPending()
utils.Success(ctx, agents)
}
// Approve 审核通过 Agent
func (c *AgentController) Approve(ctx *gin.Context) {
id, err := strconv.ParseUint(ctx.Param("id"), 10, 32)
if err != nil {
utils.BadRequest(ctx, "无效的 ID")
return
}
agent, err := c.agentService.Approve(uint(id))
if err != nil {
utils.BadRequest(ctx, err.Error())
return
}
utils.Success(ctx, agent)
}
// Reject 拒绝 Agent
func (c *AgentController) Reject(ctx *gin.Context) {
id, err := strconv.ParseUint(ctx.Param("id"), 10, 32)
if err != nil {
utils.BadRequest(ctx, "无效的 ID")
return
}
if err := c.agentService.Reject(uint(id)); err != nil {
utils.ServerError(ctx, err.Error())
return
}
utils.SuccessMsg(ctx, "已拒绝")
}
// Update 更新 Agent
func (c *AgentController) Update(ctx *gin.Context) {
id, err := strconv.ParseUint(ctx.Param("id"), 10, 32)
@@ -86,11 +60,36 @@ func (c *AgentController) Update(ctx *gin.Context) {
return
}
// 获取旧状态
oldAgent := c.agentService.GetByID(uint(id))
if oldAgent == nil {
utils.NotFound(ctx, "Agent 不存在")
return
}
wasEnabled := oldAgent.Enabled
if err := c.agentService.Update(uint(id), req.Name, req.Description, req.Enabled); err != nil {
utils.ServerError(ctx, err.Error())
return
}
// 如果启用状态发生变化，通知 Agent
if wasEnabled != req.Enabled {
if req.Enabled {
// 启用：发送任务列表
c.wsManager.SendToAgent(uint(id), services.WSTypeEnabled, map[string]interface{}{
"message": "Agent 已启用",
})
// 发送任务列表
c.wsManager.BroadcastTasks(uint(id))
} else {
// 禁用：发送禁用消息，Agent 收到后清空任务
c.wsManager.SendToAgent(uint(id), services.WSTypeDisabled, map[string]interface{}{
"message": "Agent 已禁用",
})
}
}
utils.SuccessMsg(ctx, "更新成功")
}
@@ -143,49 +142,19 @@ func (c *AgentController) Register(ctx *gin.Context) {
}
ip := ctx.ClientIP()
agent, err := c.agentService.Register(&req, ip)
agent, token, err := c.agentService.Register(&req, ip)
if err != nil {
utils.ServerError(ctx, err.Error())
utils.BadRequest(ctx, err.Error())
return
}
utils.Success(ctx, gin.H{
"agent_id": agent.ID,
"status": agent.Status,
"message": "注册成功，等待审核",
"token": token,
"message": "注册成功",
})
}
// CheckStatus Agent 检查状态（用于轮询等待审核结果）
func (c *AgentController) CheckStatus(ctx *gin.Context) {
var req struct {
Name string `json:"name"`
}
if err := ctx.ShouldBindJSON(&req); err != nil {
utils.BadRequest(ctx, "参数错误")
return
}
ip := ctx.ClientIP()
agent, err := c.agentService.CheckPendingAgent(req.Name, ip)
if err != nil {
utils.NotFound(ctx, err.Error())
return
}
response := gin.H{
"agent_id": agent.ID,
"status": agent.Status,
}
// 如果已审核通过，返回 Token
if agent.Status != "pending" && agent.Token != "" {
response["token"] = agent.Token
}
utils.Success(ctx, response)
}
// Heartbeat Agent 心跳
func (c *AgentController) Heartbeat(ctx *gin.Context) {
token := c.getAgentToken(ctx)
@@ -348,3 +317,260 @@ func (c *AgentController) ForceUpdate(ctx *gin.Context) {
utils.SuccessMsg(ctx, "已标记强制更新，Agent 下次心跳时将自动更新")
}
// ========== WebSocket ==========
// WSConnect Agent WebSocket 连接
func (c *AgentController) WSConnect(ctx *gin.Context) {
ip := ctx.ClientIP()
// 检查 IP 限流
if allowed, reason := c.wsManager.CheckRateLimit(ip); !allowed {
logger.Warnf("[AgentWS] IP %s 被限流: %s", ip, reason)
ctx.JSON(http.StatusTooManyRequests, gin.H{"error": reason})
return
}
token := ctx.Query("token")
if token == "" {
c.wsManager.RecordConnectFail(ip)
ctx.JSON(http.StatusUnauthorized, gin.H{"error": "缺少 token"})
return
}
machineID := ctx.Query("machine_id")
isNewAgent := false
// 先尝试用 token 查找已有 Agent
agent := c.agentService.GetByToken(token)
// 如果没找到，尝试用令牌注册（会检查 machine_id 是否已存在）
if agent == nil {
var err error
agent, isNewAgent, err = c.agentService.RegisterByToken(token, machineID, ip)
if err != nil {
c.wsManager.RecordConnectFail(ip)
ctx.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
return
}
}
if !agent.Enabled {
c.wsManager.RecordConnectFail(ip)
ctx.JSON(http.StatusForbidden, gin.H{"error": "Agent 已禁用"})
return
}
conn, err := agentUpgrader.Upgrade(ctx.Writer, ctx.Request, nil)
if err != nil {
logger.Errorf("[AgentWS] 升级连接失败: %v", err)
return
}
// 连接成功，重置失败计数
c.wsManager.RecordConnectSuccess(ip)
// 注册连接
ac := c.wsManager.Register(agent.ID, conn, ip)
// 更新 Agent 状态
c.agentService.Heartbeat(token, ip, "", "", "", "", "")
// 发送连接成功消息（包含注册状态）
c.wsManager.SendToAgent(agent.ID, services.WSTypeConnected, map[string]interface{}{
"agent_id": agent.ID,
"name": agent.Name,
"is_new_agent": isNewAgent,
"machine_id": machineID,
})
// 启动读写协程
go c.wsWritePump(ac)
go c.wsReadPump(ac, agent)
}
// wsReadPump 读取消息
func (c *AgentController) wsReadPump(ac *services.AgentConnection, agent *models.Agent) {
defer func() {
c.wsManager.Unregister(agent.ID)
}()
ac.Conn.SetReadDeadline(time.Now().Add(90 * time.Second))
ac.Conn.SetPongHandler(func(string) error {
ac.Conn.SetReadDeadline(time.Now().Add(90 * time.Second))
return nil
})
for {
_, message, err := ac.Conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
logger.Warnf("[AgentWS] Agent #%d 读取错误: %v", agent.ID, err)
}
break
}
var msg services.WSMessage
if err := json.Unmarshal(message, &msg); err != nil {
continue
}
c.handleWSMessage(ac, agent, &msg)
}
}
// wsWritePump 写入消息
func (c *AgentController) wsWritePump(ac *services.AgentConnection) {
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()
for {
select {
case message, ok := <-ac.Send:
if !ok {
return
}
if err := ac.WriteMessage(message); err != nil {
return
}
case <-ticker.C:
ac.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if err := ac.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}
// handleWSMessage 处理 WebSocket 消息
func (c *AgentController) handleWSMessage(ac *services.AgentConnection, agent *models.Agent, msg *services.WSMessage) {
switch msg.Type {
case services.WSTypeHeartbeat:
c.handleHeartbeat(ac, agent, msg.Data)
case services.WSTypeTaskResult:
c.handleTaskResult(agent, msg.Data)
case services.WSTypeFetchTasks:
c.handleFetchTasks(agent)
}
}
// handleFetchTasks 处理 Agent 请求任务列表
func (c *AgentController) handleFetchTasks(agent *models.Agent) {
tasks := c.agentService.GetTasks(agent.ID)
c.wsManager.SendToAgent(agent.ID, services.WSTypeTasks, map[string]interface{}{
"tasks": tasks,
})
logger.Infof("[AgentWS] Agent #%d 请求任务列表，返回 %d 个任务", agent.ID, len(tasks))
}
// handleHeartbeat 处理心跳
func (c *AgentController) handleHeartbeat(ac *services.AgentConnection, agent *models.Agent, data json.RawMessage) {
var req struct {
Version string `json:"version"`
BuildTime string `json:"build_time"`
Hostname string `json:"hostname"`
OS string `json:"os"`
Arch string `json:"arch"`
AutoUpdate bool `json:"auto_update"`
}
json.Unmarshal(data, &req)
ac.UpdatePing()
// 更新 Agent 信息（使用连接时保存的 IP）
c.agentService.Heartbeat(agent.Token, ac.IP, req.Version, req.BuildTime, req.Hostname, req.OS, req.Arch)
// 检查是否需要更新
latestVersion := c.agentService.GetLatestVersion()
needUpdate := latestVersion != "" && req.Version != "" && req.Version != latestVersion
forceUpdate := agent.ForceUpdate
if forceUpdate && needUpdate {
c.agentService.ClearForceUpdate(agent.ID)
}
// 发送心跳响应
response := map[string]interface{}{
"agent_id": agent.ID,
"name": agent.Name,
"need_update": needUpdate,
"force_update": forceUpdate,
"latest_version": latestVersion,
}
c.wsManager.SendToAgent(agent.ID, services.WSTypeHeartbeatAck, response)
}
// handleTaskResult 处理任务结果
func (c *AgentController) handleTaskResult(agent *models.Agent, data json.RawMessage) {
var result models.AgentTaskResult
if err := json.Unmarshal(data, &result); err != nil {
return
}
result.AgentID = agent.ID
c.agentService.ReportResult(&result)
}
// NotifyTaskUpdate 通知 Agent 任务更新
func (c *AgentController) NotifyTaskUpdate(agentID uint) {
c.wsManager.BroadcastTasks(agentID)
}
// ========== 注册码管理 ==========
// ListRegCodes 获取注册码列表
func (c *AgentController) ListRegCodes(ctx *gin.Context) {
codes := c.agentService.ListRegCodes()
utils.Success(ctx, codes)
}
// CreateRegCode 创建注册码
func (c *AgentController) CreateRegCode(ctx *gin.Context) {
var req struct {
Remark string `json:"remark"`
MaxUses int `json:"max_uses"`
ExpiresAt string `json:"expires_at"` // 格式: 2006-01-02 15:04:05
}
if err := ctx.ShouldBindJSON(&req); err != nil {
utils.BadRequest(ctx, "参数错误")
return
}
var expiresAt *time.Time
if req.ExpiresAt != "" {
t, err := time.ParseInLocation("2006-01-02 15:04:05", req.ExpiresAt, time.Local)
if err != nil {
utils.BadRequest(ctx, "过期时间格式错误")
return
}
expiresAt = &t
}
code, err := c.agentService.CreateRegCode(req.Remark, req.MaxUses, expiresAt)
if err != nil {
utils.ServerError(ctx, err.Error())
return
}
utils.Success(ctx, code)
}
// DeleteRegCode 删除注册码
func (c *AgentController) DeleteRegCode(ctx *gin.Context) {
id, err := strconv.ParseUint(ctx.Param("id"), 10, 32)
if err != nil {
utils.BadRequest(ctx, "无效的 ID")
return
}
if err := c.agentService.DeleteRegCode(uint(id)); err != nil {
utils.ServerError(ctx, err.Error())
return
}
utils.SuccessMsg(ctx, "删除成功")
}
internal/controllers/task_controller.go
+55 -10
View File
@@ -12,14 +12,16 @@ import (
)
type TaskController struct {
taskService *services.TaskService
cronService *services.CronService
taskService *services.TaskService
cronService *services.CronService
agentWSManager *services.AgentWSManager
}
func NewTaskController(taskService *services.TaskService, cronService *services.CronService) *TaskController {
return &TaskController{
taskService: taskService,
cronService: cronService,
taskService: taskService,
cronService: cronService,
agentWSManager: services.GetAgentWSManager(),
}
}
@@ -57,6 +59,7 @@ func (tc *TaskController) CreateTask(c *gin.Context) {
WorkDir string `json:"work_dir"`
CleanConfig string `json:"clean_config"`
Envs string `json:"envs"`
AgentID *uint `json:"agent_id"`
}
if err := c.ShouldBindJSON(&req); err != nil {
@@ -78,8 +81,14 @@ func (tc *TaskController) CreateTask(c *gin.Context) {
// 转换为绝对路径
workDir := resolveWorkDir(req.WorkDir)
task := tc.taskService.CreateTask(req.Name, req.Command, req.Schedule, req.Timeout, workDir, req.CleanConfig, req.Envs, req.Type, req.Config)
tc.cronService.AddTask(task)
task := tc.taskService.CreateTask(req.Name, req.Command, req.Schedule, req.Timeout, workDir, req.CleanConfig, req.Envs, req.Type, req.Config, req.AgentID)
// 如果是 Agent 任务，通知 Agent；否则添加到本地 cron
if task.AgentID != nil && *task.AgentID > 0 {
tc.agentWSManager.BroadcastTasks(*task.AgentID)
} else {
tc.cronService.AddTask(task)
}
utils.Success(c, task)
}
@@ -115,6 +124,13 @@ func (tc *TaskController) UpdateTask(c *gin.Context) {
return
}
// 获取旧任务信息（用于判断 agent 变更）
oldTask := tc.taskService.GetTaskByID(id)
var oldAgentID *uint
if oldTask != nil {
oldAgentID = oldTask.AgentID
}
var req struct {
Name string `json:"name"`
Command string `json:"command"`
@@ -126,6 +142,7 @@ func (tc *TaskController) UpdateTask(c *gin.Context) {
CleanConfig string `json:"clean_config"`
Envs string `json:"envs"`
Enabled bool `json:"enabled"`
AgentID *uint `json:"agent_id"`
}
if err := c.ShouldBindJSON(&req); err != nil {
@@ -140,16 +157,32 @@ func (tc *TaskController) UpdateTask(c *gin.Context) {
}
}
task := tc.taskService.UpdateTask(id, req.Name, req.Command, req.Schedule, req.Timeout, resolveWorkDir(req.WorkDir), req.CleanConfig, req.Envs, req.Enabled, req.Type, req.Config)
task := tc.taskService.UpdateTask(id, req.Name, req.Command, req.Schedule, req.Timeout, resolveWorkDir(req.WorkDir), req.CleanConfig, req.Envs, req.Enabled, req.Type, req.Config, req.AgentID)
if task == nil {
utils.NotFound(c, "任务不存在")
return
}
if task.Enabled {
tc.cronService.AddTask(task)
} else {
// 处理任务调度
if task.AgentID != nil && *task.AgentID > 0 {
// Agent 任务：从本地 cron 移除，通知 Agent
tc.cronService.RemoveTask(task.ID)
tc.agentWSManager.BroadcastTasks(*task.AgentID)
// 如果 agent 变更了，也通知旧 agent
if oldAgentID != nil && *oldAgentID > 0 && *oldAgentID != *task.AgentID {
tc.agentWSManager.BroadcastTasks(*oldAgentID)
}
} else {
// 本地任务
if task.Enabled {
tc.cronService.AddTask(task)
} else {
tc.cronService.RemoveTask(task.ID)
}
// 如果之前是 agent 任务，通知旧 agent 移除
if oldAgentID != nil && *oldAgentID > 0 {
tc.agentWSManager.BroadcastTasks(*oldAgentID)
}
}
utils.Success(c, task)
@@ -162,6 +195,13 @@ func (tc *TaskController) DeleteTask(c *gin.Context) {
return
}
// 获取任务信息（用于通知 agent）
task := tc.taskService.GetTaskByID(id)
var agentID *uint
if task != nil {
agentID = task.AgentID
}
tc.cronService.RemoveTask(uint(id))
success := tc.taskService.DeleteTask(id)
@@ -170,5 +210,10 @@ func (tc *TaskController) DeleteTask(c *gin.Context) {
return
}
// 如果是 agent 任务，通知 agent
if agentID != nil && *agentID > 0 {
tc.agentWSManager.BroadcastTasks(*agentID)
}
utils.SuccessMsg(c, "删除成功")
}
internal/database/migrate.go
+20
View File
@@ -1,10 +1,16 @@
package database
import (
"baihu/internal/logger"
"baihu/internal/models"
)
func Migrate() error {
// 先执行自定义迁移
if err := customMigrations(); err != nil {
logger.Warnf("[Database] 自定义迁移警告: %v", err)
}
return AutoMigrate(
&models.User{},
&models.Task{},
@@ -16,5 +22,19 @@ func Migrate() error {
&models.SendStats{},
&models.Dependency{},
&models.Agent{},
&models.AgentRegCode{},
)
}
// customMigrations 自定义迁移（处理 AutoMigrate 无法自动完成的变更）
func customMigrations() error {
// 检查 ql_tokens 表是否存在，如果存在则修改 code 列大小为 64
if DB.Migrator().HasTable("ql_tokens") {
// MySQL: 修改 code 列大小
if err := DB.Exec("ALTER TABLE ql_tokens MODIFY COLUMN code VARCHAR(64)").Error; err != nil {
// 忽略错误（可能是 SQLite 或列已经是正确大小）
logger.Debugf("[Database] 修改 ql_tokens.code 列: %v", err)
}
}
return nil
}
internal/models/agent.go
+27 -5
View File
@@ -10,9 +10,10 @@ import (
type Agent struct {
ID uint `json:"id" gorm:"primaryKey"`
Name string `json:"name" gorm:"size:100;not null"` // Agent 名称
Token string `json:"token" gorm:"size:64;uniqueIndex"` // 认证 Token
Token string `json:"token" gorm:"size:64;index"` // 认证 Token（可重复使用）
MachineID string `json:"machine_id" gorm:"size:64;uniqueIndex"` // 机器识别码（唯一）
Description string `json:"description" gorm:"size:255"` // 描述
Status string `json:"status" gorm:"size:20;default:'pending'"` // 状态: pending(待审核), online, offline
Status string `json:"status" gorm:"size:20;default:'pending'"` // 状态: pending(待审核), online, offline, blocked(拉黑)
LastSeen *LocalTime `json:"last_seen"` // 最后心跳时间
IP string `json:"ip" gorm:"size:45"` // Agent IP 地址
Version string `json:"version" gorm:"size:20"` // Agent 版本
@@ -31,6 +32,24 @@ func (Agent) TableName() string {
return constant.TablePrefix + "agents"
}
// AgentRegCode 注册码
type AgentRegCode struct {
ID uint `json:"id" gorm:"primaryKey"`
Code string `json:"code" gorm:"size:64;uniqueIndex;not null"` // 令牌
Remark string `json:"remark" gorm:"size:255"` // 备注
MaxUses int `json:"max_uses" gorm:"default:0"` // 最大使用次数，0 表示无限制
UsedCount int `json:"used_count" gorm:"default:0"` // 已使用次数
ExpiresAt *LocalTime `json:"expires_at"` // 过期时间，null 表示永不过期
Enabled bool `json:"enabled" gorm:"default:true"` // 是否启用
CreatedAt LocalTime `json:"created_at"`
UpdatedAt LocalTime `json:"updated_at"`
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
}
func (AgentRegCode) TableName() string {
return constant.TablePrefix + "tokens"
}
// AgentTask Agent 任务配置（用于下发给 Agent）
type AgentTask struct {
ID uint `json:"id"`
@@ -58,7 +77,10 @@ type AgentTaskResult struct {
// AgentRegisterRequest Agent 注册请求
type AgentRegisterRequest struct {
Name string `json:"name"`
Hostname string `json:"hostname"`
Version string `json:"version"`
Name string `json:"name"`
Hostname string `json:"hostname"`
Version string `json:"version"`
BuildTime string `json:"build_time"`
Token string `json:"token"` // 注册令牌
MachineID string `json:"machine_id"` // 机器识别码
}
internal/router/router.go
+5 -5
View File
@@ -203,26 +203,26 @@ func Setup(c *Controllers) *gin.Engine {
agents := authorized.Group("/agents")
{
agents.GET("", c.Agent.List)
agents.GET("/pending", c.Agent.ListPending)
agents.GET("/version", c.Agent.GetVersion)
agents.POST("/:id/approve", c.Agent.Approve)
agents.POST("/:id/reject", c.Agent.Reject)
agents.PUT("/:id", c.Agent.Update)
agents.DELETE("/:id", c.Agent.Delete)
agents.POST("/:id/token", c.Agent.RegenerateToken)
agents.POST("/:id/update", c.Agent.ForceUpdate)
// 令牌管理
agents.GET("/regcodes", c.Agent.ListRegCodes)
agents.POST("/regcodes", c.Agent.CreateRegCode)
agents.DELETE("/regcodes/:id", c.Agent.DeleteRegCode)
}
}
// Agent API（供远程 Agent 调用）
agentAPI := api.Group("/agent")
{
agentAPI.POST("/register", c.Agent.Register)
agentAPI.POST("/status", c.Agent.CheckStatus)
agentAPI.POST("/heartbeat", c.Agent.Heartbeat)
agentAPI.GET("/tasks", c.Agent.GetTasks)
agentAPI.POST("/report", c.Agent.ReportResult)
agentAPI.GET("/download", c.Agent.Download)
agentAPI.GET("/ws", c.Agent.WSConnect) // WebSocket 连接
}
}
internal/services/agent_service.go
+154 -81
View File
@@ -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
internal/services/agent_ws_service.go
+293
View File
@@ -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()
}
internal/services/task_service.go
+4 -2
View File
@@ -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
}