Fix LIVEBOOK_PORT=0 for iframes

lib/livebook/application.ex

@@ -47,11 +47,12 @@ defmodule Livebook.Application do
         if serverless?() do
           []
         else
+          {_type, module, key} = Livebook.Config.identity_provider()
+
           iframe_server_specs() ++
-            identity_provider() ++
             [
+              {module, name: LivebookWeb.ZTA, identity_key: key},
               {DNSCluster, query: Application.get_env(:livebook, :dns_cluster_query) || :ignore},
-              # Start the Endpoint (http/https)
               # We skip the access url as we do our own logging below
               {LivebookWeb.Endpoint, log_access_url: false}
             ] ++ app_specs()
@@ -367,47 +368,16 @@ defmodule Livebook.Application do
         [
           scheme: :http,
           plug: LivebookWeb.IframeEndpoint,
-          port: port
+          port: port,
+          thousand_island_options: [supervisor_options: [name: LivebookWeb.IframeEndpoint]]
         ] ++ Keyword.take(http, [:ip])
 
-      spec = Bandit.child_spec(iframe_opts)
-      spec = update_in(spec.start, &{__MODULE__, :start_iframe, [port, &1]})
-      [spec]
+      [{Bandit, iframe_opts}]
     else
       []
     end
   end
 
-  @doc false
-  def start_iframe(port, {m, f, a}) do
-    case apply(m, f, a) do
-      {:ok, pid} ->
-        {:ok, pid}
-
-      {:error, {:shutdown, {_, _, {{_, {:error, :eaddrinuse}}, _}}}} ->
-        iframe_port_in_use(port)
-
-      {:error, {:shutdown, {_, _, {:listen_error, _, :eaddrinuse}}}} ->
-        iframe_port_in_use(port)
-
-      {:error, _} = error ->
-        error
-    end
-  end
-
-  @dialyzer {:no_return, iframe_port_in_use: 1}
-
-  defp iframe_port_in_use(port) do
-    Livebook.Config.abort!(
-      "Failed to start Livebook iframe server because port #{port} is already in use"
-    )
-  end
-
-  defp identity_provider() do
-    {_type, module, key} = Livebook.Config.identity_provider()
-    [{module, name: LivebookWeb.ZTA, identity_key: key}]
-  end
-
   defp serverless?() do
     Application.get_env(:livebook, :serverless, false)
   end

lib/livebook/utils.ex

@@ -619,21 +619,6 @@ defmodule Livebook.Utils do
   defp memory_unit(:MB), do: 1_000_000
   defp memory_unit(:KB), do: 1_000
 
-  @doc """
-  Gets the port for an existing listener.
-  """
-  @spec get_port(module, :http | :https, :inet.port_number()) :: :inet.port_number()
-  def get_port(endpoint, scheme, default) do
-    try do
-      endpoint.server_info(scheme)
-    rescue
-      _ -> default
-    else
-      {:ok, {_, port}} when is_integer(port) -> port
-      _ -> default
-    end
-  end
-
   @doc """
   Converts the given IP address into a valid hostname.
 

lib/livebook_web/endpoint.ex

@@ -140,10 +140,10 @@ defmodule LivebookWeb.Endpoint do
     base =
       case struct_url() do
         %URI{scheme: "https", port: 0} = uri ->
-          %{uri | port: Livebook.Utils.get_port(__MODULE__, :https, 433)}
+          %{uri | port: port(:https, 433)}
 
         %URI{scheme: "http", port: 0} = uri ->
-          %{uri | port: Livebook.Utils.get_port(__MODULE__, :http, 80)}
+          %{uri | port: port(:http, 80)}
 
         %URI{} = uri ->
           uri
@@ -162,4 +162,15 @@ defmodule LivebookWeb.Endpoint do
   def access_url do
     URI.to_string(access_struct_url())
   end
+
+  defp port(scheme, default) do
+    try do
+      server_info(scheme)
+    rescue
+      _ -> default
+    else
+      {:ok, {_, port}} when is_integer(port) -> port
+      _ -> default
+    end
+  end
 end

lib/livebook_web/iframe_endpoint.ex

@@ -35,8 +35,18 @@ defmodule LivebookWeb.IframeEndpoint do
     iframe_port = Livebook.Config.iframe_port()
 
     case livebook_port do
-      0 -> Livebook.Utils.get_port(__MODULE__, :http, iframe_port)
-      _ -> iframe_port
+      0 ->
+        try do
+          ThousandIsland.listener_info(__MODULE__)
+        rescue
+          _ -> iframe_port
+        else
+          {:ok, {_ip, port}} -> port
+          _ -> iframe_port
+        end
+
+      _ ->
+        iframe_port
     end
   end
 end

mix.lock

@@ -47,7 +47,7 @@
   "telemetry": {:hex, :telemetry, "1.2.1", "68fdfe8d8f05a8428483a97d7aab2f268aaff24b49e0f599faa091f1d4e7f61c", [:rebar3], [], "hexpm", "dad9ce9d8effc621708f99eac538ef1cbe05d6a874dd741de2e689c47feafed5"},
   "telemetry_metrics": {:hex, :telemetry_metrics, "0.6.1", "315d9163a1d4660aedc3fee73f33f1d355dcc76c5c3ab3d59e76e3edf80eef1f", [:mix], [{:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "7be9e0871c41732c233be71e4be11b96e56177bf15dde64a8ac9ce72ac9834c6"},
   "telemetry_poller": {:hex, :telemetry_poller, "1.0.0", "db91bb424e07f2bb6e73926fcafbfcbcb295f0193e0a00e825e589a0a47e8453", [:rebar3], [{:telemetry, "~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "b3a24eafd66c3f42da30fc3ca7dda1e9d546c12250a2d60d7b81d264fbec4f6e"},
-  "thousand_island": {:hex, :thousand_island, "1.1.0", "dcc115650adc61c5e7de12619f0cb94b2b8f050326e7f21ffbf6fdeb3d291e4c", [:mix], [{:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "7745cf71520d74e119827ff32c2da6307e822cf835bebed3b2c459cc57f32d21"},
+  "thousand_island": {:hex, :thousand_island, "1.2.0", "4f548ae771ab5f96bc7e199f9824c0c2ce6d365f8c93f5f64dbbb33988e484bf", [:mix], [{:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "521671fea179672addb6af46455fc2a77be1edda4c0ed351633e0ef37a4b3584"},
   "websock": {:hex, :websock, "0.5.1", "c496036ce95bc26d08ba086b2a827b212c67e7cabaa1c06473cd26b40ed8cf10", [:mix], [], "hexpm", "b9f785108b81cd457b06e5f5dabe5f65453d86a99118b2c0a515e1e296dc2d2c"},
   "websock_adapter": {:hex, :websock_adapter, "0.5.4", "7af8408e7ed9d56578539594d1ee7d8461e2dd5c3f57b0f2a5352d610ddde757", [:mix], [{:bandit, ">= 0.6.0", [hex: :bandit, repo: "hexpm", optional: true]}, {:plug, "~> 1.14", [hex: :plug, repo: "hexpm", optional: false]}, {:plug_cowboy, "~> 2.6", [hex: :plug_cowboy, repo: "hexpm", optional: true]}, {:websock, "~> 0.5", [hex: :websock, repo: "hexpm", optional: false]}], "hexpm", "d2c238c79c52cbe223fcdae22ca0bb5007a735b9e933870e241fce66afb4f4ab"},
 }