Add crypto.Signer support for KMS/HSM keys

Check public key type instead of private key type to support
crypto.Signer implementations (GCP KMS, AWS KMS, HSM) that
aren't concrete *rsa.PrivateKey or *ecdsa.PrivateKey types.

Changes:
- samlsp/new.go: Update defaultSigningMethodForKey()
- samlsp/session_jwt.go: Add fallback signing with crypto.Signer
- samlsp/request_tracker_jwt.go: Add fallback signing
- service_provider.go: Update GetSigningContext() validation

Author: dineshudayakumar

samlsp/middleware_test.go

@@ -2,6 +2,7 @@ package samlsp
 
 import (
 	"bytes"
+	"crypto"
 	"crypto/rsa"
 	"crypto/sha256"
 	"crypto/x509"
@@ -520,3 +521,122 @@ func TestMiddlewareHandlesInvalidResponse(t *testing.T) {
 	assert.Check(t, is.Equal("", resp.Header().Get("Location")))
 	assert.Check(t, is.Equal("", resp.Header().Get("Set-Cookie")))
 }
+
+type mockSigner struct {
+	signer crypto.Signer
+}
+
+func (m *mockSigner) Public() crypto.PublicKey {
+	return m.signer.Public()
+}
+
+func (m *mockSigner) Sign(rand io.Reader, digest []byte, opts crypto.SignerOpts) ([]byte, error) {
+	return m.signer.Sign(rand, digest, opts)
+}
+
+func newMockRSASigner(t *testing.T) crypto.Signer {
+	key := mustParsePrivateKey(golden.Get(t, "key.pem"))
+	return &mockSigner{signer: key.(crypto.Signer)}
+}
+
+func TestMiddleware_WithCryptoSignerE2E(t *testing.T) {
+	saml.TimeNow = func() time.Time {
+		rv, _ := time.Parse("Mon Jan 2 15:04:05.999999999 MST 2006", "Mon Dec 1 01:57:09.123456789 UTC 2015")
+		return rv
+	}
+	saml.Clock = dsig.NewFakeClockAt(saml.TimeNow())
+	saml.RandReader = &testRandomReader{}
+
+	cert := mustParseCertificate(golden.Get(t, "cert.pem"))
+	idpMetadata := golden.Get(t, "idp_metadata.xml")
+
+	var metadata saml.EntityDescriptor
+	if err := xml.Unmarshal(idpMetadata, &metadata); err != nil {
+		panic(err)
+	}
+
+	mockSigner := newMockRSASigner(t)
+
+	opts := Options{
+		URL:         mustParseURL("https://15661444.ngrok.io/"),
+		Key:         mockSigner,
+		Certificate: cert,
+		IDPMetadata: &metadata,
+	}
+
+	middleware, err := New(opts)
+	assert.Check(t, err)
+
+	sessionProvider := DefaultSessionProvider(opts)
+	sessionProvider.Name = "ttt"
+	sessionProvider.MaxAge = 7200 * time.Second
+
+	sessionCodec := sessionProvider.Codec.(JWTSessionCodec)
+	sessionCodec.MaxAge = 7200 * time.Second
+	sessionProvider.Codec = sessionCodec
+
+	middleware.Session = sessionProvider
+	middleware.ServiceProvider.MetadataURL.Path = "/saml2/metadata"
+	middleware.ServiceProvider.AcsURL.Path = "/saml2/acs"
+	middleware.ServiceProvider.SloURL.Path = "/saml2/slo"
+
+	t.Run("SessionEncodeDecode", func(t *testing.T) {
+		var tc JWTSessionClaims
+		if err := json.Unmarshal(golden.Get(t, "token.json"), &tc); err != nil {
+			t.Fatal(err)
+		}
+
+		encoded, err := sessionProvider.Codec.Encode(tc)
+		assert.Check(t, err)
+		assert.Assert(t, encoded != "")
+
+		decoded, err := sessionProvider.Codec.Decode(encoded)
+		assert.Check(t, err)
+		decodedClaims := decoded.(JWTSessionClaims)
+		assert.Equal(t, tc.Subject, decodedClaims.Subject)
+	})
+
+	t.Run("TrackedRequestEncodeDecode", func(t *testing.T) {
+		codec := middleware.RequestTracker.(CookieRequestTracker).Codec
+		trackedReq := TrackedRequest{
+			Index:         "test-index",
+			SAMLRequestID: "test-request-id",
+			URI:           "/test-uri",
+		}
+
+		encoded, err := codec.Encode(trackedReq)
+		assert.Check(t, err)
+		assert.Assert(t, encoded != "")
+
+		decoded, err := codec.Decode(encoded)
+		assert.Check(t, err)
+		assert.Equal(t, trackedReq.Index, decoded.Index)
+		assert.Equal(t, trackedReq.SAMLRequestID, decoded.SAMLRequestID)
+	})
+
+	t.Run("RequireAccountFlow", func(t *testing.T) {
+		handler := middleware.RequireAccount(
+			http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
+				panic("not reached")
+			}))
+
+		req, _ := http.NewRequest("GET", "/protected", nil)
+		resp := httptest.NewRecorder()
+		handler.ServeHTTP(resp, req)
+
+		assert.Check(t, is.Equal(http.StatusFound, resp.Code))
+		assert.Assert(t, resp.Header().Get("Location") != "")
+		assert.Assert(t, resp.Header().Get("Set-Cookie") != "")
+	})
+
+	t.Run("Metadata", func(t *testing.T) {
+		req, _ := http.NewRequest("GET", "/saml2/metadata", nil)
+		resp := httptest.NewRecorder()
+		middleware.ServeHTTP(resp, req)
+
+		assert.Check(t, is.Equal(http.StatusOK, resp.Code))
+		assert.Check(t, is.Equal("application/samlmetadata+xml",
+			resp.Header().Get("Content-type")))
+		golden.Assert(t, resp.Body.String(), "expected_middleware_metadata.xml")
+	})
+}

samlsp/new.go

@@ -149,13 +149,16 @@ func DefaultServiceProvider(opts Options) saml.ServiceProvider {
 }
 
 func defaultSigningMethodForKey(key crypto.Signer) string {
-	switch key.(type) {
-	case *rsa.PrivateKey:
+	if key == nil {
+		return ""
+	}
+	// Check public key type to support crypto.Signer implementations (KMS/HSM)
+	// that aren't concrete *rsa.PrivateKey or *ecdsa.PrivateKey types
+	switch key.Public().(type) {
+	case *rsa.PublicKey:
 		return dsig.RSASHA1SignatureMethod
-	case *ecdsa.PrivateKey:
+	case *ecdsa.PublicKey:
 		return dsig.ECDSASHA256SignatureMethod
-	case nil:
-		return ""
 	default:
 		panic(fmt.Sprintf("programming error: unsupported key type %T", key))
 	}

samlsp/request_tracker_jwt.go

@@ -2,6 +2,9 @@ package samlsp
 
 import (
 	"crypto"
+	"crypto/ecdsa"
+	"crypto/ed25519"
+	"crypto/rsa"
 	"fmt"
 	"time"
 
@@ -44,7 +47,18 @@ func (s JWTTrackedRequestCodec) Encode(value TrackedRequest) (string, error) {
 		SAMLAuthnRequest: true,
 	}
 	token := jwt.NewWithClaims(s.SigningMethod, claims)
-	return token.SignedString(s.Key)
+
+	// Check if key is a concrete private key type that jwt library can handle directly.
+	// For crypto.Signer implementations (KMS/HSM), use custom signing.
+	switch s.Key.(type) {
+	case *rsa.PrivateKey, *ecdsa.PrivateKey, ed25519.PrivateKey:
+		return token.SignedString(s.Key)
+	default:
+		if s.Key == nil {
+			return "", fmt.Errorf("signing key is nil")
+		}
+		return signJWTWithCryptoSigner(token, s.Key, s.SigningMethod)
+	}
 }
 
 // Decode returns a Tracked request from an encoded string.

samlsp/session_jwt.go

@@ -2,7 +2,16 @@ package samlsp
 
 import (
 	"crypto"
+	"crypto/ecdsa"
+	"crypto/ed25519"
+	"crypto/rand"
+	"crypto/rsa"
+	"encoding/asn1"
+	"encoding/base64"
 	"errors"
+	"fmt"
+	"math/big"
+	"strings"
 	"time"
 
 	"github.com/golang-jwt/jwt/v5"
@@ -77,12 +86,18 @@ func (c JWTSessionCodec) Encode(s Session) (string, error) {
 	claims := s.(JWTSessionClaims) // this will panic if you pass the wrong kind of session
 
 	token := jwt.NewWithClaims(c.SigningMethod, claims)
-	signedString, err := token.SignedString(c.Key)
-	if err != nil {
-		return "", err
-	}
 
-	return signedString, nil
+	// Check if key is a concrete private key type that jwt library can handle directly.
+	// For crypto.Signer implementations (KMS/HSM), use custom signing.
+	switch c.Key.(type) {
+	case *rsa.PrivateKey, *ecdsa.PrivateKey, ed25519.PrivateKey:
+		return token.SignedString(c.Key)
+	default:
+		if c.Key == nil {
+			return "", fmt.Errorf("signing key is nil")
+		}
+		return signJWTWithCryptoSigner(token, c.Key, c.SigningMethod)
+	}
 }
 
 // Decode parses the serialized session that may have been returned by Encode
@@ -137,3 +152,74 @@ func (a Attributes) Get(key string) string {
 	}
 	return v[0]
 }
+
+// signJWTWithCryptoSigner signs a JWT token using the crypto.Signer interface.
+// This allows KMS/HSM keys that implement crypto.Signer to sign JWTs.
+func signJWTWithCryptoSigner(token *jwt.Token, signer crypto.Signer, method jwt.SigningMethod) (string, error) {
+	// Get the signing string (header.payload)
+	signingString, err := token.SigningString()
+	if err != nil {
+		return "", err
+	}
+
+	// Determine hash algorithm based on signing method
+	var hashFunc crypto.Hash
+	switch method.Alg() {
+	case "RS256", "ES256", "PS256":
+		hashFunc = crypto.SHA256
+	case "RS384", "ES384", "PS384":
+		hashFunc = crypto.SHA384
+	case "RS512", "ES512", "PS512":
+		hashFunc = crypto.SHA512
+	default:
+		return "", fmt.Errorf("unsupported signing algorithm: %s", method.Alg())
+	}
+
+	// Hash the signing string
+	hasher := hashFunc.New()
+	hasher.Write([]byte(signingString))
+	digest := hasher.Sum(nil)
+
+	// Sign using crypto.Signer
+	sig, err := signer.Sign(rand.Reader, digest, hashFunc)
+	if err != nil {
+		return "", fmt.Errorf("signing with crypto.Signer: %w", err)
+	}
+
+	// For ECDSA, the signature from crypto.Signer is ASN.1 DER encoded,
+	// but JWT expects raw R||S format
+	if _, ok := signer.Public().(*ecdsa.PublicKey); ok {
+		sig, err = convertECDSASignatureToJWT(sig, signer.Public().(*ecdsa.PublicKey))
+		if err != nil {
+			return "", err
+		}
+	}
+
+	// Encode signature and return complete JWT
+	return strings.Join([]string{signingString, base64.RawURLEncoding.EncodeToString(sig)}, "."), nil
+}
+
+// convertECDSASignatureToJWT converts ASN.1 DER encoded ECDSA signature to JWT format (R||S)
+func convertECDSASignatureToJWT(derSig []byte, pubKey *ecdsa.PublicKey) ([]byte, error) {
+	// Parse ASN.1 DER signature
+	var sig struct {
+		R, S *big.Int
+	}
+	if _, err := asn1.Unmarshal(derSig, &sig); err != nil {
+		return nil, fmt.Errorf("parsing ECDSA signature: %w", err)
+	}
+
+	// Calculate key size in bytes
+	keyBytes := (pubKey.Curve.Params().BitSize + 7) / 8
+
+	// Create R||S format
+	rBytes := sig.R.Bytes()
+	sBytes := sig.S.Bytes()
+
+	// Pad to key size
+	result := make([]byte, 2*keyBytes)
+	copy(result[keyBytes-len(rBytes):keyBytes], rBytes)
+	copy(result[2*keyBytes-len(sBytes):], sBytes)
+
+	return result, nil
+}

service_provider.go

@@ -567,21 +567,25 @@ func GetSigningContext(sp *ServiceProvider) (*dsig.SigningContext, error) {
 	// 	keyPair.Certificate = append(keyPair.Certificate, cert.Raw)
 	// }
 
+	// Validate that the key type matches the signature method.
+	// We check the public key type to support crypto.Signer implementations
+	// (like KMS/HSM signers) that aren't literal *rsa.PrivateKey or *ecdsa.PrivateKey.
+	pubKey := sp.Key.Public()
 	switch sp.SignatureMethod {
 	case dsig.RSASHA1SignatureMethod,
 		dsig.RSASHA256SignatureMethod,
 		dsig.RSASHA384SignatureMethod,
 		dsig.RSASHA512SignatureMethod:
-		if _, ok := sp.Key.(*rsa.PrivateKey); !ok {
-			return nil, fmt.Errorf("signature method %s requires a key of type rsa.PrivateKey, not %T", sp.SignatureMethod, sp.Key)
+		if _, ok := pubKey.(*rsa.PublicKey); !ok {
+			return nil, fmt.Errorf("signature method %s requires an RSA key, got %T", sp.SignatureMethod, pubKey)
 		}
 
 	case dsig.ECDSASHA1SignatureMethod,
 		dsig.ECDSASHA256SignatureMethod,
 		dsig.ECDSASHA384SignatureMethod,
 		dsig.ECDSASHA512SignatureMethod:
-		if _, ok := sp.Key.(*ecdsa.PrivateKey); !ok {
-			return nil, fmt.Errorf("signature method %s requires a key of type ecdsa.PrivateKey, not %T", sp.SignatureMethod, sp.Key)
+		if _, ok := pubKey.(*ecdsa.PublicKey); !ok {
+			return nil, fmt.Errorf("signature method %s requires a key of type *ecdsa.PrivateKey, not %T", sp.SignatureMethod, pubKey)
 		}
 	default:
 		return nil, fmt.Errorf("invalid signing method %s", sp.SignatureMethod)