From d603f592052f3fb798b5d389b95a8e8a9a124f7b Mon Sep 17 00:00:00 2001 From: Justin Erenkrantz Date: Sat, 29 Aug 2026 04:47:23 +0000 Subject: [PATCH] Propagate Lambda invoke request IDs as gRPC metadata Co-authored-by: c1-squire-dev[bot] --- pkg/lambda/grpc/client.go | 43 ++++++++++--- pkg/lambda/grpc/client_test.go | 112 +++++++++++++++++++++++++++++++++ pkg/lambda/grpc/util.go | 3 + 3 files changed, 150 insertions(+), 8 deletions(-) create mode 100644 pkg/lambda/grpc/client_test.go diff --git a/pkg/lambda/grpc/client.go b/pkg/lambda/grpc/client.go index 21e9828f5..7680bca0c 100644 --- a/pkg/lambda/grpc/client.go +++ b/pkg/lambda/grpc/client.go @@ -8,6 +8,7 @@ import ( "time" "github.com/aws/aws-sdk-go-v2/aws" + awsmiddleware "github.com/aws/aws-sdk-go-v2/aws/middleware" "github.com/aws/aws-sdk-go-v2/service/lambda" "github.com/aws/aws-sdk-go-v2/service/lambda/types" "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap" @@ -24,6 +25,8 @@ type lambdaTransport struct { functionName string } +const lambdaInvokeRequestIDMetadataKey = "x-amzn-requestid" + func (l *lambdaTransport) RoundTrip(ctx context.Context, req *Request) (*Response, error) { payload, frameOnly, err := req.marshalPayload() if err != nil { @@ -82,7 +85,17 @@ func (l *lambdaTransport) RoundTrip(ctx context.Context, req *Request) (*Respons return nil, fmt.Errorf("lambda_transport: failed to unmarshal response: %w", err) } - return resp, err + if requestID, ok := awsmiddleware.GetRequestIDMetadata(invokeResp.ResultMetadata); ok && requestID != "" { + headers := resp.Headers() + headers.Set(lambdaInvokeRequestIDMetadataKey, requestID) + respHeaders, err := MarshalMetadata(headers) + if err != nil { + return nil, fmt.Errorf("lambda_transport: failed to encode response metadata: %w", err) + } + resp.msg.SetHeaders(respHeaders) + } + + return resp, nil } // NewLambdaClientTransport returns a new client transport that invokes a lambda function. @@ -136,6 +149,8 @@ func (c *clientConn) Invoke(ctx context.Context, method string, args any, reply return err } + populateResponseMetadata(tresp, opts) + if st.Code() != codes.OK { return st.Err() } @@ -145,21 +160,33 @@ func (c *clientConn) Invoke(ctx context.Context, method string, args any, reply return err } - // TODO(morgabra): call opts here, some are probably important (e.g. PerRPCCredsCallOption, etc) + return nil +} + +func populateResponseMetadata(resp *Response, opts []grpc.CallOption) { + var headers metadata.MD + var trailers metadata.MD + for _, opt := range opts { switch o := opt.(type) { case grpc.HeaderCallOption: - for k, v := range tresp.Headers() { - o.HeaderAddr.Append(k, v...) + if o.HeaderAddr == nil { + continue + } + if headers == nil { + headers = resp.Headers() } + *o.HeaderAddr = headers case grpc.TrailerCallOption: - for k, v := range tresp.Trailers() { - o.TrailerAddr.Append(k, v...) + if o.TrailerAddr == nil { + continue } + if trailers == nil { + trailers = resp.Trailers() + } + *o.TrailerAddr = trailers } } - - return nil } func (c *clientConn) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) { diff --git a/pkg/lambda/grpc/client_test.go b/pkg/lambda/grpc/client_test.go new file mode 100644 index 000000000..ee985fa02 --- /dev/null +++ b/pkg/lambda/grpc/client_test.go @@ -0,0 +1,112 @@ +package grpc + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/lambda" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/types/known/anypb" + "google.golang.org/protobuf/types/known/structpb" + + pbtransport "github.com/conductorone/baton-sdk/pb/c1/transport/v1" +) + +type staticClientTransport struct { + response *Response + err error +} + +func (t *staticClientTransport) RoundTrip(context.Context, *Request) (*Response, error) { + return t.response, t.err +} + +func testResponse(t *testing.T, code codes.Code, headers, trailers metadata.MD) *Response { + t.Helper() + + response, err := anypb.New(&structpb.Struct{}) + require.NoError(t, err) + responseStatus, err := anypb.New(status.New(code, "response status").Proto()) + require.NoError(t, err) + responseHeaders, err := MarshalMetadata(headers) + require.NoError(t, err) + responseTrailers, err := MarshalMetadata(trailers) + require.NoError(t, err) + + return &Response{ + msg: pbtransport.Response_builder{ + Resp: response, + Status: responseStatus, + Headers: responseHeaders, + Trailers: responseTrailers, + }.Build(), + } +} + +func TestLambdaClientConnPropagatesInvokeRequestIDAndResponseMetadata(t *testing.T) { + t.Parallel() + + transportResponse := testResponse(t, codes.OK, metadata.Pairs("x-service-header", "server-value"), metadata.Pairs("x-service-trailer", "trailer-value")) + payload, err := json.Marshal(transportResponse) + require.NoError(t, err) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("X-Amzn-Requestid", "invoke-request-id") + _, _ = w.Write(payload) + })) + defer server.Close() + + lambdaClient := lambda.NewFromConfig(aws.Config{ + Region: "us-east-1", + BaseEndpoint: aws.String(server.URL), + Credentials: credentials.NewStaticCredentialsProvider("access-key", "secret-key", ""), + }) + transport, err := NewLambdaClientTransport(context.Background(), lambdaClient, "test-function") + require.NoError(t, err) + + var headers metadata.MD + var trailers metadata.MD + err = NewClientConn(transport).Invoke( + context.Background(), + "/test.Service/Method", + &structpb.Struct{}, + &structpb.Struct{}, + grpc.Header(&headers), + grpc.Trailer(&trailers), + ) + require.NoError(t, err) + require.Equal(t, []string{"invoke-request-id"}, headers.Get(lambdaInvokeRequestIDMetadataKey)) + require.Equal(t, []string{"server-value"}, headers.Get("x-service-header")) + require.Equal(t, []string{"trailer-value"}, trailers.Get("x-service-trailer")) +} + +func TestClientConnReturnsResponseMetadataWithStatusError(t *testing.T) { + t.Parallel() + + transport := &staticClientTransport{ + response: testResponse(t, codes.PermissionDenied, metadata.Pairs("x-service-header", "server-value"), metadata.Pairs("x-service-trailer", "trailer-value")), + } + var headers metadata.MD + var trailers metadata.MD + + err := NewClientConn(transport).Invoke( + context.Background(), + "/test.Service/Method", + &structpb.Struct{}, + &structpb.Struct{}, + grpc.Header(&headers), + grpc.Trailer(&trailers), + ) + require.Equal(t, codes.PermissionDenied, status.Code(err)) + require.Equal(t, []string{"server-value"}, headers.Get("x-service-header")) + require.Equal(t, []string{"trailer-value"}, trailers.Get("x-service-trailer")) +} diff --git a/pkg/lambda/grpc/util.go b/pkg/lambda/grpc/util.go index ce75a7b58..ea6fa01ba 100644 --- a/pkg/lambda/grpc/util.go +++ b/pkg/lambda/grpc/util.go @@ -153,6 +153,9 @@ func MarshalMetadata(md metadata.MD) (*structpb.Struct, error) { // Only keys with []string values are converted. // Empty string values are ignored. func UnmarshalMetadata(s *structpb.Struct) metadata.MD { + if s == nil { + return metadata.MD{} + } md := make(metadata.MD, len(s.Fields)) for k, v := range s.Fields { lv := v.GetListValue()