feat: add node resource profiling and admin device revocation

This commit is contained in:
kory
2026-09-28 09:00:31 +08:00
parent a7b23b760a
commit 3776c310e4
17 changed files with 589 additions and 36 deletions
+54 -7
View File
@@ -12,20 +12,24 @@ import (
"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
Name string
Conn *websocket.Conn
WriteMu sync.Mutex
Pending sync.Map
LastSeen atomic.Int64
LastProfile *sysinfo.Profile
}
type Result struct {
@@ -46,6 +50,11 @@ type FileResult struct {
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 := ""
@@ -61,6 +70,14 @@ func (m *Manager) ServeWS(w http.ResponseWriter, r *http.Request) {
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 {
@@ -115,7 +132,10 @@ func (m *Manager) readLoop(n *Node) {
}
n.LastSeen.Store(time.Now().UnixNano())
_ = n.Conn.SetReadDeadline(time.Now().Add(75 * time.Second))
if msg.Event == "pong" || msg.Event == "ping" {
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 != "" {
@@ -130,6 +150,7 @@ func (m *Manager) readLoop(n *Node) {
}
}
}
func (m *Manager) Online(name string) bool {
m.mu.RLock()
defer m.mu.RUnlock()
@@ -137,6 +158,22 @@ func (m *Manager) Online(name string) bool {
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()
@@ -147,7 +184,16 @@ func (m *Manager) Devices() []map[string]any {
if time.Since(lastSeen) >= 70*time.Second {
status = "offline"
}
out = append(out, map[string]any{"name": name, "status": status, "is_local": false, "last_heartbeat": lastSeen.Unix()})
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
}
@@ -180,6 +226,7 @@ func (m *Manager) call(ctx context.Context, name, action string, payload any) (M
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 {
+15 -12
View File
@@ -1,18 +1,21 @@
package cluster
import "github.com/wjyhk/vps-commander/internal/sysinfo"
type Message struct {
ID string `json:"id,omitempty"`
Event string `json:"event,omitempty"`
Action string `json:"action,omitempty"`
Device string `json:"device,omitempty"`
Status string `json:"status,omitempty"`
Payload any `json:"payload,omitempty"`
ExitCode int `json:"exit_code,omitempty"`
Stdout string `json:"stdout,omitempty"`
Stderr string `json:"stderr,omitempty"`
DurationMS int64 `json:"duration_ms,omitempty"`
Error string `json:"error,omitempty"`
TS int64 `json:"ts,omitempty"`
ID string `json:"id,omitempty"`
Event string `json:"event,omitempty"`
Action string `json:"action,omitempty"`
Device string `json:"device,omitempty"`
Status string `json:"status,omitempty"`
Payload any `json:"payload,omitempty"`
Profile *sysinfo.Profile `json:"profile,omitempty"`
ExitCode int `json:"exit_code,omitempty"`
Stdout string `json:"stdout,omitempty"`
Stderr string `json:"stderr,omitempty"`
DurationMS int64 `json:"duration_ms,omitempty"`
Error string `json:"error,omitempty"`
TS int64 `json:"ts,omitempty"`
}
type ExecPayload struct {