From 1b06712d900780dce1defe26cf21479a6af3494f Mon Sep 17 00:00:00 2001
From: c9s <yoan@maicoin.com>
Date: Sun, 16 Jan 2022 23:45:33 +0800
Subject: [PATCH] util: improve mask key function and add tests

---
 pkg/util/string.go      | 15 ++++++++++++---
 pkg/util/string_test.go | 42 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 3 deletions(-)
 create mode 100644 pkg/util/string_test.go

diff --git a/pkg/util/string.go b/pkg/util/string.go
index 8bf153c46..268233acb 100644
--- a/pkg/util/string.go
+++ b/pkg/util/string.go
@@ -13,8 +13,17 @@ func StringSliceContains(slice []string, needle string) bool {
 }
 
 func MaskKey(key string) string {
-	maskKey := key[0:5]
-	maskKey += strings.Repeat("*", len(key)-1-5-5)
-	maskKey += key[len(key)-5-1:]
+	if len(key) == 0 {
+		return "{empty}"
+	}
+
+	h := len(key) / 3
+	if h > 5 {
+		h = 5
+	}
+
+	maskKey := key[0:h]
+	maskKey += strings.Repeat("*", len(key)-h*2)
+	maskKey += key[len(key)-h:]
 	return maskKey
 }
diff --git a/pkg/util/string_test.go b/pkg/util/string_test.go
new file mode 100644
index 000000000..7d55f425a
--- /dev/null
+++ b/pkg/util/string_test.go
@@ -0,0 +1,42 @@
+package util
+
+import "testing"
+
+func TestMaskKey(t *testing.T) {
+	type args struct {
+		key string
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			name: "key length more than 5",
+			args: args{key: "abcdefghijklmnopqr"},
+			want: "abcde********nopqr",
+		},
+		{
+			name: "key length less than 10",
+			args: args{key: "12345678"},
+			want: "12****78",
+		},
+		{
+			name: "even",
+			args: args{key: "1234567"},
+			want: "12***67",
+		},
+		{
+			name: "empty",
+			args: args{key: ""},
+			want: "{empty}",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := MaskKey(tt.args.key); got != tt.want {
+				t.Errorf("MaskKey() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}