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

fix: remote agent exec url_path_prefix error

Browse Source
This commit is contained in:
engigu
2026-01-14 10:04:29 +08:00
parent b7641ab718
commit 36ffad9933
5 changed files with 52 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -10,6 +10,7 @@ baihu
# Data & Logs
data/
agent/logs/
!data/agent/
data/agent/*
!data/agent/version.txt
agent/agent.go
+11 -1
View File
@@ -202,9 +202,19 @@ func (a *Agent) connectWS() error {
wsURL = strings.Replace(wsURL, "https://", "wss://", 1)
wsURL = fmt.Sprintf("%s/api/agent/ws?token=%s&machine_id=%s", wsURL, url.QueryEscape(a.config.Token), url.QueryEscape(a.machineID))
log.Infof("正在连接 WebSocket: %s", wsURL)
log.Infof("Token: %s..., MachineID: %s...", a.config.Token[:8], a.machineID[:16])
dialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second}
conn, _, err := dialer.Dial(wsURL, nil)
conn, resp, err := dialer.Dial(wsURL, nil)
if err != nil {
if resp != nil {
bodyBytes, _ := io.ReadAll(resp.Body)
log.Errorf("WebSocket 握手失败: HTTP %d, Body: %s", resp.StatusCode, string(bodyBytes))
resp.Body.Close()
} else {
log.Errorf("WebSocket 连接失败: %v", err)
}
return err
}
agent/config.example.ini
+4
View File
@@ -1,6 +1,10 @@
[agent]
# 主服务器地址（http/https，Agent 会自动转换为 WebSocket 连接）
# 如果主服务配置了url_prefix, 这里要也要加上路径
server_url = http://192.168.1.100:8052
# 比如 url_prefix=/baihu
; server_url = http://192.168.1.100:8052/baihu
# Agent 名称（留空则使用主机名）
name =
# 注册令牌（首次注册时填写，注册成功后会自动替换为认证 Token）
internal/controllers/agent_controller.go
+23 -1
View File
@@ -337,8 +337,19 @@ func (c *AgentController) ForceUpdate(ctx *gin.Context) {
// WSConnect Agent WebSocket 连接
func (c *AgentController) WSConnect(ctx *gin.Context) {
// 添加 panic 恢复
defer func() {
if r := recover(); r != nil {
logger.Errorf("[AgentWS] WSConnect panic: %v", r)
ctx.JSON(http.StatusInternalServerError, gin.H{"error": "服务器内部错误"})
}
}()
ip := ctx.ClientIP()
// 打印请求信息用于调试
logger.Infof("[AgentWS] 收到连接请求: IP=%s, URL=%s", ip, ctx.Request.URL.String())
// 检查 IP 限流
if allowed, reason := c.wsManager.CheckRateLimit(ip); !allowed {
logger.Warnf("[AgentWS] IP %s 被限流: %s", ip, reason)
@@ -349,36 +360,45 @@ func (c *AgentController) WSConnect(ctx *gin.Context) {
token := ctx.Query("token")
if token == "" {
c.wsManager.RecordConnectFail(ip)
logger.Warnf("[AgentWS] 连接失败: 缺少 token, IP=%s", ip)
ctx.JSON(http.StatusUnauthorized, gin.H{"error": "缺少 token"})
return
}
machineID := ctx.Query("machine_id")
logger.Infof("[AgentWS] Token: %s..., MachineID: %s...", token[:8], machineID[:16])
isNewAgent := false
// 先尝试用 token 查找已有 Agent
agent := c.agentService.GetByToken(token)
logger.Infof("[AgentWS] GetByToken 结果: agent=%v", agent != nil)
// 如果没找到，尝试用令牌注册（会检查 machine_id 是否已存在）
if agent == nil {
logger.Infof("[AgentWS] 尝试注册新 Agent")
var err error
agent, isNewAgent, err = c.agentService.RegisterByToken(token, machineID, ip)
if err != nil {
c.wsManager.RecordConnectFail(ip)
logger.Warnf("[AgentWS] 注册失败: %v, IP=%s, token=%s", err, ip, token[:8]+"...")
ctx.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
return
}
logger.Infof("[AgentWS] 注册成功: Agent #%d, isNew=%v", agent.ID, isNewAgent)
}
if !agent.Enabled {
c.wsManager.RecordConnectFail(ip)
logger.Warnf("[AgentWS] Agent #%d 已禁用, IP=%s", agent.ID, ip)
ctx.JSON(http.StatusForbidden, gin.H{"error": "Agent 已禁用"})
return
}
logger.Infof("[AgentWS] 准备升级连接: Agent #%d, IP=%s", agent.ID, ip)
conn, err := agentUpgrader.Upgrade(ctx.Writer, ctx.Request, nil)
if err != nil {
logger.Errorf("[AgentWS] 升级连接失败: %v", err)
logger.Errorf("[AgentWS] 升级连接失败: %v, Agent #%d, IP=%s", err, agent.ID, ip)
return
}
@@ -399,6 +419,8 @@ func (c *AgentController) WSConnect(ctx *gin.Context) {
"machine_id": machineID,
})
logger.Infof("[AgentWS] Agent #%d 连接成功", agent.ID)
// 启动读写协程
go c.wsWritePump(ac)
go c.wsReadPump(ac, agent)
internal/router/router.go
+13 -13
View File
@@ -217,16 +217,16 @@ func Setup(c *Controllers) *gin.Engine {
agents.DELETE("/tokens/:id", c.Agent.DeleteToken)
}
}
}
// Agent API（供远程 Agent 调用）
agentAPI := api.Group("/agent")
{
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 连接
}
// Agent API（供远程 Agent 调用，不使用 /v1 版本号）
agentAPI := root.Group("/api/agent")
{
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 连接
}
// SPA fallback - serve index.html (no cache for HTML)
@@ -237,19 +237,19 @@ func Setup(c *Controllers) *gin.Engine {
ctx.Status(404)
return
}
data, err := static.ReadFile("index.html")
if err != nil {
ctx.String(500, "index.html not found")
return
}
html := string(data)
// 注入配置变量供前端使用（API 调用和路由）
configScript := `<script>window.__BASE_URL__ = "` + urlPrefix + `"; window.__API_VERSION__ = "/api/v1";</script>`
html = strings.Replace(html, "</head>", configScript+"</head>", 1)
ctx.Header("Cache-Control", "no-cache, no-store, must-revalidate")
ctx.Data(200, "text/html; charset=utf-8", []byte(html))
})