Skip to content

Commit e9f4379

Browse files
leeqvipleeqvip
committed
perf: do some optimization
1 parent 8c7a563 commit e9f4379

File tree

3 files changed

+86
-47
lines changed

3 files changed

+86
-47
lines changed

‎key_test.go‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,12 @@ import (
66
"github.com/stretchr/testify/assert"
77
)
88

9+
const (
10+
TestBlockSize10 = 10
11+
TestBlockSize20 = 20
12+
TestBlockSize30 = 30
13+
)
14+
915
func TestSHA1(t *testing.T) {
1016
testData := []byte("test")
1117
h := sha1.New()
@@ -16,14 +22,17 @@ func TestSHA1(t *testing.T) {
1622
}
1723

1824
func TestKeyGenerator(t *testing.T) {
25+
testData := []byte("test")
26+
result10 := KeyGenerator(testData, TestBlockSize10)
27+
result20 := KeyGenerator(testData, TestBlockSize20)
1928
testCases := []struct {
2029
data []byte
2130
blockSize int
2231
expected []byte
2332
}{
24-
{[]byte("test"), 10, KeyGenerator([]byte("test"), 10)},
25-
{[]byte("test"), 20, KeyGenerator([]byte("test"), 20)},
26-
{[]byte("test"), 30, []byte("test")},
33+
{testData, TestBlockSize10, result10},
34+
{testData, TestBlockSize20, result20},
35+
{testData, TestBlockSize30, []byte("test")},
2736
}
2837

2938
for _, tc := range testCases {

‎padding.go‎

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -12,28 +12,30 @@ const PKCS5_PADDING = "PKCS5"
1212
const PKCS7_PADDING = "PKCS7"
1313
const ZEROS_PADDING = "ZEROS"
1414

15+
var paddingFunctions = map[string]func([]byte, int) []byte{
16+
PKCS5_PADDING: PKCS5Padding,
17+
PKCS7_PADDING: PKCS7Padding,
18+
ZEROS_PADDING: ZerosPadding,
19+
}
20+
21+
var unpaddingFunctions = map[string]func([]byte) ([]byte, error){
22+
PKCS5_PADDING: PKCS5Unpadding,
23+
PKCS7_PADDING: PKCS7UnPadding,
24+
ZEROS_PADDING: ZerosUnPadding,
25+
}
26+
1527
// Applies the specified padding scheme to the input data.
1628
func Padding(padding string, src []byte, blockSize int) []byte {
17-
switch padding {
18-
case PKCS5_PADDING:
19-
src = PKCS5Padding(src, blockSize)
20-
case PKCS7_PADDING:
21-
src = PKCS7Padding(src, blockSize)
22-
case ZEROS_PADDING:
23-
src = ZerosPadding(src, blockSize)
29+
if fn, ok := paddingFunctions[padding]; ok {
30+
return fn(src, blockSize)
2431
}
2532
return src
2633
}
2734

2835
// Removes the specified padding from the input data.
2936
func UnPadding(padding string, src []byte) ([]byte, error) {
30-
switch padding {
31-
case PKCS5_PADDING:
32-
return PKCS5Unpadding(src)
33-
case PKCS7_PADDING:
34-
return PKCS7UnPadding(src)
35-
case ZEROS_PADDING:
36-
return ZerosUnPadding(src)
37+
if fn, ok := unpaddingFunctions[padding]; ok {
38+
return fn(src)
3739
}
3840
return src, nil
3941
}
@@ -80,9 +82,14 @@ func ZerosPadding(src []byte, blockSize int) []byte {
8082

8183
// Removes zero padding from the input data.
8284
func ZerosUnPadding(src []byte) ([]byte, error) {
83-
for i := len(src) - 1; ; i-- {
85+
length := len(src)
86+
if length == 0 {
87+
return src, nil
88+
}
89+
for i := length - 1; i >= 0; i-- {
8490
if src[i] != 0 {
8591
return src[:i+1], nil
8692
}
8793
}
94+
return []byte{}, nil
8895
}

‎padding_test.go‎

Lines changed: 52 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,28 @@ import (
66
"github.com/stretchr/testify/assert"
77
)
88

9+
const TestBlockSize = 16
10+
11+
func runTestCases(t *testing.T, testCases []struct {
12+
name string
13+
args struct{ src []byte }
14+
want []byte
15+
wantErr error
16+
}, testFunc func([]byte) ([]byte, error)) {
17+
for _, tt := range testCases {
18+
t.Run(tt.name, func(t *testing.T) {
19+
result, err := testFunc(tt.args.src)
20+
assert.Equal(t, tt.wantErr, err)
21+
assert.Equal(t, tt.want, result)
22+
})
23+
}
24+
}
25+
926
func TestPKCS7UnPadding(t *testing.T) {
1027
type args struct {
1128
src []byte
1229
}
13-
tests := []struct {
30+
testCases := []struct {
1431
name string
1532
args args
1633
want []byte
@@ -35,46 +52,57 @@ func TestPKCS7UnPadding(t *testing.T) {
3552
wantErr: ErrUnPadding,
3653
},
3754
}
38-
for _, tt := range tests {
39-
t.Run(tt.name, func(t *testing.T) {
40-
result, err := PKCS7UnPadding(tt.args.src)
41-
t.Log(string(result))
42-
assert.Equal(t, tt.wantErr, err)
43-
assert.Equal(t, tt.want, result)
44-
})
45-
}
55+
// Convert testCases to match the expected type of runTestCases
56+
convertedTestCases := make([]struct {
57+
name string
58+
args struct{ src []byte }
59+
want []byte
60+
wantErr error
61+
}, len(testCases))
62+
for i, tc := range testCases {
63+
convertedTestCases[i] = struct {
64+
name string
65+
args struct{ src []byte }
66+
want []byte
67+
wantErr error
68+
}{
69+
name: tc.name,
70+
args: struct{ src []byte }{src: tc.args.src},
71+
want: tc.want,
72+
wantErr: tc.wantErr,
73+
}
74+
}
75+
runTestCases(t, convertedTestCases, PKCS7UnPadding)
4676
}
4777

4878
func TestPadding(t *testing.T) {
49-
blockSize := 16
5079
tests := []struct {
5180
name string
5281
padding string
5382
src []byte
5483
wantLen int
5584
}{
56-
{name: "PKCS5", padding: PKCS5_PADDING, src: []byte("test"), wantLen: 16},
57-
{name: "PKCS7", padding: PKCS7_PADDING, src: []byte("test"), wantLen: 16},
58-
{name: "Zeros", padding: ZEROS_PADDING, src: []byte("test"), wantLen: 16},
85+
{name: "PKCS5", padding: PKCS5_PADDING, src: []byte("test"), wantLen: TestBlockSize},
86+
{name: "PKCS7", padding: PKCS7_PADDING, src: []byte("test"), wantLen: TestBlockSize},
87+
{name: "Zeros", padding: ZEROS_PADDING, src: []byte("test"), wantLen: TestBlockSize},
5988
}
6089
for _, tt := range tests {
6190
t.Run(tt.name, func(t *testing.T) {
62-
got := Padding(tt.padding, tt.src, blockSize)
91+
got := Padding(tt.padding, tt.src, TestBlockSize)
6392
assert.Equal(t, tt.wantLen, len(got))
6493
})
6594
}
6695
}
6796

6897
func TestUnPadding(t *testing.T) {
69-
blockSize := 16
7098
padding := PKCS7_PADDING
7199
tests := []struct {
72100
name string
73101
src []byte
74102
want []byte
75103
wantErr bool
76104
}{
77-
{name: "valid padding", src: PKCS7Padding([]byte("test"), blockSize), want: []byte("test"), wantErr: false},
105+
{name: "valid padding", src: PKCS7Padding([]byte("test"), TestBlockSize), want: []byte("test"), wantErr: false},
78106
{name: "empty src", src: []byte{}, want: []byte{}, wantErr: true},
79107
}
80108
for _, tt := range tests {
@@ -91,25 +119,23 @@ func TestUnPadding(t *testing.T) {
91119
}
92120

93121
func TestPKCS5Padding(t *testing.T) {
94-
blockSize := 16
95122
tests := []struct {
96123
name string
97124
src []byte
98125
wantLen int
99126
}{
100-
{name: "test PKCS5", src: []byte("test"), wantLen: 16},
127+
{name: "test PKCS5", src: []byte("test"), wantLen: TestBlockSize},
101128
}
102129
for _, tt := range tests {
103130
t.Run(tt.name, func(t *testing.T) {
104-
got := PKCS5Padding(tt.src, blockSize)
131+
got := PKCS5Padding(tt.src, TestBlockSize)
105132
assert.Equal(t, tt.wantLen, len(got))
106133
})
107134
}
108135
}
109136

110137
func TestPKCS5Unpadding(t *testing.T) {
111-
blockSize := 16
112-
padded := PKCS5Padding([]byte("test"), blockSize)
138+
padded := PKCS5Padding([]byte("test"), TestBlockSize)
113139
tests := []struct {
114140
name string
115141
src []byte
@@ -132,47 +158,44 @@ func TestPKCS5Unpadding(t *testing.T) {
132158
}
133159

134160
func TestPKCS7Padding(t *testing.T) {
135-
blockSize := 16
136161
tests := []struct {
137162
name string
138163
src []byte
139164
wantLen int
140165
}{
141-
{name: "test PKCS7", src: []byte("test"), wantLen: 16},
166+
{name: "test PKCS7", src: []byte("test"), wantLen: TestBlockSize},
142167
}
143168
for _, tt := range tests {
144169
t.Run(tt.name, func(t *testing.T) {
145-
got := PKCS7Padding(tt.src, blockSize)
170+
got := PKCS7Padding(tt.src, TestBlockSize)
146171
assert.Equal(t, tt.wantLen, len(got))
147172
})
148173
}
149174
}
150175

151176
func TestZerosPadding(t *testing.T) {
152-
blockSize := 16
153177
tests := []struct {
154178
name string
155179
src []byte
156180
wantLen int
157181
}{
158-
{name: "test Zeros", src: []byte("test"), wantLen: 16},
182+
{name: "test Zeros", src: []byte("test"), wantLen: TestBlockSize},
159183
}
160184
for _, tt := range tests {
161185
t.Run(tt.name, func(t *testing.T) {
162-
got := ZerosPadding(tt.src, blockSize)
186+
got := ZerosPadding(tt.src, TestBlockSize)
163187
assert.Equal(t, tt.wantLen, len(got))
164188
})
165189
}
166190
}
167191

168192
func TestZerosUnPadding(t *testing.T) {
169-
blockSize := 16
170193
tests := []struct {
171194
name string
172195
src []byte
173196
want []byte
174197
}{
175-
{name: "test ZerosUnPadding", src: ZerosPadding([]byte("test"), blockSize), want: []byte("test")},
198+
{name: "test ZerosUnPadding", src: ZerosPadding([]byte("test"), TestBlockSize), want: []byte("test")},
176199
}
177200
for _, tt := range tests {
178201
t.Run(tt.name, func(t *testing.T) {

0 commit comments

Comments
 (0)