137 lines
3.3 KiB
Go
137 lines
3.3 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"net/http"
|
|
"net/netip"
|
|
"net/url"
|
|
"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")
|
|
tr := http.DefaultTransport.(*http.Transport).Clone()
|
|
tr.Proxy = apiProxy
|
|
hc := &http.Client{Transport: tr, 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)
|
|
}
|
|
|
|
// http.ProxyFromEnvironment never proxies loopback, but a sandboxed shell can
|
|
// reach byre's loopback API only through its proxy.
|
|
func apiProxy(r *http.Request) (*url.URL, error) {
|
|
if os.Getenv("BYRE_NO_PROXY") == "1" {
|
|
return nil, nil
|
|
}
|
|
p := os.Getenv("HTTP_PROXY")
|
|
if p == "" {
|
|
p = os.Getenv("http_proxy")
|
|
}
|
|
if p == "" || noProxy(r.URL) {
|
|
return nil, nil
|
|
}
|
|
if !strings.Contains(p, "://") {
|
|
p = "http://" + p
|
|
}
|
|
return url.Parse(p)
|
|
}
|
|
|
|
func noProxy(u *url.URL) bool {
|
|
list := os.Getenv("NO_PROXY")
|
|
if list == "" {
|
|
list = os.Getenv("no_proxy")
|
|
}
|
|
host := u.Hostname()
|
|
ip, ipErr := netip.ParseAddr(host)
|
|
for _, e := range strings.Split(list, ",") {
|
|
e = strings.TrimSpace(e)
|
|
if e == "" {
|
|
continue
|
|
}
|
|
if pfx, err := netip.ParsePrefix(e); err == nil {
|
|
if ipErr == nil && pfx.Contains(ip) {
|
|
return true
|
|
}
|
|
continue
|
|
}
|
|
if e == "*" || e == u.Host || e == host || strings.HasSuffix(host, "."+strings.TrimPrefix(e, ".")) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func execFail(code int, msg string) {
|
|
fmt.Fprintln(os.Stderr, "byre client-exec: "+msg)
|
|
os.Exit(code)
|
|
}
|