feat: add deadline and trace context to iamv1.Caller

As a workaround for CEL-Go currently not supporting threading the user
context through expression and function evaluation, this patch adds
support for the caller to thread request deadline and trace context
through to the CEL functions, so that downstream gRPC calls can inherit
the request deadline and trace context.

This might all be removable if/when CEL attains built-in support for
"async functions", as per google/cel-go#368

Author: odsod

cmd/iamctl/go.sum

@@ -38,7 +38,6 @@ cloud.google.com/go/pubsub v1.1.0/go.mod h1:EwwdRX2sKPjnvnqCa270oGRyludottCI76h+
 cloud.google.com/go/pubsub v1.2.0/go.mod h1:jhfEVHT8odbXTkndysNHCcx0awwzvfOlguIAii9o8iA=
 cloud.google.com/go/pubsub v1.3.1/go.mod h1:i+ucay31+CNRpDW4Lu78I4xXG+O1r/MAHgjpRVR+TSU=
 cloud.google.com/go/spanner v1.17.0/go.mod h1:+17t2ixFwRG4lWRwE+5kipDR9Ef07Jkmc8z0IbMDKUs=
-cloud.google.com/go/spanner v1.20.0/go.mod h1:ajR/W06cMHQu7nqQ4irRGplPNoWgejGJlEhlB8xBTKk=
 cloud.google.com/go/spanner v1.21.0 h1:NWLJnTTPwKu5OB/3SwL/VkJ9rIpvNPjalWz0p6vywnk=
 cloud.google.com/go/spanner v1.21.0/go.mod h1:P1Pl0zyIIdhovaFueBrOjSQ6jKQDfl5bVemE+gdEJog=
 cloud.google.com/go/storage v1.0.0/go.mod h1:IhtSnM/ZTZV8YYJWCY8RULGVqBDmpoyjwiyrjsg+URw=

cmd/iamctl/internal/examplecmd/exampleservercmd/server.go

@@ -23,6 +23,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 func newServer(spannerClient *spanner.Client) (*iamexample.Authorization, error) {
@@ -126,6 +127,9 @@ type googleIdentityTokenCallerResolver struct{}
 func (googleIdentityTokenCallerResolver) ResolveCaller(ctx context.Context) (*iamv1.Caller, error) {
 	const authorizationKey = "authorization"
 	var result iamv1.Caller
+	if deadline, ok := ctx.Deadline(); ok {
+		result.Deadline = timestamppb.New(deadline)
+	}
 	token, ok := iamtoken.FromIncomingContext(ctx, authorizationKey)
 	if !ok {
 		return &result, nil

iamcaller/chain.go

@@ -36,6 +36,14 @@ func (c chainResolver) ResolveCaller(ctx context.Context) (*iamv1.Caller, error)
 		for key, value := range nextCaller.Metadata {
 			Add(&result, key, value)
 		}
+		// TODO: Remove this when CEL-Go supports async functions with context arguments.
+		if result.Deadline == nil && nextCaller.Deadline != nil {
+			result.Deadline = nextCaller.Deadline
+		}
+		// TODO: Remove this when CEL-Go supports async functions with context arguments.
+		if result.TraceContext == "" && nextCaller.TraceContext != "" {
+			result.TraceContext = nextCaller.TraceContext
+		}
 	}
 	return &result, nil
 }

iamcaller/chain_test.go

@@ -4,9 +4,11 @@ import (
 	"context"
 	"errors"
 	"testing"
+	"time"
 
 	iamv1 "go.einride.tech/iam/proto/gen/einride/iam/v1"
 	"google.golang.org/protobuf/testing/protocmp"
+	"google.golang.org/protobuf/types/known/timestamppb"
 	"gotest.tools/v3/assert"
 )
 
@@ -83,6 +85,34 @@ func TestChainResolvers(t *testing.T) {
 		assert.DeepEqual(t, expected, actual, protocmp.Transform())
 	})
 
+	t.Run("deadline", func(t *testing.T) {
+		expected := &iamv1.Caller{
+			Deadline: timestamppb.New(time.Unix(1234, 0).UTC()),
+			Members:  []string{"test:bar", "test:foo"},
+			Metadata: map[string]*iamv1.Caller_Metadata{
+				"key1": {Members: []string{"test:foo"}},
+				"key2": {Members: []string{"test:bar"}},
+			},
+		}
+		actual, err := ChainResolvers(constant(expected)).ResolveCaller(context.Background())
+		assert.NilError(t, err)
+		assert.DeepEqual(t, expected, actual, protocmp.Transform())
+	})
+
+	t.Run("trace-context", func(t *testing.T) {
+		expected := &iamv1.Caller{
+			TraceContext: "mock-trace-context",
+			Members:      []string{"test:bar", "test:foo"},
+			Metadata: map[string]*iamv1.Caller_Metadata{
+				"key1": {Members: []string{"test:foo"}},
+				"key2": {Members: []string{"test:bar"}},
+			},
+		}
+		actual, err := ChainResolvers(constant(expected)).ResolveCaller(context.Background())
+		assert.NilError(t, err)
+		assert.DeepEqual(t, expected, actual, protocmp.Transform())
+	})
+
 	t.Run("error", func(t *testing.T) {
 		actual, err := ChainResolvers(
 			constant(&iamv1.Caller{

iamexample/caller.go

@@ -6,6 +6,7 @@ import (
 	"go.einride.tech/iam/iamcaller"
 	iamv1 "go.einride.tech/iam/proto/gen/einride/iam/v1"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 // MemberHeader is the gRPC header used by the example server to determine IAM members of the caller.
@@ -23,13 +24,14 @@ type memberHeaderResolver struct{}
 // ResolveCaller implements iamcaller.Resolver.
 func (m *memberHeaderResolver) ResolveCaller(ctx context.Context) (*iamv1.Caller, error) {
 	var result iamv1.Caller
-	md, ok := metadata.FromIncomingContext(ctx)
-	if !ok {
-		return &result, nil
+	if md, ok := metadata.FromIncomingContext(ctx); ok {
+		iamcaller.Add(&result, MemberHeader, &iamv1.Caller_Metadata{
+			Members: md.Get(MemberHeader),
+		})
+	}
+	if deadline, ok := ctx.Deadline(); ok {
+		result.Deadline = timestamppb.New(deadline)
 	}
-	iamcaller.Add(&result, MemberHeader, &iamv1.Caller_Metadata{
-		Members: md.Get(MemberHeader),
-	})
 	return &result, nil
 }
 

proto/gen/einride/iam/v1/caller.pb.go

@@ -19,6 +19,12 @@ message Caller {
   repeated string members = 1;
   // Caller identity from gRPC metadata key/value pairs.
   map<string, Metadata> metadata = 2;
+  // Deadline for the caller's request.
+  // TODO: Remove this when cel-go supports async functions with context threading.
+  google.protobuf.Timestamp deadline = 3;
+  // Trace context for the caller's request.
+  // TODO: Remove this when cel-go supports async functions with context threading.
+  string trace_context = 4;
   // Caller identity for a gRPC metadata key/value pair.
   message Metadata {
     // The IAM members from the metadata value.