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{

iamcel/test.go

@@ -2,7 +2,6 @@ package iamcel
 
 import (
 	"context"
-	"time"
 
 	"github.com/google/cel-go/checker/decls"
 	"github.com/google/cel-go/common/types"
@@ -57,8 +56,12 @@ func NewTestFunctionImplementation(
 				return types.NewErr("%s: no permission configured for resource '%s'", codes.PermissionDenied, resource)
 			}
 			// TODO: When cel-go supports async functions, use the caller context here.
-			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-			defer cancel()
+			ctx := context.Background()
+			if caller.GetContext().GetDeadline() != nil {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithDeadline(ctx, caller.GetContext().GetDeadline().AsTime())
+				defer cancel()
+			}
 			if result, err := tester.TestPermissions(ctx, caller, map[string]string{resource: permission}); err != nil {
 				if s, ok := status.FromError(err); ok {
 					return types.NewErr("%s: %s", s.Code(), s.Message())

iamcel/testall.go

@@ -3,7 +3,6 @@ package iamcel
 import (
 	"context"
 	"reflect"
-	"time"
 
 	"github.com/google/cel-go/checker/decls"
 	"github.com/google/cel-go/common/types"
@@ -70,8 +69,12 @@ func NewTestAllFunctionImplementation(
 				resourcePermissions[resource] = permission
 			}
 			// TODO: When cel-go supports async functions, use the caller context here.
-			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-			defer cancel()
+			ctx := context.Background()
+			if caller.GetContext().GetDeadline() != nil {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithDeadline(ctx, caller.GetContext().GetDeadline().AsTime())
+				defer cancel()
+			}
 			if result, err := tester.TestPermissions(ctx, caller, resourcePermissions); err != nil {
 				if s, ok := status.FromError(err); ok {
 					return types.NewErr("%s: %s", s.Code(), s.Message())

iamcel/testany.go

@@ -3,7 +3,6 @@ package iamcel
 import (
 	"context"
 	"reflect"
-	"time"
 
 	"github.com/google/cel-go/checker/decls"
 	"github.com/google/cel-go/common/types"
@@ -70,8 +69,12 @@ func NewTestAnyFunctionImplementation(
 				resourcePermissions[resource] = permission
 			}
 			// TODO: When cel-go supports async functions, use the caller context here.
-			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-			defer cancel()
+			ctx := context.Background()
+			if caller.GetContext().GetDeadline() != nil {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithDeadline(ctx, caller.GetContext().GetDeadline().AsTime())
+				defer cancel()
+			}
 			if result, err := tester.TestPermissions(ctx, caller, resourcePermissions); err != nil {
 				if s, ok := status.FromError(err); ok {
 					return types.NewErr("%s: %s", s.Code(), s.Message())

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
 }