From ac60aba2fbe26b83d8c3de8d43b80d7539dfc2e6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Tue, 5 Oct 2021 00:44:27 +0200
Subject: [PATCH] Make it possible to configure custom plug for all requests
 (#576)

* Make it possible to configure custom plug for all requests

* Apply review comments

* Use fetch_env!
---
 config/config.exs                         |  7 +++++++
 lib/livebook_web/endpoint.ex              |  4 ++++
 lib/livebook_web/plugs/configured_plug.ex | 18 ++++++++++++++++++
 3 files changed, 29 insertions(+)
 create mode 100644 lib/livebook_web/plugs/configured_plug.ex

diff --git a/config/config.exs b/config/config.exs
index a2085c534..57a3f6832 100644
--- a/config/config.exs
+++ b/config/config.exs
@@ -26,6 +26,13 @@ config :livebook, :authentication_mode, :token
 # configured arguments.
 config :livebook, :default_runtime, {Livebook.Runtime.ElixirStandalone, []}
 
+# A list of custom plugs in the following format:
+#
+#    [{plug_module :: module(), opts :: keyword()}]
+#
+# The plugs are called directly before the Livebook router.
+config :livebook, :plugs, []
+
 # Import environment specific config. This must remain at the bottom
 # of this file so it overrides the configuration defined above.
 import_config "#{Mix.env()}.exs"
diff --git a/lib/livebook_web/endpoint.ex b/lib/livebook_web/endpoint.ex
index 7c3e0d8b1..5ec83ea86 100644
--- a/lib/livebook_web/endpoint.ex
+++ b/lib/livebook_web/endpoint.ex
@@ -71,6 +71,10 @@ defmodule LivebookWeb.Endpoint do
   plug Plug.MethodOverride
   plug Plug.Head
   plug Plug.Session, @session_options
+
+  # Run custom plugs from the app configuration
+  plug LivebookWeb.ConfiguredPlug
+
   plug LivebookWeb.Router
 
   def access_url() do
diff --git a/lib/livebook_web/plugs/configured_plug.ex b/lib/livebook_web/plugs/configured_plug.ex
new file mode 100644
index 000000000..ecb296b97
--- /dev/null
+++ b/lib/livebook_web/plugs/configured_plug.ex
@@ -0,0 +1,18 @@
+defmodule LivebookWeb.ConfiguredPlug do
+  @moduledoc false
+
+  # Runs plugs configured for the :livebook application
+
+  @behaviour Plug
+
+  @impl true
+  def init(opts), do: opts
+
+  @impl true
+  def call(conn, _opts) do
+    case Application.fetch_env!(:livebook, :plugs) do
+      [] -> conn
+      plugs -> Plug.run(conn, plugs)
+    end
+  end
+end