byre/internal/server/clients.go

package server

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"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) {
	id := r.Header.Get("Byre-Client")
	host := r.Header.Get("Byre-Host")
	if id == "" || host == "" {
		writeError(w, http.StatusBadRequest, errors.New("missing Byre-Client or Byre-Host"))
		return
	}

	ws, err := websocket.Accept(w, r, nil)
	if err != nil {
		slog.Warn("accept client websocket", "host", host, "err", err)
		return
	}
	defer ws.CloseNow()

	s.clients.Serve(r.Context(), ws, id, actor(r), host, r.Header.Get("Byre-Commit"))
}

func (s *Server) SetClientSync(w http.ResponseWriter, r *http.Request) {
	id := r.Header.Get("Byre-Client")
	host := r.Header.Get("Byre-Host")
	if id == "" || host == "" {
		writeError(w, http.StatusBadRequest, errors.New("missing Byre-Client or Byre-Host"))
		return
	}
	sync := api.ClientSync{}
	err := json.NewDecoder(r.Body).Decode(&sync)
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}
	if !sync.Ok {
		if sync.Error == nil {
			writeError(w, http.StatusBadRequest, errors.New("a failed sync needs an error"))
			return
		}
		slog.Warn("client icloud sync failing", "id", id, "host", host, "error", *sync.Error)
	}
	s.clients.Synced(id, actor(r), host, r.Header.Get("Byre-Commit"), sync)
	s.clientsChanged()
	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 {
		writeError(w, http.StatusNotFound, fmt.Errorf("no transfer %s to send", transferID))
		return
	}

	w.Header().Set("Content-Type", "application/gzip")
	cmd := exec.CommandContext(r.Context(), "tar", "-czf", "-", "-C", t.Dir, "--", t.Name)
	cmd.Stdout = w
	stderr := &bytes.Buffer{}
	cmd.Stderr = stderr
	err := cmd.Run()
	if err != nil {
		slog.Error("stream client transfer", "transfer", transferID, "dir", t.Dir, "name", t.Name, "err", err, "stderr", stderr.String())
		panic(http.ErrAbortHandler)
	}
}

func (s *Server) PutClientTransfer(w http.ResponseWriter, r *http.Request, transferID string) {
	t, ok := s.clients.Transfer(transferID)
	if !ok || !t.Receive {
		writeError(w, http.StatusNotFound, fmt.Errorf("no transfer %s to receive", transferID))
		return
	}

	f, err := os.OpenFile(filepath.Join(t.Dir, t.Name), os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if err != nil {
		writeError(w, http.StatusNotFound, fmt.Errorf("transfer %s: %w", transferID, err))
		return
	}
	_, err = io.Copy(f, r.Body)
	cerr := f.Close()
	if err == nil {
		err = cerr
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, fmt.Errorf("receive transfer %s: %w", transferID, err))
		return
	}
	w.WriteHeader(http.StatusNoContent)
}