jwe/CompactSerialize: improve performance. (#67)

lasiar · web-flow · commit aa386df3933e · 2023-11-27T13:45:49.000-08:00
Avoid fmt.Sprintf, which is unnecessarily slow. Instead allocate a `[]byte` and write directly to that.
diff --git 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
@@ -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
@@ -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.
