client-exec CLI

internal/server/clients.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -11,10 +12,12 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"time"
 
 	"github.com/coder/websocket"
 
 	"g.fc.run/ai/byre/internal/api"
+	"g.fc.run/ai/byre/internal/clients"
 )
 
 func (s *Server) ConnectClient(w http.ResponseWriter, r *http.Request) {
@@ -60,6 +63,57 @@ func (s *Server) SetClientSync(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusNoContent)
 }
 
+const execGrace = 30 * time.Second
+
+func (s *Server) ExecClient(w http.ResponseWriter, r *http.Request) {
+	if !s.bearer(r) {
+		writeError(w, http.StatusUnauthorized, errors.New("exec needs the loopback token as a bearer token"))
+		return
+	}
+	req := api.ClientExec{}
+	if !readJSON(w, r, &req) {
+		return
+	}
+	if req.TimeoutMs <= 0 {
+		writeError(w, http.StatusBadRequest, errors.New("timeout_ms must be positive"))
+		return
+	}
+
+	name := ""
+	if req.Client != nil {
+		name = *req.Client
+	}
+	if name == "" {
+		recent := s.clients.Recent()
+		if len(recent) == 0 {
+			writeError(w, http.StatusServiceUnavailable, errors.New("no client Macs are connected"))
+			return
+		}
+		name = recent[0].ID
+	}
+
+	timeout := time.Duration(req.TimeoutMs) * time.Millisecond
+	ctx, cancel := context.WithTimeout(r.Context(), timeout+execGrace)
+	defer cancel()
+	start := time.Now()
+	reply, err := s.clients.Call(ctx, name, api.ClientRequest{Op: api.Exec, Command: &req.Command, TimeoutMs: &req.TimeoutMs})
+	if err != nil {
+		slog.Info("client exec", "client", name, "command", req.Command, "err", err, "duration", time.Since(start))
+		status := http.StatusInternalServerError
+		if errors.Is(err, clients.ErrUnavailable) {
+			status = http.StatusServiceUnavailable
+		}
+		writeError(w, status, err)
+		return
+	}
+	if reply.ExitCode == nil || reply.Stdout == nil || reply.Stderr == nil {
+		writeError(w, http.StatusInternalServerError, fmt.Errorf("%s's reply has no exit_code, stdout or stderr", name))
+		return
+	}
+	slog.Info("client exec", "client", name, "command", req.Command, "exit", *reply.ExitCode, "duration", time.Since(start))
+	writeJSON(w, http.StatusOK, api.ClientExecResult{Client: name, ExitCode: *reply.ExitCode, Stdout: *reply.Stdout, Stderr: *reply.Stderr})
+}
+
 func (s *Server) GetClientTransfer(w http.ResponseWriter, r *http.Request, transferID string) {
 	t, ok := s.clients.Transfer(transferID)
 	if !ok || t.Receive {

internal/server/server.go

@@ -35,6 +35,7 @@ type Server struct {
 	terminals *terminal.Terminals
 	tailnet   *local.Client
 	hostname  string
+	token     string
 	startedAt time.Time
 
 	usersMu sync.Mutex
@@ -43,7 +44,7 @@ type Server struct {
 
 type ctxKey struct{}
 
-func New(dataDir string, portFirst int, ts *tsnet.Server) (*Server, error) {
+func New(dataDir string, portFirst int, ts *tsnet.Server, token string) (*Server, error) {
 	err := os.MkdirAll(dataDir, 0o700)
 	if err != nil {
 		return nil, err
@@ -62,6 +63,7 @@ func New(dataDir string, portFirst int, ts *tsnet.Server) (*Server, error) {
 		hub:       newHub(),
 		tailnet:   lc,
 		hostname:  ts.Hostname,
+		token:     token,
 		startedAt: time.Now().UTC(),
 		users:     map[string]string{},
 	}
@@ -130,11 +132,12 @@ func (s *Server) Handler() http.Handler {
 	return http.NewCrossOriginProtection().Handler(mux)
 }
 
-func (s *Server) LoopbackHandler(token string) http.Handler {
+func (s *Server) LoopbackHandler() http.Handler {
 	mux := http.NewServeMux()
 	mux.HandleFunc("GET /api/tailnet", s.GetTailnet)
+	mux.HandleFunc("POST /api/clients/exec", s.ExecClient)
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if subtle.ConstantTimeCompare([]byte(r.Header.Get("Authorization")), []byte("Bearer "+token)) != 1 {
+		if !s.bearer(r) {
 			http.Error(w, "unauthorized", http.StatusUnauthorized)
 			return
 		}
@@ -143,6 +146,10 @@ func (s *Server) LoopbackHandler(token string) http.Handler {
 	})
 }
 
+func (s *Server) bearer(r *http.Request) bool {
+	return subtle.ConstantTimeCompare([]byte(r.Header.Get("Authorization")), []byte("Bearer "+s.token)) == 1
+}
+
 func (s *Server) GetTailnet(w http.ResponseWriter, r *http.Request) {
 	st, err := s.tailnet.StatusWithoutPeers(r.Context())
 	if err != nil {