Update context proto resolution to use type provider (#1013)

google · Aug 30, 2024 · 6ba2c03 · 6ba2c03
1 parent 79c5dae
commit 6ba2c03
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 10 deletions.
diff --git a/policy/BUILD.bazel b/policy/BUILD.bazel
@@ -42,8 +42,6 @@ go_library(
         "//ext:go_default_library",
         "@in_gopkg_yaml_v3//:go_default_library",
         "@org_golang_google_protobuf//proto:go_default_library",
-        "@org_golang_google_protobuf//reflect/protoreflect:go_default_library",
-        "@org_golang_google_protobuf//reflect/protoregistry:go_default_library",
     ],
 )
 

diff --git a/policy/config.go b/policy/config.go
@@ -20,10 +20,10 @@ import (
 	"math"
 	"strconv"
 
-	"google.golang.org/protobuf/reflect/protoreflect"
-	"google.golang.org/protobuf/reflect/protoregistry"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/google/cel-go/cel"
+	"github.com/google/cel-go/common/types/ref"
 	"github.com/google/cel-go/ext"
 )
 
@@ -130,11 +130,13 @@ func (vd *VariableDecl) AsEnvOption(baseEnv *cel.Env) (cel.EnvOption, error) {
 		if _, found := baseEnv.CELTypeProvider().FindStructType(vd.ContextProto); !found {
 			return nil, fmt.Errorf("could not find context proto type name: %s", vd.ContextProto)
 		}
-		messageType, err := protoregistry.GlobalTypes.FindMessageByName(protoreflect.FullName(vd.ContextProto))
-		if err == protoregistry.NotFound {
-			return nil, fmt.Errorf("could not find context proto type name: %s", vd.ContextProto)
+		// Attempt to instantiate the proto in order to reflect to its descriptor
+		msg := baseEnv.CELTypeProvider().NewValue(vd.ContextProto, map[string]ref.Val{})
+		pbMsg, ok := msg.Value().(proto.Message)
+		if !ok {
+			return nil, fmt.Errorf("type name was not a protobuf: %T", msg.Value())
 		}
-		return cel.DeclareContextProto(messageType.Descriptor()), nil
+		return cel.DeclareContextProto(pbMsg.ProtoReflect().Descriptor()), nil
 	}
 	return nil, errors.New("invalid variable, must set 'name' or 'context_proto' field")
 }

diff --git a/policy/go.mod b/policy/go.mod
@@ -1,9 +1,10 @@
 module github.com/google/cel-go/policy
 
-go 1.20
+go 1.22
 
 require (
 	github.com/google/cel-go v0.21.0
+	google.golang.org/protobuf v1.34.2
 	gopkg.in/yaml.v3 v3.0.1
 )
 
@@ -14,7 +15,6 @@ require (
 	golang.org/x/text v0.17.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240826202546-f6391c0de4c7 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240826202546-f6391c0de4c7 // indirect
-	google.golang.org/protobuf v1.34.2 // indirect
 )
 
 replace github.com/google/cel-go => ../.
diff --git a/policy/go.sum b/policy/go.sum
@@ -4,6 +4,7 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stoewer/go-strcase v1.3.0 h1:g0eASXYtp+yvN9fK8sH94oCIk0fau9uV1/ZdJ0AVEzs=