feat: initial commit of vps-commander release v1.0.0
This commit is contained in:
kory
@@ -0,0 +1,234 @@
|
||||
package cluster
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
type Manager struct {
|
||||
Secret string
|
||||
mu sync.RWMutex
|
||||
nodes map[string]*Node
|
||||
}
|
||||
|
||||
type Node struct {
|
||||
Name string
|
||||
Conn *websocket.Conn
|
||||
WriteMu sync.Mutex
|
||||
Pending sync.Map
|
||||
LastSeen atomic.Int64
|
||||
}
|
||||
|
||||
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) 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
|
||||
}
|
||||
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.Event == "pong" || msg.Event == "ping" {
|
||||
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) 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"
|
||||
}
|
||||
out = append(out, map[string]any{"name": name, "status": status, "is_local": false, "last_heartbeat": lastSeen.Unix()})
|
||||
}
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package cluster
|
||||
|
||||
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"`
|
||||
}
|
||||
|
||||
type ExecPayload struct {
|
||||
Command string `json:"command"`
|
||||
Workdir string `json:"workdir,omitempty"`
|
||||
Timeout int `json:"timeout,omitempty"`
|
||||
}
|
||||
|
||||
type FileReadPayload struct {
|
||||
Path string `json:"path"`
|
||||
Offset int64 `json:"offset,omitempty"`
|
||||
Limit int64 `json:"limit,omitempty"`
|
||||
}
|
||||
|
||||
type FileWritePayload struct {
|
||||
Path string `json:"path"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
Reference in New Issue
Block a user