[v11] Add pprof diagnostics endpoints to tbot (#26106)

lib/tbot/config/config.go

@@ -149,6 +149,10 @@ type CLIConf struct {
 	// RemainingArgs is the remaining string arguments for commands that
 	// require them.
 	RemainingArgs []string
+
+	// DiagAddr is the address the diagnostics http service should listen on.
+	// If not set, no diagnostics listener is created.
+	DiagAddr string
 }
 
 // OnboardingConfig contains values only required on first connect.
@@ -214,6 +218,9 @@ type BotConfig struct {
 	CertificateTTL  time.Duration `yaml:"certificate_ttl"`
 	RenewalInterval time.Duration `yaml:"renewal_interval"`
 	Oneshot         bool          `yaml:"oneshot"`
+	// DiagAddr is the address the diagnostics http service should listen on.
+	// If not set, no diagnostics listener is created.
+	DiagAddr string `yaml:"diag_addr,omitempty"`
 }
 
 func (conf *BotConfig) CheckAndSetDefaults() error {
@@ -425,6 +432,13 @@ func FromCLIConf(cf *CLIConf) (*BotConfig, error) {
 		config.Onboarding.SetToken(cf.Token)
 	}
 
+	if cf.DiagAddr != "" {
+		if config.DiagAddr != "" {
+			log.Warnf("CLI parameters are overriding diagnostics address configured in %s", cf.ConfigPath)
+		}
+		config.DiagAddr = cf.DiagAddr
+	}
+
 	if err := config.CheckAndSetDefaults(); err != nil {
 		return nil, trace.Wrap(err, "validating merged bot config")
 	}

lib/tbot/config/config_test.go

@@ -56,6 +56,8 @@ func TestConfigCLIOnlySample(t *testing.T) {
 		Token:          "foo",
 		CAPins:         []string{"abc123"},
 		AuthServer:     "auth.example.com",
+		DiagAddr:       "127.0.0.1:1337",
+		Debug:          true,
 	}
 	cfg, err := FromCLIConf(&cf)
 	require.NoError(t, err)
@@ -91,6 +93,8 @@ func TestConfigCLIOnlySample(t *testing.T) {
 	require.True(t, ok)
 
 	require.Equal(t, cf.DestinationDir, destImplReal.Path)
+	require.Equal(t, cf.Debug, cfg.Debug)
+	require.Equal(t, cf.DiagAddr, cfg.DiagAddr)
 }
 
 func TestConfigFile(t *testing.T) {
@@ -130,6 +134,9 @@ func TestConfigFile(t *testing.T) {
 	destImplReal, ok := destImpl.(*DestinationDirectory)
 	require.True(t, ok)
 	require.Equal(t, "/tmp/foo", destImplReal.Path)
+
+	require.True(t, cfg.Debug)
+	require.Equal(t, "127.0.0.1:1337", cfg.DiagAddr)
 }
 
 func TestLoadTokenFromFile(t *testing.T) {
@@ -149,6 +156,8 @@ func TestLoadTokenFromFile(t *testing.T) {
 const exampleConfigFile = `
 auth_server: auth.example.com
 renewal_interval: 5m
+debug: true
+diag_addr: 127.0.0.1:1337
 onboarding:
   token: %s
   ca_pins:

lib/tbot/tbot.go

@@ -22,6 +22,8 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"errors"
+	"net/http"
+	"net/http/pprof"
 	"sync"
 	"time"
 
@@ -254,6 +256,38 @@ func (b *Bot) Run(ctx context.Context) error {
 		cancel()
 		return nil
 	})
+	if b.cfg.Debug && b.cfg.DiagAddr != "" {
+		eg.Go(func() error {
+			b.log.WithField("addr", b.cfg.DiagAddr).Info(
+				"DiagAddr configured, diagnostics service will be started.",
+			)
+			mux := http.NewServeMux()
+			mux.HandleFunc("/debug/pprof/", pprof.Index)
+			mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+			mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+			mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+			mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+			mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusNotFound)
+				msg := "404 - Not Found\n\nI'm a little tbot,\nshort and stout,\nthe page you seek,\nis not about."
+				_, _ = w.Write([]byte(msg))
+			}))
+			srv := http.Server{
+				Addr:    b.cfg.DiagAddr,
+				Handler: mux,
+			}
+			go func() {
+				<-egCtx.Done()
+				if err := srv.Close(); err != nil {
+					b.log.WithError(err).Warn("Failed to close HTTP server.")
+				}
+			}()
+			if err := srv.ListenAndServe(); err != http.ErrServerClosed {
+				return err
+			}
+			return nil
+		})
+	}
 
 	return eg.Wait()
 }

tool/tbot/main.go

@@ -83,6 +83,7 @@ func Run(args []string, stdout io.Writer) error {
 	startCmd.Flag("renewal-interval", "Interval at which short-lived certificates are renewed; must be less than the certificate TTL.").DurationVar(&cf.RenewalInterval)
 	startCmd.Flag("join-method", "Method to use to join the cluster. "+joinMethodList).Default(config.DefaultJoinMethod).EnumVar(&cf.JoinMethod, config.SupportedJoinMethods...)
 	startCmd.Flag("oneshot", "If set, quit after the first renewal.").BoolVar(&cf.Oneshot)
+	startCmd.Flag("diag-addr", "If set and the bot is in debug mode, a diagnostics service will listen on specified address.").StringVar(&cf.DiagAddr)
 
 	initCmd := app.Command("init", "Initialize a certificate destination directory for writes from a separate bot user.")
 	initCmd.Flag("destination-dir", "Directory to write short-lived machine certificates to.").StringVar(&cf.DestinationDir)