feat: Allow replacing default algorithms with AWS KMS versions (#7)

* feat: Allow replacing default algorithms with AWS KMS versions

* Fix a little

README.md

@@ -15,9 +15,25 @@ And require the gem in your code.
 ```ruby
 require `jwt-aws-kms`
 ```
+## Supported algorithms
+
+The gem supports the following AWS KMS algorithms:
+
+| Algorithm Name | Description                                      | JWA Name |
+|----------------|--------------------------------------------------|-------------------------|
+| RSASSA_PKCS1_V1_5_SHA_256 | RSASSA PKCS1 v1.5 using SHA-256       | RS256                   |
+| RSASSA_PKCS1_V1_5_SHA_384 | RSASSA PKCS1 v1.5 using SHA-384       | RS384                   |
+| RSASSA_PKCS1_V1_5_SHA_512 | RSASSA PKCS1 v1.5 using SHA-512       | RS512                   |
+| RSASSA_PSS_SHA_256 | RSASSA PSS using SHA-256                     | PS256                   |
+| RSASSA_PSS_SHA_384 | RSASSA PSS using SHA-384                     | PS384                   |
+| RSASSA_PSS_SHA_512 | RSASSA PSS using SHA-512                     | PS512                   |
+| ECDSA_SHA_256 | ECDSA using P-256 and SHA-256                     | ES256                   |
+| ECDSA_SHA_384 | ECDSA using P-384 and SHA-384                     | ES384                   |
+| ECDSA_SHA_512 | ECDSA using P-521 and SHA-512                     | ES512                   |
 
 ## Usage
 
+### Basic usage
 ```ruby
 
 # Create a key, for example with the ruby AWS SDK
@@ -28,22 +44,15 @@ algo = ::JWT::Aws::KMS.for(algorithm: "HS512")
 token = JWT.encode(payload, key.key_metadata.key_id, algo)
 decoded_token = JWT.decode(token, key.key_metadata.key_id, true, algorithm: algo)
 ```
+### Replace default algorithms
 
-## Supported algorithms
+You can swap the default algorithms in the JWT gem to AWS backed ones by calling `::JWT::Aws::KMS.replace_defaults!`.
 
-The gem supports the following AWS KMS algorithms:
+```ruby
+::JWT::Aws::KMS.replace_defaults! # Called in a initializer of some kind
 
-| Algorithm Name | Description                                      | JWA Name |
-|----------------|--------------------------------------------------|-------------------------|
-| RSASSA_PKCS1_V1_5_SHA_256 | RSASSA PKCS1 v1.5 using SHA-256       | RS256                   |
-| RSASSA_PKCS1_V1_5_SHA_384 | RSASSA PKCS1 v1.5 using SHA-384       | RS384                   |
-| RSASSA_PKCS1_V1_5_SHA_512 | RSASSA PKCS1 v1.5 using SHA-512       | RS512                   |
-| RSASSA_PSS_SHA_256 | RSASSA PSS using SHA-256                     | PS256                   |
-| RSASSA_PSS_SHA_384 | RSASSA PSS using SHA-384                     | PS384                   |
-| RSASSA_PSS_SHA_512 | RSASSA PSS using SHA-512                     | PS512                   |
-| ECDSA_SHA_256 | ECDSA using P-256 and SHA-256                     | ES256                   |
-| ECDSA_SHA_384 | ECDSA using P-384 and SHA-384                     | ES384                   |
-| ECDSA_SHA_512 | ECDSA using P-521 and SHA-512                     | ES512                   |
+token = JWT.encode(payload, "e25c502b-a383-44ac-a778-0d97e8688cb7", "HS512") # Encode payload with KMS key e25c502b-a383-44ac-a778-0d97e8688cb7
+```
 
 ## Development
 

lib/jwt/aws/kms.rb

@@ -25,6 +25,14 @@ def self.for(algorithm:)
           raise ArgumentError, "Algorithm #{algorithm} not supported"
         end.new(algorithm: algorithm)
       end
+
+      def self.replace_defaults!
+        [HmacKey, SignVerifyKey].each do |type|
+          type::MAPPINGS.each_key do |algorithm|
+            type.register_algorithm(type.new(algorithm: algorithm))
+          end
+        end
+      end
     end
   end
 end

spec/jwt/aws/kms_spec.rb

@@ -103,4 +103,21 @@
       expect(decoded_token).to eq([{ "pay" => "load" }, { "alg" => "HS512" }])
     end
   end
+
+  describe ".replace_defaults!" do
+    before do
+      described_class.replace_defaults!
+    end
+
+    it "replaces the default algorithms with AWS KMS backed ones" do
+      expect(JWT::JWA.resolve("RS512")).to be_a(JWT::Aws::KMS::SignVerifyKey)
+    end
+
+    it "allows utilizing the AWS KMS key using the algo name" do
+      key = Aws::KMS::Client.new.create_key(key_spec: "HMAC_512", key_usage: "GENERATE_VERIFY_MAC")
+      token = JWT.encode(payload, key.key_metadata.key_id, "HS512")
+      decoded_token = JWT.decode(token, key.key_metadata.key_id, true, algorithms: "HS512")
+      expect(decoded_token).to eq([{ "pay" => "load" }, { "alg" => "HS512" }])
+    end
+  end
 end