Handle module plugs in proxy

lib/livebook/proxy/handler.ex

@@ -1,7 +1,7 @@
 defmodule Livebook.Proxy.Handler do
   @moduledoc false
 
-  # Handles request `conn` with the configured function.
+  # Handles request `conn` with the configured plug.
   #
   # The handler forwards all actual communication to the parent
   # `Livebook.Proxy.Server` via `Livebook.Proxy.Adapter`.
@@ -15,8 +15,11 @@ defmodule Livebook.Proxy.Handler do
   @doc """
   Returns a child spec to setup the handler supervision tree.
 
-  Expects the `:listen` option to be provided, and be a function with
-  the signature `Plug.Conn.t() -> Plug.Conn.t()`.
+  Expects the `:listen` option to be provided, and be one of:
+
+    * a function plug: `Plug.Conn.t() -> Plug.Conn.t()`
+
+    * a module plug: `module` atom or `{module, options}` tuple.
   """
   @spec child_spec(keyword()) :: Supervisor.child_spec()
   def child_spec(opts) do
@@ -39,8 +42,27 @@ defmodule Livebook.Proxy.Handler do
     conn =
       struct!(Plug.Conn, %{conn_attrs | adapter: {Livebook.Proxy.Adapter, {parent_pid, ref}}})
 
-    listen = :persistent_term.get({__MODULE__, :listen})
-    listen.(conn)
+    case :persistent_term.get({__MODULE__, :listen}) do
+      fun when is_function(fun, 1) ->
+        fun.(conn)
+
+      mod when is_atom(mod) ->
+        mod.call(conn, mod.init([]))
+
+      {mod, opts} when is_atom(mod) ->
+        mod.call(conn, mod.init(opts))
+
+      other ->
+        raise """
+        expected plug to be one of:
+
+          * fun(conn)
+          * module
+          * {module, options}
+
+        got: #{inspect(other)}
+        """
+    end
   end
 
   @doc """