Set Value Case

yonesko · Dec 7, 2023 · f2ca8aa · f2ca8aa
1 parent 3fbfef2
commit f2ca8aa
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 22 deletions.
diff --git a/redact.go b/redact.go
@@ -9,21 +9,25 @@ import (
 	"google.golang.org/protobuf/types/descriptorpb"
 )
 
-var StdRedactor = Redactor{
-	RedactingHandler: func(parent protoreflect.Value, fd protoreflect.FieldDescriptor) error {
+var (
+	clearFunc = func(parent protoreflect.Value, fd protoreflect.FieldDescriptor) error {
 		parent.Message().Clear(fd)
 		return nil
-	},
-}
+	}
+)
 
 type Redactor struct {
-	RedactingHandler func(parent protoreflect.Value, field protoreflect.FieldDescriptor) error
+	SensitiveFieldAnnotation *protoimpl.ExtensionInfo
+	RedactingHandler         func(parent protoreflect.Value, field protoreflect.FieldDescriptor) error
 }
 
-func (r Redactor) Redact(msg proto.Message, sensitiveFieldAnnotation *protoimpl.ExtensionInfo) error {
+func (r Redactor) Redact(msg proto.Message) error {
+	if r.SensitiveFieldAnnotation == nil || r.RedactingHandler == nil {
+		return nil
+	}
 	return protorange.Range(msg.ProtoReflect(), func(p protopath.Values) error {
 		fd := p.Path.Index(-1).FieldDescriptor()
-		if isFieldSensetive(fd, p.Index(-1).Value, sensitiveFieldAnnotation) {
+		if isFieldSensetive(fd, p.Index(-1).Value, r.SensitiveFieldAnnotation) {
 			parent := p.Index(-2)
 			if parent.Value.IsValid() {
 				err := r.RedactingHandler(parent.Value, fd)
@@ -38,7 +42,7 @@ func (r Redactor) Redact(msg proto.Message, sensitiveFieldAnnotation *protoimpl.
 }
 
 func Redact(msg proto.Message, sensitiveFieldAnnotation *protoimpl.ExtensionInfo) error {
-	return StdRedactor.Redact(msg, sensitiveFieldAnnotation)
+	return Redactor{RedactingHandler: clearFunc, SensitiveFieldAnnotation: sensitiveFieldAnnotation}.Redact(msg)
 }
 
 func isFieldSensetive(fieldDescriptor protoreflect.FieldDescriptor, value protoreflect.Value, sensitiveFieldAnnotation *protoimpl.ExtensionInfo) bool {

diff --git a/redact_test.go b/redact_test.go
@@ -11,10 +11,13 @@ import (
 
 func TestRedactProto_Clear(t *testing.T) {
 	t.Parallel()
-	redactor := Redactor{RedactingHandler: func(parent protoreflect.Value, fd protoreflect.FieldDescriptor) error {
-		parent.Message().Clear(fd)
-		return nil
-	}}
+	redactor := Redactor{
+		RedactingHandler: func(parent protoreflect.Value, fd protoreflect.FieldDescriptor) error {
+			parent.Message().Clear(fd)
+			return nil
+		},
+		SensitiveFieldAnnotation: testproto.E_SensitiveData,
+	}
 
 	type args struct {
 		message proto.Message
@@ -130,7 +133,7 @@ func TestRedactProto_Clear(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			err := redactor.Redact(tt.args.message, testproto.E_SensitiveData)
+			err := redactor.Redact(tt.args.message)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("SensetiveFields() error = %v, wantErr %v", err, tt.wantErr)
 				return
@@ -142,14 +145,17 @@ func TestRedactProto_Clear(t *testing.T) {
 }
 func TestRedactProto_SetStringClearOther(t *testing.T) {
 	t.Parallel()
-	redactor := Redactor{RedactingHandler: func(parent protoreflect.Value, field protoreflect.FieldDescriptor) error {
-		if parent.IsValid() && field.Kind() == protoreflect.StringKind {
-			parent.Message().Set(field, protoreflect.ValueOfString("REDACTED"))
-		} else {
-			parent.Message().Clear(field)
-		}
-		return nil
-	}}
+	redactor := Redactor{
+		RedactingHandler: func(parent protoreflect.Value, field protoreflect.FieldDescriptor) error {
+			if parent.IsValid() && field.Kind() == protoreflect.StringKind {
+				parent.Message().Set(field, protoreflect.ValueOfString("REDACTED"))
+			} else {
+				parent.Message().Clear(field)
+			}
+			return nil
+		},
+		SensitiveFieldAnnotation: testproto.E_SensitiveData,
+	}
 	type args struct {
 		message proto.Message
 	}
@@ -226,7 +232,7 @@ func TestRedactProto_SetStringClearOther(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			err := redactor.Redact(tt.args.message, testproto.E_SensitiveData)
+			err := redactor.Redact(tt.args.message)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("SensetiveFields() error = %v, wantErr %v", err, tt.wantErr)
 				return