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") + } +}