Propagate Lambda invoke request IDs as gRPC metadata

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) {

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

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