|
|
|
|
@@ -60,9 +60,13 @@ type Message struct {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
type AgentConnection struct {
|
|
|
|
|
AgentID string
|
|
|
|
|
Conn *websocket.Conn
|
|
|
|
|
mu sync.Mutex
|
|
|
|
|
AgentID string
|
|
|
|
|
Conn *websocket.Conn
|
|
|
|
|
mu sync.Mutex
|
|
|
|
|
// Latency tracking — updated each ping/pong cycle.
|
|
|
|
|
latencyMu sync.Mutex
|
|
|
|
|
pingSentAt time.Time
|
|
|
|
|
LatencyMs *int // nil until first pong received
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (c *AgentConnection) SendJSON(v interface{}) error {
|
|
|
|
|
@@ -177,6 +181,35 @@ func (h *WSHub) runPingLoopRaw(conn *websocket.Conn) {
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// runPingLoopAgent is like runPingLoopRaw but also records RTT on each pong.
|
|
|
|
|
func (h *WSHub) runPingLoopAgent(ac *AgentConnection) {
|
|
|
|
|
interval := h.pingInterval()
|
|
|
|
|
ticker := time.NewTicker(interval)
|
|
|
|
|
defer ticker.Stop()
|
|
|
|
|
|
|
|
|
|
conn := ac.Conn
|
|
|
|
|
_ = conn.SetReadDeadline(time.Now().Add(interval * 2))
|
|
|
|
|
conn.SetPongHandler(func(string) error {
|
|
|
|
|
// Measure RTT.
|
|
|
|
|
ac.latencyMu.Lock()
|
|
|
|
|
if !ac.pingSentAt.IsZero() {
|
|
|
|
|
ms := int(time.Since(ac.pingSentAt).Milliseconds())
|
|
|
|
|
ac.LatencyMs = &ms
|
|
|
|
|
}
|
|
|
|
|
ac.latencyMu.Unlock()
|
|
|
|
|
return conn.SetReadDeadline(time.Now().Add(interval * 2))
|
|
|
|
|
})
|
|
|
|
|
|
|
|
|
|
for range ticker.C {
|
|
|
|
|
ac.latencyMu.Lock()
|
|
|
|
|
ac.pingSentAt = time.Now()
|
|
|
|
|
ac.latencyMu.Unlock()
|
|
|
|
|
if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
|
|
|
|
|
return
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (h *WSHub) runPingLoopDash(dc *DashboardConn) {
|
|
|
|
|
interval := h.pingInterval()
|
|
|
|
|
ticker := time.NewTicker(interval)
|
|
|
|
|
@@ -263,13 +296,17 @@ func (h *WSHub) getAgentConn(agentID string) *AgentConnection {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
clientIP := r.Header.Get("X-Forwarded-For")
|
|
|
|
|
if clientIP == "" {
|
|
|
|
|
clientIP = r.RemoteAddr
|
|
|
|
|
}
|
|
|
|
|
log.Printf("[WS] Agent connection attempt from %s (origin=%s)", clientIP, r.Header.Get("Origin"))
|
|
|
|
|
conn, err := upgrader.Upgrade(w, r, nil)
|
|
|
|
|
if err != nil {
|
|
|
|
|
log.Printf("WebSocket upgrade error: %v", err)
|
|
|
|
|
log.Printf("[WS] Agent upgrade failed from %s: %v", clientIP, err)
|
|
|
|
|
return
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
go h.runPingLoopRaw(conn)
|
|
|
|
|
log.Printf("[WS] Agent WebSocket upgraded OK from %s", clientIP)
|
|
|
|
|
|
|
|
|
|
agentID := ""
|
|
|
|
|
defer func() {
|
|
|
|
|
@@ -351,17 +388,18 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
continue
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Verify fleet secret. If the server has one configured, the agent must match.
|
|
|
|
|
h.mu.RLock()
|
|
|
|
|
requiredSecret := h.fleetSecret
|
|
|
|
|
h.mu.RUnlock()
|
|
|
|
|
if requiredSecret != "" && !secureStringEqual(auth.FleetSecret, requiredSecret) {
|
|
|
|
|
conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
|
|
|
|
|
"success": false, "error": "invalid fleet secret — re-forge this agent",
|
|
|
|
|
})})
|
|
|
|
|
log.Printf("[auth] Agent rejected: bad fleet secret (host=%s id=%s)", auth.Hostname, auth.AgentID)
|
|
|
|
|
return
|
|
|
|
|
}
|
|
|
|
|
// Verify fleet secret. If the server has one configured, the agent must match.
|
|
|
|
|
h.mu.RLock()
|
|
|
|
|
requiredSecret := h.fleetSecret
|
|
|
|
|
h.mu.RUnlock()
|
|
|
|
|
log.Printf("[WS] Agent auth: id=%s host=%s secret_prefix=%.8s", auth.AgentID, auth.Hostname, auth.FleetSecret)
|
|
|
|
|
if requiredSecret != "" && !secureStringEqual(auth.FleetSecret, requiredSecret) {
|
|
|
|
|
conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
|
|
|
|
|
"success": false, "error": "invalid fleet secret — re-forge this agent",
|
|
|
|
|
})})
|
|
|
|
|
log.Printf("[auth] Agent rejected: bad fleet secret (host=%s id=%s)", auth.Hostname, auth.AgentID)
|
|
|
|
|
return
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
agentID = auth.AgentID
|
|
|
|
|
if agentID == "" {
|
|
|
|
|
@@ -410,8 +448,6 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Build backup pool.Config list from what the agent sent at auth.
|
|
|
|
|
// These are registered on the proxy so reconnect() rotates through
|
|
|
|
|
// them automatically — not just at initial connect.
|
|
|
|
|
var backupCfgs []pool.Config
|
|
|
|
|
for _, bp := range backupPools {
|
|
|
|
|
if bp.Host == "" || bp.Port <= 0 {
|
|
|
|
|
@@ -431,9 +467,14 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
backupCfgs = append(backupCfgs, bpc)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if _, err := h.poolManager.EnsurePoolWithBackups(&poolCfg, backupCfgs); err != nil {
|
|
|
|
|
log.Printf("[WS] All pools failed for agent %s (%d backups tried) — agent will mine when pool reconnects", agentID, len(backupCfgs))
|
|
|
|
|
}
|
|
|
|
|
// Connect to pool in background — do NOT block the auth_response.
|
|
|
|
|
// The agent can start and the pool proxy will be ready by the time
|
|
|
|
|
// the first share is submitted.
|
|
|
|
|
go func(pc pool.Config, bcs []pool.Config, aid string) {
|
|
|
|
|
if _, err := h.poolManager.EnsurePoolWithBackups(&pc, bcs); err != nil {
|
|
|
|
|
log.Printf("[WS] Pool init for agent %s failed (will retry): %v", aid, err)
|
|
|
|
|
}
|
|
|
|
|
}(poolCfg, backupCfgs, agentID)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if h.aiHandler != nil && forgeCfg.AIEnabled {
|
|
|
|
|
@@ -461,6 +502,7 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
Platform: auth.Platform,
|
|
|
|
|
Arch: auth.Arch,
|
|
|
|
|
OSVersion: auth.OSVersion,
|
|
|
|
|
Hostname: auth.Hostname,
|
|
|
|
|
Capabilities: &caps,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
@@ -496,14 +538,22 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
oldConn.Close()
|
|
|
|
|
h.mu.Lock()
|
|
|
|
|
}
|
|
|
|
|
h.agents[agentID] = &AgentConnection{AgentID: agentID, Conn: conn}
|
|
|
|
|
ac := &AgentConnection{AgentID: agentID, Conn: conn}
|
|
|
|
|
h.agents[agentID] = ac
|
|
|
|
|
h.mu.Unlock()
|
|
|
|
|
|
|
|
|
|
// Start the RTT-aware ping loop now that we have an AgentConnection.
|
|
|
|
|
go h.runPingLoopAgent(ac)
|
|
|
|
|
|
|
|
|
|
conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
|
|
|
|
|
"success": true,
|
|
|
|
|
"agent_id": agentID,
|
|
|
|
|
})})
|
|
|
|
|
|
|
|
|
|
// Enrich agent with hostname before broadcasting so the dashboard
|
|
|
|
|
// immediately shows the correct machine-specific display name.
|
|
|
|
|
agent.Hostname = auth.Hostname
|
|
|
|
|
|
|
|
|
|
h.broadcastDashboard(Message{
|
|
|
|
|
Type: "agent_online",
|
|
|
|
|
Payload: mustMarshal(agent),
|
|
|
|
|
@@ -659,6 +709,14 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
if len(stats.Services) > 0 {
|
|
|
|
|
broadcast["services"] = stats.Services
|
|
|
|
|
}
|
|
|
|
|
// Attach latest RTT latency from the ping loop.
|
|
|
|
|
if ac := h.getAgentConn(agentID); ac != nil {
|
|
|
|
|
ac.latencyMu.Lock()
|
|
|
|
|
if ac.LatencyMs != nil {
|
|
|
|
|
broadcast["latency_ms"] = *ac.LatencyMs
|
|
|
|
|
}
|
|
|
|
|
ac.latencyMu.Unlock()
|
|
|
|
|
}
|
|
|
|
|
h.broadcastDashboard(Message{Type: "stats_update", Payload: mustMarshal(broadcast)})
|
|
|
|
|
|
|
|
|
|
case "submit_share":
|
|
|
|
|
@@ -751,11 +809,18 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
var proxy *pool.Proxy
|
|
|
|
|
if h.poolManager != nil {
|
|
|
|
|
poolCfg := h.agentPoolConfig(agentID)
|
|
|
|
|
// Only use GetPool (non-blocking). If the pool hasn't connected yet
|
|
|
|
|
// (background EnsurePoolWithBackups from auth is still dialing), kick
|
|
|
|
|
// off another async attempt rather than blocking the WS read loop.
|
|
|
|
|
proxy = h.poolManager.GetPool(&poolCfg)
|
|
|
|
|
if proxy == nil {
|
|
|
|
|
if p, err := h.poolManager.EnsurePool(&poolCfg); err == nil {
|
|
|
|
|
proxy = p
|
|
|
|
|
}
|
|
|
|
|
go func(pc pool.Config) {
|
|
|
|
|
if p, err := h.poolManager.EnsurePool(&pc); err != nil {
|
|
|
|
|
log.Printf("[WS] get_job EnsurePool for %s failed: %v", pc.Host, err)
|
|
|
|
|
} else {
|
|
|
|
|
_ = p
|
|
|
|
|
}
|
|
|
|
|
}(poolCfg)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
if proxy != nil {
|
|
|
|
|
@@ -763,10 +828,10 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
if job != nil {
|
|
|
|
|
conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(job)})
|
|
|
|
|
} else {
|
|
|
|
|
conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "no job available"})})
|
|
|
|
|
conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "no job available — pool connecting"})})
|
|
|
|
|
}
|
|
|
|
|
} else {
|
|
|
|
|
conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "pool not connected"})})
|
|
|
|
|
conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "pool connecting — retry shortly"})})
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
case "log_tail":
|
|
|
|
|
@@ -920,6 +985,29 @@ func (h *WSHub) SendToAgent(agentID string, msg Message) error {
|
|
|
|
|
return agent.SendJSON(msg)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// RemoveAgent forcibly disconnects an agent and removes it from the live map.
|
|
|
|
|
// It then broadcasts agent_deleted to all dashboard clients so the UI removes
|
|
|
|
|
// the agent immediately without waiting for the disconnect goroutine to fire.
|
|
|
|
|
func (h *WSHub) RemoveAgent(agentID string) {
|
|
|
|
|
h.mu.Lock()
|
|
|
|
|
if ac, ok := h.agents[agentID]; ok {
|
|
|
|
|
// Nil the map entry BEFORE closing so the agent goroutine's deferred
|
|
|
|
|
// cleanup (which checks cur.Conn == conn) falls into the else branch
|
|
|
|
|
// and skips SetAgentOffline — avoiding a write to an already-deleted row.
|
|
|
|
|
delete(h.agents, agentID)
|
|
|
|
|
delete(h.agentConfigs, agentID)
|
|
|
|
|
delete(h.agentLogs, agentID)
|
|
|
|
|
delete(h.agentCapabilities, agentID)
|
|
|
|
|
ac.Conn.Close()
|
|
|
|
|
}
|
|
|
|
|
h.mu.Unlock()
|
|
|
|
|
// Broadcast deletion so every connected dashboard removes the agent immediately.
|
|
|
|
|
h.broadcastDashboard(Message{
|
|
|
|
|
Type: "agent_deleted",
|
|
|
|
|
Payload: mustMarshal(map[string]string{"agent_id": agentID}),
|
|
|
|
|
})
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// SendAgentCommand sends a remote command to an agent.
|
|
|
|
|
func (h *WSHub) SendAgentCommand(agentID, action string, args map[string]interface{}) error {
|
|
|
|
|
payload := map[string]interface{}{"action": action}
|
|
|
|
|
@@ -970,17 +1058,23 @@ func (h *WSHub) BroadcastAIActivity(entry interface{}) {
|
|
|
|
|
h.broadcastDashboard(Message{Type: "ai_activity", Payload: mustMarshal(entry)})
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// agentDisplayName returns a display name that is unique per physical machine.
|
|
|
|
|
// Hostname is preferred because it's machine-specific — many agents deployed from
|
|
|
|
|
// the same binary would otherwise share the same baked-in worker name, making
|
|
|
|
|
// a large fleet impossible to differentiate.
|
|
|
|
|
func agentDisplayName(workerName, worker, hostname, agentID string) string {
|
|
|
|
|
if workerName != "" {
|
|
|
|
|
return workerName
|
|
|
|
|
}
|
|
|
|
|
if worker != "" {
|
|
|
|
|
return worker
|
|
|
|
|
}
|
|
|
|
|
if hostname != "" {
|
|
|
|
|
return hostname
|
|
|
|
|
}
|
|
|
|
|
return shortAgentID(agentID)
|
|
|
|
|
// No hostname reported — make the worker name unique with a short agent ID suffix.
|
|
|
|
|
base := workerName
|
|
|
|
|
if base == "" {
|
|
|
|
|
base = worker
|
|
|
|
|
}
|
|
|
|
|
if base == "" {
|
|
|
|
|
base = "agent"
|
|
|
|
|
}
|
|
|
|
|
return base + "-" + shortAgentID(agentID)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func shortAgentID(id string) string {
|
|
|
|
|
|