From 1fab2f54bff97006267aa8598fb7ef5c6188e8a9 Mon Sep 17 00:00:00 2001
From: ByeongUk Choi <cbw0916@gmail.com>
Date: Thu, 24 Mar 2022 02:24:58 +0900
Subject: [PATCH] Add LIVEBOOK_FORCE_SSL_HOST env (#1064)

Closes #983.
---
 README.md                    |  3 +++
 lib/livebook.ex              |  4 ++++
 lib/livebook/config.ex       |  7 +++++++
 lib/livebook_web/endpoint.ex | 13 +++++++++++++
 4 files changed, 27 insertions(+)

diff --git a/README.md b/README.md
index cc7ca5a05..3f8a9f299 100644
--- a/README.md
+++ b/README.md
@@ -184,6 +184,9 @@ The following environment variables configure Livebook:
     Enabled by default unless `LIVEBOOK_PASSWORD` is set. Set it to "false" to
     disable it.
 
+  * LIVEBOOK_FORCE_SSL_HOST - set a host to redirect to if the request is not over HTTP.
+    Note it does not apply when accessing Livebook via localhost. Defaults to nil.
+
 <!-- Environment variables -->
 
 If running Livebook as a Docker image or an Elixir release, [the environment
diff --git a/lib/livebook.ex b/lib/livebook.ex
index c58f59338..1d2ba8fd6 100644
--- a/lib/livebook.ex
+++ b/lib/livebook.ex
@@ -58,6 +58,10 @@ defmodule Livebook do
       config :livebook, :data_path, data_path
     end
 
+    if force_ssl_host = Livebook.Config.force_ssl_host!("LIVEBOOK_FORCE_SSL_HOST") do
+      config :livebook, :force_ssl_host, force_ssl_host
+    end
+
     config :livebook,
            :cookie,
            Livebook.Config.cookie!("LIVEBOOK_COOKIE") ||
diff --git a/lib/livebook/config.ex b/lib/livebook/config.ex
index 84b7bcbd9..b00891c6b 100644
--- a/lib/livebook/config.ex
+++ b/lib/livebook/config.ex
@@ -197,6 +197,13 @@ defmodule Livebook.Config do
     System.get_env(env, "1") in ~w(true 1)
   end
 
+  @doc """
+  Parses force ssl host setting from env.
+  """
+  def force_ssl_host!(env) do
+    System.get_env(env)
+  end
+
   @doc """
   Parses and validates default runtime from env.
   """
diff --git a/lib/livebook_web/endpoint.ex b/lib/livebook_web/endpoint.ex
index 95e9facf4..1ccb5df11 100644
--- a/lib/livebook_web/endpoint.ex
+++ b/lib/livebook_web/endpoint.ex
@@ -54,6 +54,8 @@ defmodule LivebookWeb.Endpoint do
     file_provider: AssetsMemoryProvider,
     gzip: true
 
+  plug :force_ssl
+
   # Code reloading can be explicitly enabled under the
   # :code_reloader configuration of your endpoint.
   if code_reloading? do
@@ -83,6 +85,17 @@ defmodule LivebookWeb.Endpoint do
 
   plug LivebookWeb.Router
 
+  @plug_ssl Plug.SSL.init(host: {Application, :get_env, [:livebook, :force_ssl_host, nil]})
+  def force_ssl(conn, _opts) do
+    force_ssl_host = Application.get_env(:livebook, :force_ssl_host, nil)
+
+    if force_ssl_host do
+      Plug.SSL.call(conn, @plug_ssl)
+    else
+      conn
+    end
+  end
+
   def access_struct_url() do
     base =
       case struct_url() do