From ca54111c2c137965a077418c329f99ebffe1c20a Mon Sep 17 00:00:00 2001
From: Tom Strassner <tstrassner@mitre.org>
Date: Fri, 20 Dec 2024 11:09:22 -0600
Subject: [PATCH] Overwrite OAuth URIs in metadata proxy response (#43)

---
 .../dtr_smart_app_suite.rb                     |  2 +-
 .../endpoints/mock_authorization.rb            | 12 ++++++++++--
 lib/davinci_dtr_test_kit/endpoints/mock_ehr.rb | 18 +++++++++++++++++-
 lib/davinci_dtr_test_kit/urls.rb               |  1 +
 4 files changed, 29 insertions(+), 4 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 7fb7ca3..7a0d4f3 100644
--- a/lib/davinci_dtr_test_kit/dtr_smart_app_suite.rb
+++ b/lib/davinci_dtr_test_kit/dtr_smart_app_suite.rb
@@ -67,7 +67,7 @@ class DTRSmartAppSuite < Inferno::TestSuite
     suite_endpoint :post, NEXT_PATH, MockPayer::NextQuestionEndpoint
 
     # EHR
-    route(:get, '/fhir/metadata', MockEHR.method(:metadata))
+    route(:get, METADATA_PATH, MockEHR.method(:metadata))
     suite_endpoint :post, QUESTIONNAIRE_RESPONSE_PATH, MockEHR::QuestionnaireResponseEndpoint
     suite_endpoint :get, FHIR_RESOURCE_PATH, MockEHR::FHIRGetEndpoint
     suite_endpoint :get, FHIR_SEARCH_PATH, MockEHR::FHIRGetEndpoint
diff --git a/lib/davinci_dtr_test_kit/endpoints/mock_authorization.rb b/lib/davinci_dtr_test_kit/endpoints/mock_authorization.rb
index 0a7469b..4460dc0 100644
--- a/lib/davinci_dtr_test_kit/endpoints/mock_authorization.rb
+++ b/lib/davinci_dtr_test_kit/endpoints/mock_authorization.rb
@@ -62,8 +62,8 @@ def ehr_openid_config(env)
       base_url = env_base_url(env, OPENID_CONFIG_PATH)
       response_body = {
         issuer: base_url + FHIR_BASE_PATH,
-        authorization_endpoint: base_url + EHR_AUTHORIZE_PATH,
-        token_endpoint: base_url + EHR_TOKEN_PATH,
+        authorization_endpoint: authorization_endpoint(base_url),
+        token_endpoint: token_endpoint(base_url),
         jwks_uri: base_url + JKWS_PATH,
         response_types_supported: ['id_token'],
         subject_types_supported: ['public'],
@@ -79,5 +79,13 @@ def env_base_url(env, endpoint_path)
       path.gsub!(%r{#{endpoint_path}(/)?}, '')
       "#{protocol}://#{host + path}"
     end
+
+    def authorization_endpoint(base_url)
+      base_url + EHR_AUTHORIZE_PATH
+    end
+
+    def token_endpoint(base_url)
+      base_url + EHR_TOKEN_PATH
+    end
   end
 end
diff --git a/lib/davinci_dtr_test_kit/endpoints/mock_ehr.rb b/lib/davinci_dtr_test_kit/endpoints/mock_ehr.rb
index abdfafc..21a45ad 100644
--- a/lib/davinci_dtr_test_kit/endpoints/mock_ehr.rb
+++ b/lib/davinci_dtr_test_kit/endpoints/mock_ehr.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative '../endpoints/mock_authorization'
+
 module DaVinciDTRTestKit
   module MockEHR
     RESOURCE_SERVER_BASE = ENV.fetch('FHIR_REFERENCE_SERVER')
@@ -13,9 +15,23 @@ def resource_server_client
       client
     end
 
-    def metadata(_env)
+    def metadata(env)
       cs = resource_server_client.capability_statement
       if cs.present?
+        # Overwrite the OAuth URIs returned by the reference server to point to the suite endpoints instead
+        oauth_uris_url = 'http://fhir-registry.smarthealthit.org/StructureDefinition/oauth-uris'
+        base_url = MockAuthorization.env_base_url(env, METADATA_PATH)
+        sec_ext = cs.rest.first&.security&.extension&.delete_if { |e| e.url == oauth_uris_url }
+        sec_ext&.push(
+          FHIR::Extension.new(
+            url: oauth_uris_url,
+            extension: [
+              FHIR::Extension.new(url: 'authorize', valueUri: MockAuthorization.authorization_endpoint(base_url)),
+              FHIR::Extension.new(url: 'token', valueUri: MockAuthorization.token_endpoint(base_url))
+            ]
+          )
+        )
+
         [200, { 'Content-Type' => 'application/fhir+json', 'Access-Control-Allow-Origin' => '*' }, [cs.to_json]]
       else
         [500, { 'Access-Control-Allow-Origin' => '*' }, ['Unexpected error occurred while fetching metadata']]
diff --git a/lib/davinci_dtr_test_kit/urls.rb b/lib/davinci_dtr_test_kit/urls.rb
index 04170af..aa36bf7 100644
--- a/lib/davinci_dtr_test_kit/urls.rb
+++ b/lib/davinci_dtr_test_kit/urls.rb
@@ -2,6 +2,7 @@
 
 module DaVinciDTRTestKit
   FHIR_BASE_PATH = '/fhir'
+  METADATA_PATH = "#{FHIR_BASE_PATH}/metadata".freeze
   SMART_CONFIG_PATH = "#{FHIR_BASE_PATH}/.well-known/smart-configuration".freeze
   OPENID_CONFIG_PATH = "#{FHIR_BASE_PATH}/.well-known/openid-configuration".freeze
   JKWS_PATH = "#{FHIR_BASE_PATH}/.well-known/jwks.json".freeze