improve performance of parsing jwks

provider/assume-role/github/jwk/base64.go

@@ -0,0 +1,34 @@
+package jwk
+
+import (
+	"encoding/base64"
+	"fmt"
+)
+
+type base64Context struct {
+	src []byte
+	dst []byte
+	err error
+}
+
+// newBase64Context pre-allocates base64 decoding buffers.
+func newBase64Context(n int) base64Context {
+	src := make([]byte, n)
+	dst := make([]byte, base64.RawURLEncoding.DecodedLen(n))
+	return base64Context{
+		src: src,
+		dst: dst,
+	}
+}
+
+// decode decodes s as base64 raw url encoding.
+// the returned slice is valid until next call.
+func (ctx *base64Context) decode(s string, name string) []byte {
+	src := ctx.src[:len(s)]
+	copy(src, s)
+	n, err := base64.RawURLEncoding.Decode(ctx.dst, src)
+	if err != nil && ctx.err != nil {
+		ctx.err = fmt.Errorf("jwk: failed to parse the parameter %s: %w", name, err)
+	}
+	return ctx.dst[:n]
+}

provider/assume-role/github/jwk/ecdsa.go

@@ -3,7 +3,6 @@ package jwk
 import (
 	"crypto/ecdsa"
 	"crypto/elliptic"
-	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -79,25 +78,26 @@ func (key *ecdsaPrivateKey) decode() error {
 		return fmt.Errorf("jwk: unknown elliptic curve: %q", key.Crv)
 	}
 
-	dataX, err := base64.RawURLEncoding.DecodeString(key.X)
-	if err != nil {
-		return fmt.Errorf("jwk: failed to parse parameter x: %w", err)
-	}
-	key.privateKey.X = new(big.Int).SetBytes(dataX)
+	ctx := key.getContext()
+	key.privateKey.X = new(big.Int).SetBytes(ctx.decode(key.X, "x"))
+	key.privateKey.Y = new(big.Int).SetBytes(ctx.decode(key.Y, "y"))
+	key.privateKey.D = new(big.Int).SetBytes(ctx.decode(key.D, "d"))
 
-	dataY, err := base64.RawURLEncoding.DecodeString(key.Y)
-	if err != nil {
-		return fmt.Errorf("jwk: failed to parse parameter y: %w", err)
-	}
-	key.privateKey.Y = new(big.Int).SetBytes(dataY)
+	return ctx.err
+}
 
-	dataD, err := base64.RawURLEncoding.DecodeString(key.D)
-	if err != nil {
-		return fmt.Errorf("jwk: failed to parse parameter d: %w", err)
+func (key *ecdsaPrivateKey) getContext() base64Context {
+	var size int
+	if len(key.X) > size {
+		size = len(key.X)
 	}
-	key.privateKey.D = new(big.Int).SetBytes(dataD)
-
-	return nil
+	if len(key.Y) > size {
+		size = len(key.Y)
+	}
+	if len(key.D) > size {
+		size = len(key.D)
+	}
+	return newBase64Context(size)
 }
 
 // RFC7518 6.2.1. Parameters for Elliptic Curve Public Keys
@@ -162,17 +162,20 @@ func (key *ecdsaPublicKey) decode() error {
 		return fmt.Errorf("jwk: unknown elliptic curve: %q", key.Crv)
 	}
 
-	dataX, err := base64.RawURLEncoding.DecodeString(key.X)
-	if err != nil {
-		return fmt.Errorf("jwk: failed to parse parameter x: %w", err)
-	}
-	key.publicKey.X = new(big.Int).SetBytes(dataX)
+	ctx := key.getContext()
+	key.publicKey.X = new(big.Int).SetBytes(ctx.decode(key.X, "x"))
+	key.publicKey.Y = new(big.Int).SetBytes(ctx.decode(key.Y, "y"))
 
-	dataY, err := base64.RawURLEncoding.DecodeString(key.Y)
-	if err != nil {
-		return fmt.Errorf("jwk: failed to parse parameter y: %w", err)
-	}
-	key.publicKey.Y = new(big.Int).SetBytes(dataY)
+	return ctx.err
+}
 
-	return nil
+func (key *ecdsaPublicKey) getContext() base64Context {
+	var size int
+	if len(key.X) > size {
+		size = len(key.X)
+	}
+	if len(key.Y) > size {
+		size = len(key.Y)
+	}
+	return newBase64Context(size)
 }

provider/assume-role/github/jwk/ed25519.go

@@ -2,7 +2,6 @@ package jwk
 
 import (
 	"crypto/ed25519"
-	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -61,31 +60,36 @@ func (key *ed25519PrivateKey) PublicKey() interface{} {
 }
 
 func (key *ed25519PrivateKey) decode() error {
-	if base64.RawURLEncoding.DecodedLen(len(key.X)) != ed25519.PublicKeySize {
-		return fmt.Errorf("jwk: the parameter x has invalid size")
-	}
-	if base64.RawURLEncoding.DecodedLen(len(key.D)) != ed25519.PrivateKeySize-ed25519.PublicKeySize {
-		return fmt.Errorf("jwk: the parameter d has invalid size")
-	}
+	ctx := key.getContext()
 
-	publicKey := make([]byte, ed25519.PublicKeySize)
 	privateKey := make([]byte, ed25519.PrivateKeySize)
-
-	var err error
-	_, err = base64.RawURLEncoding.Decode(publicKey, []byte(key.X))
-	if err != nil {
-		return fmt.Errorf("jwk: failed to parse parameter x: %w", err)
+	data := ctx.decode(key.D, "d")
+	if len(data) != ed25519.PrivateKeySize-ed25519.PublicKeySize {
+		return fmt.Errorf("jwk: the parameter d has invalid size")
 	}
-	_, err = base64.RawURLEncoding.Decode(privateKey, []byte(key.D))
-	if err != nil {
-		return fmt.Errorf("jwk: failed to parse parameter d: %w", err)
+	copy(privateKey, data)
+
+	publicKey := ctx.decode(key.X, "x")
+	if len(publicKey) != ed25519.PublicKeySize {
+		return fmt.Errorf("jwk: the parameter x has invalid size")
 	}
 	copy(privateKey[32:], publicKey)
 
 	key.publicKey = ed25519.PublicKey(publicKey)
 	key.privateKey = ed25519.PrivateKey(privateKey)
 
-	return nil
+	return ctx.err
+}
+
+func (key *ed25519PrivateKey) getContext() base64Context {
+	var size int
+	if len(key.X) > size {
+		size = len(key.X)
+	}
+	if len(key.D) > size {
+		size = len(key.D)
+	}
+	return newBase64Context(size)
 }
 
 // RFC8037 2. Key Type "OKP"
@@ -132,15 +136,15 @@ func parseEd25519PublicKey(data []byte) (Key, error) {
 
 // decode decodes the encoded values into publicKey.
 func (key *ed25519PublicKey) decode() error {
-	if base64.RawURLEncoding.DecodedLen(len(key.X)) != ed25519.PublicKeySize {
+	ctx := key.getContext()
+	data := ctx.decode(key.X, "x")
+	if len(data) != ed25519.PublicKeySize {
 		return fmt.Errorf("jwk: the parameter x has invalid size")
 	}
-	x := []byte(key.X)
-	data := make([]byte, ed25519.PublicKeySize)
-	_, err := base64.RawURLEncoding.Decode(data, x)
-	if err != nil {
-		return fmt.Errorf("jwk: failed to parse parameter x: %w", err)
-	}
 	key.publicKey = ed25519.PublicKey(data)
-	return nil
+	return ctx.err
+}
+
+func (key *ed25519PublicKey) getContext() base64Context {
+	return newBase64Context(len(key.X))
 }

provider/assume-role/github/jwk/jwk.go

@@ -169,29 +169,29 @@ func (key *commonKey) decode() error {
 // ParseKey parses a JWK.
 func ParseKey(data []byte) (Key, error) {
 	var hint struct {
-		Kty string          `json:"kty"`
-		Crv string          `json:"crv"`
-		D   json.RawMessage `json:"d"`
+		Kty string   `json:"kty"`
+		Crv string   `json:"crv"`
+		D   hasValue `json:"d"`
 	}
 
 	if err := json.Unmarshal(data, &hint); err != nil {
 		return nil, err
 	}
 	switch hint.Kty {
 	case "EC":
-		if len(hint.D) > 0 {
+		if hint.D {
 			return parseEcdsaPrivateKey(data)
 		} else {
 			return parseEcdsaPublicKey(data)
 		}
 	case "RSA":
-		if len(hint.D) > 0 {
+		if hint.D {
 			return parseRSAPrivateKey(data)
 		} else {
 			return parseRSAPublicKey(data)
 		}
 	case "OKP":
-		if len(hint.D) > 0 {
+		if hint.D {
 			return parseOkpPrivateKey(data, hint.Crv)
 		} else {
 			return parseOkpPublicKey(data, hint.Crv)
@@ -202,3 +202,12 @@ func ParseKey(data []byte) (Key, error) {
 		return nil, fmt.Errorf("jwk: unknown key type: %s", hint.Kty)
 	}
 }
+
+// hasValue checks a JSON string has the specific key.
+// it don't parse the value.
+type hasValue bool
+
+func (v *hasValue) UnmarshalJSON(data []byte) error {
+	*v = true
+	return nil
+}