From f5f33bf9db42fcc224610f0a8f7ca67a08b491f2 Mon Sep 17 00:00:00 2001
From: Zoey Lan <zoey.lan@shopify.com>
Date: Tue, 2 Apr 2024 15:37:38 -0600
Subject: [PATCH] Retrying block if encountered 401 error

---
 .../controller_concerns/token_exchange.rb     | 13 ++-
 .../token_exchange_test.rb                    | 81 +++++++++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/lib/shopify_app/controller_concerns/token_exchange.rb b/lib/shopify_app/controller_concerns/token_exchange.rb
index dc5522779..96013cec2 100644
--- a/lib/shopify_app/controller_concerns/token_exchange.rb
+++ b/lib/shopify_app/controller_concerns/token_exchange.rb
@@ -10,14 +10,24 @@ def activate_shopify_session
       end
 
       if ShopifyApp.configuration.check_session_expiry_date && current_shopify_session.expired?
-        @current_shopify_session = nil
         retrieve_session_from_token_exchange
       end
 
+      attempts = 0
       begin
         ShopifyApp::Logger.debug("Activating Shopify session")
         ShopifyAPI::Context.activate_session(current_shopify_session)
         yield
+      rescue ShopifyAPI::Errors::HttpResponseError => error
+        if error.code == 401 && attempts.zero?
+          ShopifyApp::Logger.debug("Encountered 401 error, exchanging token and retrying with new access token")
+          attempts += 1
+          retrieve_session_from_token_exchange
+          retry
+        else
+          ShopifyApp::Logger.debug("Encountered error: #{error.code} - #{error.message}, re-raising")
+          raise
+        end
       ensure
         ShopifyApp::Logger.debug("Deactivating session")
         ShopifyAPI::Context.deactivate_session
@@ -46,6 +56,7 @@ def current_shopify_domain
     private
 
     def retrieve_session_from_token_exchange
+      @current_shopify_session = nil
       # 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
diff --git a/test/shopify_app/controller_concerns/token_exchange_test.rb b/test/shopify_app/controller_concerns/token_exchange_test.rb
index a02414c77..e233f8c60 100644
--- a/test/shopify_app/controller_concerns/token_exchange_test.rb
+++ b/test/shopify_app/controller_concerns/token_exchange_test.rb
@@ -4,6 +4,10 @@
 require "action_controller"
 require "action_controller/base"
 
+class ApiClass
+  def self.perform; end
+end
+
 class TokenExchangeController < ActionController::Base
   include ShopifyApp::TokenExchange
 
@@ -12,6 +16,11 @@ class TokenExchangeController < ActionController::Base
   def index
     render(plain: "OK")
   end
+
+  def make_api_call
+    ApiClass.perform
+    render(plain: "OK")
+  end
 end
 
 class TokenExchangeControllerTest < ActionController::TestCase
@@ -190,14 +199,86 @@ class TokenExchangeControllerTest < ActionController::TestCase
     end
   end
 
+  test "Handles 401 error by exchanging token before retrying with new access token" do
+    ShopifyApp::SessionRepository.store_shop_session(@offline_session)
+    ShopifyAPI::Utils::SessionUtils.stubs(:current_session_id).returns(@offline_session_id)
+
+    ApiClass.expects(:perform).times(2)
+      .raises(http_response_401_error)
+      .then.returns(true)
+
+    ShopifyAPI::Auth::TokenExchange.expects(:exchange_token).with(
+      shop: @shop,
+      session_token: @session_token,
+      requested_token_type: ShopifyAPI::Auth::TokenExchange::RequestedTokenType::OFFLINE_ACCESS_TOKEN,
+    ).returns(@offline_session)
+
+    with_application_test_routes do
+      get :make_api_call, params: { shop: @shop }
+    end
+  end
+
+  test "Only retry once when encountering 401 error and raises the second error" do
+    ShopifyApp::SessionRepository.store_shop_session(@offline_session)
+    ShopifyAPI::Utils::SessionUtils.stubs(:current_session_id).returns(@offline_session_id)
+    ShopifyAPI::Auth::TokenExchange.stubs(:exchange_token).returns(@offline_session)
+
+    ApiClass.expects(:perform).times(2).raises(http_response_401_error)
+
+    with_application_test_routes do
+      actual_error = assert_raises(ShopifyAPI::Errors::HttpResponseError) do
+        get :make_api_call, params: { shop: @shop }
+      end
+
+      assert_equal http_response_401_error.code, actual_error.code
+    end
+  end
+
+  test "Raises HttpResponseError without retrying if not a 401 error" do
+    ShopifyApp::SessionRepository.store_shop_session(@offline_session)
+    ShopifyAPI::Utils::SessionUtils.stubs(:current_session_id).returns(@offline_session_id)
+    ShopifyAPI::Auth::TokenExchange.stubs(:exchange_token).returns(@offline_session)
+
+    ApiClass.expects(:perform).raises(http_response_500_error)
+
+    with_application_test_routes do
+      actual_error = assert_raises(ShopifyAPI::Errors::HttpResponseError) do
+        get :make_api_call, params: { shop: @shop }
+      end
+
+      assert_equal http_response_500_error.code, actual_error.code
+    end
+  end
+
   private
 
   def with_application_test_routes
     with_routing do |set|
       set.draw do
         get "/" => "token_exchange#index"
+        get "/make_api_call" => "token_exchange#make_api_call"
       end
       yield
     end
   end
+
+  def http_response_401_error
+    ShopifyAPI::Errors::HttpResponseError.new(
+      response: ShopifyAPI::Clients::HttpResponse.new(
+        code: 401,
+        headers: {},
+        body: "Invalid API key or access token (unrecognized login or wrong password)",
+      ),
+    )
+  end
+
+  def http_response_500_error
+    ShopifyAPI::Errors::HttpResponseError.new(
+      response: ShopifyAPI::Clients::HttpResponse.new(
+        code: 500,
+        headers: {},
+        body: "Internal Server Error",
+      ),
+    )
+  end
 end