diff --git a/cmd/byre/serve.go b/cmd/byre/serve.go index 05155e2..52d0d03 100644 --- a/cmd/byre/serve.go +++ b/cmd/byre/serve.go @@ -19,6 +19,7 @@ import ( "syscall" "time" + "g.fc.run/ai/byre/internal/peer" "g.fc.run/ai/byre/internal/server" "tailscale.com/tsnet" ) @@ -117,7 +118,7 @@ func serveCmd(args []string) { base := func(net.Listener) context.Context { return ctx } loopback := &http.Server{Addr: *listenAddr, Handler: srv.LoopbackHandler(), BaseContext: base} - tailnet := &http.Server{Handler: srv.Handler(), BaseContext: base} + tailnet := &http.Server{Handler: srv.Handler(), BaseContext: base, ConnContext: peer.ConnContext} go func() { <-ctx.Done() diff --git a/go.mod b/go.mod index eda809b..f103169 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( github.com/getkin/kin-openapi v0.142.0 github.com/go-git/go-git/v5 v5.19.3 github.com/oapi-codegen/runtime v1.7.0 + golang.org/x/sys v0.48.0 howett.net/plist v1.0.0 modernc.org/sqlite v1.60.1 tailscale.com v1.104.2 @@ -69,7 +70,6 @@ require ( golang.org/x/net v0.59.0 // indirect golang.org/x/oauth2 v0.37.0 // indirect golang.org/x/sync v0.23.0 // indirect - golang.org/x/sys v0.48.0 // indirect golang.org/x/term v0.46.0 // indirect golang.org/x/text v0.42.0 // indirect golang.org/x/time v0.16.0 // indirect diff --git a/internal/agents/health.go b/internal/agents/health.go index d101cf2..664b113 100644 --- a/internal/agents/health.go +++ b/internal/agents/health.go @@ -1,24 +1,21 @@ package agents import ( - "bufio" - "bytes" "context" "encoding/json" "errors" "fmt" "log/slog" "os" - "os/exec" "path/filepath" "slices" - "strconv" "strings" "sync" "syscall" "time" "g.fc.run/ai/byre/internal/api" + "golang.org/x/sys/unix" ) const ( @@ -112,7 +109,7 @@ func (m *Manager) Health(ctx context.Context) (Health, error) { return Health{}, err } - children, err := claudeChildren(ctx, os.Getpid()) + children, err := claudeChildren(os.Getpid()) if err != nil { return Health{}, err } @@ -226,39 +223,29 @@ func pidAlive(pid int) bool { return err == nil || errors.Is(err, syscall.EPERM) } -func claudeChildren(ctx context.Context, parent int) ([]int, error) { - cctx, cancel := context.WithTimeout(ctx, 5*time.Second) - defer cancel() - - out, err := exec.CommandContext(cctx, "ps", "-A", "-o", "pid=,ppid=,comm=").Output() +func claudeChildren(parent int) ([]int, error) { + procs, err := unix.SysctlKinfoProcSlice("kern.proc.all") if err != nil { return nil, fmt.Errorf("list processes: %w", err) } - return parseChildren(out, parent), nil -} - -func parseChildren(out []byte, parent int) []int { pids := []int{} - sc := bufio.NewScanner(bytes.NewReader(out)) - for sc.Scan() { - fields := strings.Fields(sc.Text()) - if len(fields) < 3 { + for _, p := range procs { + if int(p.Eproc.Ppid) != parent { continue } - pid, err := strconv.Atoi(fields[0]) - if err != nil { + pid := int(p.Proc.P_pid) + // p_comm holds the resolved binary's name (claude's is its version), so match the path it was started by. + args, err := unix.SysctlRaw("kern.procargs2", pid) + if err != nil || len(args) < 4 { continue } - ppid, err := strconv.Atoi(fields[1]) - if err != nil || ppid != parent { - continue - } - if filepath.Base(strings.Join(fields[2:], " ")) != "claude" { + path, _, _ := strings.Cut(string(args[4:]), "\x00") + if filepath.Base(path) != "claude" { continue } pids = append(pids, pid) } - return pids + return pids, nil } func problems(h Health) []string { diff --git a/internal/agents/preview.go b/internal/agents/preview.go index 80f5184..444052d 100644 --- a/internal/agents/preview.go +++ b/internal/agents/preview.go @@ -17,6 +17,7 @@ import ( "g.fc.run/ai/byre/internal/api" "g.fc.run/ai/byre/internal/mcp" + "g.fc.run/ai/byre/internal/peer" "g.fc.run/ai/byre/internal/store" ) @@ -169,7 +170,7 @@ func (m *Manager) servePreview(port int) error { http.Error(w, "forbidden preview host", http.StatusForbidden) return } - who, err := lc.WhoIs(r.Context(), r.RemoteAddr) + who, err := peer.WhoIs(r, lc) if err != nil { http.Error(w, err.Error(), http.StatusForbidden) return @@ -183,7 +184,7 @@ func (m *Manager) servePreview(port int) error { proxy.ServeHTTP(w, r) }) go func() { - err := (&http.Server{Handler: handler}).Serve(tracked) + err := (&http.Server{Handler: handler, ConnContext: peer.ConnContext}).Serve(tracked) if !errors.Is(err, net.ErrClosed) { slog.Error("serve preview", "port", port, "err", err) } diff --git a/internal/peer/peer.go b/internal/peer/peer.go new file mode 100644 index 0000000..90e2168 --- /dev/null +++ b/internal/peer/peer.go @@ -0,0 +1,37 @@ +package peer + +import ( + "context" + "net" + "net/http" + "sync" + + "tailscale.com/client/local" + "tailscale.com/client/tailscale/apitype" +) + +type key struct{} + +type peer struct { + mu sync.Mutex + who *apitype.WhoIsResponse +} + +func ConnContext(ctx context.Context, _ net.Conn) context.Context { + return context.WithValue(ctx, key{}, &peer{}) +} + +func WhoIs(r *http.Request, lc *local.Client) (*apitype.WhoIsResponse, error) { + p := r.Context().Value(key{}).(*peer) + p.mu.Lock() + defer p.mu.Unlock() + if p.who != nil { + return p.who, nil + } + who, err := lc.WhoIs(r.Context(), r.RemoteAddr) + if err != nil { + return nil, err + } + p.who = who + return who, nil +} diff --git a/internal/server/server.go b/internal/server/server.go index a68a986..dc18b8c 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -19,6 +19,7 @@ import ( "g.fc.run/ai/byre/internal/clients" "g.fc.run/ai/byre/internal/git" "g.fc.run/ai/byre/internal/ids" + "g.fc.run/ai/byre/internal/peer" "g.fc.run/ai/byre/internal/schedule" "g.fc.run/ai/byre/internal/store" "g.fc.run/ai/byre/internal/terminal" @@ -172,7 +173,7 @@ func (s *Server) GetTailnet(w http.ResponseWriter, r *http.Request) { func (s *Server) identify(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { identified := trace.Start(r.Context(), "whois") - who, err := s.tailnet.WhoIs(r.Context(), r.RemoteAddr) + who, err := peer.WhoIs(r, s.tailnet) identified() if err != nil { writeError(w, http.StatusUnauthorized, fmt.Errorf("identify %s on the tailnet: %w", r.RemoteAddr, err))