From 76f0153c963637d560cd659898c2bc616da904c1 Mon Sep 17 00:00:00 2001
From: Dj Gilcrease <devlon.gilcrease@pindrop.com>
Date: Tue, 13 Aug 2024 09:18:47 -0700
Subject: [PATCH] feat: make it so you can optionally have DiscardUnknown not
 apply to enums

---
 encoding/protojson/decode.go      | 10 +++--
 encoding/protojson/decode_test.go | 66 +++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 3 deletions(-)

diff --git a/encoding/protojson/decode.go b/encoding/protojson/decode.go
index 8f9e592f..f184ea8b 100644
--- a/encoding/protojson/decode.go
+++ b/encoding/protojson/decode.go
@@ -41,6 +41,10 @@ type UnmarshalOptions struct {
 	// If DiscardUnknown is set, unknown fields and enum name values are ignored.
 	DiscardUnknown bool
 
+	// If ErrorOnUnknownEnumValue is set, unknown enum values will return an error.
+	// this has no effect if DiscardUnknown is false.
+	ErrorOnUnknownEnumValue bool
+
 	// Resolver is used for looking up types when unmarshaling
 	// google.protobuf.Any messages or extension fields.
 	// If nil, this defaults to using protoregistry.GlobalTypes.
@@ -343,7 +347,7 @@ func (d decoder) unmarshalScalar(fd protoreflect.FieldDescriptor) (protoreflect.
 		}
 
 	case protoreflect.EnumKind:
-		if v, ok := unmarshalEnum(tok, fd, d.opts.DiscardUnknown); ok {
+		if v, ok := unmarshalEnum(tok, fd, d.opts.DiscardUnknown, d.opts.ErrorOnUnknownEnumValue); ok {
 			return v, nil
 		}
 
@@ -488,7 +492,7 @@ func unmarshalBytes(tok json.Token) (protoreflect.Value, bool) {
 	return protoreflect.ValueOfBytes(b), true
 }
 
-func unmarshalEnum(tok json.Token, fd protoreflect.FieldDescriptor, discardUnknown bool) (protoreflect.Value, bool) {
+func unmarshalEnum(tok json.Token, fd protoreflect.FieldDescriptor, discardUnknown, errorOnUnknown bool) (protoreflect.Value, bool) {
 	switch tok.Kind() {
 	case json.String:
 		// Lookup EnumNumber based on name.
@@ -496,7 +500,7 @@ func unmarshalEnum(tok json.Token, fd protoreflect.FieldDescriptor, discardUnkno
 		if enumVal := fd.Enum().Values().ByName(protoreflect.Name(s)); enumVal != nil {
 			return protoreflect.ValueOfEnum(enumVal.Number()), true
 		}
-		if discardUnknown {
+		if discardUnknown && !errorOnUnknown {
 			return protoreflect.Value{}, true
 		}
 
diff --git a/encoding/protojson/decode_test.go b/encoding/protojson/decode_test.go
index faf5eb54..c4b509f8 100644
--- a/encoding/protojson/decode_test.go
+++ b/encoding/protojson/decode_test.go
@@ -2673,6 +2673,72 @@ func TestUnmarshal(t *testing.T) {
 				10: 101,
 			},
 		},
+	}, /* test ErrorOnUnknownEnumValue undoes DiscardUnknown for enums, but not others */ {
+		desc: "DiscardUnknown: Any, ErrorOnUnknownEnumValue: true",
+		umo: protojson.UnmarshalOptions{
+			DiscardUnknown:          true,
+			ErrorOnUnknownEnumValue: true,
+		},
+		inputMessage: &anypb.Any{},
+		inputText: `{
+  "@type": "foo/pb2.Nested",
+  "unknown": "none"
+}`,
+		wantMessage: &anypb.Any{
+			TypeUrl: "foo/pb2.Nested",
+		},
+	}, {
+		desc: "DiscardUnknown: Any with Empty, ErrorOnUnknownEnumValue: true",
+		umo: protojson.UnmarshalOptions{
+			DiscardUnknown:          true,
+			ErrorOnUnknownEnumValue: true,
+		},
+		inputMessage: &anypb.Any{},
+		inputText: `{
+  "@type": "type.googleapis.com/google.protobuf.Empty",
+  "value": {"unknown": 47}
+}`,
+		wantMessage: &anypb.Any{
+			TypeUrl: "type.googleapis.com/google.protobuf.Empty",
+		},
+	}, {
+		desc:         "DiscardUnknown: true, ErrorOnUnknownEnumValue: true - unknown enum name",
+		inputMessage: &pb3.Enums{},
+		inputText: `{
+  "sEnum": "UNNAMED"
+}`,
+		umo: protojson.UnmarshalOptions{
+			DiscardUnknown:          true,
+			ErrorOnUnknownEnumValue: true,
+		},
+		wantErr: `invalid value for enum field sEnum: "UNNAMED"`, // weak_message2 is unknown since the package containing it is not imported
+	}, {
+		desc:         "DiscardUnknown: true, ErrorOnUnknownEnumValue: true - repeated enum unknown name",
+		inputMessage: &pb2.Enums{},
+		inputText: `{
+  "rptEnum"      : ["TEN", 1, 42, "UNNAMED"]
+}`,
+		umo: protojson.UnmarshalOptions{
+			DiscardUnknown:          true,
+			ErrorOnUnknownEnumValue: true,
+		},
+		wantErr: `invalid value for enum field rptEnum: "UNNAMED"`, // weak_message2 is unknown since the package containing it is not imported
+	}, {
+		desc:         "DiscardUnknown: true, ErrorOnUnknownEnumValue: true - enum map value unknown name",
+		inputMessage: &pb3.Maps{},
+		inputText: `{
+  "uint64ToEnum": {
+    "1" : "ONE",
+	"2" : 2,
+	"10": 101,
+	"3": "UNNAMED"
+  }
+}`,
+		umo: protojson.UnmarshalOptions{
+			DiscardUnknown:          true,
+			ErrorOnUnknownEnumValue: true,
+		},
+		wantErr: `invalid value for enum field value: "UNNAMED"`, // weak_message2 is unknown since the package containing it is not imported
 	}, {
 		desc:         "weak fields",
 		inputMessage: &testpb.TestWeak{},