Add redacted values in marshalobject

thriftrw · Mar 15, 2024 · 899cabf · 899cabf
1 parent 11e60c8
commit 899cabf
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 20 deletions.
diff --git 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
@@ -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,20 +173,28 @@ 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
 		got  any
 		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
diff --git a/gen/internal/tests/structs/structs.go b/gen/internal/tests/structs/structs.go
diff --git 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 {