From 20c5ec74f610fce821c2221b7cae0d842eb5cf21 Mon Sep 17 00:00:00 2001
From: Kirill Platonov <mail@kirillplatonov.com>
Date: Fri, 7 Jan 2022 19:13:59 +0300
Subject: [PATCH] Add AppBridgeMiddleware

---
 lib/shopify_app.rb                            |  1 +
 lib/shopify_app/engine.rb                     |  3 ++-
 .../middleware/app_bridge_middleware.rb       | 20 +++++++++++++++++
 .../middleware/app_bridge_middleware_test.rb  | 22 +++++++++++++++++++
 4 files changed, 45 insertions(+), 1 deletion(-)
 create mode 100644 lib/shopify_app/middleware/app_bridge_middleware.rb
 create mode 100644 test/shopify_app/middleware/app_bridge_middleware_test.rb

diff --git a/lib/shopify_app.rb b/lib/shopify_app.rb
index f40d2e4a7..cb7323a37 100644
--- a/lib/shopify_app.rb
+++ b/lib/shopify_app.rb
@@ -46,6 +46,7 @@ def self.use_webpacker?
   require 'shopify_app/managers/scripttags_manager'
 
   # middleware
+  require 'shopify_app/middleware/app_bridge_middleware'
   require 'shopify_app/middleware/jwt_middleware'
   require 'shopify_app/middleware/same_site_cookie_middleware'
 
diff --git a/lib/shopify_app/engine.rb b/lib/shopify_app/engine.rb
index c03e2a845..63941681d 100644
--- a/lib/shopify_app/engine.rb
+++ b/lib/shopify_app/engine.rb
@@ -27,9 +27,10 @@ class Engine < Rails::Engine
 
     initializer "shopify_app.middleware" do |app|
       app.config.middleware.insert_after(::Rack::Runtime, ShopifyApp::SameSiteCookieMiddleware)
+      app.config.middleware.insert_after(ShopifyApp::SameSiteCookieMiddleware, ShopifyApp::AppBridgeMiddleware)
 
       if ShopifyApp.configuration.allow_jwt_authentication
-        app.config.middleware.insert_after(ShopifyApp::SameSiteCookieMiddleware, ShopifyApp::JWTMiddleware)
+        app.config.middleware.insert_after(ShopifyApp::AppBridgeMiddleware, ShopifyApp::JWTMiddleware)
       end
     end
 
diff --git a/lib/shopify_app/middleware/app_bridge_middleware.rb b/lib/shopify_app/middleware/app_bridge_middleware.rb
new file mode 100644
index 000000000..4c39d89e6
--- /dev/null
+++ b/lib/shopify_app/middleware/app_bridge_middleware.rb
@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+module ShopifyApp
+  class AppBridgeMiddleware
+    def initialize(app)
+      @app = app
+    end
+
+    def call(env)
+      request = Rack::Request.new(env)
+
+      if request.params.has_key?("shop") && !request.params.has_key?("host")
+        shop = request.params["shop"]
+        host = Base64.urlsafe_encode64("#{shop}/admin", padding: false)
+        request.update_param("host", host)
+      end
+
+      @app.call(env)
+    end
+  end
+end
diff --git a/test/shopify_app/middleware/app_bridge_middleware_test.rb b/test/shopify_app/middleware/app_bridge_middleware_test.rb
new file mode 100644
index 000000000..dda4b8c05
--- /dev/null
+++ b/test/shopify_app/middleware/app_bridge_middleware_test.rb
@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+require 'test_helper'
+
+class ShopifyApp::AppBridgeMiddlewareTest < ActiveSupport::TestCase
+  def simple_app
+    lambda { |env|
+      [200, { "Content-Type" => "text/plain" }, ["OK"]]
+    }
+  end
+
+  def app
+    Rack::Lint.new(ShopifyApp::AppBridgeMiddleware.new(simple_app))
+  end
+
+  test "adds missing host params" do
+    env = Rack::MockRequest.env_for('https://example.com', params: { shop: "test-shop.myshopify.com" })
+
+    app.call(env)
+
+    assert_equal "dGVzdC1zaG9wLm15c2hvcGlmeS5jb20vYWRtaW4", env["rack.request.query_hash"]["host"]
+  end
+end