Run byre and the app on their own embedded Tailscale nodes

internal/agents/preview.go

@@ -7,15 +7,17 @@ import (
 	"fmt"
 	"log/slog"
 	"net"
+	"net/http"
+	"net/http/httputil"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 	"unicode/utf8"
 
 	"g.fc.run/ai/byre/internal/api"
 	"g.fc.run/ai/byre/internal/mcp"
 	"g.fc.run/ai/byre/internal/store"
-	"g.fc.run/ai/byre/internal/tailscale"
 )
 
 const maxPreviewName = 40
@@ -34,7 +36,7 @@ func (m *Manager) exposePortTool(l *live) mcp.Tool {
 		Name: "expose_port",
 		Description: "Show the human a server you are running, such as a dev server, as a link above the task's tabs that opens it in a new window. " +
 			"Start the server yourself first (in the background), listening on 127.0.0.1 (not 0.0.0.0 or localhost) on a port in this task's block. " +
-			"byre then serves it over HTTPS on the same port of byre's host through tailscale serve, and returns the URL. " +
+			"byre then serves it on the same port of byre's tailnet node, and the human's app opens it through a local link. " +
 			"name labels it in the tab (\"web\", \"storybook\"); exposing a name or port again replaces the earlier one. " +
 			"The server stops when your process does (a byre restart, compaction, the task closing); the human sees it as unreachable until you start it again.",
 		InputSchema: json.RawMessage(`{
@@ -79,26 +81,138 @@ func (m *Manager) exposePort(ctx context.Context, l *live, raw json.RawMessage)
 	}
 	conn.Close()
 
-	err = tailscale.Serve(ctx, args.Port)
+	err = m.servePreview(args.Port)
 	if err != nil {
 		return "", err
 	}
 
 	dropped, err := m.store.SetPreview(ctx, api.Preview{TaskId: task.Id, Name: name, Port: args.Port, CreatedAt: time.Now().UTC()}, agentActor)
 	if errors.Is(err, store.ErrNotFound) {
+		m.stopPreview(args.Port)
 		return "", errors.New("this task is no longer active")
 	}
 	if err != nil {
 		return "", err
 	}
 	for _, port := range dropped {
-		err = tailscale.Off(ctx, port)
-		if err != nil {
-			slog.Error("stop serving replaced preview", "task", task.Id, "port", port, "err", err)
-		}
+		m.stopPreview(port)
 	}
 
-	return fmt.Sprintf("Exposed %q: https://<byre's host>:%d/ proxies to %s, and the human sees a link to it above the task's tabs.", name, args.Port, addr), nil
+	return fmt.Sprintf("Exposed %q: port %d of byre's tailnet node proxies to %s, and the human sees a link to it above the task's tabs.", name, args.Port, addr), nil
+}
+
+func (m *Manager) servePreview(port int) error {
+	m.stopPreview(port)
+	ln, err := m.tailnet.Listen("tcp", ":"+strconv.Itoa(port))
+	if err != nil {
+		return fmt.Errorf("listen on the tailnet on port %d: %w", port, err)
+	}
+	lc, err := m.tailnet.LocalClient()
+	if err != nil {
+		ln.Close()
+		return err
+	}
+	tracked := &previewListener{Listener: ln, conns: map[net.Conn]bool{}}
+	m.previewMu.Lock()
+	m.previewListeners[port] = tracked
+	m.previewMu.Unlock()
+
+	target := net.JoinHostPort("127.0.0.1", strconv.Itoa(port))
+	proxy := &httputil.ReverseProxy{
+		Rewrite: func(pr *httputil.ProxyRequest) {
+			pr.Out.URL.Scheme = "http"
+			pr.Out.URL.Host = target
+		},
+		FlushInterval: -1,
+	}
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		who, err := lc.WhoIs(r.Context(), r.RemoteAddr)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusForbidden)
+			return
+		}
+		r.Header.Set("Tailscale-User-Login", who.UserProfile.LoginName)
+		r.Header.Set("Tailscale-User-Name", who.UserProfile.DisplayName)
+		proxy.ServeHTTP(w, r)
+	})
+	go func() {
+		err := (&http.Server{Handler: handler}).Serve(tracked)
+		if !errors.Is(err, net.ErrClosed) {
+			slog.Error("serve preview", "port", port, "err", err)
+		}
+	}()
+	return nil
+}
+
+func (m *Manager) stopPreview(port int) {
+	m.previewMu.Lock()
+	ln := m.previewListeners[port]
+	delete(m.previewListeners, port)
+	m.previewMu.Unlock()
+	if ln != nil {
+		ln.closeAll()
+	}
+}
+
+type previewListener struct {
+	net.Listener
+	mu     sync.Mutex
+	conns  map[net.Conn]bool
+	closed bool
+}
+
+func (l *previewListener) Accept() (net.Conn, error) {
+	c, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if l.closed {
+		c.Close()
+		return nil, net.ErrClosed
+	}
+	tc := &previewConn{Conn: c, l: l}
+	l.conns[tc] = true
+	return tc, nil
+}
+
+func (l *previewListener) closeAll() {
+	l.Listener.Close()
+	l.mu.Lock()
+	l.closed = true
+	conns := l.conns
+	l.conns = map[net.Conn]bool{}
+	l.mu.Unlock()
+	for c := range conns {
+		c.Close()
+	}
+}
+
+type previewConn struct {
+	net.Conn
+	l *previewListener
+}
+
+func (c *previewConn) Close() error {
+	c.l.mu.Lock()
+	delete(c.l.conns, c)
+	c.l.mu.Unlock()
+	return c.Conn.Close()
+}
+
+func (m *Manager) recoverPreviews(ctx context.Context) error {
+	previews, err := m.store.ActivePreviews(ctx)
+	if err != nil {
+		return err
+	}
+	for _, p := range previews {
+		err = m.servePreview(p.Port)
+		if err != nil {
+			slog.Error("serve preview", "task", p.TaskId, "port", p.Port, "err", err)
+		}
+	}
+	return nil
 }
 
 func (m *Manager) Unexpose(ctx context.Context, taskID string) {
@@ -108,11 +222,7 @@ func (m *Manager) Unexpose(ctx context.Context, taskID string) {
 		return
 	}
 	for _, p := range previews {
-		err = tailscale.Off(ctx, p.Port)
-		if err != nil {
-			slog.Error("stop serving preview", "task", taskID, "port", p.Port, "err", err)
-			continue
-		}
+		m.stopPreview(p.Port)
 		slog.Info("stopped serving preview", "task", taskID, "name", p.Name, "port", p.Port)
 	}
 }