Machine ID: Support overriding Proxy address from ProxyPing when usin…

…g TLS routing (#48373)

* Start working on ability to use explicit proxy address

* Convert more call-sites to use new helper

* Convert final callsite

* Simplify implementation

* Fix teests

* Remove additional check

* Fix dialling auth client

* Update env var name

lib/tbot/service_application_tunnel.go

@@ -133,7 +133,10 @@ func (s *ApplicationTunnelService) buildLocalProxyConfig(ctx context.Context) (l
 	if err != nil {
 		return alpnproxy.LocalProxyConfig{}, trace.Wrap(err, "pinging proxy")
 	}
-	proxyAddr := proxyPing.Proxy.SSH.PublicAddr
+	proxyAddr, err := proxyPing.proxyWebAddr()
+	if err != nil {
+		return alpnproxy.LocalProxyConfig{}, trace.Wrap(err, "determining proxy web addr")
+	}
 
 	s.log.DebugContext(ctx, "Issuing initial certificate for local proxy.")
 	appCert, app, err := s.issueCert(ctx, roles)

lib/tbot/service_database_tunnel.go

@@ -93,7 +93,10 @@ func (s *DatabaseTunnelService) buildLocalProxyConfig(ctx context.Context) (lpCf
 	if err != nil {
 		return alpnproxy.LocalProxyConfig{}, trace.Wrap(err, "pinging proxy")
 	}
-	proxyAddr := proxyPing.Proxy.SSH.PublicAddr
+	proxyAddr, err := proxyPing.proxyWebAddr()
+	if err != nil {
+		return alpnproxy.LocalProxyConfig{}, trace.Wrap(err, "determining proxy web address")
+	}
 
 	// Fetch information about the database and then issue the initial
 	// certificate. We issue the initial certificate to allow us to fail faster.

lib/tbot/service_identity_output.go

@@ -28,7 +28,6 @@ import (
 	"github.com/gravitational/trace"
 
 	"github.com/gravitational/teleport/api/client/proto"
-	"github.com/gravitational/teleport/api/client/webclient"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/lib/auth/authclient"
 	"github.com/gravitational/teleport/lib/config/openssh"
@@ -239,7 +238,7 @@ type alpnTester interface {
 func renderSSHConfig(
 	ctx context.Context,
 	log *slog.Logger,
-	proxyPing *webclient.PingResponse,
+	proxyPing *proxyPingResponse,
 	clusterNames []string,
 	dest bot.Destination,
 	certAuthGetter certAuthGetter,
@@ -253,11 +252,16 @@ func renderSSHConfig(
 	)
 	defer span.End()
 
-	proxyHost, proxyPort, err := utils.SplitHostPort(proxyPing.Proxy.SSH.PublicAddr)
+	proxyAddr, err := proxyPing.proxyWebAddr()
+	if err != nil {
+		return trace.Wrap(err, "determining proxy web addr")
+	}
+
+	proxyHost, proxyPort, err := utils.SplitHostPort(proxyAddr)
 	if err != nil {
 		return trace.BadParameter(
 			"proxy %+v has no usable public address: %v",
-			proxyPing.Proxy.SSH.PublicAddr, err,
+			proxyAddr, err,
 		)
 	}
 
@@ -310,7 +314,7 @@ func renderSSHConfig(
 	connUpgradeRequired := false
 	if proxyPing.Proxy.TLSRoutingEnabled {
 		connUpgradeRequired, err = alpnTester.isUpgradeRequired(
-			ctx, proxyPing.Proxy.SSH.PublicAddr, botCfg.Insecure,
+			ctx, proxyAddr, botCfg.Insecure,
 		)
 		if err != nil {
 			return trace.Wrap(err, "determining if ALPN upgrade is required")

lib/tbot/service_identity_output_test.go

@@ -157,12 +157,14 @@ func Test_renderSSHConfig(t *testing.T) {
 			err := renderSSHConfig(
 				context.Background(),
 				utils.NewSlogLoggerForTests(),
-				&webclient.PingResponse{
-					ClusterName: mockClusterName,
-					Proxy: webclient.ProxySettings{
-						TLSRoutingEnabled: tc.TLSRouting,
-						SSH: webclient.SSHProxySettings{
-							PublicAddr: mockProxyAddr,
+				&proxyPingResponse{
+					PingResponse: &webclient.PingResponse{
+						ClusterName: mockClusterName,
+						Proxy: webclient.ProxySettings{
+							TLSRoutingEnabled: tc.TLSRouting,
+							SSH: webclient.SSHProxySettings{
+								PublicAddr: mockProxyAddr,
+							},
 						},
 					},
 				},

lib/tbot/service_kubernetes_output.go

@@ -32,7 +32,6 @@ import (
 
 	apiclient "github.com/gravitational/teleport/api/client"
 	"github.com/gravitational/teleport/api/client/proto"
-	"github.com/gravitational/teleport/api/client/webclient"
 	"github.com/gravitational/teleport/api/defaults"
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/lib/auth/authclient"
@@ -418,7 +417,7 @@ func getKubeCluster(ctx context.Context, clt *authclient.Client, name string) (t
 
 // selectKubeConnectionMethod determines the address and SNI that should be
 // put into the kubeconfig file.
-func selectKubeConnectionMethod(proxyPong *webclient.PingResponse) (
+func selectKubeConnectionMethod(proxyPong *proxyPingResponse) (
 	clusterAddr string, sni string, err error,
 ) {
 	// First we check for TLS routing. If this is enabled, we use the Proxy's
@@ -427,8 +426,11 @@ func selectKubeConnectionMethod(proxyPong *webclient.PingResponse) (
 	// Even if KubePublicAddr is specified, we still use the general
 	// PublicAddr when using TLS routing.
 	if proxyPong.Proxy.TLSRoutingEnabled {
-		addr := proxyPong.Proxy.SSH.PublicAddr
-		host, _, err := net.SplitHostPort(proxyPong.Proxy.SSH.PublicAddr)
+		addr, err := proxyPong.proxyWebAddr()
+		if err != nil {
+			return "", "", trace.Wrap(err)
+		}
+		host, _, err := net.SplitHostPort(addr)
 		if err != nil {
 			return "", "", trace.Wrap(err, "parsing proxy public_addr")
 		}

lib/tbot/service_kubernetes_output_test.go

@@ -279,7 +279,9 @@ func Test_selectKubeConnectionMethod(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			addr, sni, err := selectKubeConnectionMethod(tt.proxyPing)
+			addr, sni, err := selectKubeConnectionMethod(&proxyPingResponse{
+				PingResponse: tt.proxyPing,
+			})
 			require.NoError(t, err)
 			require.Equal(t, tt.wantAddr, addr)
 			require.Equal(t, tt.wantSNI, sni)

lib/tbot/service_ssh_multiplexer.go

@@ -273,11 +273,15 @@ func (s *SSHMultiplexerService) setup(ctx context.Context) (
 	if err != nil {
 		return nil, nil, "", nil, trace.Wrap(err)
 	}
-	proxyAddr := proxyPing.Proxy.SSH.PublicAddr
-	proxyHost, _, err = net.SplitHostPort(proxyPing.Proxy.SSH.PublicAddr)
+	proxyAddr, err := proxyPing.proxyWebAddr()
+	if err != nil {
+		return nil, nil, "", nil, trace.Wrap(err, "determining proxy web addr")
+	}
+	proxyHost, _, err = net.SplitHostPort(proxyAddr)
 	if err != nil {
 		return nil, nil, "", nil, trace.Wrap(err)
 	}
+
 	connUpgradeRequired := false
 	if proxyPing.Proxy.TLSRoutingEnabled {
 		connUpgradeRequired, err = s.alpnUpgradeCache.isUpgradeRequired(

lib/tbot/tbot.go

@@ -153,17 +153,32 @@ func (b *Bot) Run(ctx context.Context) (err error) {
 		return trace.Wrap(err)
 	}
 
-	addr, _ := b.cfg.Address()
-	resolver, err := reversetunnelclient.CachingResolver(
-		ctx,
-		reversetunnelclient.WebClientResolver(&webclient.Config{
-			Context:   ctx,
-			ProxyAddr: addr,
-			Insecure:  b.cfg.Insecure,
-		}),
-		nil /* clock */)
-	if err != nil {
-		return trace.Wrap(err)
+	addr, addrKind := b.cfg.Address()
+	var resolver reversetunnelclient.Resolver
+	if shouldUseProxyAddr() {
+		if addrKind != config.AddressKindProxy {
+			return trace.BadParameter("TBOT_USE_PROXY_ADDR requires that a proxy address is set using --proxy-server or proxy_server")
+		}
+		// If the user has indicated they want tbot to prefer using the proxy
+		// address they have configured, we use a static resolver set to this
+		// address. We also assume that they have TLS routing/multiplexing
+		// enabled, since otherwise we'd need them to manually configure an
+		// an entry for each kind of address.
+		resolver = reversetunnelclient.StaticResolver(
+			addr, types.ProxyListenerMode_Multiplex,
+		)
+	} else {
+		resolver, err = reversetunnelclient.CachingResolver(
+			ctx,
+			reversetunnelclient.WebClientResolver(&webclient.Config{
+				Context:   ctx,
+				ProxyAddr: addr,
+				Insecure:  b.cfg.Insecure,
+			}),
+			nil /* clock */)
+		if err != nil {
+			return trace.Wrap(err)
+		}
 	}
 
 	// Create an error group to manage all the services lifetimes.
@@ -527,7 +542,8 @@ func (b *Bot) preRunChecks(ctx context.Context) (_ func() error, err error) {
 	ctx, span := tracer.Start(ctx, "Bot/preRunChecks")
 	defer func() { apitracing.EndSpan(span, err) }()
 
-	switch _, addrKind := b.cfg.Address(); addrKind {
+	_, addrKind := b.cfg.Address()
+	switch addrKind {
 	case config.AddressKindUnspecified:
 		return nil, trace.BadParameter(
 			"either a proxy or auth address must be set using --proxy, --auth-server or configuration",
@@ -702,10 +718,10 @@ type proxyPingCache struct {
 	log           *slog.Logger
 
 	mu          sync.RWMutex
-	cachedValue *webclient.PingResponse
+	cachedValue *proxyPingResponse
 }
 
-func (p *proxyPingCache) ping(ctx context.Context) (*webclient.PingResponse, error) {
+func (p *proxyPingCache) ping(ctx context.Context) (*proxyPingResponse, error) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 	if p.cachedValue != nil {
@@ -739,11 +755,50 @@ func (p *proxyPingCache) ping(ctx context.Context) (*webclient.PingResponse, err
 		return nil, trace.Wrap(err)
 	}
 	p.log.DebugContext(ctx, "Successfully pinged proxy.", "pong", res)
-	p.cachedValue = res
+	p.cachedValue = &proxyPingResponse{
+		PingResponse:        res,
+		configuredProxyAddr: p.botCfg.ProxyServer,
+	}
 
 	return p.cachedValue, nil
 }
 
+type proxyPingResponse struct {
+	*webclient.PingResponse
+	configuredProxyAddr string
+}
+
+// useProxyAddrEnv is an environment variable which can be set to
+// force `tbot` to prefer using the proxy address explicitly provided by the
+// user over the one fetched from the proxy ping. This is only intended to work
+// in cases where TLS routing is enabled, and is intended to support cases where
+// the Proxy is accessible from multiple addresses, and the one included in the
+// ProxyPing is incorrect.
+const useProxyAddrEnv = "TBOT_USE_PROXY_ADDR"
+
+// shouldUseProxyAddr returns true if the TBOT_USE_PROXY_ADDR environment
+// variable is set to "1". More generally, this indicates that the user wishes
+// for tbot to prefer using the proxy address that has been explicitly provided
+// by the user rather than the one fetched via a discovery process (e.g ping).
+func shouldUseProxyAddr() bool {
+	return os.Getenv(useProxyAddrEnv) == "1"
+}
+
+// proxyWebAddr returns the address to use to connect to the proxy web port.
+// In TLS routing mode, this address should be used for most/all connections.
+// This function takes into account the TBOT_USE_PROXY_ADDR environment
+// variable, which can be used to force the use of the proxy address explicitly
+// provided by the user rather than use the one fetched from the proxy ping.
+func (p *proxyPingResponse) proxyWebAddr() (string, error) {
+	if shouldUseProxyAddr() {
+		if p.configuredProxyAddr == "" {
+			return "", trace.BadParameter("TBOT_USE_PROXY_ADDR set but no explicit proxy address configured")
+		}
+		return p.configuredProxyAddr, nil
+	}
+	return p.Proxy.SSH.PublicAddr, nil
+}
+
 type alpnProxyConnUpgradeRequiredCache struct {
 	botCfg *config.BotConfig
 	log    *slog.Logger