From 47087d3d5ea9543fd8994e010e07c243c94434f3 Mon Sep 17 00:00:00 2001
From: Tom Strassner <tstrassner@mitre.org>
Date: Wed, 23 Oct 2024 17:24:57 -0500
Subject: [PATCH] Implement ID Token, JWKS endpoint

---
 .../dtr_smart_app_suite.rb                    |  4 +-
 lib/davinci_dtr_test_kit/mock_auth_server.rb  | 47 +++++++++++++++++--
 lib/davinci_dtr_test_kit/urls.rb              |  1 +
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/lib/davinci_dtr_test_kit/dtr_smart_app_suite.rb b/lib/davinci_dtr_test_kit/dtr_smart_app_suite.rb
index 060363c..d21ca0b 100644
--- a/lib/davinci_dtr_test_kit/dtr_smart_app_suite.rb
+++ b/lib/davinci_dtr_test_kit/dtr_smart_app_suite.rb
@@ -50,12 +50,14 @@ class DTRSmartAppSuite < Inferno::TestSuite
     end
 
     allow_cors QUESTIONNAIRE_PACKAGE_PATH, QUESTIONNAIRE_RESPONSE_PATH, FHIR_RESOURCE_PATH, FHIR_SEARCH_PATH,
-               EHR_AUTHORIZE_PATH, EHR_TOKEN_PATH
+               EHR_AUTHORIZE_PATH, EHR_TOKEN_PATH, JKWS_PATH
 
     route(:get, '/fhir/metadata', method(:metadata_handler))
 
     route(:get, SMART_CONFIG_PATH, method(:ehr_smart_config))
 
+    route(:get, JKWS_PATH, method(:auth_server_jwks))
+
     record_response_route :get, EHR_AUTHORIZE_PATH, 'dtr_smart_app_ehr_authorize', method(:ehr_authorize),
                           resumes: ->(_) { false } do |request|
       DTRSmartAppSuite.extract_client_id_from_query_params(request)
diff --git a/lib/davinci_dtr_test_kit/mock_auth_server.rb b/lib/davinci_dtr_test_kit/mock_auth_server.rb
index 0bb0591..528efe0 100644
--- a/lib/davinci_dtr_test_kit/mock_auth_server.rb
+++ b/lib/davinci_dtr_test_kit/mock_auth_server.rb
@@ -1,7 +1,29 @@
+# frozen_string_literal: true
+
 require_relative 'urls'
 
 module DaVinciDTRTestKit
   module MockAuthServer
+    AUTHORIZED_PRACTITIONER_ID = 'pra1234' # Must exist on the FHIR_REFERENCE_SERVER (env var)
+
+    RSA_PRIVATE_KEY = OpenSSL::PKey::RSA.generate(2048)
+    RSA_PUBLIC_KEY = RSA_PRIVATE_KEY.public_key
+
+    def auth_server_jwks(_env)
+      response_body = {
+        keys: [
+          {
+            alg: 'RSA',
+            mod: Base64.urlsafe_encode64(RSA_PUBLIC_KEY.n.to_s(2), padding: false),
+            exp: Base64.urlsafe_encode64(RSA_PUBLIC_KEY.e.to_s(2), padding: false),
+            use: 'sig'
+          }
+        ]
+      }.to_json
+
+      [200, { 'Content-Type' => 'application/json', 'Access-Control-Allow-Origin' => '*' }, [response_body]]
+    end
+
     def ehr_smart_config(env)
       protocol = env['rack.url_scheme']
       host = env['HTTP_HOST']
@@ -10,12 +32,13 @@ def ehr_smart_config(env)
       base_url = "#{protocol}://#{host + path}"
       response_body =
         {
+          jwks_uri: base_url + JKWS_PATH,
           authorization_endpoint: base_url + EHR_AUTHORIZE_PATH,
           token_endpoint: base_url + EHR_TOKEN_PATH,
           token_endpoint_auth_methods_supported: ['private_key_jwt'],
           token_endpoint_auth_signing_alg_values_supported: ['RS256'],
           grant_types_supported: ['authorization_code'],
-          scopes_supported: ['launch', 'patient/*.rs', 'user/*.rs', 'offline_access'],
+          scopes_supported: ['launch', 'patient/*.rs', 'user/*.rs', 'offline_access', 'openid', 'fhirUser'],
           response_types_supported: ['code'],
           code_challenge_methods_supported: ['S256'],
           capabilities: [
@@ -54,10 +77,26 @@ def ehr_authorize(request, _test = nil, _test_result = nil)
 
     def ehr_token_response(request, _test = nil, test_result = nil)
       client_id = extract_client_id_from_token_request(request)
-      token = JWT.encode({ inferno_client_id: client_id }, nil, 'none')
-      response = { access_token: token, token_type: 'bearer', expires_in: 3600 }
-      test_input = JSON.parse(test_result.input_json)
+      access_token = JWT.encode({ inferno_client_id: client_id }, nil, 'none')
+
+      # No point in mocking an identity provider, just always use known Practitioner as the authorized user
+      suite_base_url = request.url.split(EHR_TOKEN_PATH).first
+      id_token = JWT.encode(
+        {
+          fhirUser: "#{suite_base_url}/fhir/Practitioner/#{AUTHORIZED_PRACTITIONER_ID}",
+          iss: suite_base_url,
+          sub: AUTHORIZED_PRACTITIONER_ID,
+          aud: client_id,
+          exp: Time.now.to_i + (24 * 60 * 60), # 24 hrs
+          iat: Time.now.to_i
+        },
+        RSA_PRIVATE_KEY,
+        'RS256'
+      )
+
+      response = { access_token:, id_token:, token_type: 'bearer', expires_in: 3600 }
 
+      test_input = JSON.parse(test_result.input_json)
       fhir_context_input = test_input.find { |input| input['name'] == 'smart_fhir_context' }
       fhir_context_input_value = fhir_context_input['value'] if fhir_context_input.present?
       begin
diff --git a/lib/davinci_dtr_test_kit/urls.rb b/lib/davinci_dtr_test_kit/urls.rb
index a42f75a..33583b0 100644
--- a/lib/davinci_dtr_test_kit/urls.rb
+++ b/lib/davinci_dtr_test_kit/urls.rb
@@ -2,6 +2,7 @@
 
 module DaVinciDTRTestKit
   SMART_CONFIG_PATH = '/fhir/.well-known/smart-configuration'
+  JKWS_PATH = '/fhir/.well-known/jwks.json'
   EHR_AUTHORIZE_PATH = '/mock_ehr_auth/authorize'
   EHR_TOKEN_PATH = '/mock_ehr_auth/token'
   PAYER_TOKEN_PATH = '/mock_payer_auth/token'