feat(agent): add http debug routes for magicsock

coadler · coadler · commit 3138432ade2d · 2023-04-25T23:18:45.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -60,7 +60,7 @@ type Options struct {
 	ReconnectingPTYTimeout time.Duration
 	EnvironmentVariables   map[string]string
 	Logger                 slog.Logger
-	AgentPorts             map[int]string
+	IgnorePorts            map[int]string
 	SSHMaxTimeout          time.Duration
 	TailnetListenPort      uint16
 }
@@ -76,7 +76,12 @@ type Client interface {
 	PatchStartupLogs(ctx context.Context, req agentsdk.PatchStartupLogs) error
 }
 
-func New(options Options) io.Closer {
+type Agent interface {
+	MagicsockServeHTTPDebug(w http.ResponseWriter, r *http.Request)
+	io.Closer
+}
+
+func New(options Options) Agent {
 	if options.ReconnectingPTYTimeout == 0 {
 		options.ReconnectingPTYTimeout = 5 * time.Minute
 	}
@@ -112,7 +117,7 @@ func New(options Options) io.Closer {
 		tempDir:                options.TempDir,
 		lifecycleUpdate:        make(chan struct{}, 1),
 		lifecycleReported:      make(chan codersdk.WorkspaceAgentLifecycle, 1),
-		ignorePorts:            options.AgentPorts,
+		ignorePorts:            options.IgnorePorts,
 		connStatsChan:          make(chan *agentsdk.Stats, 1),
 		sshMaxTimeout:          options.SSHMaxTimeout,
 	}
@@ -1267,6 +1272,20 @@ func (a *agent) isClosed() bool {
 	}
 }
 
+func (a *agent) MagicsockServeHTTPDebug(w http.ResponseWriter, r *http.Request) {
+	a.closeMutex.Lock()
+	net := a.network
+	a.closeMutex.Unlock()
+
+	if net == nil {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("network is not ready yet"))
+		return
+	}
+
+	net.MagicsockServeHTTPDebug(w, r)
+}
+
 func (a *agent) Close() error {
 	a.closeMutex.Lock()
 	defer a.closeMutex.Unlock()
diff --git a/cli/agent.go b/cli/agent.go
@@ -16,6 +16,7 @@ import (
 	"time"
 
 	"cloud.google.com/go/compute/metadata"
+	"github.com/go-chi/chi/v5"
 	"golang.org/x/xerrors"
 	"gopkg.in/natefinch/lumberjack.v2"
 	"tailscale.com/util/clientmetric"
@@ -37,7 +38,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 		noReap            bool
 		sshMaxTimeout     time.Duration
 		tailnetListenPort int64
-		prometheusAddress string
+		metricsAddress    string
 	)
 	cmd := &clibase.Cmd{
 		Use:   "agent",
@@ -48,7 +49,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 			ctx, cancel := context.WithCancel(inv.Context())
 			defer cancel()
 
-			agentPorts := map[int]string{}
+			ignorePorts := map[int]string{}
 
 			isLinux := runtime.GOOS == "linux"
 
@@ -125,14 +126,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 			defer pprofSrvClose()
 			// Do a best effort here. If this fails, it's not a big deal.
 			if port, err := urlPort(pprofAddress); err == nil {
-				agentPorts[port] = "pprof"
-			}
-
-			prometheusSrvClose := ServeHandler(ctx, logger, prometheusMetricsHandler(), prometheusAddress, "prometheus")
-			defer prometheusSrvClose()
-			// Do a best effort here. If this fails, it's not a big deal.
-			if port, err := urlPort(prometheusAddress); err == nil {
-				agentPorts[port] = "prometheus"
+				ignorePorts[port] = "pprof"
 			}
 
 			// exchangeToken returns a session token.
@@ -196,7 +190,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 				return xerrors.Errorf("add executable to $PATH: %w", err)
 			}
 
-			closer := agent.New(agent.Options{
+			agnt := agent.New(agent.Options{
 				Client:            client,
 				Logger:            logger,
 				LogDir:            logDir,
@@ -215,11 +209,19 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 				EnvironmentVariables: map[string]string{
 					"GIT_ASKPASS": executablePath,
 				},
-				AgentPorts:    agentPorts,
+				IgnorePorts:   ignorePorts,
 				SSHMaxTimeout: sshMaxTimeout,
 			})
+
+			metricsSrvClose := ServeHandler(ctx, logger, metricsHandler(agnt), metricsAddress, "metrics")
+			defer metricsSrvClose()
+			// Do a best effort here. If this fails, it's not a big deal.
+			if port, err := urlPort(metricsAddress); err == nil {
+				ignorePorts[port] = "metrics"
+			}
+
 			<-ctx.Done()
-			return closer.Close()
+			return agnt.Close()
 		},
 	}
 
@@ -270,7 +272,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 			Flag:        "prometheus-address",
 			Default:     "127.0.0.1:2112",
 			Env:         "CODER_AGENT_PROMETHEUS_ADDRESS",
-			Value:       clibase.StringOf(&prometheusAddress),
+			Value:       clibase.StringOf(&metricsAddress),
 			Description: "The bind address to serve Prometheus metrics.",
 		},
 	}
@@ -360,11 +362,17 @@ func urlPort(u string) (int, error) {
 	return -1, xerrors.Errorf("invalid port: %s", u)
 }
 
-func prometheusMetricsHandler() http.Handler {
-	// We don't have any other internal metrics so far, so it's safe to expose metrics this way.
-	// Based on: https://github.com/tailscale/tailscale/blob/280255acae604796a1113861f5a84e6fa2dc6121/ipn/localapi/localapi.go#L489
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+func metricsHandler(a agent.Agent) http.Handler {
+	r := chi.NewRouter()
+
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "text/plain")
 		clientmetric.WritePrometheusExpositionFormat(w)
 	})
+
+	r.Route("/debug", func(r chi.Router) {
+		r.Get("/magicsock", a.MagicsockServeHTTPDebug)
+	})
+
+	return r
 }
diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -828,6 +828,10 @@ func (c *Conn) SetConnStatsCallback(maxPeriod time.Duration, maxConns int, dump
 	c.tunDevice.SetStatistics(connStats)
 }
 
+func (c *Conn) MagicsockServeHTTPDebug(w http.ResponseWriter, r *http.Request) {
+	c.magicConn.ServeHTTPDebug(w, r)
+}
+
 type listenKey struct {
 	network string
 	host    string
