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,11 @@ 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.Context = &iamv1.Caller_Context{
+			Deadline: timestamppb.New(deadline),
+		}
+	}
 	token, ok := iamtoken.FromIncomingContext(ctx, authorizationKey)
 	if !ok {
 		return &result, nil

iamcaller/chain.go

@@ -36,6 +36,10 @@ 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.Context == nil && nextCaller.Context != nil {
+			result.Context = nextCaller.Context
+		}
 	}
 	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,23 @@ func TestChainResolvers(t *testing.T) {
 		assert.DeepEqual(t, expected, actual, protocmp.Transform())
 	})
 
+	t.Run("context", func(t *testing.T) {
+		expected := &iamv1.Caller{
+			Context: &iamv1.Caller_Context{
+				Deadline: timestamppb.New(time.Unix(1234, 0).UTC()),
+				Trace:    "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,16 @@ 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.Context = &iamv1.Caller_Context{
+			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,11 +19,22 @@ message Caller {
   repeated string members = 1;
   // Caller identity from gRPC metadata key/value pairs.
   map<string, Metadata> metadata = 2;
+  // Caller context.
+  // TODO: Remove this when cel-go supports async functions with context threading.
+  Context context = 3;
   // Caller identity for a gRPC metadata key/value pair.
   message Metadata {
     // The IAM members from the metadata value.
     repeated string members = 1;
     // The identity token from the metadata value.
     IdentityToken identity_token = 2;
   }
+  // Caller context for downstream network calls.
+  // TODO: Remove this when cel-go supports async functions with context threading.
+  message Context {
+    // Deadline for the caller's request.
+    google.protobuf.Timestamp deadline = 1;
+    // Trace context for the caller's request.
+    string trace = 2;
+  }
 }