Add token exchange concern

Gemfile.lock

@@ -256,6 +256,7 @@ GEM
 PLATFORMS
   arm64-darwin-21
   arm64-darwin-22
+  arm64-darwin-23
   x86_64-darwin-19
   x86_64-darwin-20
   x86_64-darwin-21

app/controllers/concerns/shopify_app/ensure_has_session.rb

@@ -6,14 +6,20 @@ module EnsureHasSession
 
     included do
       include ShopifyApp::Localization
-      include ShopifyApp::LoginProtection
+
+      if ShopifyApp.configuration.use_new_embedded_auth_strategy?
+        include ShopifyApp::TokenExchange
+        around_action :activate_shopify_session
+      else
+        include ShopifyApp::LoginProtection
+        before_action :login_again_if_different_user_or_shop
+        around_action :activate_shopify_session
+        after_action :add_top_level_redirection_headers
+      end
+
       include ShopifyApp::CsrfProtection
       include ShopifyApp::EmbeddedApp
       include ShopifyApp::EnsureBilling
-
-      before_action :login_again_if_different_user_or_shop
-      around_action :activate_shopify_session
-      after_action :add_top_level_redirection_headers
     end
   end
 end

app/controllers/concerns/shopify_app/ensure_installed.rb

@@ -16,8 +16,13 @@ module EnsureInstalled
       end
 
       before_action :check_shop_domain
-      before_action :check_shop_known
-      before_action :validate_non_embedded_session
+
+      unless ShopifyApp.configuration.use_new_embedded_auth_strategy?
+        # TODO: Add support to use new embedded auth strategy here when invalid
+        # session token can be handled by AppBridge app reload
+        before_action :check_shop_known
+        before_action :validate_non_embedded_session
+      end
     end
 
     def current_shopify_domain

app/controllers/concerns/shopify_app/shop_access_scopes_verification.rb

@@ -6,7 +6,11 @@ module ShopAccessScopesVerification
     include ShopifyApp::RedirectForEmbedded
 
     included do
-      before_action :login_on_scope_changes
+      # Embedded auth strategy uses Shopify managed install to ensure latest access scopes,
+      # This will be handled automatically through token exchange
+      unless ShopifyApp.configuration.use_new_embedded_auth_strategy?
+        before_action :login_on_scope_changes
+      end
     end
 
     protected

lib/shopify_app.rb

@@ -52,6 +52,7 @@ def self.use_webpacker?
   require "shopify_app/controller_concerns/payload_verification"
   require "shopify_app/controller_concerns/app_proxy_verification"
   require "shopify_app/controller_concerns/webhook_verification"
+  require "shopify_app/controller_concerns/token_exchange"
 
   # jobs
   require "shopify_app/jobs/webhooks_manager_job"

lib/shopify_app/configuration.rb

@@ -126,6 +126,10 @@ def user_access_scopes
     def use_new_embedded_auth_strategy?
       wip_new_embedded_auth_strategy && embedded_app?
     end
+
+    def online_token_configured?
+      !ShopifyApp.configuration.user_session_repository.blank? && ShopifyApp::SessionRepository.user_storage.present?
+    end
   end
 
   class BillingConfiguration

lib/shopify_app/controller_concerns/login_protection.rb

@@ -263,7 +263,7 @@ def shop_session
     end
 
     def online_token_configured?
-      !ShopifyApp.configuration.user_session_repository.blank? && ShopifyApp::SessionRepository.user_storage.present?
+      ShopifyApp.configuration.online_token_configured?
     end
 
     def user_session_expected?

lib/shopify_app/controller_concerns/token_exchange.rb

@@ -0,0 +1,142 @@
+# frozen_string_literal: true
+
+module ShopifyApp
+  module TokenExchange
+    extend ActiveSupport::Concern
+
+    def activate_shopify_session
+      if current_shopify_session.blank?
+        retrieve_session_from_token_exchange
+      end
+
+      if ShopifyApp.configuration.check_session_expiry_date && current_shopify_session.expired?
+        @current_shopify_session = nil
+        retrieve_session_from_token_exchange
+      end
+
+      begin
+        ShopifyApp::Logger.debug("Activating Shopify session")
+        ShopifyAPI::Context.activate_session(current_shopify_session)
+        yield
+      ensure
+        ShopifyApp::Logger.debug("Deactivating session")
+        ShopifyAPI::Context.deactivate_session
+      end
+    end
+
+    def current_shopify_session
+      @current_shopify_session ||= begin
+        session_id = ShopifyAPI::Utils::SessionUtils.current_session_id(
+          request.headers["HTTP_AUTHORIZATION"],
+          nil,
+          online_token_configured?,
+        )
+        return nil unless session_id
+
+        ShopifyApp::SessionRepository.load_session(session_id)
+      end
+    end
+
+    def current_shopify_domain
+      return if params[:shop].blank?
+
+      ShopifyApp::Utils.sanitize_shop_domain(params[:shop])
+    end
+
+    private
+
+    def retrieve_session_from_token_exchange
+      # TODO: Right now JWT Middleware only updates env['jwt.shopify_domain'] from request headers tokens,
+      # which won't work for new installs.
+      # we need to update the middleware to also update the env['jwt.shopify_domain'] from the query params
+      domain = ShopifyApp::JWT.new(session_token).shopify_domain
+
+      ShopifyApp::Logger.info("Performing Token Exchange for [#{domain}] - (Offline)")
+      session = exchange_token(
+        shop: domain, # TODO: use jwt_shopify_domain ?
+        session_token: session_token,
+        requested_token_type: ShopifyAPI::Auth::TokenExchange::RequestedTokenType::OFFLINE_ACCESS_TOKEN,
+      )
+
+      if session && online_token_configured?
+        ShopifyApp::Logger.info("Performing Token Exchange for [#{domain}] - (Online)")
+        exchange_token(
+          shop: domain, # TODO: use jwt_shopify_domain ?
+          session_token: session_token,
+          requested_token_type: ShopifyAPI::Auth::TokenExchange::RequestedTokenType::ONLINE_ACCESS_TOKEN,
+        )
+      end
+
+      # TODO: Refactor and add post authenticate tasks
+      # ShopifyApp.configuration.post_authenticate_tasks.perform(session) if session
+    end
+
+    def exchange_token(shop:, session_token:, requested_token_type:)
+      if session_token.blank?
+        # respond_to_invalid_session_token
+        return
+      end
+
+      begin
+        session = ShopifyAPI::Auth::TokenExchange.exchange_token(
+          shop: shop,
+          session_token: session_token,
+          requested_token_type: requested_token_type,
+        )
+      rescue ShopifyAPI::Errors::InvalidJwtTokenError
+        # respond_to_invalid_session_token
+        return
+      rescue ShopifyAPI::Errors::HttpResponseError => error
+        ShopifyApp::Logger.error(
+          "A #{error.code} error (#{error.class}) occurred during the token exchange. Response: #{error.response.body}",
+        )
+        raise
+      rescue => error
+        ShopifyApp::Logger.error("An error occurred during the token exchange: #{error.message}")
+        raise
+      end
+
+      if session
+        begin
+          ShopifyApp::SessionRepository.store_session(session)
+        rescue ActiveRecord::RecordNotUnique
+          ShopifyApp::Logger.debug("Session not stored due to concurrent token exchange calls")
+        end
+      end
+
+      session
+    end
+
+    def session_token
+      @session_token ||= id_token_header
+    end
+
+    def id_token_header
+      request.headers["HTTP_AUTHORIZATION"]&.match(/^Bearer (.+)$/)&.[](1)
+    end
+
+    def respond_to_invalid_session_token
+      # TODO: Implement this method to handle invalid session tokens
+
+      # if request.xhr?
+      # response.set_header("X-Shopify-Retry-Invalid-Session-Request", 1)
+      # unauthorized_response = { message: :unauthorized }
+      # render(json: { errors: [unauthorized_response] }, status: :unauthorized)
+      # else
+      # patch_session_token_url = "#{ShopifyAPI::Context.host}/patch_session_token"
+      # patch_session_token_params = request.query_parameters.except(:id_token)
+
+      # bounce_url = "#{ShopifyAPI::Context.host}#{request.path}?#{patch_session_token_params.to_query}"
+
+      # # App Bridge will trigger a fetch to the URL in shopify-reload, with a new session token in headers
+      # patch_session_token_params["shopify-reload"] = bounce_url
+
+      # redirect_to("#{patch_session_token_url}?#{patch_session_token_params.to_query}", allow_other_host: true)
+      # end
+    end
+
+    def online_token_configured?
+      ShopifyApp.configuration.online_token_configured?
+    end
+  end
+end

test/controllers/concerns/ensure_has_session_test.rb

@@ -3,20 +3,36 @@
 require "test_helper"
 
 class EnsureHasSessionTest < ActionController::TestCase
-  class EnsureHasSessionTestController < ActionController::Base
-    include ShopifyApp::EnsureHasSession
+  test "includes all the needed concerns" do
+    ShopifyApp.configuration.stubs(:use_new_embedded_auth_strategy?).returns(false)
 
-    def index
-    end
+    controller = define_controller
+    assert_common_concerns_are_included(controller)
+    assert controller.include?(ShopifyApp::LoginProtection), "ShopifyApp::LoginProtection"
+    refute controller.include?(ShopifyApp::TokenExchange), "ShopifyApp::TokenExchange"
   end
 
-  tests EnsureHasSessionTestController
+  test "includes TokenExchange when use_new_embedded_auth_strategy is true" do
+    ShopifyApp.configuration.stubs(:use_new_embedded_auth_strategy?).returns(true)
 
-  test "includes all the needed concerns" do
-    EnsureHasSessionTestController.include?(ShopifyApp::Localization)
-    EnsureHasSessionTestController.include?(ShopifyApp::LoginProtection)
-    EnsureHasSessionTestController.include?(ShopifyApp::CsrfProtection)
-    EnsureHasSessionTestController.include?(ShopifyApp::EmbeddedApp)
-    EnsureHasSessionTestController.include?(ShopifyApp::EnsureBilling)
+    controller = define_controller
+    assert_common_concerns_are_included(controller)
+    assert controller.include?(ShopifyApp::TokenExchange), "ShopifyApp::TokenExchange"
+    refute controller.include?(ShopifyApp::LoginProtection), "ShopifyApp::LoginProtection"
+  end
+
+  private
+
+  def assert_common_concerns_are_included(controller)
+    assert controller.include?(ShopifyApp::Localization), "ShopifyApp::Localization"
+    assert controller.include?(ShopifyApp::CsrfProtection), "ShopifyApp::CsrfProtection"
+    assert controller.include?(ShopifyApp::EmbeddedApp), "ShopifyApp::EmbeddedApp"
+    assert controller.include?(ShopifyApp::EnsureBilling), "ShopifyApp::EnsureBilling"
+  end
+
+  def define_controller
+    Class.new(ActionController::Base) do
+      include ShopifyApp::EnsureHasSession
+    end
   end
 end

test/shopify_app/configuration_test.rb

@@ -254,4 +254,21 @@ class ConfigurationTest < ActiveSupport::TestCase
 
     refute ShopifyApp.configuration.use_new_embedded_auth_strategy?
   end
+
+  test "#online_token_configured? is true when user_session_repository is set" do
+    ShopifyApp.configure do |config|
+      config.user_session_repository = "ShopifyApp::InMemoryUserSessionStore"
+    end
+
+    assert ShopifyApp.configuration.online_token_configured?
+  end
+
+  test "#online_token_configured? is false when user storage is nil" do
+    ShopifyApp.configure do |config|
+      config.user_session_repository = "ShopifyApp::InMemoryUserSessionStore"
+    end
+    ShopifyApp::SessionRepository.user_storage = nil
+
+    refute ShopifyApp.configuration.online_token_configured?
+  end
 end