[interceptor]: Add payload codec interceptor and wire into the proxy

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,

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

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

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

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