From eb3b483e6b24d5f4417cd5eaf74d6b0415628b28 Mon Sep 17 00:00:00 2001 From: "David Muto (pseudomuto)" Date: Thu, 2 Jul 2026 12:53:33 -0400 Subject: [PATCH] [interceptor]: Add payload codec interceptor and wire into the proxy The proxy forwarded WorkflowService traffic upstream with no hook to transform payloads in flight. This adds the machinery to run codec chains (compression, encryption, etc.) over payloads crossing the proxy. The new internal/interceptor package provides `Payloads`, a gRPC unary client interceptor that encodes outbound request payloads through an ordered codec chain and decodes inbound responses in reverse, so one chain undoes its own transformations. --- cmd/proxy/serve.go | 2 + internal/interceptor/doc.go | 7 ++ internal/interceptor/fx.go | 38 ++++++ internal/interceptor/fx_test.go | 63 ++++++++++ internal/interceptor/payload.go | 116 +++++++++++++++++ internal/interceptor/payload_test.go | 181 +++++++++++++++++++++++++++ internal/proxy/fx.go | 29 ++++- internal/proxy/fx_test.go | 140 ++++++++++++++++++++- internal/proxy/server.go | 46 +++++-- internal/proxy/server_test.go | 94 +++++++++++++- 10 files changed, 696 insertions(+), 20 deletions(-) create mode 100644 internal/interceptor/doc.go create mode 100644 internal/interceptor/fx.go create mode 100644 internal/interceptor/fx_test.go create mode 100644 internal/interceptor/payload.go create mode 100644 internal/interceptor/payload_test.go diff --git a/cmd/proxy/serve.go b/cmd/proxy/serve.go index c10a948..260ff8d 100644 --- a/cmd/proxy/serve.go +++ b/cmd/proxy/serve.go @@ -9,6 +9,7 @@ import ( "go.uber.org/fx" "github.com/temporalio/temporal-proxy/internal/config" + "github.com/temporalio/temporal-proxy/internal/interceptor" "github.com/temporalio/temporal-proxy/internal/proxy" "github.com/temporalio/temporal-proxy/internal/router" "github.com/temporalio/temporal-proxy/internal/server" @@ -52,6 +53,7 @@ func serve() *cli.Command { ), config.Module, connect.Module, + interceptor.Module, proxy.Module, router.Module, server.Module, diff --git a/internal/interceptor/doc.go b/internal/interceptor/doc.go new file mode 100644 index 0000000..9882cac --- /dev/null +++ b/internal/interceptor/doc.go @@ -0,0 +1,7 @@ +// Package interceptor provides gRPC client interceptors for the Temporal proxy. +// +// Use Payloads to build a [google.golang.org/grpc.UnaryClientInterceptor] that +// runs a chain of PayloadCodec implementations over every payload passing +// through the proxy: codecs encode in the order given on outbound requests and +// decode in reverse order on inbound responses. +package interceptor diff --git a/internal/interceptor/fx.go b/internal/interceptor/fx.go new file mode 100644 index 0000000..d1269c4 --- /dev/null +++ b/internal/interceptor/fx.go @@ -0,0 +1,38 @@ +package interceptor + +import ( + "fmt" + + "go.uber.org/fx" + "google.golang.org/grpc" + + "github.com/temporalio/temporal-proxy/internal/proxy" +) + +// Module is the fx module that provides the proxy's outbound client +// interceptors. It builds a single payload interceptor from the optional +// [PayloadCodec] chain in [InterceptorParams] and provides it as a +// []grpc.UnaryClientInterceptor for the proxy to chain onto its upstream +// connection. +var Module = fx.Options(fx.Provide( + fx.Annotate(func(p InterceptorParams) ([]grpc.UnaryClientInterceptor, error) { + payloads, err := Payloads(p.Codecs...) + if err != nil { + return nil, fmt.Errorf("failed to construct payloads interceptor: %w", err) + } + + return []grpc.UnaryClientInterceptor{payloads}, nil + }, proxy.UnaryInterceptorsTag), + fx.Annotate(func(_ InterceptorParams) ([]grpc.StreamClientInterceptor, error) { + return nil, nil + }, proxy.StreamInterceptorsTag), +)) + +// InterceptorParams collects the fx-provided dependencies used to build the +// proxy's interceptors. Codecs is optional; when none are supplied the payload +// interceptor is a passthrough that leaves payloads unchanged. +type InterceptorParams struct { + fx.In + + Codecs []PayloadCodec `optional:"true"` +} diff --git a/internal/interceptor/fx_test.go b/internal/interceptor/fx_test.go new file mode 100644 index 0000000..9906c22 --- /dev/null +++ b/internal/interceptor/fx_test.go @@ -0,0 +1,63 @@ +package interceptor_test + +import ( + "testing" + + "github.com/stretchr/testify/require" + workflowservice "go.temporal.io/api/workflowservice/v1" + "go.uber.org/fx" + "google.golang.org/grpc" + + "github.com/temporalio/temporal-proxy/internal/interceptor" +) + +// unaryGroup collects the interceptors Module contributes. The group name +// mirrors proxy.UnaryInterceptorsTag, where the proxy consumes them. +type unaryGroup struct { + fx.In + + Interceptors []grpc.UnaryClientInterceptor `group:"proxy_unary_interceptors"` +} + +func TestModule(t *testing.T) { + t.Parallel() + + t.Run("contributes an interceptor that applies the supplied codecs", func(t *testing.T) { + t.Parallel() + + var got unaryGroup + app := fx.New( + interceptor.Module, + fx.Supply([]interceptor.PayloadCodec{tagCodec{tag: 0x01}}), + fx.Populate(&got), + fx.NopLogger, + ) + require.NoError(t, app.Err()) + require.Len(t, got.Interceptors, 1) + + // Running the contributed interceptor over a request encodes its payload, + // proving the supplied codec was wired through to Payloads. + req := &workflowservice.StartWorkflowExecutionRequest{Input: payload('h', 'i')} + _, err := invoke(t, got.Interceptors[0], req, &workflowservice.QueryWorkflowResponse{}) + require.NoError(t, err) + require.Equal(t, []byte{'h', 'i', 0x01}, req.GetInput().GetPayloads()[0].GetData()) + }) + + t.Run("contributes a passthrough interceptor when no codecs are supplied", func(t *testing.T) { + t.Parallel() + + var got unaryGroup + app := fx.New( + interceptor.Module, + fx.Populate(&got), + fx.NopLogger, + ) + require.NoError(t, app.Err()) + require.Len(t, got.Interceptors, 1) + + req := &workflowservice.StartWorkflowExecutionRequest{Input: payload('h', 'i')} + _, err := invoke(t, got.Interceptors[0], req, &workflowservice.QueryWorkflowResponse{}) + require.NoError(t, err) + require.Equal(t, []byte{'h', 'i'}, req.GetInput().GetPayloads()[0].GetData()) + }) +} diff --git a/internal/interceptor/payload.go b/internal/interceptor/payload.go new file mode 100644 index 0000000..bfcd6e6 --- /dev/null +++ b/internal/interceptor/payload.go @@ -0,0 +1,116 @@ +package interceptor + +import ( + "context" + "fmt" + "runtime" + + "go.temporal.io/api/common/v1" + "go.temporal.io/api/proxy" + "google.golang.org/grpc" +) + +type ( + // PayloadCodec transforms a single Temporal payload. Implementations + // typically rewrite the payload's bytes (for example compressing or + // encrypting them) in Encode and reverse that transformation in Decode. + // Encode and Decode must be inverses: Decode(Encode(p)) should yield the + // original payload. + // + // The interceptor visits payloads concurrently, so implementations must be + // safe for use by multiple goroutines. + PayloadCodec interface { + Encode(context.Context, *common.Payload) (*common.Payload, error) + Decode(context.Context, *common.Payload) (*common.Payload, error) + } + + // payloadCodecChain applies an ordered series of codecs to payloads. + payloadCodecChain []PayloadCodec +) + +// Payloads returns a gRPC unary client interceptor that runs the given codecs +// over every payload in outbound requests and inbound responses. +// +// Outbound payloads are encoded by applying the codecs in the order given. +// Inbound payloads are decoded by applying the codecs in reverse order, so the +// same chain used on the way out undoes its own transformations on the way back. +// +// Search attributes are intentionally skipped so they remain readable/indexable. +// +// These codecs transform the Temporal payloads carried inside WorkflowService +// messages, not the gRPC frames themselves. Unlike gRPC transport compression +// (grpc.UseCompressor), which a peer decodes on arrival, these transformations +// persist: the upstream server stores the encoded bytes. That is what enables +// at-rest compression and encryption, and lets an ordered chain express +// transforms (e.g. compress-then-encrypt) that fixed transport options cannot. +func Payloads(codecs ...PayloadCodec) (grpc.UnaryClientInterceptor, error) { + if len(codecs) == 0 { + return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error { + return invoker(ctx, method, req, reply, cc, opts...) + }, nil + } + + cores := runtime.NumCPU() + chain := payloadCodecChain(codecs) + + pv, err := proxy.NewPayloadVisitorInterceptor(proxy.PayloadVisitorInterceptorOptions{ + Outbound: &proxy.VisitPayloadsOptions{ + ConcurrencyLimit: cores, + SkipSearchAttributes: true, + Visitor: func(ctx *proxy.VisitPayloadsContext, p []*common.Payload) ([]*common.Payload, error) { + return chain.encode(ctx.Context, p) + }, + }, + Inbound: &proxy.VisitPayloadsOptions{ + ConcurrencyLimit: cores, + SkipSearchAttributes: true, + Visitor: func(ctx *proxy.VisitPayloadsContext, p []*common.Payload) ([]*common.Payload, error) { + return chain.decode(ctx.Context, p) + }, + }, + }) + if err != nil { + return nil, fmt.Errorf("failed to construct PayloadVisitorInterceptor: %w", err) + } + + return pv, nil +} + +// encode applies each codec's Encode to every payload in forward order. +func (codecs payloadCodecChain) encode(ctx context.Context, p []*common.Payload) ([]*common.Payload, error) { + var err error + + res := make([]*common.Payload, len(p)) + for i := range p { + res[i] = p[i] + + for j := range codecs { + res[i], err = codecs[j].Encode(ctx, res[i]) + if err != nil { + return nil, fmt.Errorf("failed to encode payload: %w", err) + } + } + } + + return res, nil +} + +// decode applies each codec's Decode to every payload in reverse order, +// undoing the transformations applied by encode. +func (codecs payloadCodecChain) decode(ctx context.Context, p []*common.Payload) ([]*common.Payload, error) { + var err error + + res := make([]*common.Payload, len(p)) + for i := range p { + res[i] = p[i] + + for j := len(codecs) - 1; j >= 0; j-- { + res[i], err = codecs[j].Decode(ctx, res[i]) + if err != nil { + return nil, fmt.Errorf("failed to decode payload: %w", err) + } + } + } + + return res, nil +} diff --git a/internal/interceptor/payload_test.go b/internal/interceptor/payload_test.go new file mode 100644 index 0000000..4a8b589 --- /dev/null +++ b/internal/interceptor/payload_test.go @@ -0,0 +1,181 @@ +package interceptor_test + +import ( + "context" + "errors" + "slices" + "testing" + + "github.com/stretchr/testify/require" + common "go.temporal.io/api/common/v1" + workflowservice "go.temporal.io/api/workflowservice/v1" + "google.golang.org/grpc" + + "github.com/temporalio/temporal-proxy/internal/interceptor" +) + +// tagCodec is a test PayloadCodec that appends a single tag byte on Encode and +// strips it on Decode. Decode fails if the trailing byte is not its own tag, +// which surfaces incorrect ordering (decode must run in reverse of encode). +type tagCodec struct { + tag byte + encErr error + decErr error +} + +func TestPayloadsOutbound(t *testing.T) { + t.Parallel() + + sentinel := errors.New("boom") + + tests := []struct { + name string + codecs []interceptor.PayloadCodec + wantErr error + wantCalled bool + wantData []byte + }{ + { + name: "encodes in order", + codecs: []interceptor.PayloadCodec{tagCodec{tag: 0x01}, tagCodec{tag: 0x02}}, + wantCalled: true, + wantData: []byte{'h', 'i', 0x01, 0x02}, + }, + { + name: "encode error skips invoker", + codecs: []interceptor.PayloadCodec{tagCodec{tag: 0x01}, tagCodec{tag: 0x02, encErr: sentinel}}, + wantErr: sentinel, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + in, err := interceptor.Payloads(tt.codecs...) + require.NoError(t, err) + + req := &workflowservice.StartWorkflowExecutionRequest{Input: payload('h', 'i')} + + called, err := invoke(t, in, req, &workflowservice.QueryWorkflowResponse{}) + require.Equal(t, tt.wantCalled, called) + + if tt.wantErr != nil { + require.ErrorIs(t, err, tt.wantErr) + return + } + + require.NoError(t, err) + require.Equal(t, tt.wantData, req.GetInput().GetPayloads()[0].GetData()) + }) + } +} + +func TestPayloadsInbound(t *testing.T) { + t.Parallel() + + sentinel := errors.New("boom") + + tests := []struct { + name string + codecs []interceptor.PayloadCodec + respData []byte + wantErr error + wantData []byte + }{ + { + name: "decodes in reverse", + codecs: []interceptor.PayloadCodec{tagCodec{tag: 0x01}, tagCodec{tag: 0x02}}, + respData: []byte{'h', 'i', 0x01, 0x02}, + wantData: []byte{'h', 'i'}, + }, + { + name: "decode error propagates", + codecs: []interceptor.PayloadCodec{tagCodec{tag: 0x01, decErr: sentinel}}, + respData: []byte{'h', 'i', 0x01}, + wantErr: sentinel, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + in, err := interceptor.Payloads(tt.codecs...) + require.NoError(t, err) + + resp := &workflowservice.QueryWorkflowResponse{QueryResult: payload(tt.respData...)} + + // Inbound decoding runs after the RPC, so the invoker is always reached. + called, err := invoke(t, in, &workflowservice.StartWorkflowExecutionRequest{}, resp) + require.True(t, called) + + if tt.wantErr != nil { + require.ErrorIs(t, err, tt.wantErr) + return + } + + require.NoError(t, err) + require.Equal(t, tt.wantData, resp.GetQueryResult().GetPayloads()[0].GetData()) + }) + } +} + +func TestPayloadsNoCodecs(t *testing.T) { + t.Parallel() + + // With no codecs, Payloads short-circuits to a plain pass-through: it invokes + // the next handler and leaves request and response payloads untouched, without + // constructing a payload visitor. + in, err := interceptor.Payloads() + require.NoError(t, err) + + req := &workflowservice.StartWorkflowExecutionRequest{Input: payload('h', 'i')} + resp := &workflowservice.QueryWorkflowResponse{QueryResult: payload('y', 'o')} + + called, err := invoke(t, in, req, resp) + require.NoError(t, err) + require.True(t, called) + require.Equal(t, []byte{'h', 'i'}, req.GetInput().GetPayloads()[0].GetData()) + require.Equal(t, []byte{'y', 'o'}, resp.GetQueryResult().GetPayloads()[0].GetData()) +} + +func (c tagCodec) Encode(_ context.Context, p *common.Payload) (*common.Payload, error) { + if c.encErr != nil { + return nil, c.encErr + } + + return &common.Payload{Data: append(slices.Clone(p.Data), c.tag)}, nil +} + +func (c tagCodec) Decode(_ context.Context, p *common.Payload) (*common.Payload, error) { + if c.decErr != nil { + return nil, c.decErr + } + + data := p.Data + if len(data) == 0 || data[len(data)-1] != c.tag { + return nil, errors.New("tagCodec: unexpected trailing byte") + } + + return &common.Payload{Data: slices.Clone(data[:len(data)-1])}, nil +} + +// payload builds a single-payload Payloads message from the given bytes. +func payload(data ...byte) *common.Payloads { + return &common.Payloads{Payloads: []*common.Payload{{Data: data}}} +} + +// invoke runs interceptor over req/resp with a no-op invoker, reporting whether +// the invoker was reached. +func invoke(t *testing.T, in grpc.UnaryClientInterceptor, req, resp any) (called bool, err error) { + t.Helper() + + invoker := func(context.Context, string, any, any, *grpc.ClientConn, ...grpc.CallOption) error { + called = true + return nil + } + + err = in(t.Context(), "/temporal.api.workflowservice.v1.WorkflowService/Method", req, resp, nil, invoker) + return called, err +} diff --git a/internal/proxy/fx.go b/internal/proxy/fx.go index 5460e9c..26775e5 100644 --- a/internal/proxy/fx.go +++ b/internal/proxy/fx.go @@ -5,12 +5,27 @@ import ( "fmt" "go.uber.org/fx" + "google.golang.org/grpc" "github.com/temporalio/temporal-proxy/internal/config" "github.com/temporalio/temporal-proxy/internal/transport/creds" "github.com/temporalio/temporal-proxy/pkg/logger" ) +var ( + // UnaryInterceptorsTag annotates a provider that returns a + // []grpc.UnaryClientInterceptor so its elements are contributed into the + // group [ProxyParams] consumes. flatten spreads the slice into individual + // group members. The group name must match the ProxyParams.UnaryInterceptors + // struct tag. + UnaryInterceptorsTag = fx.ResultTags(`group:"proxy_unary_interceptors,flatten"`) + + // StreamInterceptorsTag is the stream-interceptor counterpart to + // [UnaryInterceptorsTag]; see it for details. The group name must match the + // ProxyParams.StreamInterceptors struct tag. + StreamInterceptorsTag = fx.ResultTags(`group:"proxy_stream_interceptors,flatten"`) +) + // Module is the fx module that constructs the proxy [Server] from [ProxyParams] // and binds its lifecycle to the application. var Module = fx.Options(fx.Invoke(func(p ProxyParams) error { @@ -18,7 +33,11 @@ var Module = fx.Options(fx.Invoke(func(p ProxyParams) error { return fmt.Errorf("invalid upstream configuration: %w", err) } - opts := []Option{WithCredentials(p.creds())} + opts := []Option{ + WithCredentials(p.creds()), + WithUnaryInterceptor(p.UnaryInterceptors...), + WithStreamInterceptor(p.StreamInterceptors...), + } if p.Logger != nil { opts = append(opts, WithLogger(p.Logger)) } @@ -49,7 +68,9 @@ 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 and Config are required; Logger is optional and -// falls back to the default used by [New] when not supplied. +// falls back to the default used by [New] when not supplied. UnaryInterceptors +// and StreamInterceptors are optional and, when provided, are chained onto the +// outbound connection to the upstream frontend. type ProxyParams struct { fx.In Lifecycle fx.Lifecycle @@ -60,7 +81,9 @@ type ProxyParams struct { Config *config.Config // Optional values - Logger logger.Logger `optional:"true"` + Logger logger.Logger `optional:"true"` + UnaryInterceptors []grpc.UnaryClientInterceptor `group:"proxy_unary_interceptors"` + StreamInterceptors []grpc.StreamClientInterceptor `group:"proxy_stream_interceptors"` } // creds derives the credentials used to dial the upstream frontend from the diff --git a/internal/proxy/fx_test.go b/internal/proxy/fx_test.go index b01d605..f0b09bb 100644 --- a/internal/proxy/fx_test.go +++ b/internal/proxy/fx_test.go @@ -2,15 +2,21 @@ package proxy_test import ( "context" + "net" "testing" "time" "github.com/stretchr/testify/require" + common "go.temporal.io/api/common/v1" + "go.temporal.io/api/workflowservice/v1" "go.uber.org/fx" "google.golang.org/grpc" + "google.golang.org/grpc/codes" "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" "github.com/temporalio/temporal-proxy/internal/config" + "github.com/temporalio/temporal-proxy/internal/interceptor" "github.com/temporalio/temporal-proxy/internal/proxy" "github.com/temporalio/temporal-proxy/pkg/logger" "github.com/temporalio/temporal-proxy/pkg/validation" @@ -22,7 +28,7 @@ func TestModule(t *testing.T) { t.Run("wires defaults and runs the lifecycle", func(t *testing.T) { t.Parallel() - const upstream = "127.0.0.1:47233" + upstream := freeUpstream(t) app := newProxyApp(t, &config.Config{ Upstream: config.Upstream{Listen: config.ListenConfig{HostPort: upstream}}, @@ -35,7 +41,7 @@ func TestModule(t *testing.T) { t.Run("uses the supplied logger", func(t *testing.T) { t.Parallel() - const upstream = "127.0.0.1:57233" + upstream := freeUpstream(t) log := logger.NewTestLogger() app := newProxyApp( @@ -50,6 +56,102 @@ func TestModule(t *testing.T) { require.True(t, log.Contains("Starting the server"), "expected the injected logger to be used") }) + t.Run("chains provided unary interceptors onto the upstream connection", func(t *testing.T) { + t.Parallel() + + upstream := freeUpstream(t) + + fired := make(chan string, 4) + in := grpc.UnaryClientInterceptor(func(_ context.Context, method string, _, _ any, _ *grpc.ClientConn, _ grpc.UnaryInvoker, _ ...grpc.CallOption) error { + fired <- method + return status.Error(codes.Unavailable, "short-circuit") + }) + + app := newProxyApp( + t, + &config.Config{Upstream: config.Upstream{Listen: config.ListenConfig{HostPort: upstream}}}, + fx.Provide(fx.Annotate( + func() []grpc.UnaryClientInterceptor { return []grpc.UnaryClientInterceptor{in} }, + proxy.UnaryInterceptorsTag, + )), + ) + require.NoError(t, app.Err()) + + startCtx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + require.NoError(t, app.Start(startCtx)) + + conn := dialUnix(t, upstream) + defer func() { _ = conn.Close() }() + + // The interceptor short-circuits the forwarded call, so it errors, but only + // after fx has passed it through the value group and it has run on the + // upstream connection. It fires synchronously before the call returns. + _, err := workflowservice.NewWorkflowServiceClient(conn).GetSystemInfo( + startCtx, + &workflowservice.GetSystemInfoRequest{}, + grpc.WaitForReady(true), + ) + require.Error(t, err) + + stopCtx, stopCancel := context.WithTimeout(t.Context(), 5*time.Second) + defer stopCancel() + require.NoError(t, app.Stop(stopCtx)) + + // Non-blocking: if the interceptor was not wired through, fail loudly here + // rather than block forever on an empty channel. + select { + case method := <-fired: + require.Equal(t, "/temporal.api.workflowservice.v1.WorkflowService/GetSystemInfo", method) + default: + t.Fatal("expected the supplied interceptor to run on the forwarded call") + } + }) + + t.Run("runs the payload interceptor contributed by interceptor.Module", func(t *testing.T) { + t.Parallel() + + upstream := freeUpstream(t) + + encoded := make(chan struct{}, 1) + app := newProxyApp( + t, + &config.Config{Upstream: config.Upstream{Listen: config.ListenConfig{HostPort: upstream}}}, + interceptor.Module, + fx.Supply([]interceptor.PayloadCodec{recordingCodec{encoded: encoded}}), + ) + require.NoError(t, app.Err()) + + startCtx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + require.NoError(t, app.Start(startCtx)) + + conn := dialUnix(t, upstream) + defer func() { _ = conn.Close() }() + + // StartWorkflowExecution carries an input payload. Forwarding it upstream + // runs the payload interceptor's outbound codec chain before the (down) + // upstream is dialed, so the call errors but the codec has already fired. + _, err := workflowservice.NewWorkflowServiceClient(conn).StartWorkflowExecution( + startCtx, + &workflowservice.StartWorkflowExecutionRequest{ + Input: &common.Payloads{Payloads: []*common.Payload{{Data: []byte("hi")}}}, + }, + grpc.WaitForReady(true), + ) + require.Error(t, err) + + stopCtx, stopCancel := context.WithTimeout(t.Context(), 5*time.Second) + defer stopCancel() + require.NoError(t, app.Stop(stopCtx)) + + select { + case <-encoded: + default: + t.Fatal("expected the interceptor.Module codec to encode the forwarded payload") + } + }) + t.Run("rejects invalid upstream configuration before construction", func(t *testing.T) { t.Parallel() @@ -66,6 +168,40 @@ func TestModule(t *testing.T) { }) } +// recordingCodec is an interceptor.PayloadCodec that signals encoded when Encode +// runs and otherwise leaves payloads unchanged. The non-blocking send keeps it +// safe for the concurrent visits the interceptor performs. +type recordingCodec struct { + encoded chan<- struct{} +} + +func (c recordingCodec) Encode(_ context.Context, p *common.Payload) (*common.Payload, error) { + select { + case c.encoded <- struct{}{}: + default: + } + + return p, nil +} + +func (c recordingCodec) Decode(_ context.Context, p *common.Payload) (*common.Payload, error) { + return p, nil +} + +// freeUpstream returns a loopback host:port backed by a free ephemeral port. +// The proxy never binds or connects to it in these tests; a fresh port simply +// gives each test a unique derived socket path so they run in parallel without +// colliding. +func freeUpstream(t *testing.T) string { + t.Helper() + + lis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + require.NoError(t, lis.Close()) + + return lis.Addr().String() +} + func newProxyApp(t *testing.T, cfg *config.Config, opts ...fx.Option) *fx.App { t.Helper() diff --git a/internal/proxy/server.go b/internal/proxy/server.go index 9c45655..0c88587 100644 --- a/internal/proxy/server.go +++ b/internal/proxy/server.go @@ -35,14 +35,16 @@ type ( path string // path to unix socket } - // Options configures a [Server] at construction time. - Options struct { - creds Credentials - logger logger.Logger + // options configures a [Server] at construction time. + options struct { + creds Credentials + logger logger.Logger + unaryInterceptors []grpc.UnaryClientInterceptor + streamInterceptors []grpc.StreamClientInterceptor } // Option configures a [Server] via [New]. - Option func(*Options) + Option func(*options) ) // New constructs a [Server] that forwards WorkflowService traffic to the @@ -50,7 +52,7 @@ type ( // is derived from hostPort (see [Server.Start]). With no options it dials the // upstream with insecure credentials and logs via a CLI logger. func New(hostPort string, opts ...Option) (*Server, error) { - pops := &Options{ + pops := &options{ creds: creds.NewInsecure(), logger: logger.Default(), } @@ -63,7 +65,17 @@ func New(hostPort string, opts ...Option) (*Server, error) { return nil, fmt.Errorf("failed to generate outbound credentials: %w", err) } - conn, err := grpc.NewClient(hostPort, upstreamCreds) + dialOpts := make([]grpc.DialOption, 0, 3) + dialOpts = append(dialOpts, upstreamCreds) + if len(pops.unaryInterceptors) > 0 { + dialOpts = append(dialOpts, grpc.WithChainUnaryInterceptor(pops.unaryInterceptors...)) + } + + if len(pops.streamInterceptors) > 0 { + dialOpts = append(dialOpts, grpc.WithChainStreamInterceptor(pops.streamInterceptors...)) + } + + conn, err := grpc.NewClient(hostPort, dialOpts...) if err != nil { return nil, fmt.Errorf("failed to dial: %s, %w", hostPort, err) } @@ -103,12 +115,28 @@ func New(hostPort string, opts ...Option) (*Server, error) { // WithCredentials sets the transport credentials used to dial the upstream // frontend. func WithCredentials(creds Credentials) Option { - return Option(func(o *Options) { o.creds = creds }) + return Option(func(o *options) { o.creds = creds }) } // WithLogger sets the logger used by the proxy. func WithLogger(log logger.Logger) Option { - return Option(func(o *Options) { o.logger = log }) + return Option(func(o *options) { o.logger = log }) +} + +// WithUnaryInterceptor adds unary client interceptors to the outbound +// connection to the upstream frontend. Interceptors are chained in the order +// supplied, accumulating across calls, and run on every unary RPC the proxy +// forwards upstream. +func WithUnaryInterceptor(in ...grpc.UnaryClientInterceptor) Option { + return Option(func(o *options) { o.unaryInterceptors = append(o.unaryInterceptors, in...) }) +} + +// WithStreamInterceptor adds stream client interceptors to the outbound +// connection to the upstream frontend. Interceptors are chained in the order +// supplied, accumulating across calls, and run on every streaming RPC the proxy +// forwards upstream. +func WithStreamInterceptor(in ...grpc.StreamClientInterceptor) Option { + return Option(func(o *options) { o.streamInterceptors = append(o.streamInterceptors, in...) }) } // Start binds the local unix socket and serves until the proxy is stopped or diff --git a/internal/proxy/server_test.go b/internal/proxy/server_test.go index 10f5ae7..7b3044b 100644 --- a/internal/proxy/server_test.go +++ b/internal/proxy/server_test.go @@ -9,9 +9,12 @@ import ( "time" "github.com/stretchr/testify/require" + "go.temporal.io/api/workflowservice/v1" "google.golang.org/grpc" + "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" "github.com/temporalio/temporal-proxy/internal/proxy" "github.com/temporalio/temporal-proxy/internal/transport/socket" @@ -28,7 +31,7 @@ func TestNew(t *testing.T) { t.Run("returns a server with default options", func(t *testing.T) { t.Parallel() - svr, err := proxy.New("127.0.0.1:7233") + svr, err := proxy.New(freeUpstream(t)) require.NoError(t, err) require.NotNil(t, svr) }) @@ -37,7 +40,7 @@ func TestNew(t *testing.T) { t.Parallel() svr, err := proxy.New( - "127.0.0.1:7233", + freeUpstream(t), proxy.WithCredentials(failingCredentials{err: errors.New("boom")}), ) require.Error(t, err) @@ -50,10 +53,10 @@ func TestNew(t *testing.T) { func TestServerStartAndStop(t *testing.T) { t.Parallel() - // A unique upstream host gives this test its own socket path so it can run in + // A free ephemeral port gives this test its own socket path so it can run in // parallel with the others. The upstream is never dialed: the health service // the proxy serves locally answers the Check below. - const upstream = "127.0.0.1:17233" + upstream := freeUpstream(t) log := logger.NewTestLogger() svr, err := proxy.New(upstream, proxy.WithLogger(log)) @@ -92,7 +95,7 @@ func TestServerStartAndStop(t *testing.T) { func TestStartRemovesStaleSocket(t *testing.T) { t.Parallel() - const upstream = "127.0.0.1:27233" + upstream := freeUpstream(t) path, err := socket.UnixPath(upstream) require.NoError(t, err) @@ -128,7 +131,7 @@ func TestStartRemovesStaleSocket(t *testing.T) { func TestStartReturnsErrorWhenStaleSocketCannotBeRemoved(t *testing.T) { t.Parallel() - const upstream = "127.0.0.1:37233" + upstream := freeUpstream(t) path, err := socket.UnixPath(upstream) require.NoError(t, err) @@ -147,6 +150,70 @@ func TestStartReturnsErrorWhenStaleSocketCannotBeRemoved(t *testing.T) { require.ErrorContains(t, err, "failed to remove stale socket") } +func TestWithUnaryInterceptor(t *testing.T) { + t.Parallel() + + upstream := freeUpstream(t) + + // first records and forwards; second records and short-circuits with an error + // so the never-dialed upstream is never actually contacted. Supplying them via + // two separate options proves the interceptors accumulate and chain in order. + fired := make(chan string, 8) + first := func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error { + fired <- "first:" + method + return invoker(ctx, method, req, reply, cc, opts...) + } + second := func(_ context.Context, _ string, _, _ any, _ *grpc.ClientConn, _ grpc.UnaryInvoker, _ ...grpc.CallOption) error { + fired <- "second" + return status.Error(codes.Unavailable, "short-circuit") + } + + svr, err := proxy.New(upstream, proxy.WithUnaryInterceptor(first), proxy.WithUnaryInterceptor(second)) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + errCh := make(chan error, 1) + go func() { errCh <- svr.Start(ctx) }() + + conn := dialUnix(t, upstream) + defer func() { _ = conn.Close() }() + + // The forwarded call reaches the upstream client interceptors and is + // short-circuited before any upstream dial, so it returns an error. + _, callErr := workflowservice.NewWorkflowServiceClient(conn).GetSystemInfo( + t.Context(), + &workflowservice.GetSystemInfoRequest{}, + grpc.WaitForReady(true), + ) + require.Error(t, callErr) + + require.NoError(t, svr.Stop(t.Context())) + require.NoError(t, <-errCh) + + // Both interceptors fire synchronously before the call returns, so read + // without blocking: a regression that stops them firing fails loudly here + // instead of hanging on an empty channel. + require.Equal(t, "first:/temporal.api.workflowservice.v1.WorkflowService/GetSystemInfo", recvOr(t, fired)) + require.Equal(t, "second", recvOr(t, fired)) +} + +func TestWithStreamInterceptor(t *testing.T) { + t.Parallel() + + // WorkflowService exposes no streaming RPCs, so a stream interceptor cannot be + // exercised end-to-end through the proxy. This confirms the option is accepted + // and chained onto the dial options without error at construction. + streamIn := func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) { + return streamer(ctx, desc, cc, method, opts...) + } + + svr, err := proxy.New(freeUpstream(t), proxy.WithStreamInterceptor(streamIn)) + require.NoError(t, err) + require.NotNil(t, svr) +} + func (f failingCredentials) DialOption() (grpc.DialOption, error) { return nil, f.err } @@ -166,3 +233,18 @@ func dialUnix(t *testing.T, upstream string) *grpc.ClientConn { require.NoError(t, err) return conn } + +// recvOr returns a value already buffered on ch, or fails the test if none is +// present. Use it for values a synchronous call is expected to have produced by +// the time it returns, so a missing value fails loudly instead of blocking. +func recvOr[T any](t *testing.T, ch <-chan T) T { + t.Helper() + + select { + case v := <-ch: + return v + default: + t.Fatal("expected a value on the channel, but none was ready") + panic("unreachable") + } +}