feat: add RS256 signing keys backed by AWS KMS (#2571)

Adds support for RSA signing keys backed by AWS KMS, which are the
cheapest type of key.

You specify `aws:kms:arn` as a claim in the private key's JWK and it all
flows from there. It uses the ambient credentials of the process to talk
to KMS.

---------

Co-authored-by: Chris Stockton <chris.stockton@supabase.io>

Author: hf

go.mod

@@ -35,6 +35,21 @@ require (
 
 require (
 	github.com/ProjectZKM/Ziren/crates/go-runtime/zkvm_runtime v0.0.0-20251001021608-1fe7b43fc4d6 // indirect
+	github.com/aws/aws-sdk-go-v2 v1.41.7 // indirect
+	github.com/aws/aws-sdk-go-v2/config v1.32.18 // indirect
+	github.com/aws/aws-sdk-go-v2/credentials v1.19.17 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.23 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.23 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.23 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.24 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.9 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.23 // indirect
+	github.com/aws/aws-sdk-go-v2/service/kms v1.52.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/signin v1.0.11 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.30.17 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sts v1.42.1 // indirect
+	github.com/aws/smithy-go v1.25.1 // indirect
 	github.com/bits-and-blooms/bitset v1.20.0 // indirect
 	github.com/cenkalti/backoff/v5 v5.0.3 // indirect
 	github.com/consensys/gnark-crypto v0.18.1 // indirect

go.sum

@@ -17,6 +17,36 @@ github.com/ajstarks/svgo v0.0.0-20211024235047-1546f124cd8b h1:slYM766cy2nI3BwyR
 github.com/ajstarks/svgo v0.0.0-20211024235047-1546f124cd8b/go.mod h1:1KcenG0jGWcpt8ov532z81sp/kMMUG485J2InIOyADM=
 github.com/apapsch/go-jsonmerge/v2 v2.0.0 h1:axGnT1gRIfimI7gJifB699GoE/oq+F2MU7Dml6nw9rQ=
 github.com/apapsch/go-jsonmerge/v2 v2.0.0/go.mod h1:lvDnEdqiQrp0O42VQGgmlKpxL1AP2+08jFMw88y4klk=
+github.com/aws/aws-sdk-go-v2 v1.41.7 h1:DWpAJt66FmnnaRIOT/8ASTucrvuDPZASqhhLey6tLY8=
+github.com/aws/aws-sdk-go-v2 v1.41.7/go.mod h1:4LAfZOPHNVNQEckOACQx60Y8pSRjIkNZQz1w92xpMJc=
+github.com/aws/aws-sdk-go-v2/config v1.32.18 h1:Hcia46bxhGgF3BaSnG8nSNCWmqTK6bj9xN9/FJ3WK6Q=
+github.com/aws/aws-sdk-go-v2/config v1.32.18/go.mod h1:zEjCAYmxqDadH1WX8CdBvmLKhUEUVFgKRQG38zjDmrY=
+github.com/aws/aws-sdk-go-v2/credentials v1.19.17 h1:gP2nkGsS+KMvF/jfFz2Vv2qiiOqWKyPACSzPsqHgoW8=
+github.com/aws/aws-sdk-go-v2/credentials v1.19.17/go.mod h1:Bsew3S/moG5iT77giPj1q8wb/s0RE5/QfH+ASjYtuQc=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.23 h1:UuSfcORqNSz/ey3VPRS8TcVH2Ikf0/sC+Hdj400QI6U=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.23/go.mod h1:+G/OSGiOFnSOkYloKj/9M35s74LgVAdJBSD5lsFfqKg=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.23 h1:GpT/TrnBYuE5gan2cZbTtvP+JlHsutdmlV2YfEyNde0=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.23/go.mod h1:xYWD6BS9ywC5bS3sz9Xh04whO/hzK2plt2Zkyrp4JuA=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.23 h1:bpd8vxhlQi2r1hiueOw02f/duEPTMK59Q4QMAoTTtTo=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.23/go.mod h1:15DfR2nw+CRHIk0tqNyifu3G1YdAOy68RftkhMDDwYk=
+github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.24 h1:OQqn11BtaYv1WLUowvcA30MpzIu8Ti4pcLPIIyoKZrA=
+github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.24/go.mod h1:X5ZJyfwVrWA96GzPmUCWFQaEARPR7gCrpq2E92PJwAE=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.9 h1:FLudkZLt5ci0ozzgkVo8BJGwvqNaZbTWb3UcucAateA=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.9/go.mod h1:w7wZ/s9qK7c8g4al+UyoF1Sp/Z45UwMGcqIzLWVQHWk=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.23 h1:pbrxO/kuIwgEsOPLkaHu0O+m4fNgLU8B3vxQ+72jTPw=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.23/go.mod h1:/CMNUqoj46HpS3MNRDEDIwcgEnrtZlKRaHNaHxIFpNA=
+github.com/aws/aws-sdk-go-v2/service/kms v1.52.0 h1:QNtg+Mtj1zmepk568+UKBD5DFfqh+ESTUUqQT27JkQc=
+github.com/aws/aws-sdk-go-v2/service/kms v1.52.0/go.mod h1:Y0+uxvxz6ib4KktRdK0V4X45Vcs/JyYoz8H71pO8xeI=
+github.com/aws/aws-sdk-go-v2/service/signin v1.0.11 h1:TdJ+HdzOBhU8+iVAOGUTU63VXopcumCOF1paFulHWZc=
+github.com/aws/aws-sdk-go-v2/service/signin v1.0.11/go.mod h1:R82ZRExE/nheo0N+T8zHPcLRTcH8MGsnR3BiVGX0TwI=
+github.com/aws/aws-sdk-go-v2/service/sso v1.30.17 h1:7byT8HUWrgoRp6sXjxtZwgOKfhss5fW6SkLBtqzgRoE=
+github.com/aws/aws-sdk-go-v2/service/sso v1.30.17/go.mod h1:xNWknVi4Ezm1vg1QsB/5EWpAJURq22uqd38U8qKvOJc=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.0 h1:nDARhv/oF55bcxF7rCI/4PDxOKnVXVWwDuDwCs2I2SQ=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.36.0/go.mod h1:4vIRDq+CJB2xFAXZ+YgGUTiEft7oAQlhIs71xcSeuVg=
+github.com/aws/aws-sdk-go-v2/service/sts v1.42.1 h1:F/M5Y9I3nwr2IEpshZgh1GeHpOItExNM9L1euNuh/fk=
+github.com/aws/aws-sdk-go-v2/service/sts v1.42.1/go.mod h1:mTNxImtovCOEEuD65mKW7DCsL+2gjEH+RPEAexAzAio=
+github.com/aws/smithy-go v1.25.1 h1:J8ERsGSU7d+aCmdQur5Txg6bVoYelvQJgtZehD12GkI=
+github.com/aws/smithy-go v1.25.1/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc=
 github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
 github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
 github.com/badoux/checkmail v0.0.0-20170203135005-d0a759655d62 h1:vMqcPzLT1/mbYew0gM6EJy4/sCNy9lY9rmlFO+pPwhY=

hack/kms-rsa-to-jwk.js

@@ -0,0 +1,73 @@
+#!/usr/bin/env node
+
+import { execFileSync } from "node:child_process";
+import { webcrypto } from "node:crypto";
+
+const keyId = process.argv[2];
+const compact = process.argv[3] === '--compact';
+
+if (!keyId) {
+  console.error("Usage: kms-rsa-to-jwk.js <key-arn> [--compact]");
+  process.exit(1);
+}
+
+// arn:partition:kms:region:account:key/uuid
+const arnParts = keyId.split(":");
+if (arnParts.length < 6 || arnParts[2] !== "kms") {
+  throw new Error(`Invalid KMS ARN: ${keyId}`);
+}
+
+const region = arnParts[3];
+
+const publicKeyB64 = execFileSync(
+  "aws",
+  [
+    "kms",
+    "get-public-key",
+    "--key-id",
+    keyId,
+    "--query",
+    "PublicKey",
+    "--output",
+    "text",
+    "--region",
+    region,
+  ],
+  { encoding: "utf8" },
+).trim();
+
+const spki = Buffer.from(publicKeyB64, "base64");
+
+const key = await webcrypto.subtle.importKey(
+  "spki",
+  spki,
+  {
+    name: "RSASSA-PKCS1-v1_5",
+    hash: "SHA-256",
+  },
+  true,
+  ["verify"],
+);
+
+const jwk = await webcrypto.subtle.exportKey("jwk", key);
+
+console.log(
+  JSON.stringify(
+    {
+      ...jwk,
+      ext: undefined,
+      use: "sig",
+      key_ops: ['sign', 'verify'],
+      'aws:kms:arn': keyId,
+
+      //kty: jwk.kty,
+      //use: "sig",
+      //alg: "RS256",
+      //kid: keyId,
+      //n: jwk.n,
+      //e: jwk.e,
+    },
+    null,
+    compact ? 0 : 2,
+  ),
+);

internal/api/auth.go

@@ -82,7 +82,7 @@ func (a *API) parseJWTClaims(bearer string, r *http.Request) (context.Context, e
 	token, err := p.ParseWithClaims(bearer, &AccessTokenClaims{}, func(token *jwt.Token) (interface{}, error) {
 		if kid, ok := token.Header["kid"]; ok {
 			if kidStr, ok := kid.(string); ok {
-				key, err := conf.FindPublicKeyByKid(kidStr, &config.JWT)
+				key, err := conf.FindPublicKeyByKid(ctx, kidStr, &config.JWT)
 				if err != nil {
 					return nil, err
 				}

internal/api/auth_test.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -167,7 +168,7 @@ func (ts *AuthTestSuite) TestParseJWTClaims() {
 			jwk, err := conf.GetSigningJwk(&ts.Config.JWT)
 			require.NoError(ts.T(), err)
 			signingMethod := conf.GetSigningAlg(jwk)
-			signingKey, err := conf.GetSigningKey(jwk)
+			signingKey, err := ts.Config.JWT.SigningKey(context.Background())
 			require.NoError(ts.T(), err)
 
 			userJwtToken := jwt.NewWithClaims(signingMethod, userClaims)

internal/api/e2e_test.go

@@ -980,7 +980,7 @@ func TestE2EHooks(t *testing.T) {
 					) (any, error) {
 						if kid, ok := token.Header["kid"]; ok {
 							if kidStr, ok := kid.(string); ok {
-								return conf.FindPublicKeyByKid(kidStr, &globalCfg.JWT)
+								return conf.FindPublicKeyByKid(context.Background(), kidStr, &globalCfg.JWT)
 							}
 						}
 						if alg, ok := token.Header["alg"]; ok {
@@ -1055,7 +1055,7 @@ func TestE2EHooks(t *testing.T) {
 				func(token *jwt.Token) (any, error) {
 					if kid, ok := token.Header["kid"]; ok {
 						if kidStr, ok := kid.(string); ok {
-							return conf.FindPublicKeyByKid(kidStr, &globalCfg.JWT)
+							return conf.FindPublicKeyByKid(context.Background(), kidStr, &globalCfg.JWT)
 						}
 					}
 					if alg, ok := token.Header["alg"]; ok {

internal/api/oauthserver/handlers.go

@@ -435,7 +435,7 @@ func (s *Server) handleAuthorizationCodeGrant(ctx context.Context, w http.Respon
 			nonce = *authorization.Nonce
 		}
 
-		idToken, err := tokenService.GenerateIDToken(tokens.GenerateIDTokenParams{
+		idToken, err := tokenService.GenerateIDToken(ctx, tokens.GenerateIDTokenParams{
 			User:     user,
 			ClientID: client.ID,
 			Nonce:    nonce,

internal/conf/awskmsjwk/kms.go

@@ -0,0 +1,11 @@
+package awskmsjwk
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go-v2/service/kms"
+)
+
+type KMSAPI interface {
+	Sign(ctx context.Context, in *kms.SignInput, optFns ...func(*kms.Options)) (*kms.SignOutput, error)
+}

internal/conf/awskmsjwk/rs256.go

@@ -0,0 +1,64 @@
+package awskmsjwk
+
+import (
+	"context"
+	"crypto/sha256"
+	"errors"
+
+	"github.com/aws/aws-sdk-go-v2/service/kms"
+	kmstypes "github.com/aws/aws-sdk-go-v2/service/kms/types"
+	"github.com/golang-jwt/jwt/v5"
+	"github.com/sirupsen/logrus"
+)
+
+var ErrNotRS256Key = errors.New("awskmsjwk: key needs to be *RS256Key")
+
+type RS256Key struct {
+	Ctx context.Context
+	KMS KMSAPI
+
+	KeyID string
+	Raw   any
+}
+
+type signingMethodKMSRS256 struct{}
+
+var SigningMethodRS256KMS jwt.SigningMethod = &signingMethodKMSRS256{}
+
+func (m *signingMethodKMSRS256) Alg() string {
+	return jwt.SigningMethodRS256.Alg() // "RS256"
+}
+
+func (m *signingMethodKMSRS256) Sign(signingString string, key any) ([]byte, error) {
+	k, ok := key.(*RS256Key)
+	if !ok {
+		return nil, ErrNotRS256Key
+	}
+
+	// JWT RS256 signs SHA256(base64url(header) + "." + base64url(payload)).
+	// Use DIGEST so large JWTs do not hit KMS RAW message size limits.
+	digest := sha256.Sum256([]byte(signingString))
+
+	out, err := k.KMS.Sign(k.Ctx, &kms.SignInput{
+		KeyId:            &k.KeyID,
+		Message:          digest[:],
+		MessageType:      kmstypes.MessageTypeDigest,
+		SigningAlgorithm: kmstypes.SigningAlgorithmSpecRsassaPkcs1V15Sha256,
+	})
+	if err != nil {
+		logrus.WithError(err).Errorf("Unable to sign RS256 JWT with AWS KMS key %q", k.KeyID)
+
+		return nil, err
+	}
+
+	return out.Signature, nil
+}
+
+func (m *signingMethodKMSRS256) Verify(signingString string, sig []byte, key any) error {
+	k, ok := key.(*RS256Key)
+	if !ok {
+		return ErrNotRS256Key
+	}
+
+	return jwt.SigningMethodRS256.Verify(signingString, sig, k.Raw)
+}

internal/conf/configuration.go

@@ -2,6 +2,7 @@ package conf
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -141,6 +142,8 @@ type JWTConfiguration struct {
 	KeyID            string         `json:"key_id" split_words:"true"`
 	Keys             JwtKeysDecoder `json:"keys"`
 	ValidMethods     []string       `json:"-" split_words:"true"`
+
+	SigningKey func(context.Context) (any, error) `json:"-"`
 }
 
 type MFAFactorTypeConfiguration struct {
@@ -1043,13 +1046,22 @@ func (config *GlobalConfiguration) ApplyDefaults() error {
 		if err := config.applyDefaultsJWT([]byte(config.JWT.Secret)); err != nil {
 			return err
 		}
+	} else {
+		jwk, err := GetSigningJwk(&config.JWT)
+		if err != nil {
+			return err
+		}
+		sk, err := getSigningKey(jwk)
+		if err != nil {
+			return err
+		}
+		config.JWT.SigningKey = sk
 	}
 
 	if config.JWT.ValidMethods == nil {
 		config.JWT.ValidMethods = []string{}
 		for _, key := range config.JWT.Keys {
-			alg := GetSigningAlg(key.PublicKey)
-			config.JWT.ValidMethods = append(config.JWT.ValidMethods, alg.Alg())
+			config.JWT.ValidMethods = append(config.JWT.ValidMethods, key.PublicKey.Algorithm().String())
 		}
 
 	}
@@ -1193,6 +1205,16 @@ func (config *GlobalConfiguration) applyDefaultsJWTPrivateKey(privKey jwk.Key) e
 		PublicKey:  pubKey,
 		PrivateKey: privKey,
 	}
+
+	var key any
+	if err := privKey.Raw(&key); err != nil {
+		return err
+	}
+
+	config.JWT.SigningKey = func(ctx context.Context) (any, error) {
+		return key, nil
+	}
+
 	return nil
 }