From 899cabfde231246ee680a8473b74e40ea5aaaa18 Mon Sep 17 00:00:00 2001
From: Moises Vega <moises@uber.com>
Date: Fri, 15 Mar 2024 15:58:58 -0700
Subject: [PATCH] Add redacted values in marshalobject

---
 gen/field.go                                | 26 +++++++++++++++------
 gen/field_test.go                           | 18 ++++++++++----
 gen/internal/tests/exceptions/exceptions.go |  2 +-
 gen/internal/tests/structs/structs.go       |  2 +-
 gen/zap_test.go                             |  6 -----
 5 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/gen/field.go b/gen/field.go
index f2faec09..ce6629c4 100644
--- a/gen/field.go
+++ b/gen/field.go
@@ -792,18 +792,28 @@ func (f fieldGroupGenerator) Zap(g Generator) error {
 			if <$v> == nil {
 				return nil
 			}
+			< $redactedContent := redactedContent ->
 			<range .Fields>
+				<- $requiresRedaction := requiresRedaction . ->
 				<- if not (zapOptOut .) ->
 					<- $fval := printf "%s.%s" $v (goName .) ->
 					<- if .Required ->
-						<zapEncodeBegin .Type ->
-							<$enc>.Add<zapEncoder .Type>("<fieldLabel .>", <zapMarshaler .Type $fval>)
-						<- zapEncodeEnd .Type>
+						<- if $requiresRedaction ->
+								<$enc>.AddString("<fieldLabel .>", "<$redactedContent>")
+						<- else ->
+							<- zapEncodeBegin .Type ->
+								<$enc>.Add<zapEncoder .Type>("<fieldLabel .>", <zapMarshaler .Type $fval>)
+							<- zapEncodeEnd .Type ->
+						<- end ->
 					<- else ->
 						if <$fval> != nil {
-							<zapEncodeBegin .Type ->
-								<$enc>.Add<zapEncoder .Type>("<fieldLabel .>", <zapMarshalerPtr .Type $fval>)
-							<- zapEncodeEnd .Type>
+							<- if $requiresRedaction ->
+								<$enc>.AddString("<fieldLabel .>", "<$redactedContent>")
+							<- else ->
+								<- zapEncodeBegin .Type ->
+									<$enc>.Add<zapEncoder .Type>("<fieldLabel .>", <zapMarshalerPtr .Type $fval>)
+								<- zapEncodeEnd .Type ->
+							<- end >
 						}
 					<- end>
 				<- end>
@@ -813,6 +823,8 @@ func (f fieldGroupGenerator) Zap(g Generator) error {
 		`, f,
 		TemplateFunc("zapOptOut", zapOptOut),
 		TemplateFunc("fieldLabel", entityLabel),
+		TemplateFunc("requiresRedaction", requiresRedaction),
+		TemplateFunc("redactedContent", redactedContent),
 	)
 }
 
@@ -890,7 +902,7 @@ func verifyUniqueFieldLabels(fs compile.FieldGroup) error {
 }
 
 // RedactedLabel provides a mechanism to redact certain struct fields from
-// the outputs of String() and Error() methods.
+// the outputs of String(), Error() and MarshalLogObject() methods.
 //
 //	struct Contact {
 //	    1: required string name
diff --git a/gen/field_test.go b/gen/field_test.go
index 90779c52..837e7702 100644
--- a/gen/field_test.go
+++ b/gen/field_test.go
@@ -163,7 +163,7 @@ func TestHasRedactedAnnotation(t *testing.T) {
 	}
 }
 
-func TestSanitizeRedacted(t *testing.T) {
+func TestRedactedAnnotation(t *testing.T) {
 	age := int32(21)
 	pi := ts.PersonalInfo{
 		Age:  toPtr(age),
@@ -173,11 +173,17 @@ func TestSanitizeRedacted(t *testing.T) {
 		Key:      "s",
 		UserName: toPtr("john doe"),
 	}
-	piEncoder := zapcore.NewMapObjectEncoder()
-	require.NoError(t, pi.MarshalLogObject(piEncoder))
+	enc := zapcore.NewMapObjectEncoder()
+	require.NoError(t, pi.MarshalLogObject(enc))
+	require.Len(t, enc.Fields, 2)
+	_, ok := enc.Fields["race"]
+	require.True(t, ok)
 
-	redactedExceptionEncoder := zapcore.NewMapObjectEncoder()
-	require.NoError(t, redactedException.MarshalLogObject(redactedExceptionEncoder))
+	eEncoder := zapcore.NewMapObjectEncoder()
+	require.NoError(t, redactedException.MarshalLogObject(eEncoder))
+	require.Len(t, eEncoder.Fields, 2)
+	_, ok = eEncoder.Fields["userName"]
+	require.True(t, ok)
 
 	tests := []struct {
 		name string
@@ -185,8 +191,10 @@ func TestSanitizeRedacted(t *testing.T) {
 		want any
 	}{
 		{name: "struct/string", got: pi.String(), want: "PersonalInfo{Age: 21, Race: <redacted>}"},
+		{name: "struct/MarshalLogObject", got: enc.Fields["race"], want: _redactedContent},
 		{name: "exception/string", got: redactedException.String(), want: "DoesNotExistException{Key: s, UserName: <redacted>}"},
 		{name: "exception/error", got: redactedException.Error(), want: "DoesNotExistException{Key: s, UserName: <redacted>}"},
+		{name: "exception/MarshalLogObject", got: eEncoder.Fields["userName"], want: _redactedContent},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
diff --git a/gen/internal/tests/exceptions/exceptions.go b/gen/internal/tests/exceptions/exceptions.go
index e1596a45..bfe5beeb 100644
--- a/gen/internal/tests/exceptions/exceptions.go
+++ b/gen/internal/tests/exceptions/exceptions.go
@@ -319,7 +319,7 @@ func (v *DoesNotExistException) MarshalLogObject(enc zapcore.ObjectEncoder) (err
 		enc.AddString("Error", *v.Error2)
 	}
 	if v.UserName != nil {
-		enc.AddString("userName", *v.UserName)
+		enc.AddString("userName", "<redacted>")
 	}
 	return err
 }
diff --git a/gen/internal/tests/structs/structs.go b/gen/internal/tests/structs/structs.go
index 14166d84..1c22b6a9 100644
--- a/gen/internal/tests/structs/structs.go
+++ b/gen/internal/tests/structs/structs.go
@@ -4989,7 +4989,7 @@ func (v *PersonalInfo) MarshalLogObject(enc zapcore.ObjectEncoder) (err error) {
 		enc.AddInt32("age", *v.Age)
 	}
 	if v.Race != nil {
-		enc.AddString("race", *v.Race)
+		enc.AddString("race", "<redacted>")
 	}
 	return err
 }
diff --git a/gen/zap_test.go b/gen/zap_test.go
index 007b8925..a412aed7 100644
--- a/gen/zap_test.go
+++ b/gen/zap_test.go
@@ -858,11 +858,6 @@ func TestZapOptOut(t *testing.T) {
 	foo := &compile.FieldSpec{
 		Name: "foo",
 	}
-	redacted := &compile.FieldSpec{
-		Name:        "redacted",
-		Annotations: compile.Annotations{RedactedLabel: ""},
-	}
-
 	nolog := &compile.FieldSpec{
 		Name:        "nolog",
 		Annotations: compile.Annotations{NoZapLabel: ""},
@@ -873,7 +868,6 @@ func TestZapOptOut(t *testing.T) {
 		want bool
 	}{
 		{name: "no annotation", spec: foo, want: false},
-		{name: "redacted annotation", spec: redacted, want: false},
 		{name: "nolog annotation", spec: nolog, want: true},
 	}
 	for _, tt := range tests {