Reduce allocations in concatkdf Read (#1562)

* Reduce allocations in concatkdf Read

Use h.Sum(k.buf) instead of append(k.buf, h.Sum(nil)...) to avoid
intermediate slice allocation. Replace binary.Write with
BigEndian.PutUint32 + h.Write to avoid its internal buffer alloc.

Read: -7% time (225→208 ns/op), -10% mem (324→292 B/op), -1 alloc (6→5)

* Fix benchmark lint: use b.Loop()

Author: lestrrat

jwe/internal/concatkdf/bench_test.go

@@ -0,0 +1,28 @@
+package concatkdf_test
+
+import (
+	"crypto"
+	"testing"
+
+	"github.com/lestrrat-go/jwx/v3/jwe/internal/concatkdf"
+)
+
+func BenchmarkRead(b *testing.B) {
+	hash := crypto.SHA256
+	alg := []byte("A128CBC-HS256")
+	z := make([]byte, 32)
+	apu := []byte("Alice")
+	apv := []byte("Bob")
+	pubinfo := make([]byte, 4)
+	privinfo := []byte{}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for b.Loop() {
+		kdf := concatkdf.New(hash, alg, z, apu, apv, pubinfo, privinfo)
+		out := make([]byte, 32)
+		if _, err := kdf.Read(out); err != nil {
+			b.Fatal(err)
+		}
+	}
+}

jwe/internal/concatkdf/concatkdf.go

@@ -42,11 +42,13 @@ func New(hash crypto.Hash, alg, Z, apu, apv, pubinfo, privinfo []byte) *KDF {
 func (k *KDF) Read(out []byte) (int, error) {
 	var round uint32 = 1
 	h := k.hash.New()
+	var roundBuf [4]byte
 
 	for len(out) > len(k.buf) {
 		h.Reset()
 
-		if err := binary.Write(h, binary.BigEndian, round); err != nil {
+		binary.BigEndian.PutUint32(roundBuf[:], round)
+		if _, err := h.Write(roundBuf[:]); err != nil {
 			return 0, fmt.Errorf(`failed to write round using kdf: %w`, err)
 		}
 		if _, err := h.Write(k.z); err != nil {
@@ -56,7 +58,7 @@ func (k *KDF) Read(out []byte) (int, error) {
 			return 0, fmt.Errorf(`failed to write other info using kdf: %w`, err)
 		}
 
-		k.buf = append(k.buf, h.Sum(nil)...)
+		k.buf = h.Sum(k.buf)
 		round++
 	}