Merge pull request #2 from anakinj/rsa-algos

feat: Add RSA algos

README.md

@@ -23,15 +23,17 @@ require `jwt/kms`
 # Create a key, for example with the ruby AWS SDK
 key = Aws::KMS::Client.new.create_key(key_spec: "HMAC_512", key_usage: "GENERATE_VERIFY_MAC")
 
-algo = ::JWT::KMS.by(key_id: key.key_metadata.key_id)
+algo = ::JWT::KMS.for(algorithm: "HS512")
 
-token = JWT.encode(payload, nil, algo)
-decoded_token = JWT.decode(token, "Not relevant", true, algorithm: algo)
+token = JWT.encode(payload, key.key_metadata.key_id, algo)
+decoded_token = JWT.decode(token, key.key_metadata.key_id, true, algorithm: algo)
 ```
 
 
 ## Development
 
+[Localstack](https://www.localstack.cloud/) can be used to simulate the AWS KMS environment.
+
 ```
 docker run \
   --rm -it \

lib/jwt/kms.rb

@@ -5,6 +5,7 @@
 
 require_relative "kms/version"
 require_relative "kms/hmac_key"
+require_relative "kms/sign_verify_key"
 
 module JWT
   # :nodoc:
@@ -13,19 +14,14 @@ def self.client
       @client ||= Aws::KMS::Client.new
     end
 
-    def self.by(key_id:)
-      from_description(KMS.client.describe_key(key_id: key_id))
-    end
-
-    def self.from_description(description)
-      case description.key_metadata.key_usage
-      when "GENERATE_VERIFY_MAC"
-        HmacKey.new(key_id: description.key_metadata.key_id, key_spec: description.key_metadata.key_spec)
-      when "SIGN_VERIFY"
-        SignVerifyKey.new(key_id: description.key_metadata.key_id, key_spec: description.key_metadata.key_spec)
+    def self.for(algorithm:)
+      if HmacKey::MAPPINGS.key?(algorithm)
+        HmacKey
+      elsif SignVerifyKey::MAPPINGS.key?(algorithm)
+        SignVerifyKey
       else
-        raise ArgumentError, "Keys with key_usage #{description.key_metadata.key_usage} not supported"
-      end
+        raise ArgumentError, "Algorithm #{algorithm} not supported"
+      end.new(algorithm: algorithm)
     end
   end
 end

lib/jwt/kms/hmac_key.rb

@@ -8,42 +8,30 @@ class HmacKey
       include JWT::JWA::SigningAlgorithm
 
       MAPPINGS = {
-        "HMAC_256" => { alg: "HS256", mac_algorithm: "HMAC_SHA_256" },
-        "HMAC_384" => { alg: "HS384", mac_algorithm: "HMAC_SHA_384" },
-        "HMAC_512" => { alg: "HS512", mac_algorithm: "HMAC_SHA_512" }
+        "HS256" => "HMAC_SHA_256",
+        "HS384" => "HMAC_SHA_384",
+        "HS512" => "HMAC_SHA_512"
       }.freeze
 
-      def initialize(key_id:, key_spec: nil)
-        @key_id = key_id
-        @key_spec = key_spec
+      def initialize(algorithm:)
+        @alg = algorithm
       end
 
-      def alg
-        MAPPINGS.dig(key_spec, :alg)
+      def sign(data:, signing_key:, **)
+        KMS.client.generate_mac(key_id: signing_key, mac_algorithm: mac_algorithm, message: data).mac
       end
 
-      def sign(data:, **)
-        KMS.client.generate_mac(key_id: key_id, mac_algorithm: mac_algorithm, message: data).mac
-      end
-
-      def verify(data:, signature:, **)
-        KMS.client.verify_mac(key_id: key_id, mac_algorithm: mac_algorithm, message: data, mac: signature).mac_valid
+      def verify(data:, verification_key:, signature:, **)
+        KMS.client.verify_mac(key_id: verification_key, mac_algorithm: mac_algorithm, message: data,
+                              mac: signature).mac_valid
       end
 
       private
 
       attr_reader :key_id
 
-      def key_spec
-        @key_spec ||= description.key_spec
-      end
-
       def mac_algorithm
-        MAPPINGS.dig(key_spec, :mac_algorithm)
-      end
-
-      def description
-        @description ||= KMS.client.describe_key(key_id: key_id)
+        MAPPINGS.fetch(alg, nil)
       end
     end
   end

lib/jwt/kms/sign_verify_key.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module JWT
+  module KMS
+    # Represent a AWS asymmetric key
+    # https://docs.aws.amazon.com/kms/latest/developerguide/symmetric-asymmetric.html
+    class SignVerifyKey
+      include JWT::JWA::SigningAlgorithm
+
+      MAPPINGS = {
+        "RS256" => "RSASSA_PKCS1_V1_5_SHA_256",
+        "RS384" => "RSASSA_PKCS1_V1_5_SHA_384",
+        "RS512" => "RSASSA_PKCS1_V1_5_SHA_512",
+        "PS256" => "RSASSA_PSS_SHA_256",
+        "PS384" => "RSASSA_PSS_SHA_384",
+        "PS512" => "RSASSA_PSS_SHA_512",
+        "ES256" => "ECDSA_SHA_256",
+        "ES384" => "ECDSA_SHA_384",
+        "ES512" => "ECDSA_SHA_512"
+      }.freeze
+
+      def initialize(algorithm:)
+        @alg = algorithm
+      end
+
+      def sign(data:, signing_key:, **)
+        KMS.client.sign(key_id: signing_key, signing_algorithm: signing_algorithm,
+                        message: data).signature
+      end
+
+      def verify(data:, verification_key:, signature:, **)
+        KMS.client.verify(key_id: verification_key, signing_algorithm: signing_algorithm,
+                          message: data, signature: signature).signature_valid
+      end
+
+      private
+
+      attr_reader :key_id
+
+      def signing_algorithm
+        MAPPINGS.fetch(alg, nil)
+      end
+    end
+  end
+end

spec/jwt/kms_spec.rb

@@ -1,40 +1,106 @@
 # frozen_string_literal: true
 
 RSpec.describe JWT::KMS do
-  let(:algo_instance) { described_class.by(key_id: key_id) }
   let(:payload) { { "pay" => "load" } }
 
-  describe "algorithm is given directly" do
+  describe ".for" do
+    subject(:algo_instance) { described_class.for(algorithm: key_algorithm) }
+
     let(:key_id) { kms_key.key_metadata.key_id }
+    let(:key_algorithm) { nil }
 
-    context "when id to HMAC_256 key is given" do
-      let(:kms_key) { Aws::KMS::Client.new.create_key(key_spec: "HMAC_256", key_usage: "GENERATE_VERIFY_MAC") }
+    shared_examples "a AWS KMS algorithm" do |key_spec, key_usage, key_algorithm|
+      let(:key_algorithm) { key_algorithm }
+      let(:kms_key) { Aws::KMS::Client.new.create_key(key_spec: key_spec, key_usage: key_usage) }
 
-      it "encodes and decodes" do
-        token = JWT.encode(payload, nil, algo_instance)
-        expect(JWT.decode(token, "Not relevant", true, algorithm: algo_instance))
-          .to eq([payload, { "alg" => "HS256" }])
+      it "encodes and decodes as #{key_algorithm}" do
+        token = JWT.encode(payload, key_id, algo_instance)
+        expect(JWT.decode(token, key_id, true, algorithm: algo_instance))
+          .to eq([payload, { "alg" => key_algorithm }])
       end
     end
 
-    context "when id to HMAC_384 key is given" do
-      let(:kms_key) { Aws::KMS::Client.new.create_key(key_spec: "HMAC_384", key_usage: "GENERATE_VERIFY_MAC") }
+    context "when key_id is referring a HMAC_256 key" do
+      it_behaves_like "a AWS KMS algorithm", "HMAC_256", "GENERATE_VERIFY_MAC", "HS256"
+    end
+
+    context "when key_id is referring a HMAC_384 key" do
+      it_behaves_like "a AWS KMS algorithm", "HMAC_384", "GENERATE_VERIFY_MAC", "HS384"
+    end
+
+    context "when key_id is referring a HMAC_512 key" do
+      it_behaves_like "a AWS KMS algorithm", "HMAC_512", "GENERATE_VERIFY_MAC", "HS512"
+    end
+
+    context "when key_id is referring a RSA_2048 key" do
+      it_behaves_like "a AWS KMS algorithm", "RSA_2048", "SIGN_VERIFY", "RS256"
+      it_behaves_like "a AWS KMS algorithm", "RSA_2048", "SIGN_VERIFY", "RS384"
+      it_behaves_like "a AWS KMS algorithm", "RSA_2048", "SIGN_VERIFY", "RS512"
+
+      it_behaves_like "a AWS KMS algorithm", "RSA_2048", "SIGN_VERIFY", "PS256"
+      it_behaves_like "a AWS KMS algorithm", "RSA_2048", "SIGN_VERIFY", "PS384"
+      it_behaves_like "a AWS KMS algorithm", "RSA_2048", "SIGN_VERIFY", "PS512"
+    end
+
+    context "when key_id is referring a RSA_3072 key" do
+      it_behaves_like "a AWS KMS algorithm", "RSA_3072", "SIGN_VERIFY", "RS256"
+      it_behaves_like "a AWS KMS algorithm", "RSA_3072", "SIGN_VERIFY", "RS384"
+      it_behaves_like "a AWS KMS algorithm", "RSA_3072", "SIGN_VERIFY", "RS512"
+
+      it_behaves_like "a AWS KMS algorithm", "RSA_3072", "SIGN_VERIFY", "PS256"
+      it_behaves_like "a AWS KMS algorithm", "RSA_3072", "SIGN_VERIFY", "PS384"
+      it_behaves_like "a AWS KMS algorithm", "RSA_3072", "SIGN_VERIFY", "PS512"
+    end
+
+    context "when key_id is referring a RSA_4096 key" do
+      it_behaves_like "a AWS KMS algorithm", "RSA_4096", "SIGN_VERIFY", "RS256"
+      it_behaves_like "a AWS KMS algorithm", "RSA_4096", "SIGN_VERIFY", "RS384"
+      it_behaves_like "a AWS KMS algorithm", "RSA_4096", "SIGN_VERIFY", "RS512"
+
+      it_behaves_like "a AWS KMS algorithm", "RSA_4096", "SIGN_VERIFY", "PS256"
+      it_behaves_like "a AWS KMS algorithm", "RSA_4096", "SIGN_VERIFY", "PS384"
+      it_behaves_like "a AWS KMS algorithm", "RSA_4096", "SIGN_VERIFY", "PS512"
+    end
+
+    context "when key_id is referring a ECC_NIST_P256 key" do
+      it_behaves_like "a AWS KMS algorithm", "ECC_NIST_P256", "SIGN_VERIFY", "ES256"
+    end
+
+    context "when key_id is referring a ECC_NIST_P384 key" do
+      it_behaves_like "a AWS KMS algorithm", "ECC_NIST_P384", "SIGN_VERIFY", "ES384"
+    end
+
+    context "when key_id is referring a ECC_NIST_P521 key" do
+      it_behaves_like "a AWS KMS algorithm", "ECC_NIST_P521", "SIGN_VERIFY", "ES512"
+    end
+
+    context "when algorithm is not supported" do
+      let(:key_algorithm) { "HS666" }
 
-      it "encodes and decodes" do
-        token = JWT.encode(payload, nil, algo_instance)
-        expect(JWT.decode(token, "Not relevant", true, algorithm: algo_instance))
-          .to eq([payload, { "alg" => "HS384" }])
+      it "raises an ArgumentError" do
+        expect { algo_instance }.to raise_error(ArgumentError)
       end
     end
 
-    context "when id to HMAC_512 key is given" do
-      let(:kms_key) { Aws::KMS::Client.new.create_key(key_spec: "HMAC_512", key_usage: "GENERATE_VERIFY_MAC") }
+    context "when algorithm key is not found" do
+      let(:key_algorithm) { "HS256" }
 
-      it "encodes and decodes" do
-        token = JWT.encode(payload, nil, algo_instance)
-        expect(JWT.decode(token, "Not relevant", true, algorithm: algo_instance))
-          .to eq([payload, { "alg" => "HS512" }])
+      it "raises native AWS component error" do
+        expect { JWT.encode(payload, "not-found", algo_instance) }.to raise_error(Aws::KMS::Errors::NotFoundException)
       end
     end
   end
+
+  describe "readme example" do
+    it "works as documented" do
+      key = Aws::KMS::Client.new.create_key(key_spec: "HMAC_512", key_usage: "GENERATE_VERIFY_MAC")
+
+      algo = described_class.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)
+
+      expect(decoded_token).to eq([{ "pay" => "load" }, { "alg" => "HS512" }])
+    end
+  end
 end