Merge pull request #951 from Shopify/jwt-callback

Implement JWT callback endpoint

app/controllers/shopify_app/callback_controller.rb

@@ -6,10 +6,22 @@ class CallbackController < ActionController::Base
     include ShopifyApp::LoginProtection
 
     def callback
-      if auth_hash
-        login_shop
+      unless auth_hash
+        return respond_with_error
+      end
+
+      if jwt_request? && !valid_jwt_auth?
+        return respond_with_error
+      end
+
+      if jwt_request?
+        set_shopify_session
+        head(:ok)
+      else
+        reset_session_options
+        set_shopify_session
 
-        if ShopifyApp::SessionRepository.user_storage.present? && user_session.blank?
+        if redirect_for_user_token?
           return redirect_to(login_url_with_optional_shop)
         end
 
@@ -18,17 +30,30 @@ def callback
         perform_after_authenticate_job
 
         redirect_to(return_address)
+      end
+    end
+
+    private
+
+    def respond_with_error
+      if jwt_request?
+        head(:unauthorized)
       else
         flash[:error] = I18n.t('could_not_log_in')
         redirect_to(login_url_with_optional_shop)
       end
     end
 
-    private
+    def redirect_for_user_token?
+      ShopifyApp::SessionRepository.user_storage.present? && user_session.blank?
+    end
 
-    def login_shop
-      reset_session_options
-      set_shopify_session
+    def jwt_request?
+      jwt_shopify_domain || jwt_shopify_user_id
+    end
+
+    def valid_jwt_auth?
+      auth_hash && jwt_shopify_domain == shop_name && jwt_shopify_user_id == associated_user_id
     end
 
     def auth_hash
@@ -45,6 +70,10 @@ def associated_user
       auth_hash['extra']['associated_user']
     end
 
+    def associated_user_id
+      associated_user && associated_user['id']
+    end
+
     def token
       auth_hash['credentials']['token']
     end

lib/shopify_app.rb

@@ -43,6 +43,7 @@ def self.use_webpacker?
   require 'shopify_app/managers/scripttags_manager'
 
   # middleware
+  require 'shopify_app/middleware/jwt_middleware'
   require 'shopify_app/middleware/same_site_cookie_middleware'
 
   # session

lib/shopify_app/controller_concerns/login_protection.rb

@@ -83,17 +83,11 @@ def login_again_if_different_user_or_shop
     protected
 
     def jwt_shopify_domain
-      return unless jwt
-      @jwt_shopify_domain ||= JWT.new(jwt).shopify_domain
+      request.env['jwt.shopify_domain']
     end
 
     def jwt_shopify_user_id
-      return unless jwt
-      @jwt_user_id ||= JWT.new(jwt).shopify_user_id
-    end
-
-    def jwt
-      @jwt ||= authenticate_with_http_token { |token| token }
+      request.env['jwt.shopify_user_id']
     end
 
     def redirect_to_login

lib/shopify_app/engine.rb

@@ -16,6 +16,10 @@ class Engine < Rails::Engine
 
     initializer "shopify_app.middleware" do |app|
       app.config.middleware.insert_after(::Rack::Runtime, ShopifyApp::SameSiteCookieMiddleware)
+
+      if ShopifyApp.configuration.allow_jwt_authentication
+        app.config.middleware.insert_after(ShopifyApp::SameSiteCookieMiddleware, ShopifyApp::JWTMiddleware)
+      end
     end
   end
 end

lib/shopify_app/middleware/jwt_middleware.rb

@@ -0,0 +1,41 @@
+module ShopifyApp
+  class JWTMiddleware
+    TOKEN_REGEX = /^Bearer\s+(.*?)$/
+
+    def initialize(app)
+      @app = app
+    end
+
+    def call(env)
+      return call_next(env) unless authorization_header(env)
+
+      token = extract_token(env)
+      return call_next(env) unless token
+
+      set_env_variables(token, env)
+      call_next(env)
+    end
+
+    private
+
+    def call_next(env)
+      @app.call(env)
+    end
+
+    def authorization_header(env)
+      env['HTTP_AUTHORIZATION']
+    end
+
+    def extract_token(env)
+      match = authorization_header(env).match(TOKEN_REGEX)
+      match && match[1]
+    end
+
+    def set_env_variables(token, env)
+      jwt = ShopifyApp::JWT.new(token)
+
+      env['jwt.shopify_domain'] = jwt.shopify_domain
+      env['jwt.shopify_user_id'] = jwt.shopify_user_id
+    end
+  end
+end

shopify_app.gemspec

@@ -15,7 +15,7 @@ Gem::Specification.new do |s|
 
   s.add_runtime_dependency('browser_sniffer', '~> 1.2.2')
   s.add_runtime_dependency('rails', '> 5.2.1')
-  s.add_runtime_dependency('shopify_api', '~> 9.0.2')
+  s.add_runtime_dependency('shopify_api', '~> 9.1.0')
   s.add_runtime_dependency('omniauth-shopify-oauth2', '~> 2.2.2')
   s.add_runtime_dependency('jwt', '~> 2.2.1')
   s.add_runtime_dependency('redirect_safely', '~> 1.0')

test/controllers/callback_controller_test.rb

@@ -11,7 +11,7 @@ def perform; end
 module ShopifyApp
   class CallbackControllerTest < ActionController::TestCase
     TEST_SHOPIFY_DOMAIN = "shop.myshopify.com"
-    TEST_ASSOCIATED_USER = { "shopify_user_id" => 'test-shopify-user' }
+    TEST_ASSOCIATED_USER = { 'id' => 'test-shopify-user' }
     TEST_SESSION = "this.is.a.user.session"
 
     setup do
@@ -99,6 +99,42 @@ class CallbackControllerTest < ActionController::TestCase
       get :callback, params: { shop: 'shop' }
     end
 
+    test '#jwt_callback persists the user token' do
+      mock_shopify_user_omniauth
+      mock_shopify_jwt
+      session = mock_user_session
+
+      ShopifyApp::SessionRepository.expects(:store_user_session).with(session, TEST_ASSOCIATED_USER)
+
+      get :callback
+      assert_response :ok
+    end
+
+    test '#jwt_callback returns unauthorized if no omniauth data' do
+      mock_shopify_jwt
+
+      get :callback
+      assert_response :unauthorized
+    end
+
+    test '#jwt_callback returns unauthorized if the jwt user does not match omniauth user' do
+      mock_shopify_user_omniauth
+      mock_shopify_jwt
+      request.env['jwt.shopify_user_id'] = 'bad-user'
+
+      get :callback
+      assert_response :unauthorized
+    end
+
+    test '#jwt_callback returns unauthorized if the jwt shop does not match omniauth shop' do
+      mock_shopify_user_omniauth
+      mock_shopify_jwt
+      request.env['jwt.shopify_domain'] = 'bad-shop'
+
+      get :callback
+      assert_response :unauthorized
+    end
+
     test '#install_webhooks uses the shop token for shop strategy' do
       shop_session = ShopifyAPI::Session.new(domain: 'shop', token: '1234', api_version: '2019-1')
       ShopifyApp::SessionRepository.expects(:retrieve_shop_session).returns(shop_session)
@@ -227,6 +263,19 @@ class CallbackControllerTest < ActionController::TestCase
 
     private
 
+    def mock_shopify_jwt
+      request.env['jwt.shopify_domain'] = TEST_SHOPIFY_DOMAIN
+      request.env['jwt.shopify_user_id'] = TEST_ASSOCIATED_USER['id']
+    end
+
+    def mock_user_session
+      ShopifyAPI::Session.new(
+        token: '1234',
+        domain: TEST_SHOPIFY_DOMAIN,
+        api_version: ShopifyApp.configuration.api_version,
+      )
+    end
+
     def mock_shopify_omniauth
       ShopifyApp::SessionRepository.shop_storage = ShopifyApp::InMemoryShopSessionStore
       ShopifyApp::SessionRepository.user_storage = nil

test/shopify_app/controller_concerns/login_protection_test.rb

@@ -75,28 +75,24 @@ class LoginProtectionControllerTest < ActionController::TestCase
     ShopifyApp.configuration.allow_jwt_authentication = true
     domain = 'https://test.myshopify.io'
     token = 'admin_api_token'
-    payload = {
-      'dest' => 'shopify_domain',
-      'sub' => 'shopify_user',
-    }
-
-    jwt = JWT.encode(payload, nil, 'none')
-    jwt_mock = Struct.new(:shopify_user_id).new(payload['sub'])
-    ShopifyApp::JWT.stubs(:new).with(jwt).returns(jwt_mock)
+    dest = 'shopify_domain'
+    sub = 'shopify_user'
 
     expected_session = ShopifyAPI::Session.new(
       domain: domain,
       token: token,
       api_version: '2020-01',
     )
+
     ShopifyApp::SessionRepository.expects(:retrieve_user_session_by_shopify_user_id)
-      .with(payload['sub']).returns(expected_session)
+      .with(sub).returns(expected_session)
     ShopifyApp::SessionRepository.expects(:retrieve_user_session).never
     ShopifyApp::SessionRepository.expects(:retrieve_shop_session_by_shopify_domain).never
     ShopifyApp::SessionRepository.expects(:retrieve_shop_session).never
 
     with_application_test_routes do
-      request.env['HTTP_AUTHORIZATION'] = "Bearer #{jwt}"
+      request.env['jwt.shopify_domain'] = dest
+      request.env['jwt.shopify_user_id'] = sub
       get :index
 
       assert_equal expected_session, @controller.current_shopify_session
@@ -107,13 +103,7 @@ class LoginProtectionControllerTest < ActionController::TestCase
     ShopifyApp.configuration.allow_jwt_authentication = true
     domain = 'https://test.myshopify.io'
     token = 'admin_api_token'
-    payload = {
-      'dest' => 'test.shopify.com',
-    }
-
-    jwt = JWT.encode(payload, nil, 'none')
-    jwt_mock = Struct.new(:shopify_domain, :shopify_user_id).new(payload['dest'], nil)
-    ShopifyApp::JWT.stubs(:new).with(jwt).returns(jwt_mock)
+    dest = 'test.shopify.com'
 
     expected_session = ShopifyAPI::Session.new(
       domain: domain,
@@ -124,11 +114,11 @@ class LoginProtectionControllerTest < ActionController::TestCase
     ShopifyApp::SessionRepository.expects(:retrieve_user_session_by_shopify_user_id).never
     ShopifyApp::SessionRepository.expects(:retrieve_user_session).never
     ShopifyApp::SessionRepository.expects(:retrieve_shop_session_by_shopify_domain)
-      .with(payload['dest']).returns(expected_session)
+      .with(dest).returns(expected_session)
     ShopifyApp::SessionRepository.expects(:retrieve_shop_session).never
 
     with_application_test_routes do
-      request.env['HTTP_AUTHORIZATION'] = "Bearer #{jwt}"
+      request.env['jwt.shopify_domain'] = dest
       get :index
 
       assert_equal expected_session, @controller.current_shopify_session

test/shopify_app/middleware/jwt_middleware_test.rb

@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+require 'test_helper'
+
+class ShopifyApp::JWTMiddlewareTest < ActiveSupport::TestCase
+  def simple_app
+    lambda { |env|
+      [200, { "Content-Type" => "text/yaml" }, env['jwt.shopify_domain'] || '']
+    }
+  end
+
+  def app
+    Rack::Lint.new(ShopifyApp::JWTMiddleware.new(simple_app))
+  end
+
+  test 'does not change env if no authorization header' do
+    env = Rack::MockRequest.env_for('https://example.com')
+
+    app.call(env)
+
+    assert_nil env['jwt.shopify_domain']
+  end
+
+  test 'does not change env if no bearer token' do
+    env = Rack::MockRequest.env_for('https://example.com')
+    env['HTTP_AUTHORIZATION'] = 'something'
+
+    app.call(env)
+
+    assert_nil env['jwt.shopify_domain']
+  end
+
+  test 'does not add the shop to the env if nil shop value' do
+    jwt_mock = Struct.new(:shopify_domain, :shopify_user_id).new(nil, 1)
+    ShopifyApp::JWT.stubs(:new).with('abc').returns(jwt_mock)
+
+    env = Rack::MockRequest.env_for('https://example.com')
+    env['HTTP_AUTHORIZATION'] = 'Bearer abc'
+
+    app.call(env)
+
+    assert_nil env['jwt.shopify_domain']
+    assert_equal 1, env['jwt.shopify_user_id']
+  end
+
+  test 'does not add the user to the env if nil user value' do
+    jwt_mock = Struct.new(:shopify_domain, :shopify_user_id).new('example.myshopify.com', nil)
+    ShopifyApp::JWT.stubs(:new).with('abc').returns(jwt_mock)
+
+    env = Rack::MockRequest.env_for('https://example.com')
+    env['HTTP_AUTHORIZATION'] = 'Bearer abc'
+
+    app.call(env)
+
+    assert_equal 'example.myshopify.com', env['jwt.shopify_domain']
+    assert_nil env['jwt.shopify_user_id']
+  end
+
+  test 'sets shopify_domain and shopify_user_id if non-nil values' do
+    jwt_mock = Struct.new(:shopify_domain, :shopify_user_id).new('example.myshopify.com', 1)
+    ShopifyApp::JWT.stubs(:new).with('abc').returns(jwt_mock)
+
+    env = Rack::MockRequest.env_for('https://example.com')
+    env['HTTP_AUTHORIZATION'] = 'Bearer abc'
+
+    app.call(env)
+
+    assert_equal 'example.myshopify.com', env['jwt.shopify_domain']
+    assert_equal 1, env['jwt.shopify_user_id']
+  end
+
+  test 'sets the jwt values before calling the next middleware' do
+    jwt_mock = Struct.new(:shopify_domain, :shopify_user_id).new('example.myshopify.com', 1)
+    ShopifyApp::JWT.stubs(:new).with('abc').returns(jwt_mock)
+
+    env = Rack::MockRequest.env_for('https://example.com')
+    env['HTTP_AUTHORIZATION'] = 'Bearer abc'
+
+    _, _, body = ShopifyApp::JWTMiddleware.new(simple_app).call(env)
+
+    assert_equal 'example.myshopify.com', body
+  end
+end