EnsureBilling and EnsureAuthenticatedLinks works with token exchange

CHANGELOG.md

@@ -8,6 +8,7 @@ Unreleased
   * `ShopifyApp::JWTMiddleware` returns early if the app is not embedded to avoid unnecessary JWT verification
   * `LoginProtection` now uses `WithShopifyIdToken` concern to retrieve the Shopify Id token, thus accepting the session token from the URL param `id_token`
 * Marking `ShopifyApp::JWT` to be deprecated in version 23.0.0 [1832](https://github.com/Shopify/shopify_app/pull/1832), use `ShopifyAPI::Auth::JwtPayload` instead.
+* Fix infinite redirect loop caused by handling errors from Billing API [1833](https://github.com/Shopify/shopify_app/pull/1833)
 
 22.1.0 (April 9,2024)
 ----------

app/controllers/concerns/shopify_app/ensure_installed.rb

@@ -19,7 +19,6 @@ module EnsureInstalled
 
       if ShopifyApp.configuration.use_new_embedded_auth_strategy?
         include ShopifyApp::TokenExchange
-        include ShopifyApp::EmbeddedApp
         around_action :activate_shopify_session
       else
         before_action :check_shop_known

lib/shopify_app/controller_concerns/embedded_app.rb

@@ -5,6 +5,7 @@ module EmbeddedApp
     extend ActiveSupport::Concern
 
     include ShopifyApp::FrameAncestors
+    include ShopifyApp::SanitizedParams
 
     included do
       layout :embedded_app_layout
@@ -13,6 +14,20 @@ module EmbeddedApp
 
     protected
 
+    def redirect_to_embed_app_in_admin
+      ShopifyApp::Logger.debug("Redirecting to embed app in admin")
+
+      host = if params[:host]
+        params[:host]
+      elsif params[:shop]
+        Base64.encode64("#{sanitized_shop_name}/admin")
+      else
+        raise ShopifyApp::ShopifyDomainNotFound, "Host or shop param is missing"
+      end
+
+      redirect_to(ShopifyAPI::Auth.embedded_app_url(host), allow_other_host: true)
+    end
+
     def use_embedded_app_layout?
       ShopifyApp.configuration.embedded_app?
     end

lib/shopify_app/controller_concerns/ensure_billing.rb

@@ -27,7 +27,7 @@ def check_billing(session = current_shopify_session)
 
       unless has_payment
         if request.xhr?
-          add_top_level_redirection_headers(url: confirmation_url, ignore_response_code: true)
+          RedirectForEmbedded.add_app_bridge_redirect_url_header(confirmation_url, response)
           ShopifyApp::Logger.debug("Responding with 401 unauthorized")
           head(:unauthorized)
         elsif ShopifyApp.configuration.embedded_app?
@@ -45,8 +45,16 @@ def billing_required?
     end
 
     def handle_billing_error(error)
-      logger.info("#{error.message}: #{error.errors}")
-      redirect_to_login
+      ShopifyApp::Logger.warn("Encountered billing error - #{error.message}: #{error.errors}\n" \
+        "Redirecting to login page")
+
+      login_url = ShopifyApp.configuration.login_url
+      if request.xhr?
+        RedirectForEmbedded.add_app_bridge_redirect_url_header(login_url, response)
+        head(:unauthorized)
+      else
+        fullpage_redirect_to(login_url)
+      end
     end
 
     def has_active_payment?(session)

lib/shopify_app/controller_concerns/login_protection.rb

@@ -79,13 +79,6 @@ def signal_access_token_required
       response.set_header(ACCESS_TOKEN_REQUIRED_HEADER, "true")
     end
 
-    def jwt_expire_at
-      expire_at = request.env["jwt.expire_at"]
-      return unless expire_at
-
-      expire_at - 5.seconds # 5s gap to start fetching new token in advance
-    end
-
     def add_top_level_redirection_headers(url: nil, ignore_response_code: false)
       if request.xhr? && (ignore_response_code || response.code.to_i == 401)
         ShopifyApp::Logger.debug("Adding top level redirection headers")
@@ -104,21 +97,12 @@ def add_top_level_redirection_headers(url: nil, ignore_response_code: false)
         url ||= login_url_with_optional_shop
 
         ShopifyApp::Logger.debug("Setting Reauthorize-Url to #{url}")
-        response.set_header("X-Shopify-API-Request-Failure-Reauthorize", "1")
-        response.set_header("X-Shopify-API-Request-Failure-Reauthorize-Url", url)
+        RedirectForEmbedded.add_app_bridge_redirect_url_header(url, response)
       end
     end
 
     protected
 
-    def jwt_shopify_domain
-      request.env["jwt.shopify_domain"]
-    end
-
-    def jwt_shopify_user_id
-      request.env["jwt.shopify_user_id"]
-    end
-
     def host
       params[:host]
     end

lib/shopify_app/controller_concerns/redirect_for_embedded.rb

@@ -4,6 +4,11 @@ module ShopifyApp
   module RedirectForEmbedded
     include ShopifyApp::SanitizedParams
 
+    def self.add_app_bridge_redirect_url_header(url, response)
+      response.set_header("X-Shopify-API-Request-Failure-Reauthorize", "1")
+      response.set_header("X-Shopify-API-Request-Failure-Reauthorize-Url", url)
+    end
+
     private
 
     def embedded_redirect_url?

lib/shopify_app/controller_concerns/token_exchange.rb

@@ -4,6 +4,8 @@ module ShopifyApp
   module TokenExchange
     extend ActiveSupport::Concern
     include ShopifyApp::AdminAPI::WithTokenRefetch
+    include ShopifyApp::SanitizedParams
+    include ShopifyApp::EmbeddedApp
 
     included do
       include ShopifyApp::WithShopifyIdToken
@@ -46,9 +48,7 @@ def current_shopify_session_id
     end
 
     def current_shopify_domain
-      return if params[:shop].blank?
-
-      ShopifyApp::Utils.sanitize_shop_domain(params[:shop])
+      sanitized_shop_name || current_shopify_session&.shop
     end
 
     private
@@ -59,17 +59,24 @@ def retrieve_session_from_token_exchange
     end
 
     def respond_to_invalid_shopify_id_token
-      return redirect_to_bounce_page if request.headers["HTTP_AUTHORIZATION"].blank?
-
-      ShopifyApp::Logger.debug("Responding to invalid Shopify ID token with unauthorized response")
-      response.set_header("X-Shopify-Retry-Invalid-Session-Request", 1)
-      unauthorized_response = { message: :unauthorized }
-      render(json: { errors: [unauthorized_response] }, status: :unauthorized)
+      if request.headers["HTTP_AUTHORIZATION"].blank?
+        if missing_embedded_param?
+          redirect_to_embed_app_in_admin
+        else
+          redirect_to_bounce_page
+        end
+      else
+        ShopifyApp::Logger.debug("Responding to invalid Shopify ID token with unauthorized response")
+        response.set_header("X-Shopify-Retry-Invalid-Session-Request", 1)
+        unauthorized_response = { message: :unauthorized }
+        render(json: { errors: [unauthorized_response] }, status: :unauthorized)
+      end
     end
 
     def redirect_to_bounce_page
       ShopifyApp::Logger.debug("Redirecting to bounce page for patching Shopify ID token")
-      patch_shopify_id_token_url = "#{ShopifyApp.configuration.root_url}/patch_shopify_id_token"
+      patch_shopify_id_token_url =
+        "#{ShopifyAPI::Context.host}#{ShopifyApp.configuration.root_url}/patch_shopify_id_token"
       patch_shopify_id_token_params = request.query_parameters.except(:id_token)
 
       bounce_url = "#{request.path}?#{patch_shopify_id_token_params.to_query}"
@@ -83,8 +90,22 @@ def redirect_to_bounce_page
       )
     end
 
+    def missing_embedded_param?
+      !params[:embedded].present? || params[:embedded] != "1"
+    end
+
     def online_token_configured?
       ShopifyApp.configuration.online_token_configured?
     end
+
+    def fullpage_redirect_to(url)
+      raise ShopifyApp::ShopifyDomainNotFound if current_shopify_domain.nil?
+
+      render(
+        "shopify_app/shared/redirect",
+        layout: false,
+        locals: { url: url, current_shopify_domain: current_shopify_domain },
+      )
+    end
   end
 end

lib/shopify_app/controller_concerns/with_shopify_id_token.rb

@@ -8,6 +8,21 @@ def shopify_id_token
       @shopify_id_token ||= id_token_from_request_env || id_token_from_authorization_header || id_token_from_url_param
     end
 
+    def jwt_shopify_domain
+      request.env["jwt.shopify_domain"]
+    end
+
+    def jwt_shopify_user_id
+      request.env["jwt.shopify_user_id"]
+    end
+
+    def jwt_expire_at
+      expire_at = request.env["jwt.expire_at"]
+      return unless expire_at
+
+      expire_at - 5.seconds # 5s gap to start fetching new token in advance
+    end
+
     private
 
     def id_token_from_request_env

test/controllers/concerns/embedded_app_test.rb

@@ -24,13 +24,18 @@ class EmbeddedAppTestController < ApplicationTestController
     include ShopifyApp::EmbeddedApp
 
     def index; end
+
+    def redirect_to_embed
+      redirect_to_embed_app_in_admin
+    end
   end
 
   tests EmbeddedAppTestController
 
   setup do
     Rails.application.routes.draw do
       get "/embedded_app", to: "embedded_app_test/embedded_app_test#index"
+      get "/redirect_to_embed", to: "embedded_app_test/embedded_app_test#redirect_to_embed"
     end
   end
 
@@ -63,4 +68,29 @@ def index; end
     assert_not_includes @controller.response.headers, "P3P"
     assert_includes @controller.response.headers, "X-Frame-Options"
   end
+
+  test "#redirect_to_embed_app_in_admin redirects to the embed app in the admin when the host param is present" do
+    ShopifyApp.configuration.embedded_app = true
+
+    shop = "my-shop.myshopify.com"
+    host = Base64.encode64("#{shop}/admin")
+    get :redirect_to_embed, params: { host: host }
+    assert_redirected_to "https://#{shop}/admin/apps/#{ShopifyApp.configuration.api_key}"
+  end
+
+  test "#redirect_to_embed_app_in_admin redirects to the embed app in the admin when the shop param is present" do
+    ShopifyApp.configuration.embedded_app = true
+
+    shop = "my-shop.myshopify.com"
+    get :redirect_to_embed, params: { shop: shop }
+    assert_redirected_to "https://#{shop}/admin/apps/#{ShopifyApp.configuration.api_key}"
+  end
+
+  test "raises an exception when neither the host nor shop param is present" do
+    ShopifyApp.configuration.embedded_app = true
+
+    assert_raises ShopifyApp::ShopifyDomainNotFound do
+      get :redirect_to_embed
+    end
+  end
 end

test/controllers/concerns/ensure_billing_test.rb

@@ -240,6 +240,40 @@ def index
     refute response.headers["X-Shopify-API-Request-Failure-Reauthorize-Url"].present?
   end
 
+  test "Add app bridge redirect headers when handling billing error for XHR requests" do
+    ShopifyApp.configuration.billing = ShopifyApp::BillingConfiguration.new(
+      charge_name: TEST_CHARGE_NAME,
+      amount: 5,
+      interval: ShopifyApp::BillingConfiguration::INTERVAL_ONE_TIME,
+    )
+    @controller.stubs(:run_query).raises(ShopifyApp::BillingError.new("Billing error", { errors: "not good" }))
+
+    ShopifyApp::Logger.expects(:warn).with("Encountered billing error - Billing error: {:errors=>\"not good\"}"\
+      "\nRedirecting to login page")
+
+    get :index, xhr: true
+
+    assert_response :unauthorized
+    assert_match "1", response.headers["X-Shopify-API-Request-Failure-Reauthorize"]
+    assert_match(ShopifyApp.configuration.login_url, response.headers["X-Shopify-API-Request-Failure-Reauthorize-Url"])
+  end
+
+  test "Redirect to login when handling billing errors" do
+    ShopifyApp.configuration.billing = ShopifyApp::BillingConfiguration.new(
+      charge_name: TEST_CHARGE_NAME,
+      amount: 5,
+      interval: ShopifyApp::BillingConfiguration::INTERVAL_ONE_TIME,
+    )
+
+    @controller.stubs(:run_query).raises(ShopifyApp::BillingError.new("Billing error", { errors: "not good" }))
+
+    @controller.expects(:fullpage_redirect_to).with(ShopifyApp.configuration.login_url)
+    ShopifyApp::Logger.expects(:warn).with("Encountered billing error - Billing error: {:errors=>\"not good\"}"\
+      "\nRedirecting to login page")
+
+    get :index
+  end
+
   private
 
   def assert_client_side_redirection(url)