KMS implementation for "RS256" Algorithm, how to keep expected header?

I'm trying to add support for KMS while keeping the header {"alg": "RS256"}. We found a workaround but would like some opinion on how to make it easier to follow.

Expected Result

To be able to add a custom algorithm that defines the expected header.

Allow us to override the header when calling jwt: (you may have a better idea)

jwt.register_algorithm("KMS-RS256", AWSKMSAlgorithm())
jwt_token = jwt.encode(payload, key=key_id, algorithm="KMS-RS256", header={"alg":"RS256"})
# jwt_token is created using `AWSKMSAlgorithm` and header has `{"alg":"RS256"}`
decoded = jwt.decode(jwt_token, key=key_id, algorithms=["KMS-RS256"], header={"alg":"RS256"})
# jwt_token is decoded using `AWSKMSAlgorithm` and header is validated with `{"alg":"RS256"}`

With the current implementation, there's no option to enforce a header, when the header "alg" is set, the parameter algorithm is ignored. When the header is not set, algorithm is used as "alg" in the header.

I'm happy to create a PR allowing algorithm to be different than header: {"alg"} when both are defined.

Actual Result

The header "alg" is enforced to the same algorithm used to load the expected implementation, with no way to change it.

A workaround was to unregister the default RS256 algorithm, register my custom algorithm, decode/encode, and then undo my change. This is less than ideal as we are changing a global state, which could result in race conditions.

        jwt.unregister_algorithm("RS256")
        jwt.register_algorithm("RS256", AWSKMSAlgorithm())
        jwt_token = jwt.encode(payload, key=key_id, algorithm="RS256")
        jwt.unregister_algorithm("RS256")
        jwt.register_algorithm("RS256", RSAAlgorithm(RSAAlgorithm.SHA256))

Reproduction Steps

import jwt
from jwt.algorithms import Algorithm
import boto3

class AWSKMSAlgorithm(Algorithm):
    def __init__(self, **kwargs):
        pass

    def prepare_key(self, key):
        return key

    def sign(self, msg, key):
        kms = boto3.client("kms")

        res = kms.sign(
            KeyId=key,
            Message=msg,
            SigningAlgorithm="RSASSA_PKCS1_V1_5_SHA_256",
        )

        signature = res["Signature"]
        return signature

    def verify(self, msg, key, sig):
        kms = boto3.client("kms")

        res = kms.verify(
            KeyId=key,
            Message=msg,
            Signature=sig,
            SigningAlgorithm="RSASSA_PKCS1_V1_5_SHA_256",
        )
        return res["SignatureValid"]

    @staticmethod
    def from_jwk(jwk):
        pass

    @staticmethod
    def to_jwk(key_obj):
        pass

jwt.register_algorithm("KMS-RS256", AWSKMSAlgorithm())
jwt_token = jwt.encode(payload, key=key, algorithm="KMS-RS256")
# jwt_token has header {"alg":"KMS-RS256"} as expected

jwt_token = jwt.encode(payload, key=key, algorithm="KMS-RS256", header={"alg":"RS256"})
# AWSKMSAlgorithm is not used, the default implementation of "RS256" is used instead

System Information

$ python -m jwt.help

{
  "cryptography": {
    "version": "3.4.7"
  },
  "implementation": {
    "name": "CPython",
    "version": "3.10.4"
  },
  "platform": {
    "release": "21.5.0",
    "system": "Darwin"
  },
  "pyjwt": {
    "version": "2.4.0"
  }
}

jpadilla / pyjwt