From 20a67e08b5487cf6af4f9532e1fbbb910efd71ed Mon Sep 17 00:00:00 2001
From: Tien Nguyen <duytien.nguyen@okta.com>
Date: Thu, 21 Sep 2023 14:11:08 -0400
Subject: [PATCH] update key signer algo

---
 .generator/templates/client.mustache | 44 ++++++++++++++++++++--------
 1 file changed, 32 insertions(+), 12 deletions(-)

diff --git a/.generator/templates/client.mustache b/.generator/templates/client.mustache
index a23402c0..5c967dc3 100644
--- a/.generator/templates/client.mustache
+++ b/.generator/templates/client.mustache
@@ -4,6 +4,8 @@ package {{packageName}}
 import (
 	"bytes"
 	"context"
+	"crypto/ecdsa"
+	"crypto/rsa"
 	"crypto/x509"
 	"encoding/json"
 	"encoding/pem"
@@ -233,27 +235,45 @@ func (a *JWTAuth) Authorize() error {
 }
 
 func createKeySigner(privateKey, privateKeyID string) (jose.Signer, error) {
+	var signerOptions *jose.SignerOptions
+	if privateKeyID != "" {
+		signerOptions = (&jose.SignerOptions{}).WithHeader("kid", privateKeyID)
+	}
+
 	priv := []byte(strings.ReplaceAll(privateKey, `\n`, "\n"))
 
 	privPem, _ := pem.Decode(priv)
 	if privPem == nil {
 		return nil, errors.New("invalid private key")
 	}
-	if privPem.Type != "RSA PRIVATE KEY" {
-		return nil, fmt.Errorf("RSA private key is of the wrong type")
-	}
-
-	parsedKey, err := x509.ParsePKCS1PrivateKey(privPem.Bytes)
-	if err != nil {
-		return nil, err
+	if privPem.Type == "RSA PRIVATE KEY" {
+		parsedKey, err := x509.ParsePKCS1PrivateKey(privPem.Bytes)
+		if err != nil {
+			return nil, err
+		}
+		return jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: parsedKey}, signerOptions)
 	}
-
-	var signerOptions *jose.SignerOptions
-	if privateKeyID != "" {
-		signerOptions = (&jose.SignerOptions{}).WithHeader("kid", privateKeyID)
+	if privPem.Type == "PRIVATE KEY" {
+		parsedKey, err := x509.ParsePKCS8PrivateKey(privPem.Bytes)
+		if err != nil {
+			return nil, err
+		}
+		var alg jose.SignatureAlgorithm
+		switch parsedKey.(type) {
+		case *rsa.PrivateKey:
+			alg = jose.RS256
+		case *ecdsa.PrivateKey:
+			alg = jose.ES256 // TODO handle ES384 or ES512 ?
+		default:
+			// TODO are either of these also valid?
+			// ed25519.PrivateKey:
+			// *ecdh.PrivateKey
+			return nil, fmt.Errorf("private key %q is unknown pkcs#8 format type", privPem.Type)
+		}
+		return jose.NewSigner(jose.SigningKey{Algorithm: alg, Key: parsedKey}, signerOptions)
 	}
 
-	return jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: parsedKey}, signerOptions)
+	return nil, fmt.Errorf("private key %q is not pkcs#1 or pkcs#8 format", privPem.Type)
 }
 
 func createClientAssertion(orgURL, clientID string, privateKeySinger jose.Signer) (clientAssertion string, err error) {