[proxy]: Resolve upstreams per request from namespace and metadata

e2e/harness_test.go

@@ -115,17 +115,19 @@ func newFullApp(t *testing.T, cfg *config.Config) *fx.App {
 }
 
 // newProxyApp is a minimal fx app for the socket-level test in this package: it
-// wires proxy.Module plus protoutil.Module (which provides the Translator
-// proxy.Module requires). internal/proxy/fx_test.go has its own copy (with
-// support for extra fx.Options) that its unit tests still depend on; this is a
-// deliberate small duplication rather than an exported helper, so the two
-// packages stay decoupled.
+// wires proxy.Module plus protoutil.Module (which provides the Translator) and
+// connect.Module (which provides the Pool), both required by proxy.Module.
+// internal/proxy/fx_test.go has its own copy (with support for extra
+// fx.Options) that its unit tests still depend on; this is a deliberate small
+// duplication rather than an exported helper, so the two packages stay
+// decoupled.
 func newProxyApp(t *testing.T, cfg *config.Config) *fx.App {
 	t.Helper()
 
 	return fx.New(
 		fx.Supply(fx.Annotate(t.Context(), fx.As(new(context.Context)))),
 		fx.Supply(cfg),
+		connect.Module,
 		protoutil.Module,
 		proxy.Module,
 		fx.NopLogger,
@@ -191,7 +193,7 @@ func dialInbound(t *testing.T, addr string) *grpc.ClientConn {
 }
 
 // dialUnix returns a client connection to the proxy's unix socket for the
-// given upstream host. The socket path matches what proxy.Start binds.
+// given upstream host. The socket path matches what proxy.Listen binds.
 // internal/proxy/server_test.go has its own copy that its unit tests still
 // depend on; this is a deliberate small duplication for the socket-level test
 // in this package.

e2e/templated_upstream_test.go

@@ -0,0 +1,83 @@
+package e2e
+
+import (
+	"context"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.temporal.io/api/workflowservice/v1"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+
+	"github.com/temporalio/temporal-proxy/internal/config"
+)
+
+// TestEndToEndTemplatedUpstreamRoutesByRenderedAddress drives the full stack
+// with a single templated upstream whose hostPort renders from request
+// metadata, and proves each request reaches the upstream its metadata names.
+// This exercises per-request address rendering end to end, and because two
+// distinct rendered targets are both reached, it also proves the connection
+// pool keys by rendered target rather than collapsing onto the first dial.
+func TestEndToEndTemplatedUpstreamRoutesByRenderedAddress(t *testing.T) {
+	t.Parallel()
+
+	svcA, addrA := newFakeUpstream(t)
+	svcB, addrB := newFakeUpstream(t)
+
+	inboundAddr := freeTCPAddr(t)
+	app := newFullApp(t, &config.Config{
+		Listen:  config.ListenConfig{HostPort: inboundAddr},
+		Routing: config.Routing{DefaultUpstream: "dynamic"},
+		Upstreams: []config.Upstream{{
+			Name:   "dynamic",
+			Listen: config.ListenConfig{HostPort: `{{ index .Metadata "x-upstream" }}`},
+		}},
+	})
+	require.NoError(t, app.Err())
+
+	startCtx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
+	defer cancel()
+	require.NoError(t, app.Start(startCtx))
+	t.Cleanup(func() {
+		stopCtx, stopCancel := context.WithTimeout(t.Context(), 5*time.Second)
+		defer stopCancel()
+		_ = app.Stop(stopCtx)
+	})
+
+	conn := dialInbound(t, inboundAddr)
+	defer func() { _ = conn.Close() }()
+	client := workflowservice.NewWorkflowServiceClient(conn)
+
+	call := func(target string) {
+		ctx := metadata.AppendToOutgoingContext(startCtx, "x-upstream", target)
+		_, err := client.GetSystemInfo(ctx, &workflowservice.GetSystemInfoRequest{}, grpc.WaitForReady(true))
+		require.NoError(t, err)
+	}
+
+	call(addrA)
+	call(addrB)
+
+	require.NotNil(t, svcA.received(), "the request naming upstream A must reach A")
+	require.NotNil(t, svcB.received(), "the request naming upstream B must reach B")
+	require.Equal(t, []string{addrA}, svcA.received().Get("x-upstream"))
+	require.Equal(t, []string{addrB}, svcB.received().Get("x-upstream"))
+}
+
+// newFakeUpstream stands up a plaintext fake WorkflowService frontend and
+// returns it with its dial address.
+func newFakeUpstream(t *testing.T) (*capturingWorkflowService, string) {
+	t.Helper()
+
+	svc := &capturingWorkflowService{}
+	srv := grpc.NewServer()
+	workflowservice.RegisterWorkflowServiceServer(srv, svc)
+
+	lis, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	go func() { _ = srv.Serve(lis) }()
+	t.Cleanup(srv.Stop)
+
+	return svc, lis.Addr().String()
+}

internal/config/config_test.go

@@ -359,6 +359,18 @@ func TestUpstream_IsTemplated(t *testing.T) {
 
 	require.True(t, (&config.Upstream{Listen: config.ListenConfig{HostPort: "{{ .LocalNamespace }}.acme.cloud:7233"}}).IsTemplated())
 	require.False(t, (&config.Upstream{Listen: config.ListenConfig{HostPort: "127.0.0.1:7233"}}).IsTemplated())
+
+	// A templated TLS server name makes the upstream templated even when the
+	// hostPort is static.
+	require.True(t, (&config.Upstream{Listen: config.ListenConfig{
+		HostPort: "127.0.0.1:7233",
+		TLS:      &config.TLSConfig{ServerName: "{{ .RemoteNamespace }}.acme.cloud"},
+	}}).IsTemplated())
+
+	require.False(t, (&config.Upstream{Listen: config.ListenConfig{
+		HostPort: "127.0.0.1:7233",
+		TLS:      &config.TLSConfig{ServerName: "static.acme.cloud"},
+	}}).IsTemplated())
 }
 
 func (e *errReader) Read(_ []byte) (int, error) { return 0, e.err }

internal/config/upstream.go

@@ -79,10 +79,15 @@ func (u *Upstream) Validate() error {
 	)
 }
 
-// IsTemplated reports whether the upstream's hostPort contains a text/template
-// action and so must be resolved per request.
+// IsTemplated reports whether the upstream must be resolved per request because
+// its hostPort, or its TLS server name when one is configured, contains a
+// text/template action.
 func (u *Upstream) IsTemplated() bool {
-	return isTemplated(u.Listen.HostPort)
+	if isTemplated(u.Listen.HostPort) {
+		return true
+	}
+
+	return u.Listen.TLS != nil && isTemplated(u.Listen.TLS.ServerName)
 }
 
 // Validate checks the namespace translation rules.

internal/proxy/fx.go

@@ -3,12 +3,15 @@ package proxy
 import (
 	"context"
 	"fmt"
+	"slices"
 
 	"go.uber.org/fx"
+	"google.golang.org/grpc"
 
 	"github.com/temporalio/temporal-proxy/internal/auth"
 	"github.com/temporalio/temporal-proxy/internal/config"
 	"github.com/temporalio/temporal-proxy/internal/protoutil"
+	"github.com/temporalio/temporal-proxy/internal/transport/connect"
 	"github.com/temporalio/temporal-proxy/internal/transport/creds"
 	"github.com/temporalio/temporal-proxy/pkg/logger"
 )
@@ -17,48 +20,61 @@ import (
 // and binds its lifecycle to the application.
 var Module = fx.Options(fx.Invoke(func(p ProxyParams) error {
 	for i := range p.Config.Upstreams {
-		upstream := &p.Config.Upstreams[i]
-
-		if err := upstream.Validate(); err != nil {
+		up := &p.Config.Upstreams[i]
+		if err := up.Validate(); err != nil {
 			return fmt.Errorf("invalid upstream configuration: %w", err)
 		}
 
-		if upstream.IsTemplated() {
-			return fmt.Errorf(
-				"upstream %q has a templated hostPort %q: templated upstreams are not yet supported",
-				upstream.Name,
-				upstream.Listen.HostPort,
-			)
-		}
-
-		opts := []Option{WithCredentials(upstreamCreds(upstream))}
-
-		rules := &upstream.Namespaces.Rules
+		// Request-independent dial options: namespace translation and outbound
+		// credentials. Per-request credentials are added by the resolver.
+		var dialOpts []grpc.DialOption
+		rules := &up.Namespaces.Rules
 		if rules.Configured() {
-			opts = append(opts, WithDialOptions(translationDialOptions(p.Translator, rules.Remote, rules.Local)...))
+			dialOpts = append(dialOpts, translationDialOptions(p.Translator, rules.Remote, rules.Local)...)
 		}
 
-		cp, err := auth.CredentialProviderFor(upstream.Credentials)
+		cp, err := auth.CredentialProviderFor(up.Credentials)
 		if err != nil {
-			return fmt.Errorf("invalid credentials for upstream %q: %w", upstream.Name, err)
+			return fmt.Errorf("invalid credentials for upstream %q: %w", up.Name, err)
 		}
 		if cp != nil {
-			opts = append(opts, WithDialOptions(auth.DialOptions(cp)...))
+			dialOpts = append(dialOpts, auth.DialOptions(cp)...)
+		}
+
+		res, err := upstreamResolver(up, dialOpts)
+		if err != nil {
+			return err
 		}
 
+		conn, err := connect.NewConn(p.Pool.ConnOrCreate, res)
+		if err != nil {
+			return err
+		}
+
+		var opts []Option
 		if p.Logger != nil {
 			opts = append(opts, WithLogger(p.Logger))
 		}
 
-		svr, err := New(upstream.Listen.HostPort, opts...)
+		svr, err := New(up.Listen.HostPort, conn, opts...)
 		if err != nil {
-			return fmt.Errorf("failed to create proxy for upstream %q: %w", upstream.Name, err)
+			return fmt.Errorf("failed to create proxy for upstream %q: %w", up.Name, err)
 		}
 
 		p.Lifecycle.Append(fx.Hook{
 			OnStart: func(context.Context) error {
+				// Bind synchronously so the socket is listening before the
+				// inbound server (whose OnStart runs after this one) starts
+				// routing requests to it; then serve in the background.
+				lis, err := svr.Listen(p.Context)
+				if err != nil {
+					return fmt.Errorf("failed to start proxy for upstream %q: %w", up.Name, err)
+				}
+
 				go func() {
-					if err := svr.Start(p.Context); err != nil {
+					defer func() { _ = lis.Close() }()
+
+					if err := svr.Start(p.Context, lis); err != nil {
 						// The proxy stopped serving unexpectedly. Bring the app
 						// down rather than linger in a non-serving state; Start
 						// has already logged the cause.
@@ -76,9 +92,10 @@ var Module = fx.Options(fx.Invoke(func(p ProxyParams) error {
 }))
 
 // ProxyParams collects the fx-provided dependencies needed to construct and run
-// the proxy [Server]. Context, Config, and Translator are required; Logger is
-// optional and falls back to the default used by [New] when not supplied.
-// [protoutil.Module] provides the Translator in the assembled application.
+// the proxy [Server]. Context, Config, Translator, and Pool are required;
+// Logger is optional and falls back to the default used by [New] when not
+// supplied. [protoutil.Module] provides the Translator and [connect.Module]
+// provides the Pool in the assembled application.
 type ProxyParams struct {
 	fx.In
 	Lifecycle  fx.Lifecycle
@@ -88,28 +105,83 @@ type ProxyParams struct {
 	Context    context.Context
 	Config     *config.Config
 	Translator *protoutil.Translator
+	Pool       *connect.Pool
 
 	// Optional values
 	Logger logger.Logger `optional:"true"`
 }
 
+// upstreamResolver builds the [connect.Resolver] for an upstream. When neither
+// the hostPort nor the TLS server name is templated it returns a static
+// resolver whose connection is constructed eagerly (and reused for every
+// request); otherwise it returns a DynamicResolver that renders the target and
+// server name, and rebuilds credentials, per request. opts holds the
+// request-independent dial options (namespace translation and outbound
+// credentials).
+func upstreamResolver(upstream *config.Upstream, opts []grpc.DialOption) (connect.Resolver, error) {
+	if upstream.IsTemplated() {
+		translator := func(s string) string { return s }
+		if upstream.Namespaces.Rules.Configured() {
+			translator = upstream.Namespaces.Rules.Remote
+		}
+
+		// Share one loader across the per-request credentials so templated mTLS
+		// routing reads and parses the certificate and CA files once rather than
+		// on every request. Only ServerName varies per request. nil for non-mTLS
+		// upstreams, which have no files to load.
+		var loader *creds.CertLoader
+		if tls := upstream.Listen.TLS; tls != nil && tls.CA != "" {
+			loader = creds.NewCertLoader(tls.CA, tls.Cert, tls.Key)
+		}
+
+		return NewDynamicResolver(
+			upstream,
+			WithRemoteNamespacer(translator),
+			WithOptionsFactory(func(data RouteData) ([]grpc.DialOption, error) {
+				creds, err := upstreamCreds(upstream, data.ResolvedServerName, loader).DialOption()
+				if err != nil {
+					return nil, err
+				}
+
+				return append(slices.Clone(opts), creds), nil
+			}),
+		)
+	}
+
+	serverName := ""
+	if upstream.Listen.TLS != nil {
+		serverName = upstream.Listen.TLS.ServerName
+	}
+
+	creds, err := upstreamCreds(upstream, serverName, nil).DialOption()
+	if err != nil {
+		return nil, fmt.Errorf("failed to build credentials for upstream %q: %w", upstream.Name, err)
+	}
+
+	return connect.StaticResolver(upstream.Listen.HostPort, append(slices.Clone(opts), creds)...), nil
+}
+
 // upstreamCreds derives the credentials used to dial the upstream frontend from
 // the upstream TLS configuration: mutual TLS when a CA is set, server-verified
 // client TLS when TLS is configured without one, and insecure otherwise.
-func upstreamCreds(upstream *config.Upstream) Credentials {
+// serverName overrides the SNI/certificate-verification name; it is the
+// upstream's static configured ServerName, or a per-request value rendered
+// from a templated ServerName. loader, when non-nil, supplies the mTLS
+// certificate material so it is loaded once and reused across the per-request
+// credentials of a templated upstream; pass nil for a fixed-address upstream,
+// whose credentials are built once.
+func upstreamCreds(upstream *config.Upstream, serverName string, loader *creds.CertLoader) Credentials {
 	tls := upstream.Listen.TLS
 	if tls == nil {
 		return creds.NewInsecure()
 	}
 
 	if tls.CA != "" {
-		return creds.NewMTLS(
-			tls.CA,
-			tls.Cert,
-			tls.Key,
-			creds.MTLSOptions{ServerName: tls.ServerName},
-		)
+		return creds.NewMTLS(tls.CA, tls.Cert, tls.Key, creds.MTLSOptions{
+			ServerName: serverName,
+			Loader:     loader,
+		})
 	}
 
-	return creds.NewClientTLS()
+	return creds.NewClientTLS(serverName)
 }