Add LIVEBOOK_PROXY_HEADERS (#2604)

README.md

@@ -266,6 +266,10 @@ The following environment variables can be used to configure Livebook on boot:
   * `LIVEBOOK_PASSWORD` - sets a password that must be used to access Livebook.
     Must be at least 12 characters. Defaults to token authentication.
 
+  * `LIVEBOOK_PROXY_HEADERS` - a comma-separated list of headers that are set by
+    proxies. For example, `x-forwarded-for,x-forwarded-proto`. Configuring those
+    may be required when running Livebook behind reverse proxies.
+
   * `LIVEBOOK_PORT` - sets the port Livebook runs on. If you want to run multiple
     instances on the same domain with the same credentials but on different ports,
     you also need to set `LIVEBOOK_SECRET_KEY_BASE`. Defaults to 8080. If set to 0,

config/config.exs

@@ -35,6 +35,7 @@ config :livebook,
   force_ssl_host: nil,
   learn_notebooks: [],
   plugs: [],
+  rewrite_on: [],
   shutdown_callback: nil,
   teams_url: "https://teams.livebook.dev",
   update_instructions_url: nil,

lib/livebook.ex

@@ -199,6 +199,10 @@ defmodule Livebook do
       config :livebook, :cacertfile, cacertfile
     end
 
+    if rewrite_on = Livebook.Config.rewrite_on!("LIVEBOOK_PROXY_HEADERS") do
+      config :livebook, :rewrite_on, rewrite_on
+    end
+
     config :livebook,
            :cookie,
            Livebook.Config.cookie!("LIVEBOOK_COOKIE") ||

lib/livebook/config.ex

@@ -349,6 +349,13 @@ defmodule Livebook.Config do
     Application.fetch_env!(:livebook, :force_ssl_host)
   end
 
+  @doc """
+  Returns rewrite_on headers.
+  """
+  def rewrite_on do
+    Application.fetch_env!(:livebook, :rewrite_on)
+  end
+
   @doc """
   Returns the application cacertfile if any.
   """
@@ -543,6 +550,25 @@ defmodule Livebook.Config do
     end
   end
 
+  @doc """
+  Parses info for `Plug.RewriteOn`.
+  """
+  def rewrite_on!(env) do
+    if headers = System.get_env(env) do
+      headers
+      |> String.split(",")
+      |> Enum.map(&(&1 |> String.trim() |> rewrite_on!(env)))
+    else
+      []
+    end
+  end
+
+  defp rewrite_on!("x-forwarded-for", _env), do: :x_forwarded_for
+  defp rewrite_on!("x-forwarded-host", _env), do: :x_forwarded_host
+  defp rewrite_on!("x-forwarded-port", _env), do: :x_forwarded_port
+  defp rewrite_on!("x-forwarded-proto", _env), do: :x_forwarded_proto
+  defp rewrite_on!(header, env), do: abort!("unknown header #{inspect(header)} given to #{env}")
+
   @doc """
   Parses and validates the password from env.
   """

lib/livebook_web/endpoint.ex

@@ -95,7 +95,10 @@ defmodule LivebookWeb.Endpoint do
     end
   end
 
-  @plug_ssl Plug.SSL.init(host: {Livebook.Config, :force_ssl_host, []})
+  @plug_ssl Plug.SSL.init(
+              host: {Livebook.Config, :force_ssl_host, []},
+              rewrite_on: {Livebook.Config, :rewrite_on, []}
+            )
   def force_ssl(conn, _opts) do
     if Livebook.Config.force_ssl_host() do
       Plug.SSL.call(conn, @plug_ssl)

mix.exs

@@ -106,6 +106,7 @@ defmodule Livebook.MixProject do
       {:telemetry_poller, "~> 1.0"},
       {:jason, "~> 1.0"},
       {:bandit, "~> 1.0"},
+      {:plug, github: "elixir-plug/plug", override: true},
       {:plug_crypto, "~> 2.0"},
       {:earmark_parser, "~> 1.4"},
       {:ecto, "~> 3.10"},

mix.lock

@@ -39,7 +39,7 @@
   "phoenix_live_view": {:hex, :phoenix_live_view, "0.20.14", "70fa101aa0539e81bed4238777498f6215e9dda3461bdaa067cad6908110c364", [:mix], [{:floki, "~> 0.36", [hex: :floki, repo: "hexpm", optional: true]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:phoenix, "~> 1.6.15 or ~> 1.7.0", [hex: :phoenix, repo: "hexpm", optional: false]}, {:phoenix_html, "~> 3.3 or ~> 4.0", [hex: :phoenix_html, repo: "hexpm", optional: false]}, {:phoenix_template, "~> 1.0", [hex: :phoenix_template, repo: "hexpm", optional: false]}, {:phoenix_view, "~> 2.0", [hex: :phoenix_view, repo: "hexpm", optional: true]}, {:plug, "~> 1.15", [hex: :plug, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.2 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "82f6d006c5264f979ed5eb75593d808bbe39020f20df2e78426f4f2d570e2402"},
   "phoenix_pubsub": {:hex, :phoenix_pubsub, "2.1.3", "3168d78ba41835aecad272d5e8cd51aa87a7ac9eb836eabc42f6e57538e3731d", [:mix], [], "hexpm", "bba06bc1dcfd8cb086759f0edc94a8ba2bc8896d5331a1e2c2902bf8e36ee502"},
   "phoenix_template": {:hex, :phoenix_template, "1.0.4", "e2092c132f3b5e5b2d49c96695342eb36d0ed514c5b252a77048d5969330d639", [:mix], [{:phoenix_html, "~> 2.14.2 or ~> 3.0 or ~> 4.0", [hex: :phoenix_html, repo: "hexpm", optional: true]}], "hexpm", "2c0c81f0e5c6753faf5cca2f229c9709919aba34fab866d3bc05060c9c444206"},
-  "plug": {:hex, :plug, "1.15.3", "712976f504418f6dff0a3e554c40d705a9bcf89a7ccef92fc6a5ef8f16a30a97", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:plug_crypto, "~> 1.1.1 or ~> 1.2 or ~> 2.0", [hex: :plug_crypto, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.3 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "cc4365a3c010a56af402e0809208873d113e9c38c401cabd88027ef4f5c01fd2"},
+  "plug": {:git, "https://github.com/elixir-plug/plug.git", "0574733fb933e4a2ea78532e38e687d9cffb4858", []},
   "plug_cowboy": {:hex, :plug_cowboy, "2.7.1", "87677ffe3b765bc96a89be7960f81703223fe2e21efa42c125fcd0127dd9d6b2", [:mix], [{:cowboy, "~> 2.7", [hex: :cowboy, repo: "hexpm", optional: false]}, {:cowboy_telemetry, "~> 0.3", [hex: :cowboy_telemetry, repo: "hexpm", optional: false]}, {:plug, "~> 1.14", [hex: :plug, repo: "hexpm", optional: false]}], "hexpm", "02dbd5f9ab571b864ae39418db7811618506256f6d13b4a45037e5fe78dc5de3"},
   "plug_crypto": {:hex, :plug_crypto, "2.1.0", "f44309c2b06d249c27c8d3f65cfe08158ade08418cf540fd4f72d4d6863abb7b", [:mix], [], "hexpm", "131216a4b030b8f8ce0f26038bc4421ae60e4bb95c5cf5395e1421437824c4fa"},
   "protobuf": {:hex, :protobuf, "0.12.0", "58c0dfea5f929b96b5aa54ec02b7130688f09d2de5ddc521d696eec2a015b223", [:mix], [{:jason, "~> 1.2", [hex: :jason, repo: "hexpm", optional: true]}], "hexpm", "75fa6cbf262062073dd51be44dd0ab940500e18386a6c4e87d5819a58964dc45"},

test/livebook/config_test.exs

@@ -4,6 +4,14 @@ defmodule Livebook.ConfigTest do
   doctest Livebook.Config
   alias Livebook.Config
 
+  describe "rewrite_on!/1" do
+    test "parses headers" do
+      with_env([TEST_REWRITE_ON: "x-forwarded-for, x-forwarded-proto"], fn ->
+        assert Config.rewrite_on!("TEST_REWRITE_ON") == [:x_forwarded_for, :x_forwarded_proto]
+      end)
+    end
+  end
+
   describe "node!/1" do
     test "parses longnames" do
       with_env([TEST_LIVEBOOK_NODE: "test@::1", TEST_LIVEBOOK_DISTRIBUTION: "name"], fn ->