89 lines
2.2 KiB
Go
89 lines
2.2 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"net/http"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
|
|
"g.fc.run/ai/byre/internal/api"
|
|
)
|
|
|
|
const (
|
|
execUnavailable = 75
|
|
execTimedOut = 124
|
|
execGrace = time.Minute
|
|
)
|
|
|
|
func clientExecCmd(args []string) {
|
|
fs := flag.NewFlagSet("client-exec", flag.ExitOnError)
|
|
client := fs.String("client", "", "the client's name or id; the most recently active one when absent")
|
|
timeout := fs.Int("timeout", 30, "kill the command after this many seconds")
|
|
fs.Parse(args)
|
|
if fs.NArg() == 0 || *timeout <= 0 {
|
|
usage()
|
|
}
|
|
token := os.Getenv("BYRE_TOKEN")
|
|
if token == "" {
|
|
execFail(1, "BYRE_TOKEN is not set; byre sets it for everything it runs")
|
|
}
|
|
|
|
words := make([]string, fs.NArg())
|
|
for i, w := range fs.Args() {
|
|
words[i] = "'" + strings.ReplaceAll(w, "'", `'\''`) + "'"
|
|
}
|
|
req := api.ClientExec{Command: strings.Join(words, " "), TimeoutMs: *timeout * 1000}
|
|
if *client != "" {
|
|
req.Client = client
|
|
}
|
|
body, err := json.Marshal(req)
|
|
if err != nil {
|
|
execFail(1, err.Error())
|
|
}
|
|
hr, err := http.NewRequest(http.MethodPost, loopbackURL+"/api/clients/exec", bytes.NewReader(body))
|
|
if err != nil {
|
|
execFail(1, err.Error())
|
|
}
|
|
hr.Header.Set("Authorization", "Bearer "+token)
|
|
hr.Header.Set("Content-Type", "application/json")
|
|
hc := &http.Client{Timeout: time.Duration(*timeout)*time.Second + execGrace}
|
|
resp, err := hc.Do(hr)
|
|
if err != nil {
|
|
execFail(execUnavailable, err.Error())
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
e := api.Error{}
|
|
err = json.NewDecoder(resp.Body).Decode(&e)
|
|
if err != nil {
|
|
e.Message = resp.Status
|
|
}
|
|
code := 1
|
|
if resp.StatusCode == http.StatusServiceUnavailable {
|
|
code = execUnavailable
|
|
}
|
|
execFail(code, e.Message)
|
|
}
|
|
res := api.ClientExecResult{}
|
|
err = json.NewDecoder(resp.Body).Decode(&res)
|
|
if err != nil {
|
|
execFail(1, fmt.Sprintf("decode reply: %s", err))
|
|
}
|
|
os.Stdout.Write(res.Stdout)
|
|
os.Stderr.Write(res.Stderr)
|
|
if res.ExitCode == -1 {
|
|
execFail(execTimedOut, fmt.Sprintf("killed on %s after %ds", res.Client, *timeout))
|
|
}
|
|
os.Exit(res.ExitCode)
|
|
}
|
|
|
|
func execFail(code int, msg string) {
|
|
fmt.Fprintln(os.Stderr, "byre client-exec: "+msg)
|
|
os.Exit(code)
|
|
}
|