From aa386df3933e180ca8a2f20b9af1a1288b7b5aeb Mon Sep 17 00:00:00 2001
From: Roman Chaliy <31652715+Lasiar@users.noreply.github.com>
Date: Tue, 28 Nov 2023 00:45:49 +0300
Subject: [PATCH] jwe/CompactSerialize: improve performance. (#67)

Avoid fmt.Sprintf, which is unnecessarily slow. Instead allocate a `[]byte` and write directly to that.
---
 encoding.go | 33 +++++++++++++++++++++++++++++++++
 jwe.go      | 14 +++++++-------
 jws.go      | 13 ++++++++-----
 3 files changed, 48 insertions(+), 12 deletions(-)

diff --git a/encoding.go b/encoding.go
index 968a424..62f8b8a 100644
--- a/encoding.go
+++ b/encoding.go
@@ -189,3 +189,36 @@ func base64URLDecode(value string) ([]byte, error) {
 	value = strings.TrimRight(value, "=")
 	return base64.RawURLEncoding.DecodeString(value)
 }
+
+func base64EncodeLen(sl []byte) int {
+	return base64.RawURLEncoding.EncodedLen(len(sl))
+}
+
+func base64JoinWithDots(inputs ...[]byte) string {
+	if len(inputs) == 0 {
+		return ""
+	}
+
+	// Count of dots.
+	totalCount := len(inputs) - 1
+
+	for _, input := range inputs {
+		totalCount += base64EncodeLen(input)
+	}
+
+	out := make([]byte, totalCount)
+	startEncode := 0
+	for i, input := range inputs {
+		base64.RawURLEncoding.Encode(out[startEncode:], input)
+
+		if i == len(inputs)-1 {
+			continue
+		}
+
+		startEncode += base64EncodeLen(input)
+		out[startEncode] = '.'
+		startEncode++
+	}
+
+	return string(out)
+}
diff --git a/jwe.go b/jwe.go
index bce3045..4267ac7 100644
--- a/jwe.go
+++ b/jwe.go
@@ -252,13 +252,13 @@ func (obj JSONWebEncryption) CompactSerialize() (string, error) {
 
 	serializedProtected := mustSerializeJSON(obj.protected)
 
-	return fmt.Sprintf(
-		"%s.%s.%s.%s.%s",
-		base64.RawURLEncoding.EncodeToString(serializedProtected),
-		base64.RawURLEncoding.EncodeToString(obj.recipients[0].encryptedKey),
-		base64.RawURLEncoding.EncodeToString(obj.iv),
-		base64.RawURLEncoding.EncodeToString(obj.ciphertext),
-		base64.RawURLEncoding.EncodeToString(obj.tag)), nil
+	return base64JoinWithDots(
+		serializedProtected,
+		obj.recipients[0].encryptedKey,
+		obj.iv,
+		obj.ciphertext,
+		obj.tag,
+	), nil
 }
 
 // FullSerialize serializes an object using the full JSON serialization format.
diff --git a/jws.go b/jws.go
index 865f16a..e37007d 100644
--- a/jws.go
+++ b/jws.go
@@ -314,15 +314,18 @@ func (obj JSONWebSignature) compactSerialize(detached bool) (string, error) {
 		return "", ErrNotSupported
 	}
 
-	serializedProtected := base64.RawURLEncoding.EncodeToString(mustSerializeJSON(obj.Signatures[0].protected))
-	payload := ""
-	signature := base64.RawURLEncoding.EncodeToString(obj.Signatures[0].Signature)
+	serializedProtected := mustSerializeJSON(obj.Signatures[0].protected)
 
+	var payload []byte
 	if !detached {
-		payload = base64.RawURLEncoding.EncodeToString(obj.payload)
+		payload = obj.payload
 	}
 
-	return fmt.Sprintf("%s.%s.%s", serializedProtected, payload, signature), nil
+	return base64JoinWithDots(
+		serializedProtected,
+		payload,
+		obj.Signatures[0].Signature,
+	), nil
 }
 
 // CompactSerialize serializes an object using the compact serialization format.