vps-commander/cmd/agent/main.go

package main

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/wjyhk/vps-commander/internal/cluster"
	"github.com/wjyhk/vps-commander/internal/executor"
)

func main() {
	hub := flag.String("hub", "", "Hub WebSocket URL")
	name := flag.String("name", "", "device name")
	token := flag.String("token", os.Getenv("VPS_COMMANDER_CLUSTER_SECRET"), "cluster secret")
	insecure := flag.Bool("insecure", false, "skip TLS verification")
	flag.Parse()
	if *hub == "" || *name == "" || *token == "" {
		log.Fatal("-hub, -name and -token are required")
	}
	backoff := time.Second
	for {
		if err := run(*hub, *name, *token, *insecure); err != nil {
			log.Printf("agent connection: %v", err)
		}
		time.Sleep(backoff)
		if backoff < 30*time.Second {
			backoff *= 2
			if backoff > 30*time.Second {
				backoff = 30 * time.Second
			}
		}
	}
}

func run(hub, name, token string, insecure bool) error {
	u := hub + "?device_name=" + name
	dialer := websocket.DefaultDialer
	if insecure {
		dialer = &websocket.Dialer{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
	}
	header := http.Header{}
	header.Set("Authorization", "Bearer "+token)
	c, _, err := dialer.Dial(u, header)
	if err != nil {
		return err
	}
	defer c.Close()
	var writeMu sync.Mutex
	writeJSON := func(msg cluster.Message) error {
		writeMu.Lock()
		defer writeMu.Unlock()
		if err := c.SetWriteDeadline(time.Now().Add(10 * time.Second)); err != nil {
			return err
		}
		return c.WriteJSON(msg)
	}
	if err := c.SetReadDeadline(time.Now().Add(75 * time.Second)); err != nil {
		return err
	}
	exec := executor.Local{MaxOutput: 1024 * 1024}
	if err := writeJSON(cluster.Message{Event: "hello", Device: name, TS: time.Now().Unix()}); err != nil {
		return err
	}
	stopPing := make(chan struct{})
	defer close(stopPing)
	go func() {
		t := time.NewTicker(30 * time.Second)
		defer t.Stop()
		for {
			select {
			case <-t.C:
				_ = writeJSON(cluster.Message{Event: "ping", Device: name, TS: time.Now().Unix()})
			case <-stopPing:
				return
			}
		}
	}()
	for {
		var msg cluster.Message
		if err := c.ReadJSON(&msg); err != nil {
			return err
		}
		_ = c.SetReadDeadline(time.Now().Add(75 * time.Second))
		if msg.Event == "registered" {
			continue
		}
		if msg.Event == "ping" {
			if err := writeJSON(cluster.Message{Event: "pong", Device: name, TS: msg.TS}); err != nil {
				return err
			}
			continue
		}
		if msg.Action == "" || msg.ID == "" {
			continue
		}
		go func(msg cluster.Message) {
			if err := handle(writeJSON, exec, msg); err != nil {
				log.Printf("request %s failed: %v", msg.ID, err)
				_ = c.Close()
			}
		}(msg)
	}
}

func handle(writeJSON func(cluster.Message) error, exec executor.Local, msg cluster.Message) error {
	switch msg.Action {
	case "exec":
		var p cluster.ExecPayload
		if err := decodePayload(msg.Payload, &p); err != nil {
			return sendError(writeJSON, msg.ID, err)
		}
		timeout := p.Timeout
		if timeout <= 0 {
			timeout = 30
		}
		if timeout > 300 {
			timeout = 300
		}
		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
		res := exec.Run(ctx, p.Command, p.Workdir)
		cancel()
		return writeJSON(cluster.Message{ID: msg.ID, Status: "success", ExitCode: res.ExitCode, Stdout: res.Stdout, Stderr: res.Stderr, DurationMS: res.DurationMS})
	case "file_read":
		var p cluster.FileReadPayload
		if err := decodePayload(msg.Payload, &p); err != nil {
			return sendError(writeJSON, msg.ID, err)
		}
		if p.Offset < 0 {
			return sendError(writeJSON, msg.ID, errors.New("negative offset"))
		}
		f, err := os.Open(p.Path)
		if err != nil {
			return sendError(writeJSON, msg.ID, err)
		}
		defer f.Close()
		info, err := f.Stat()
		if err != nil {
			return sendError(writeJSON, msg.ID, err)
		}
		if _, err = f.Seek(p.Offset, io.SeekStart); err != nil {
			return sendError(writeJSON, msg.ID, err)
		}
		limit := p.Limit
		if limit <= 0 || limit > 1024*1024 {
			limit = 256 * 1024
		}
		buf := make([]byte, limit)
		n, err := io.ReadFull(f, buf)
		eof := err == io.EOF || err == io.ErrUnexpectedEOF
		if err != nil && !eof {
			return sendError(writeJSON, msg.ID, err)
		}
		return writeJSON(cluster.Message{ID: msg.ID, Status: "success", Payload: map[string]any{"content": string(buf[:n]), "bytes": n, "total_size": info.Size(), "has_more": p.Offset+int64(n) < info.Size(), "eof": eof}})
	case "file_write":
		var p cluster.FileWritePayload
		if err := decodePayload(msg.Payload, &p); err != nil {
			return sendError(writeJSON, msg.ID, err)
		}
		if len(p.Content) > 1024*1024 {
			return sendError(writeJSON, msg.ID, errors.New("content too large"))
		}
		if err := os.WriteFile(p.Path, []byte(p.Content), 0644); err != nil {
			return sendError(writeJSON, msg.ID, err)
		}
		return writeJSON(cluster.Message{ID: msg.ID, Status: "success", Payload: map[string]any{"bytes": len(p.Content)}})
	default:
		return sendError(writeJSON, msg.ID, errors.New("unsupported action"))
	}
}

func decodePayload(v any, out any) error {
	b, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return json.Unmarshal(b, out)
}
func sendError(writeJSON func(cluster.Message) error, id string, err error) error {
	return writeJSON(cluster.Message{ID: id, Status: "error", Error: err.Error()})
}