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 {