From 41347c51be7b86a352984ab9109c6c38b6320dce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Abhijit=20Kar=20=E3=83=84?= <ak24284@gmail.com>
Date: Tue, 11 Oct 2022 16:55:54 +0530
Subject: [PATCH] Infer hostname automatically in CLI for attached mode (#1474)

---
 lib/livebook/config.ex                            |  4 ----
 lib/livebook/runtime/attached.ex                  | 15 ++++++++++++++-
 lib/livebook/utils.ex                             | 13 -------------
 .../live/session_live/attached_live.ex            |  4 ++--
 4 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/lib/livebook/config.ex b/lib/livebook/config.ex
index f59d5c1fe..78099b7a8 100644
--- a/lib/livebook/config.ex
+++ b/lib/livebook/config.ex
@@ -336,10 +336,6 @@ defmodule Livebook.Config do
   defp parse_connection_config!(config) do
     {node, cookie} = split_at_last_occurrence(config, ":")
 
-    unless node =~ "@" do
-      abort!(~s{expected node to include hostname, got: #{inspect(node)}})
-    end
-
     node = String.to_atom(node)
     cookie = String.to_atom(cookie)
 
diff --git a/lib/livebook/runtime/attached.ex b/lib/livebook/runtime/attached.ex
index 8db6edf76..a556b1ee5 100644
--- a/lib/livebook/runtime/attached.ex
+++ b/lib/livebook/runtime/attached.ex
@@ -32,6 +32,10 @@ defmodule Livebook.Runtime.Attached do
   def connect(runtime) do
     %{node: node, cookie: cookie} = runtime
 
+    # We need to append the hostname on connect because
+    # net_kernel has not yet started during new/2.
+    node = append_hostname(node)
+
     # Set cookie for connecting to this specific node
     Node.set_cookie(node, cookie)
 
@@ -42,12 +46,21 @@ defmodule Livebook.Runtime.Attached do
             node_manager_opts: [parent_node: node(), capture_orphan_logs: false]
           )
 
-        {:ok, %{runtime | server_pid: server_pid}}
+        {:ok, %{runtime | node: node, server_pid: server_pid}}
 
       :pang ->
         {:error, "node #{inspect(node)} is unreachable"}
     end
   end
+
+  defp append_hostname(node) do
+    with :nomatch <- :string.find(Atom.to_string(node), "@"),
+         <<suffix::binary>> <- :string.find(Atom.to_string(:net_kernel.nodename()), "@") do
+      :"#{node}#{suffix}"
+    else
+      _ -> node
+    end
+  end
 end
 
 defimpl Livebook.Runtime, for: Livebook.Runtime.Attached do
diff --git a/lib/livebook/utils.ex b/lib/livebook/utils.ex
index 07d88db8b..7029f3d23 100644
--- a/lib/livebook/utils.ex
+++ b/lib/livebook/utils.ex
@@ -73,19 +73,6 @@ defmodule Livebook.Utils do
     end)
   end
 
-  @doc """
-  Converts the given name to node identifier.
-  """
-  @spec node_from_name(String.t()) :: atom()
-  def node_from_name(name) do
-    if name =~ "@" do
-      String.to_atom(name)
-    else
-      # Default to the same host as the current node
-      :"#{name}@#{node_host()}"
-    end
-  end
-
   @doc """
   Returns the host part of a node.
   """
diff --git a/lib/livebook_web/live/session_live/attached_live.ex b/lib/livebook_web/live/session_live/attached_live.ex
index e14a220eb..ebf69be1d 100644
--- a/lib/livebook_web/live/session_live/attached_live.ex
+++ b/lib/livebook_web/live/session_live/attached_live.ex
@@ -1,7 +1,7 @@
 defmodule LivebookWeb.SessionLive.AttachedLive do
   use LivebookWeb, :live_view
 
-  alias Livebook.{Session, Runtime, Utils}
+  alias Livebook.{Session, Runtime}
 
   @impl true
   def mount(_params, %{"session" => session, "current_runtime" => current_runtime}, socket) do
@@ -90,7 +90,7 @@ defmodule LivebookWeb.SessionLive.AttachedLive do
   end
 
   def handle_event("init", %{"data" => data}, socket) do
-    node = Utils.node_from_name(data["name"])
+    node = String.to_atom(data["name"])
     cookie = String.to_atom(data["cookie"])
 
     runtime = Runtime.Attached.new(node, cookie)