From 2cb368027951d78fc4181a856ee7aea10e48b297 Mon Sep 17 00:00:00 2001
From: Mina Mikhail <mina@fightingtheboss.com>
Date: Fri, 25 Nov 2016 00:35:23 -0500
Subject: [PATCH] Add ability to decode tokens via JWKs

---
 app/model/knock/auth_token.rb       | 31 +++++++++++++++++++++++++++--
 knock.gemspec                       |  1 +
 test/model/knock/auth_token_test.rb | 19 ++++++++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/app/model/knock/auth_token.rb b/app/model/knock/auth_token.rb
index 6b85d67..003c1b6 100644
--- a/app/model/knock/auth_token.rb
+++ b/app/model/knock/auth_token.rb
@@ -1,4 +1,7 @@
+require 'uri'
 require 'jwt'
+require 'json/jwt'
+require 'net/http'
 
 module Knock
   class AuthToken
@@ -7,7 +10,20 @@ class AuthToken
 
     def initialize payload: {}, token: nil, verify_options: {}
       if token.present?
-        @payload, _ = JWT.decode token.to_s, decode_key, true, options.merge(verify_options)
+        token_decode_key = decode_key
+
+        if token_decode_key.is_a?(JSON::JWK::Set)
+          @payload = JSON::JWT.decode(token.to_s, token_decode_key)
+
+          options.merge(verify_options).each do |key, val|
+            next unless key.to_s =~ /verify/
+
+            JWT::Verify.send(key, payload, @options) if val
+          end
+        else
+          @payload, _ = JWT.decode token.to_s, token_decode_key, true, options.merge(verify_options)
+        end
+
         @token = token
       else
         @payload = claims.merge(payload)
@@ -35,7 +51,18 @@ def secret_key
     end
 
     def decode_key
-      Knock.token_public_key || secret_key
+      if Knock.token_public_key
+        if Knock.token_public_key =~ /^#{URI::Parser.new.make_regexp(['http', 'https'])}$/
+          # This means there's a JWK or JWKs public key that we need to fetch
+          JSON::JWK::Set.new(
+            JSON.parse( Net::HTTP.get( URI(Knock.token_public_key) ) )
+          )
+        else
+          Knock.token_public_key
+        end
+      else
+        secret_key
+      end
     end
 
     def options
diff --git a/knock.gemspec b/knock.gemspec
index da76724..a8d7ce4 100644
--- a/knock.gemspec
+++ b/knock.gemspec
@@ -24,4 +24,5 @@ Gem::Specification.new do |s|
 
   s.add_development_dependency "sqlite3", "~> 1.3"
   s.add_development_dependency "timecop", "~> 0.8.0"
+  s.add_development_dependency "webmock", "~> 2.1"
 end
diff --git a/test/model/knock/auth_token_test.rb b/test/model/knock/auth_token_test.rb
index bd0730e..1a43038 100644
--- a/test/model/knock/auth_token_test.rb
+++ b/test/model/knock/auth_token_test.rb
@@ -1,6 +1,10 @@
 require 'test_helper'
 require 'jwt'
 require 'timecop'
+require 'webmock/minitest'
+
+# Disable all remote connections
+WebMock.disable_net_connect!
 
 module Knock
   class AuthTokenTest < ActiveSupport::TestCase
@@ -27,6 +31,21 @@ class AuthTokenTest < ActiveSupport::TestCase
       assert_nothing_raised { AuthToken.new(token: token) }
     end
 
+    test "decode RSA encoded tokens with JWKs from URL" do
+      rsa_private = OpenSSL::PKey::RSA.generate 2048
+      rsa_public = rsa_private.public_key
+
+      Knock.token_public_key = 'https://example.com/.well-known/jwks.json'
+      Knock.token_signature_algorithm = 'RS256'
+
+      stub_request(:get, Knock.token_public_key)
+        .to_return(body: JSON::JWK::Set.new(JSON::JWK.new(rsa_public)).to_json)
+
+      token = JSON::JWT.new({sub: "1"}).sign(JSON::JWK.new(rsa_private)).to_s
+
+      assert_nothing_raised { AuthToken.new(token: token) }
+    end
+
     test "encode tokens with RSA" do
       rsa_private = OpenSSL::PKey::RSA.generate 2048
       Knock.token_secret_signature_key = -> { rsa_private }