From 9db832f4356c15bbb853133e8504a2152d17b63f Mon Sep 17 00:00:00 2001
From: Maisem Ali <maisem@tailscale.com>
Date: Thu, 8 Aug 2024 15:37:47 -0700
Subject: [PATCH] cmd/{tta,vnet}: proxy to gokrazy UI

Signed-off-by: Maisem Ali <maisem@tailscale.com>
---
 cmd/tta/tta.go        | 44 +++++++++++++++++++++++++++++++------------
 cmd/vnet/vnet-main.go | 14 ++++++++++++++
 2 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/cmd/tta/tta.go b/cmd/tta/tta.go
index 11bfd9ead..9c4bdf212 100644
--- a/cmd/tta/tta.go
+++ b/cmd/tta/tta.go
@@ -11,6 +11,7 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"flag"
 	"io"
@@ -59,10 +60,10 @@ func serveCmd(w http.ResponseWriter, cmd string, args ...string) {
 }
 
 type localClientRoundTripper struct {
-	lc *tailscale.LocalClient
+	lc tailscale.LocalClient
 }
 
-func (rt localClientRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+func (rt *localClientRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	req = req.Clone(req.Context())
 	req.RequestURI = ""
 	return rt.lc.DoLocalRequest(req)
@@ -96,9 +97,31 @@ func main() {
 
 	log.Printf("Tailscale Test Agent running.")
 
-	var mux http.ServeMux
+	gokRP := httputil.NewSingleHostReverseProxy(must.Get(url.Parse("http://gokrazy")))
+	gokRP.Transport = &http.Transport{
+		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+			if network != "tcp" {
+				return nil, errors.New("unexpected network")
+			}
+			if addr != "gokrazy:80" {
+				return nil, errors.New("unexpected addr")
+			}
+			var d net.Dialer
+			return d.DialContext(ctx, "unix", "/run/gokrazy-http.sock")
+		},
+	}
+
+	var ttaMux http.ServeMux // agent mux
+	var serveMux http.ServeMux
+	serveMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("X-TTA-GoKrazy") == "1" {
+			gokRP.ServeHTTP(w, r)
+			return
+		}
+		ttaMux.ServeHTTP(w, r)
+	})
 	var hs http.Server
-	hs.Handler = &mux
+	hs.Handler = &serveMux
 	var (
 		stMu   sync.Mutex
 		newSet = set.Set[net.Conn]{} // conns in StateNew
@@ -121,17 +144,16 @@ func main() {
 		}
 	}
 	conns := make(chan net.Conn, 1)
-	var lc tailscale.LocalClient
-	rp := httputil.NewSingleHostReverseProxy(must.Get(url.Parse("http://local-tailscaled.sock")))
-	rp.Transport = localClientRoundTripper{&lc}
 
-	mux.Handle("/localapi/", rp)
+	lcRP := httputil.NewSingleHostReverseProxy(must.Get(url.Parse("http://local-tailscaled.sock")))
+	lcRP.Transport = new(localClientRoundTripper)
+	ttaMux.Handle("/localapi/", lcRP)
 
-	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+	ttaMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		io.WriteString(w, "TTA\n")
 		return
 	})
-	mux.HandleFunc("/up", func(w http.ResponseWriter, r *http.Request) {
+	ttaMux.HandleFunc("/up", func(w http.ResponseWriter, r *http.Request) {
 		serveCmd(w, "tailscale", "up", "--login-server=http://control.tailscale")
 	})
 	go hs.Serve(chanListener(conns))
@@ -151,8 +173,6 @@ func main() {
 			continue
 		}
 		conns <- c
-
-		time.Sleep(time.Second)
 	}
 }
 
diff --git a/cmd/vnet/vnet-main.go b/cmd/vnet/vnet-main.go
index 5fb83bce4..99eb022a8 100644
--- a/cmd/vnet/vnet-main.go
+++ b/cmd/vnet/vnet-main.go
@@ -10,11 +10,15 @@
 	"flag"
 	"log"
 	"net"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
 	"os"
 	"time"
 
 	"tailscale.com/tstest/natlab/vnet"
 	"tailscale.com/types/logger"
+	"tailscale.com/util/must"
 )
 
 var (
@@ -70,6 +74,16 @@ func main() {
 
 	s.WriteStartingBanner(os.Stdout)
 	nc := s.NodeAgentClient(node1)
+	go func() {
+		rp := httputil.NewSingleHostReverseProxy(must.Get(url.Parse("http://gokrazy")))
+		d := rp.Director
+		rp.Director = func(r *http.Request) {
+			d(r)
+			r.Header.Set("X-TTA-GoKrazy", "1")
+		}
+		rp.Transport = nc.HTTPClient.Transport
+		http.ListenAndServe(":8080", rp)
+	}()
 	go func() {
 		getStatus := func() {
 			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)