Only allow the config value to be true/false/"MARKER"

Signed-off-by: JmPotato <ghzpotato@gmail.com>
tikv · Jul 1, 2024 · bc2c7d3 · bc2c7d3
1 parent 2d8bcb0
commit bc2c7d3
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 58 deletions.
diff --git a/pkg/utils/logutil/log.go b/pkg/utils/logutil/log.go
@@ -123,27 +123,22 @@ const invalidRedactInfoLogTypeErrMsg = `the "redact-info-log" value is invalid;
 // UnmarshalJSON implements the `json.Marshaler` interface to ensure the compatibility.
 func (t *RedactInfoLogType) UnmarshalJSON(data []byte) error {
 	var s string
-	if err := json.Unmarshal(data, &s); err == nil {
-		switch strings.ToUpper(s) {
-		case "FALSE", "OFF", "":
-			*t = RedactInfoLogOFF
-		case "MARKER":
-			*t = RedactInfoLogMarker
-		default:
-			*t = RedactInfoLogON
-		}
+	err := json.Unmarshal(data, &s)
+	if err == nil && strings.ToUpper(s) == "MARKER" {
+		*t = RedactInfoLogMarker
 		return nil
 	}
 	var b bool
-	if err := json.Unmarshal(data, &b); err == nil {
-		if b {
-			*t = RedactInfoLogON
-		} else {
-			*t = RedactInfoLogOFF
-		}
-		return nil
+	err = json.Unmarshal(data, &b)
+	if err != nil {
+		return errors.New(invalidRedactInfoLogTypeErrMsg)
 	}
-	return errors.New(invalidRedactInfoLogTypeErrMsg)
+	if b {
+		*t = RedactInfoLogON
+	} else {
+		*t = RedactInfoLogOFF
+	}
+	return nil
 }
 
 // UnmarshalTOML implements the `toml.Unmarshaler` interface to ensure the compatibility.
@@ -157,18 +152,14 @@ func (t *RedactInfoLogType) UnmarshalTOML(data any) error {
 		}
 		return nil
 	case string:
-		switch strings.ToUpper(v) {
-		case "FALSE", "OFF", "":
-			*t = RedactInfoLogOFF
-		case "MARKER":
+		if strings.ToUpper(v) == "MARKER" {
 			*t = RedactInfoLogMarker
-		default:
-			*t = RedactInfoLogON
+			return nil
 		}
-		return nil
-	default:
 		return errors.New(invalidRedactInfoLogTypeErrMsg)
+	default:
 	}
+	return errors.New(invalidRedactInfoLogTypeErrMsg)
 }
 
 var (

diff --git a/pkg/utils/logutil/log_test.go b/pkg/utils/logutil/log_test.go
@@ -39,27 +39,34 @@ func TestRedactInfoLogType(t *testing.T) {
 	re := require.New(t)
 	// JSON unmarshal.
 	jsonUnmarshalTestCases := []struct {
-		jsonStr string
-		expect  RedactInfoLogType
+		jsonStr   string
+		expect    RedactInfoLogType
+		expectErr bool
 	}{
-		{`false`, RedactInfoLogOFF},
-		{`true`, RedactInfoLogON},
-		{`"MARKER"`, RedactInfoLogMarker},
-		{`"OTHER"`, RedactInfoLogON},
-		{`"OFF"`, RedactInfoLogOFF},
-		{`"ON"`, RedactInfoLogON},
-		{`"marker"`, RedactInfoLogMarker},
-		{`"off"`, RedactInfoLogOFF},
-		{`"on"`, RedactInfoLogON},
-		{`""`, RedactInfoLogOFF},
-		{`"fALSe"`, RedactInfoLogOFF},
-		{`"trUE"`, RedactInfoLogON},
+		{`false`, RedactInfoLogOFF, false},
+		{`true`, RedactInfoLogON, false},
+		{`"MARKER"`, RedactInfoLogMarker, false},
+		{`"marker"`, RedactInfoLogMarker, false},
+		{`"OTHER"`, RedactInfoLogOFF, true},
+		{`"OFF"`, RedactInfoLogOFF, true},
+		{`"ON"`, RedactInfoLogOFF, true},
+		{`"off"`, RedactInfoLogOFF, true},
+		{`"on"`, RedactInfoLogOFF, true},
+		{`""`, RedactInfoLogOFF, true},
+		{`"fALSe"`, RedactInfoLogOFF, true},
+		{`"trUE"`, RedactInfoLogOFF, true},
 	}
 	var redactType RedactInfoLogType
-	for _, tc := range jsonUnmarshalTestCases {
+	for idx, tc := range jsonUnmarshalTestCases {
+		t.Logf("test case %d: %s", idx, tc.jsonStr)
 		err := json.Unmarshal([]byte(tc.jsonStr), &redactType)
-		re.NoError(err)
-		re.Equal(tc.expect, redactType)
+		if tc.expectErr {
+			re.Error(err)
+			re.ErrorContains(err, invalidRedactInfoLogTypeErrMsg)
+		} else {
+			re.NoError(err)
+			re.Equal(tc.expect, redactType)
+		}
 	}
 	// JSON marshal.
 	jsonMarshalTestCases := []struct {
@@ -77,29 +84,35 @@ func TestRedactInfoLogType(t *testing.T) {
 	}
 	// TOML unmarshal.
 	tomlTestCases := []struct {
-		tomlStr string
-		expect  RedactInfoLogType
+		tomlStr   string
+		expect    RedactInfoLogType
+		expectErr bool
 	}{
-		{`redact-info-log = false`, RedactInfoLogOFF},
-		{`redact-info-log = true`, RedactInfoLogON},
-		{`redact-info-log = "MARKER"`, RedactInfoLogMarker},
-		{`redact-info-log = "OTHER"`, RedactInfoLogON},
-		{`redact-info-log = "OFF"`, RedactInfoLogOFF},
-		{`redact-info-log = "ON"`, RedactInfoLogON},
-		{`redact-info-log = "marker"`, RedactInfoLogMarker},
-		{`redact-info-log = "off"`, RedactInfoLogOFF},
-		{`redact-info-log = "on"`, RedactInfoLogON},
-		{`redact-info-log = ""`, RedactInfoLogOFF},
-		{`redact-info-log = "fALSe"`, RedactInfoLogOFF},
-		{`redact-info-log = "trUE"`, RedactInfoLogON},
+		{`redact-info-log = false`, RedactInfoLogOFF, false},
+		{`redact-info-log = true`, RedactInfoLogON, false},
+		{`redact-info-log = "MARKER"`, RedactInfoLogMarker, false},
+		{`redact-info-log = "marker"`, RedactInfoLogMarker, false},
+		{`redact-info-log = "OTHER"`, RedactInfoLogOFF, true},
+		{`redact-info-log = "OFF"`, RedactInfoLogOFF, true},
+		{`redact-info-log = "ON"`, RedactInfoLogOFF, true},
+		{`redact-info-log = "off"`, RedactInfoLogOFF, true},
+		{`redact-info-log = "on"`, RedactInfoLogOFF, true},
+		{`redact-info-log = ""`, RedactInfoLogOFF, true},
+		{`redact-info-log = "fALSe"`, RedactInfoLogOFF, true},
+		{`redact-info-log = "trUE"`, RedactInfoLogOFF, true},
 	}
 	var config struct {
 		RedactInfoLog RedactInfoLogType `toml:"redact-info-log"`
 	}
 	for _, tc := range tomlTestCases {
 		_, err := toml.Decode(tc.tomlStr, &config)
-		re.NoError(err)
-		re.Equal(tc.expect, config.RedactInfoLog)
+		if tc.expectErr {
+			re.Error(err)
+			re.ErrorContains(err, invalidRedactInfoLogTypeErrMsg)
+		} else {
+			re.NoError(err)
+			re.Equal(tc.expect, config.RedactInfoLog)
+		}
 	}
 }