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
einride · Jun 27, 2021 · 28da40c · 28da40c
1 parent d807e57
commit 28da40c
Show file tree

Hide file tree

Showing 10 changed files with 199 additions and 52 deletions.
diff --git a/cmd/iamctl/go.sum b/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=

diff --git a/cmd/iamctl/internal/examplecmd/exampleservercmd/server.go b/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

diff --git a/iamcaller/chain.go b/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
 }
diff --git a/iamcaller/chain_test.go b/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{

diff --git a/iamcel/test.go b/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())

diff --git a/iamcel/testall.go b/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())

diff --git a/iamcel/testany.go b/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())

diff --git a/iamexample/caller.go b/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
 }