From ca5764c7d5ebbe260329fcac3d58f9c8951cf5d8 Mon Sep 17 00:00:00 2001
From: Jaedon Farrugia <jaedonlilydale@hotmail.com>
Date: Fri, 31 May 2024 20:36:58 +1000
Subject: [PATCH] Add support for common JWKS

Adds support for common JWKS, facilitating logins from non-business
emails without sacrificing domain verification support.
---
 .../microsoft_graph/domain_verifier.rb        | 21 +++++--
 .../microsoft_graph/domain_verifier_spec.rb   | 57 ++++++++++++++-----
 2 files changed, 61 insertions(+), 17 deletions(-)

diff --git a/lib/omniauth/microsoft_graph/domain_verifier.rb b/lib/omniauth/microsoft_graph/domain_verifier.rb
index 6cbdf03..4401a5e 100644
--- a/lib/omniauth/microsoft_graph/domain_verifier.rb
+++ b/lib/omniauth/microsoft_graph/domain_verifier.rb
@@ -9,6 +9,7 @@ module MicrosoftGraph
     # https://www.descope.com/blog/post/noauth
     # https://clerk.com/docs/authentication/social-connections/microsoft#stay-secure-against-the-n-o-auth-vulnerability
     OIDC_CONFIG_URL = 'https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration'
+    COMMON_JWKS_URL = 'https://login.microsoftonline.com/common/discovery/v2.0/keys'
 
     class DomainVerificationError < OmniAuth::Error; end
 
@@ -62,13 +63,25 @@ def verify!
       def domain_verified_jwt_claim
         oidc_config = access_token.get(OIDC_CONFIG_URL).parsed
         algorithms = oidc_config['id_token_signing_alg_values_supported']
-        keys = JWT::JWK::Set.new(access_token.get(oidc_config['jwks_uri']).parsed)
-        decoded_token = JWT.decode(id_token, nil, true, algorithms: algorithms, jwks: keys)
+        jwks = get_jwks(oidc_config)
+        decoded_token = JWT.decode(id_token, nil, true, algorithms: algorithms, jwks: jwks)
+        xms_edov_valid?(decoded_token)
+      rescue JWT::VerificationError, ::OAuth2::Error
+        false
+      end
+
+      def xms_edov_valid?(decoded_token)
         # https://github.com/MicrosoftDocs/azure-docs/issues/111425#issuecomment-1761043378
         # Comments seemed to indicate the value is not consistent
         ['1', 1, 'true', true].include?(decoded_token.first['xms_edov'])
-      rescue JWT::VerificationError, ::OAuth2::Error
-        false
+      end
+
+      def get_jwks(oidc_config)
+        # Depending on the tenant, the JWKS endpoint might be different. We need to
+        # consider both the JWKS from the OIDC configuration and the common JWKS endpoint.
+        oidc_config_jwk_keys = access_token.get(oidc_config['jwks_uri']).parsed[:keys]
+        common_jwk_keys = access_token.get(COMMON_JWKS_URL).parsed[:keys]
+        JWT::JWK::Set.new(oidc_config_jwk_keys + common_jwk_keys)
       end
 
       def verification_error_message
diff --git a/spec/omniauth/microsoft_graph/domain_verifier_spec.rb b/spec/omniauth/microsoft_graph/domain_verifier_spec.rb
index 0c5c94b..777695b 100644
--- a/spec/omniauth/microsoft_graph/domain_verifier_spec.rb
+++ b/spec/omniauth/microsoft_graph/domain_verifier_spec.rb
@@ -41,34 +41,65 @@
     end
 
     context 'when the ID token indicates domain verification' do
-      # Sign a fake ID token with our own local key
-      let(:mock_key) do
-        optional_parameters = { kid: 'mock-kid', use: 'sig', alg: 'RS256' }
+      let(:mock_oidc_key) do
+        optional_parameters = { kid: 'mock_oidc_key', use: 'sig', alg: 'RS256' }
         JWT::JWK.new(OpenSSL::PKey::RSA.new(2048), optional_parameters)
       end
-      let(:id_token) do
-        payload = { email: email, xms_edov: true }
-        JWT.encode(payload, mock_key.signing_key, mock_key[:alg], kid: mock_key[:kid])
+
+      let(:mock_common_key) do
+        optional_parameters = { kid: 'mock_common_key', use: 'sig', alg: 'RS256' }
+        JWT::JWK.new(OpenSSL::PKey::RSA.new(2048), optional_parameters)
       end
 
-      # Mock the API responses to return the local key
+      # Mock the API responses to return the mock keys
       before do
         allow(access_token).to receive(:get)
           .with(OmniAuth::MicrosoftGraph::OIDC_CONFIG_URL)
           .and_return(
-            double('OAuth2::Response', parsed: {
-              'id_token_signing_alg_values_supported' => ['RS256'],
-              'jwks_uri' => 'https://example.com/jwks-keys'
-            })
+            double(
+              'OAuth2::Response',
+              parsed: {
+                'id_token_signing_alg_values_supported' => ['RS256'],
+                'jwks_uri' => 'https://example.com/jwks-keys',
+              }
+            )
           )
         allow(access_token).to receive(:get)
           .with('https://example.com/jwks-keys')
           .and_return(
-            double('OAuth2::Response', parsed: JWT::JWK::Set.new(mock_key).export)
+            double(
+              'OAuth2::Response',
+              parsed: JWT::JWK::Set.new(mock_oidc_key).export
+            )
+          )
+        allow(access_token).to receive(:get)
+          .with(OmniAuth::MicrosoftGraph::COMMON_JWKS_URL)
+          .and_return(
+            double(
+              'OAuth2::Response',
+              parsed: JWT::JWK::Set.new(mock_common_key).export,
+              body: JWT::JWK::Set.new(mock_common_key).export.to_json
+            )
           )
       end
 
-      it { is_expected.to be_truthy }
+      context 'when the kid exists in the oidc key' do
+        let(:id_token) do
+          payload = { email: email, xms_edov: true }
+          JWT.encode(payload, mock_oidc_key.signing_key, mock_oidc_key[:alg], kid: mock_oidc_key[:kid])
+        end
+
+        it { is_expected.to be_truthy }
+      end
+
+      context "when the kid exists in the common key" do
+        let(:id_token) do
+          payload = { email: email, xms_edov: true }
+          JWT.encode(payload, mock_common_key.signing_key, mock_common_key[:alg], kid: mock_common_key[:kid])
+        end
+
+        it { is_expected.to be_truthy }
+      end
     end
 
     context 'when all verification strategies fail' do