package cluster import ( "context" "crypto/subtle" "encoding/json" "errors" "fmt" "net/http" "sync" "sync/atomic" "time" "github.com/gorilla/websocket" "github.com/wjyhk/vps-commander/internal/storage" "github.com/wjyhk/vps-commander/internal/sysinfo" ) type Manager struct { Secret string Store *storage.Store mu sync.RWMutex nodes map[string]*Node } type Node struct { Name string Conn *websocket.Conn WriteMu sync.Mutex Pending sync.Map LastSeen atomic.Int64 LastProfile *sysinfo.Profile } type Result struct { ExitCode int Stdout string Stderr string DurationMS int64 } type FileResult struct { Content string `json:"content"` Bytes int64 `json:"bytes"` TotalSize int64 `json:"total_size"` HasMore bool `json:"has_more"` EOF bool `json:"eof"` } func NewManager(secret string) *Manager { return &Manager{Secret: secret, nodes: make(map[string]*Node)} } func (m *Manager) SetStore(s *storage.Store) { m.Store = s } func (m *Manager) ServeWS(w http.ResponseWriter, r *http.Request) { auth := r.Header.Get("Authorization") provided := "" if len(auth) >= 7 && auth[:7] == "Bearer " { provided = auth[7:] } if m.Secret == "" || subtle.ConstantTimeCompare([]byte(provided), []byte(m.Secret)) != 1 { http.Error(w, "unauthorized", 401) return } name := r.URL.Query().Get("device_name") if name == "" { http.Error(w, "device_name required", 400) return } if m.Store != nil { if revoked, _ := m.Store.IsRevoked(name); revoked { http.Error(w, "device has been revoked by admin", 403) return } } up := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} c, err := up.Upgrade(w, r, nil) if err != nil { return } n := &Node{Name: name, Conn: c} n.LastSeen.Store(time.Now().UnixNano()) m.mu.Lock() if old := m.nodes[name]; old != nil { _ = old.Conn.Close() } m.nodes[name] = n m.mu.Unlock() _ = n.write(Message{Event: "registered", Device: name, TS: time.Now().Unix()}) go m.heartbeat(n) go m.readLoop(n) } func (n *Node) write(msg Message) error { n.WriteMu.Lock() defer n.WriteMu.Unlock() if err := n.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second)); err != nil { return err } return n.Conn.WriteJSON(msg) } func (m *Manager) readLoop(n *Node) { defer func() { m.mu.Lock() if m.nodes[n.Name] == n { delete(m.nodes, n.Name) } m.mu.Unlock() _ = n.Conn.Close() n.Pending.Range(func(key, v any) bool { if ch, ok := v.(chan Message); ok { select { case ch <- Message{Status: "error", Error: "device disconnected"}: default: } n.Pending.Delete(key) } return true }) }() _ = n.Conn.SetReadDeadline(time.Now().Add(75 * time.Second)) for { var msg Message if err := n.Conn.ReadJSON(&msg); err != nil { return } n.LastSeen.Store(time.Now().UnixNano()) _ = n.Conn.SetReadDeadline(time.Now().Add(75 * time.Second)) if msg.Profile != nil { n.LastProfile = msg.Profile } if msg.Event == "pong" || msg.Event == "ping" || msg.Event == "hello" || msg.Event == "profile" { continue } if msg.ID != "" { if v, ok := n.Pending.LoadAndDelete(msg.ID); ok { if ch, ok := v.(chan Message); ok { select { case ch <- msg: default: } } } } } } func (m *Manager) Online(name string) bool { m.mu.RLock() defer m.mu.RUnlock() n := m.nodes[name] return n != nil && time.Since(time.Unix(0, n.LastSeen.Load())) < 70*time.Second } func (m *Manager) RevokeDevice(name, reason string) error { m.mu.Lock() n, exists := m.nodes[name] if exists { delete(m.nodes, name) _ = n.write(Message{Event: "revoked", Error: "revoked by admin"}) _ = n.Conn.Close() } m.mu.Unlock() if m.Store != nil { return m.Store.RevokeDevice(name, reason) } return nil } func (m *Manager) Devices() []map[string]any { m.mu.RLock() defer m.mu.RUnlock() out := make([]map[string]any, 0, len(m.nodes)) for name, n := range m.nodes { status := "online" lastSeen := time.Unix(0, n.LastSeen.Load()) if time.Since(lastSeen) >= 70*time.Second { status = "offline" } item := map[string]any{ "name": name, "status": status, "is_local": false, "last_heartbeat": lastSeen.Unix(), } if n.LastProfile != nil { item["profile"] = n.LastProfile } out = append(out, item) } return out } func (m *Manager) call(ctx context.Context, name, action string, payload any) (Message, error) { m.mu.RLock() n := m.nodes[name] m.mu.RUnlock() if n == nil || !m.Online(name) { return Message{}, errors.New("device offline") } id := fmt.Sprintf("task_%d", time.Now().UnixNano()) ch := make(chan Message, 1) n.Pending.Store(id, ch) if err := n.write(Message{ID: id, Action: action, Payload: payload}); err != nil { n.Pending.Delete(id) return Message{}, err } select { case msg, ok := <-ch: if !ok { return Message{}, errors.New("device disconnected") } if msg.Status == "error" { return Message{}, errors.New(msg.Error) } return msg, nil case <-ctx.Done(): n.Pending.Delete(id) return Message{}, ctx.Err() } } func (m *Manager) Exec(ctx context.Context, name string, p ExecPayload) (Result, error) { msg, err := m.call(ctx, name, "exec", p) if err != nil { return Result{}, err } return Result{msg.ExitCode, msg.Stdout, msg.Stderr, msg.DurationMS}, nil } func (m *Manager) ReadFile(ctx context.Context, name string, p FileReadPayload) (FileResult, error) { msg, err := m.call(ctx, name, "file_read", p) if err != nil { return FileResult{}, err } var v FileResult b, _ := json.Marshal(msg.Payload) if err := json.Unmarshal(b, &v); err != nil { return FileResult{}, err } return v, nil } func (m *Manager) WriteFile(ctx context.Context, name string, p FileWritePayload) (int64, error) { msg, err := m.call(ctx, name, "file_write", p) if err != nil { return 0, err } var v struct { Bytes int64 `json:"bytes"` } b, _ := json.Marshal(msg.Payload) if err := json.Unmarshal(b, &v); err != nil { return 0, err } return v.Bytes, nil } func (m *Manager) heartbeat(n *Node) { t := time.NewTicker(30 * time.Second) defer t.Stop() for range t.C { m.mu.RLock() current := m.nodes[n.Name] == n m.mu.RUnlock() if !current { return } if err := n.write(Message{Event: "ping", TS: time.Now().Unix()}); err != nil { _ = n.Conn.Close() return } } }