From f2199843b12c534606081c5d004d9f9f8c872904 Mon Sep 17 00:00:00 2001
From: Sun Yimin <emmansun@users.noreply.github.com>
Date: Thu, 21 Sep 2023 13:48:51 +0800
Subject: [PATCH] kdf: add benchmark test

---
 kdf/kdf.go      |  6 +++---
 kdf/kdf_test.go | 42 ++++++++++++++++++++++++++++++++++++------
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/kdf/kdf.go b/kdf/kdf.go
index 58d4b69..0cd2cee 100644
--- a/kdf/kdf.go
+++ b/kdf/kdf.go
@@ -15,14 +15,14 @@ func Kdf(md hash.Hash, z []byte, len int) []byte {
 	}
 	var countBytes [4]byte
 	var ct uint32 = 1
-	k := make([]byte, len)
+	var k []byte
 	for i := 0; i < int(limit); i++ {
 		binary.BigEndian.PutUint32(countBytes[:], ct)
 		md.Write(z)
 		md.Write(countBytes[:])
-		copy(k[i*md.Size():], md.Sum(nil))
+		k = md.Sum(k)
 		ct++
 		md.Reset()
 	}
-	return k
+	return k[:len]
 }
diff --git a/kdf/kdf_test.go b/kdf/kdf_test.go
index 8e09b52..0593201 100644
--- a/kdf/kdf_test.go
+++ b/kdf/kdf_test.go
@@ -2,6 +2,7 @@ package kdf
 
 import (
 	"encoding/hex"
+	"fmt"
 	"hash"
 	"math/big"
 	"reflect"
@@ -19,16 +20,19 @@ func TestKdf(t *testing.T) {
 	tests := []struct {
 		name string
 		args args
-		want []byte
+		want string
 	}{
-		{"sm3 case 1", args{sm3.New(), []byte("emmansun"), 16}, []byte{112, 137, 147, 239, 19, 136, 160, 174, 66, 69, 161, 155, 182, 192, 37, 84}},
-		{"sm3 case 2", args{sm3.New(), []byte("emmansun"), 32}, []byte{112, 137, 147, 239, 19, 136, 160, 174, 66, 69, 161, 155, 182, 192, 37, 84, 198, 50, 99, 62, 53, 109, 219, 152, 155, 235, 128, 79, 218, 150, 207, 212}},
-		{"sm3 case 3", args{sm3.New(), []byte("emmansun"), 48}, []byte{112, 137, 147, 239, 19, 136, 160, 174, 66, 69, 161, 155, 182, 192, 37, 84, 198, 50, 99, 62, 53, 109, 219, 152, 155, 235, 128, 79, 218, 150, 207, 212, 126, 186, 79, 164, 96, 231, 178, 119, 188, 107, 76, 228, 208, 126, 212, 147}},
+		{"sm3 case 1", args{sm3.New(), []byte("emmansun"), 16}, "708993ef1388a0ae4245a19bb6c02554"},
+		{"sm3 case 2", args{sm3.New(), []byte("emmansun"), 32}, "708993ef1388a0ae4245a19bb6c02554c632633e356ddb989beb804fda96cfd4"},
+		{"sm3 case 3", args{sm3.New(), []byte("emmansun"), 48}, "708993ef1388a0ae4245a19bb6c02554c632633e356ddb989beb804fda96cfd47eba4fa460e7b277bc6b4ce4d07ed493"},
+		{"sm3 case 4", args{sm3.New(), []byte("708993ef1388a0ae4245a19bb6c02554c632633e356ddb989beb804fda96cfd47eba4fa460e7b277bc6b4ce4d07ed493708993ef1388a0ae4245a19bb6c02554c632633e356ddb989beb804fda96cfd47eba4fa460e7b277bc6b4ce4d07ed493"), 48}, "49cf14649f324a07e0d5bb2a00f7f05d5f5bdd6d14dff028e071327ec031104590eddb18f98b763e18bf382ff7c3875f"},
+		{"sm3 case 5", args{sm3.New(), []byte("708993ef1388a0ae4245a19bb6c02554c632633e356ddb989beb804fda96cfd47eba4fa460e7b277bc6b4ce4d07ed493708993ef1388a0ae4245a19bb6c02554c632633e356ddb989beb804fda96cfd47eba4fa460e7b277bc6b4ce4d07ed493"), 128}, "49cf14649f324a07e0d5bb2a00f7f05d5f5bdd6d14dff028e071327ec031104590eddb18f98b763e18bf382ff7c3875f30277f3179baebd795e7853fa643fdf280d8d7b81a2ab7829f615e132ab376d32194cd315908d27090e1180ce442d9be99322523db5bfac40ac5acb03550f5c93e5b01b1d71f2630868909a6a1250edb"},
 	}
 	for _, tt := range tests {
+		wantBytes, _ := hex.DecodeString(tt.want)
 		t.Run(tt.name, func(t *testing.T) {
-			if got := Kdf(tt.args.md, tt.args.z, tt.args.len); !reflect.DeepEqual(got, tt.want) {
-				t.Errorf("Kdf() = %v, want %v", got, tt.want)
+			if got := Kdf(tt.args.md, tt.args.z, tt.args.len); !reflect.DeepEqual(got, wantBytes) {
+				t.Errorf("Kdf(%v) = %x, want %v", tt.name, got, tt.want)
 			}
 		})
 	}
@@ -62,3 +66,29 @@ func TestKdfPanic(t *testing.T) {
 		Kdf(sm3.New(), []byte("123456"), 1<<37)
 	})
 }
+
+func BenchmarkKdf(b *testing.B) {
+	tests := []struct {
+		zLen int
+		kLen int
+	}{
+		{32, 32},
+		{32, 64},
+		{32, 128},
+		{64, 32},
+		{64, 64},
+		{64, 128},
+		{440, 32},
+	}
+	sm3Hash := sm3.New()
+	z := make([]byte, 512)
+	for _, tt := range tests {
+		b.Run(fmt.Sprintf("zLen=%v-kLen=%v", tt.zLen, tt.kLen), func(b *testing.B) {
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				Kdf(sm3Hash, z[:tt.zLen], tt.kLen)
+			}
+		})
+	}
+}