Fix omniauth strategy not being set correctly for apps using session tokens

app/controllers/shopify_app/callback_controller.rb

@@ -65,6 +65,13 @@ def respond_with_error
       end
     end
 
+    # Override user_session_by_cookie from LoginProtection to bypass allow_cookie_authentication
+    # setting check because session cookies are justified at top level
+    def user_session_by_cookie
+      return unless session[:user_id].present?
+      ShopifyApp::SessionRepository.retrieve_user_session(session[:user_id])
+    end
+
     def start_user_token_flow?
       if jwt_request?
         false

app/controllers/shopify_app/sessions_controller.rb

@@ -92,19 +92,28 @@ def authenticate_with_partitioning
       end
     end
 
+    # Override shop_session_by_cookie from LoginProtection to bypass allow_cookie_authentication
+    # setting check because session cookies are justified at top level
+    def shop_session_by_cookie
+      return unless session[:shop_id].present?
+      ShopifyApp::SessionRepository.retrieve_shop_session(session[:shop_id])
+    end
+
     # rubocop:disable Lint/SuppressedException
     def set_user_tokens_option
-      if shop_session.blank?
+      current_shop_session = shop_session
+
+      if current_shop_session.blank?
         session[:user_tokens] = false
         return
       end
 
       session[:user_tokens] = ShopifyApp::SessionRepository.user_storage.present?
 
       ShopifyAPI::Session.temp(
-        domain: shop_session.domain,
-        token: shop_session.token,
-        api_version: shop_session.api_version
+        domain: current_shop_session.domain,
+        token: current_shop_session.token,
+        api_version: current_shop_session.api_version
       ) do
         ShopifyAPI::Metafield.find(:token_validity_bogus_check)
       end

test/controllers/callback_controller_test.rb

@@ -113,7 +113,9 @@ class CallbackControllerTest < ActionController::TestCase
       assert_equal 'valid-shop-id', session[:shop_id]
     end
 
-    test '#callback clears a stale shop_id if user session is for a different shop' do
+    test '#callback clears a stale shop_id if user session is for a different shop when using cookie auth' do
+      ShopifyApp.configuration.allow_jwt_authentication = false
+      ShopifyApp.configuration.allow_cookie_authentication = true
       mock_shopify_user_omniauth
       session[:shop_id] = 'valid-shop-id'
       other_shop_session = ShopifyAPI::Session.new(
@@ -311,6 +313,9 @@ class CallbackControllerTest < ActionController::TestCase
     end
 
     test "#callback redirects to the root_url when not using JWT authentication" do
+      ShopifyApp.configuration.allow_jwt_authentication = false
+      ShopifyApp.configuration.allow_cookie_authentication = true
+
       mock_shopify_omniauth
 
       get :callback, params: { shop: 'shop' }

test/controllers/sessions_controller_test.rb

@@ -12,8 +12,6 @@ class SessionsControllerTest < ActionController::TestCase
     setup do
       @routes = ShopifyApp::Engine.routes
       ShopifyApp::SessionRepository.shop_storage = ShopifyApp::InMemoryShopSessionStore
-      ShopifyApp.configuration = nil
-      ShopifyApp.configuration.embedded_app = true
 
       I18n.locale = :en
 
@@ -149,6 +147,31 @@ class SessionsControllerTest < ActionController::TestCase
       assert_match(/Shopify App — Installation/, response.body)
     end
 
+    test "#new sets session[:user_tokens] to true if online tokens are expected" do
+      session[:shop_id] = 1
+      shop_session = ShopifyAPI::Session.new(
+        domain: 'my-shop',
+        token: '1234',
+        api_version: nil,
+      )
+      ShopifyApp::SessionRepository.user_storage.stubs(:present?).returns(true)
+      ShopifyApp::SessionRepository.stubs(:retrieve_shop_session).with(session[:shop_id]).returns(shop_session)
+
+      get :new, params: { shop: 'my-shop' }
+
+      assert session[:user_tokens]
+    end
+
+    test "#new sets session[:user_tokens] to false if there is no existing offline token" do
+      session[:shop_id] = 1
+      ShopifyApp::SessionRepository.user_storage.stubs(:present?).returns(true)
+      ShopifyApp::SessionRepository.stubs(:retrieve_shop_session).with(session[:shop_id]).returns(nil)
+
+      get :new, params: { shop: 'my-shop' }
+
+      refute session[:user_tokens]
+    end
+
     ['my-shop', 'my-shop.myshopify.com', 'https://my-shop.myshopify.com',
      'http://my-shop.myshopify.com'].each do |good_url|
       test "#create should authenticate the shop for the URL (#{good_url})" do

test/dummy/config/initializers/shopify_app.rb

@@ -9,6 +9,7 @@ def self.call
       config.myshopify_domain = 'myshopify.com'
       config.api_version = :unstable
       config.allow_jwt_authentication = true
+      config.allow_cookie_authentication = false
     end
   end
 end

test/shopify_app/controller_concerns/login_protection_test.rb

@@ -48,11 +48,6 @@ class LoginProtectionControllerTest < ActionController::TestCase
                                      'AppleWebKit/537.36 (KHTML, like Gecko) Chrome/69.0.3497.100 Safari/537.36'
   end
 
-  teardown do
-    ShopifyApp.configuration.allow_jwt_authentication = false
-    ShopifyApp.configuration.allow_cookie_authentication = true
-  end
-
   test '#index sets test cookie if embedded app and user agent can partition cookies' do
     with_application_test_routes do
       request.env['HTTP_USER_AGENT'] = 'Version/12.0 Safari'
@@ -132,7 +127,8 @@ class LoginProtectionControllerTest < ActionController::TestCase
     end
   end
 
-  test "#current_shopify_session retrieves using session user_id" do
+  test "#current_shopify_session retrieves using session user_id when allow_cookie_authentication is enabled" do
+    ShopifyApp.configuration.allow_cookie_authentication = true
     with_application_test_routes do
       session[:user_id] = '145'
       get :index
@@ -141,7 +137,8 @@ class LoginProtectionControllerTest < ActionController::TestCase
     end
   end
 
-  test "#current_shopify_session retrieves using shop_id when shopify_user not present" do
+  test "#current_shopify_session retrieves using session[:shop_id] when shopify_user not present and allow_cookie_authentication" do
+    ShopifyApp.configuration.allow_cookie_authentication = true
     with_application_test_routes do
       session[:shop_id] = "shopify_id"
       get :index
@@ -162,20 +159,16 @@ class LoginProtectionControllerTest < ActionController::TestCase
     end
   end
 
-  test "#current_shopify_session retreives the session from storage" do
-    with_application_test_routes do
-      session[:shop_id] = "foobar"
-      get :index
-      ShopifyApp::SessionRepository.expects(:retrieve_shop_session).returns(session).once
-      assert @controller.current_shopify_session
-    end
-  end
-
-  test "#current_shopify_session is memoized and does not retreive session twice" do
+  test "#current_shopify_session is memoized and does not retrieve session twice" do
+    shop_session_record = ShopifyAPI::Session.new(
+      domain: 'my-shop',
+      token: '1234',
+      api_version: nil,
+    )
     with_application_test_routes do
-      session[:shop_id] = "foobar"
+      request.env['jwt.shopify_domain'] = 'foobar'
       get :index
-      ShopifyApp::SessionRepository.expects(:retrieve_shop_session).returns(session).once
+      ShopifyApp::SessionRepository.expects(:retrieve_shop_session_by_shopify_domain).returns(shop_session_record).once
       assert @controller.current_shopify_session
     end
   end
@@ -225,6 +218,7 @@ class LoginProtectionControllerTest < ActionController::TestCase
   end
 
   test "#login_again_if_different_user_or_shop removes current session and redirects to login url" do
+    ShopifyApp.configuration.allow_cookie_authentication = true
     with_application_test_routes do
       session[:shop_id] = "foobar"
       session[:user_id] = 123
@@ -242,6 +236,7 @@ class LoginProtectionControllerTest < ActionController::TestCase
   end
 
   test "#login_again_if_different_user_or_shop ignores non-String shop params so that Rails params for Shop model can be accepted" do
+    ShopifyApp.configuration.allow_cookie_authentication = true
     with_application_test_routes do
       session[:shop_id] = "foobar"
       session[:shopify_domain] = "foobar"
@@ -362,6 +357,7 @@ class LoginProtectionControllerTest < ActionController::TestCase
 
   test '#activate_shopify_session with shop_session and no user_session when \
         user_session expected returns an HTTP 401 when the request is an XHR' do
+    ShopifyApp.configuration.allow_cookie_authentication = true
     # Set up a shop_session
     with_application_test_routes do
       session[:shop_id] = 'foobar'
@@ -374,6 +370,7 @@ class LoginProtectionControllerTest < ActionController::TestCase
 
   test '#activate_shopify_session with shop_session and no user_session when \
         user_session expected redirect to login when the request is not an XHR' do
+    ShopifyApp.configuration.allow_cookie_authentication = true
     # Set up a shop_session
     with_application_test_routes do
       session[:shop_id] = 'foobar'