From 7b1addb7ebacfb886ddb47519027296d6b959d01 Mon Sep 17 00:00:00 2001 From: Alexandre de Souza Date: Thu, 3 Nov 2022 13:49:07 -0300 Subject: [PATCH] Use WebSocket and Protobuf instead of GraphQL with Livebook Enterprise (#1504) --- lib/livebook/hubs/enterprise.ex | 12 +- lib/livebook/hubs/enterprise_client.ex | 62 ------ lib/livebook/hubs/fly.ex | 16 +- lib/livebook/web_socket.ex | 96 +++++++++ lib/livebook/web_socket/client.ex | 190 ++++++++++++++++++ lib/livebook/web_socket/server.ex | 115 +++++++++++ lib/livebook_web/live/hub/edit_live.ex | 17 +- .../live/hub/new/enterprise_component.ex | 69 +++++-- mix.exs | 11 +- mix.lock | 8 +- proto/.formatter.exs | 5 + proto/.gitignore | 26 +++ proto/README.md | 11 + proto/lib/livebook_proto.ex | 10 + proto/lib/livebook_proto/error.pb.ex | 6 + proto/lib/livebook_proto/request.pb.ex | 8 + proto/lib/livebook_proto/response.pb.ex | 9 + .../lib/livebook_proto/session_request.pb.ex | 6 + .../lib/livebook_proto/session_response.pb.ex | 7 + proto/lib/livebook_proto/user.pb.ex | 7 + proto/messages.proto | 33 +++ proto/mix.exs | 33 +++ proto/mix.lock | 3 + test/livebook/hubs/enterprise_client_test.exs | 25 --- test/livebook/web_socket/client_test.exs | 42 ++++ test/livebook/web_socket/server_test.exs | 47 +++++ test/livebook/web_socket_test.exs | 47 +++++ .../hub/new/enterprise_component_test.exs | 43 +++- test/livebook_web/live/hub/new_live_test.exs | 149 -------------- test/support/integration/enterprise_server.ex | 10 +- 30 files changed, 832 insertions(+), 291 deletions(-) delete mode 100644 lib/livebook/hubs/enterprise_client.ex create mode 100644 lib/livebook/web_socket.ex create mode 100644 lib/livebook/web_socket/client.ex create mode 100644 lib/livebook/web_socket/server.ex create mode 100644 proto/.formatter.exs create mode 100644 proto/.gitignore create mode 100644 proto/README.md create mode 100644 proto/lib/livebook_proto.ex create mode 100644 proto/lib/livebook_proto/error.pb.ex create mode 100644 proto/lib/livebook_proto/request.pb.ex create mode 100644 proto/lib/livebook_proto/response.pb.ex create mode 100644 proto/lib/livebook_proto/session_request.pb.ex create mode 100644 proto/lib/livebook_proto/session_response.pb.ex create mode 100644 proto/lib/livebook_proto/user.pb.ex create mode 100644 proto/messages.proto create mode 100644 proto/mix.exs create mode 100644 proto/mix.lock delete mode 100644 test/livebook/hubs/enterprise_client_test.exs create mode 100644 test/livebook/web_socket/client_test.exs create mode 100644 test/livebook/web_socket/server_test.exs create mode 100644 test/livebook/web_socket_test.exs diff --git a/lib/livebook/hubs/enterprise.ex b/lib/livebook/hubs/enterprise.ex index dedc92d29..d4ace2e51 100644 --- a/lib/livebook/hubs/enterprise.ex +++ b/lib/livebook/hubs/enterprise.ex @@ -7,12 +7,12 @@ defmodule Livebook.Hubs.Enterprise do alias Livebook.Hubs @type t :: %__MODULE__{ - id: Livebook.Utils.id(), - url: String.t(), - token: String.t(), - external_id: String.t(), - hub_name: String.t(), - hub_color: String.t() + id: String.t() | nil, + url: String.t() | nil, + token: String.t() | nil, + external_id: String.t() | nil, + hub_name: String.t() | nil, + hub_color: String.t() | nil } embedded_schema do diff --git a/lib/livebook/hubs/enterprise_client.ex b/lib/livebook/hubs/enterprise_client.ex deleted file mode 100644 index 7038ddb67..000000000 --- a/lib/livebook/hubs/enterprise_client.ex +++ /dev/null @@ -1,62 +0,0 @@ -defmodule Livebook.Hubs.EnterpriseClient do - @moduledoc false - - alias Livebook.Utils.HTTP - - @path "/api/v1" - - def fetch_info(url, token) do - query = """ - query { - info { - id - } - } - """ - - with {:ok, %{"info" => info}} <- graphql(url, token, query) do - {:ok, info} - end - end - - def fetch_me(url, token) do - query = """ - query { - me { - id - } - } - """ - - with {:ok, %{"me" => me}} <- graphql(url, token, query) do - {:ok, me} - end - end - - defp graphql(url, token, query, input \\ %{}) do - headers = [{"Authorization", "Bearer #{token}"}] - body = {"application/json", Jason.encode!(%{query: query, variables: input})} - - case HTTP.request(:post, graphql_endpoint(url), headers: headers, body: body) do - {:ok, 200, _, body} -> - case Jason.decode!(body) do - %{"errors" => [%{"message" => "invalid_token"}]} -> - {:error, "request failed with invalid token", :invalid_token} - - %{"errors" => [%{"message" => "unauthorized"}]} -> - {:error, "request failed with unauthorized", :unauthorized} - - %{"errors" => [%{"message" => message}]} -> - {:error, "request failed with message: #{message}", :other} - - %{"data" => data} -> - {:ok, data} - end - - {:error, {:failed_connect, _}} -> - {:error, "request failed to connect", :invalid_url} - end - end - - defp graphql_endpoint(url), do: url <> @path -end diff --git a/lib/livebook/hubs/fly.ex b/lib/livebook/hubs/fly.ex index c19afaae7..fb7ff2314 100644 --- a/lib/livebook/hubs/fly.ex +++ b/lib/livebook/hubs/fly.ex @@ -7,14 +7,14 @@ defmodule Livebook.Hubs.Fly do alias Livebook.Hubs @type t :: %__MODULE__{ - id: Livebook.Utils.id(), - access_token: String.t(), - hub_name: String.t(), - hub_color: String.t(), - organization_id: String.t(), - organization_type: String.t(), - organization_name: String.t(), - application_id: String.t() + id: String.t() | nil, + access_token: String.t() | nil, + hub_name: String.t() | nil, + hub_color: String.t() | nil, + organization_id: String.t() | nil, + organization_type: String.t() | nil, + organization_name: String.t() | nil, + application_id: String.t() | nil } embedded_schema do diff --git a/lib/livebook/web_socket.ex b/lib/livebook/web_socket.ex new file mode 100644 index 000000000..0449e1360 --- /dev/null +++ b/lib/livebook/web_socket.ex @@ -0,0 +1,96 @@ +defmodule Livebook.WebSocket do + @moduledoc false + + alias Livebook.WebSocket.Client + alias LivebookProto.{Request, SessionRequest} + + defmodule Connection do + defstruct [:conn, :websocket, :ref] + + @type t :: %__MODULE__{ + conn: Client.conn(), + websocket: Client.websocket(), + ref: Client.ref() + } + end + + @type proto :: SessionRequest.t() + + @typep header :: {String.t(), String.t()} + @typep headers :: list(header()) + + @doc """ + Connects with the WebSocket server for given URL and headers. + """ + @spec connect(String.t(), headers()) :: + {:ok, Connection.t(), :connected | {atom(), proto()}} + | {:error, Connection.t(), String.t() | LivebookProto.Error.t()} + def connect(url, headers \\ []) do + with {:ok, conn, ref} <- Client.connect(url, headers) do + conn + |> Client.receive(ref) + |> handle_receive(ref) + end + end + + @doc """ + Disconnects the given WebSocket client. + """ + @spec disconnect(Connection.t()) :: :ok + def disconnect(%Connection{} = connection) do + Client.disconnect(connection.conn, connection.websocket, connection.ref) + end + + @doc """ + Sends a request to the given server. + """ + @spec send_request(Connection.t(), proto()) :: + {:ok, Connection.t()} + | {:error, Connection.t(), Client.ws_error() | Client.mint_error()} + def send_request(%Connection{} = connection, %struct{} = data) do + type = LivebookProto.request_type(struct) + message = Request.new!(type: {type, data}) + binary = {:binary, Request.encode(message)} + + case Client.send(connection.conn, connection.websocket, connection.ref, binary) do + {:ok, conn, websocket} -> + {:ok, %{connection | conn: conn, websocket: websocket}} + + {:error, %Mint.WebSocket{} = websocket, reason} -> + {:error, %{connection | websocket: websocket}, reason} + + {:error, conn, reason} -> + {:error, %{connection | conn: conn}, reason} + end + end + + @dialyzer {:nowarn_function, receive_response: 1} + + @doc """ + Receives a response from the given server. + """ + @spec receive_response(Connection.t()) :: Client.receive_fun() + def receive_response(%Connection{conn: conn, websocket: websocket, ref: ref}) do + conn + |> Client.receive(ref, websocket) + |> handle_receive(ref) + end + + defp handle_receive({:ok, conn, websocket, :connected}, ref) do + {:ok, %Connection{conn: conn, websocket: websocket, ref: ref}, :connected} + end + + defp handle_receive({:ok, conn, websocket, %Client.Response{body: response}}, ref) do + %{type: result} = LivebookProto.Response.decode(response) + {:ok, %Connection{conn: conn, websocket: websocket, ref: ref}, result} + end + + defp handle_receive({:error, conn, %Client.Response{body: nil, status: status}}, ref) do + {:error, %Connection{conn: conn, ref: ref}, Plug.Conn.Status.reason_phrase(status)} + end + + defp handle_receive({:error, conn, %Client.Response{body: response}}, ref) do + %{type: {:error, error}} = LivebookProto.Response.decode(response) + {:error, %Connection{conn: conn, ref: ref}, error} + end +end diff --git a/lib/livebook/web_socket/client.ex b/lib/livebook/web_socket/client.ex new file mode 100644 index 000000000..effd5fdc5 --- /dev/null +++ b/lib/livebook/web_socket/client.ex @@ -0,0 +1,190 @@ +defmodule Livebook.WebSocket.Client do + @moduledoc false + + alias Mint.WebSocket.UpgradeFailureError + + @ws_path "/livebook/websocket" + + @type conn :: Mint.HTTP.t() + @type websocket :: Mint.WebSocket.t() + @type frame :: :close | {:binary, binary()} + @type ref :: Mint.Types.request_ref() + @type ws_error :: Mint.WebSocket.error() + @type mint_error :: Mint.Types.error() + + defmodule Response do + defstruct [:body, :status, :headers] + + @type t :: %__MODULE__{ + body: Livebook.WebSocket.Response.t(), + status: Mint.Types.status(), + headers: Mint.Types.headers() + } + end + + defguard is_frame(value) when value == :close or elem(value, 0) == :binary + + @doc """ + Connects to the WebSocket server with given url and headers. + """ + @spec connect(String.t(), list({String.t(), String.t()})) :: + {:ok, conn(), ref()} + | {:error, mint_error()} + | {:error, conn(), ws_error()} + def connect(url, headers \\ []) do + uri = URI.parse(url) + http_scheme = parse_http_scheme(uri) + ws_scheme = parse_ws_scheme(uri) + + with {:ok, conn} <- Mint.HTTP.connect(http_scheme, uri.host, uri.port), + {:ok, conn, ref} <- Mint.WebSocket.upgrade(ws_scheme, conn, @ws_path, headers) do + {:ok, conn, ref} + end + end + + defp parse_http_scheme(uri) when uri.scheme in ["http", "ws"], do: :http + defp parse_http_scheme(uri) when uri.scheme in ["https", "wss"], do: :https + + defp parse_ws_scheme(uri) when uri.scheme in ["http", "ws"], do: :ws + defp parse_ws_scheme(uri) when uri.scheme in ["https", "wss"], do: :wss + + @doc """ + Disconnects from the given connection, WebSocket and reference. + + If there's no WebSocket connection yet, it'll only close the HTTP connection. + """ + @spec disconnect(conn(), websocket(), ref()) :: :ok + def disconnect(conn, websocket, ref) do + if websocket do + send(conn, websocket, ref, :close) + end + + Mint.HTTP.close(conn) + + :ok + end + + @doc """ + Receive the message from the given HTTP connection. + + If the WebSocket isn't connected yet, it will try to get the connection + response to start a new WebSocket connection. + """ + @spec receive(conn(), ref(), term()) :: + {:ok, conn(), Response.t() | :connect} + | {:error, conn(), Response.t()} + | {:error, conn(), :unknown} + def receive(conn, ref, websocket \\ nil, message \\ receive(do: (message -> message))) do + case Mint.WebSocket.stream(conn, message) do + {:ok, conn, responses} -> + handle_responses(conn, ref, websocket, responses) + + {:error, conn, reason, []} -> + {:error, conn, reason} + + {:error, conn, _reason, responses} -> + handle_responses(conn, ref, websocket, responses) + + :unknown -> + {:error, :unknown} + end + end + + @successful_status 100..299 + + defp handle_responses(conn, ref, nil, responses) do + result = + Enum.reduce(responses, %Response{}, fn + {:status, ^ref, status}, acc -> %{acc | status: status} + {:headers, ^ref, headers}, acc -> %{acc | headers: headers} + {:data, ^ref, body}, acc -> %{acc | body: body} + {:done, ^ref}, acc -> handle_done_response(conn, ref, acc) + end) + + case result do + %Response{} = response when response.status not in @successful_status -> + {:error, conn, response} + + result -> + result + end + end + + defp handle_responses(conn, ref, websocket, [{:data, ref, data}]) do + with {:ok, websocket, frames} <- Mint.WebSocket.decode(websocket, data) do + case handle_frames(%Response{}, frames) do + {:ok, %Response{body: %{type: {:error, _}}} = response} -> + {:error, conn, websocket, response} + + {:ok, response} -> + {:ok, conn, websocket, response} + + {:close, result} -> + disconnect(conn, websocket, ref) + {:ok, conn, websocket, result} + + {:error, response} -> + {:error, conn, websocket, response} + end + end + end + + defp handle_done_response(conn, ref, response) do + case Mint.WebSocket.new(conn, ref, response.status, response.headers) do + {:ok, conn, websocket} -> + case decode_response(websocket, response) do + {websocket, {:ok, result}} -> + {:ok, conn, websocket, result} + + {websocket, {:close, result}} -> + disconnect(conn, websocket, ref) + {:ok, conn, websocket, result} + + {websocket, {:error, reason}} -> + {:error, conn, websocket, reason} + end + + {:error, conn, %UpgradeFailureError{status_code: status, headers: headers}} -> + {:error, conn, %{response | status: status, headers: headers}} + end + end + + defp decode_response(websocket, %Response{status: 101, body: nil}) do + {websocket, {:ok, :connected}} + end + + defp decode_response(websocket, response) do + case Mint.WebSocket.decode(websocket, response.body) do + {:ok, websocket, frames} -> + {websocket, handle_frames(response, frames)} + + {:error, websocket, reason} -> + {websocket, {:error, reason}} + end + end + + defp handle_frames(response, frames) do + Enum.reduce(frames, response, fn + {:binary, binary}, acc -> + {:ok, %{acc | body: binary}} + + {:close, _code, _data}, acc -> + {:close, acc} + end) + end + + @dialyzer {:nowarn_function, send: 4} + + @doc """ + Sends a message to the given HTTP Connection and WebSocket connection. + """ + @spec send(conn(), websocket(), ref(), frame()) :: + {:ok, conn(), websocket()} + | {:error, conn() | websocket(), term()} + def send(conn, websocket, ref, frame) when is_frame(frame) do + with {:ok, websocket, data} <- Mint.WebSocket.encode(websocket, frame), + {:ok, conn} <- Mint.WebSocket.stream_request_body(conn, ref, data) do + {:ok, conn, websocket} + end + end +end diff --git a/lib/livebook/web_socket/server.ex b/lib/livebook/web_socket/server.ex new file mode 100644 index 000000000..6f9fdb3b7 --- /dev/null +++ b/lib/livebook/web_socket/server.ex @@ -0,0 +1,115 @@ +defmodule Livebook.WebSocket.Server do + @moduledoc false + use GenServer + + require Logger + + import Livebook.WebSocket.Client, only: [is_frame: 1] + + alias Livebook.WebSocket.Client + + defstruct [ + :conn, + :websocket, + :caller, + :status, + :resp_headers, + :resp_body, + :ref, + closing?: false + ] + + def start_link(opts \\ []) do + GenServer.start_link(__MODULE__, opts) + end + + @doc """ + Connects the WebSocket client. + """ + def connect(pid, url, headers \\ []) do + GenServer.call(pid, {:connect, url, headers}) + end + + @doc """ + Disconnects the WebSocket client. + """ + def disconnect(pid) do + GenServer.cast(pid, :close) + end + + @doc """ + Sends a message to the WebSocket server the message request. + """ + def send_message(socket, frame) when is_frame(frame) do + GenServer.cast(socket, {:send_message, frame}) + end + + ## GenServer callbacks + + @impl true + def init(_) do + {:ok, %__MODULE__{}} + end + + @impl true + def handle_call({:connect, url, headers}, from, state) do + case Client.connect(url, headers) do + {:ok, conn, ref} -> {:noreply, %{state | conn: conn, ref: ref, caller: from}} + {:error, _reason} = error -> {:reply, error, state} + {:error, conn, reason} -> {:reply, {:error, reason}, %{state | conn: conn}} + end + end + + @impl true + def handle_cast(:close, state) do + Client.disconnect(state.conn, state.websocket, state.ref) + + {:stop, :normal, state} + end + + def handle_cast({:send_message, frame}, state) do + case Client.send(state.conn, state.websocket, state.ref, frame) do + {:ok, conn, websocket} -> + {:noreply, %{state | conn: conn, websocket: websocket}} + + {:error, %Mint.WebSocket{} = websocket, _reason} -> + {:noreply, %{state | websocket: websocket}} + + {:error, conn, _reason} -> + {:noreply, %{state | conn: conn}} + end + end + + @impl true + def handle_info(message, state) do + case Client.receive(state.conn, state.ref, state.websocket, message) do + {:ok, conn, websocket, response} -> + state = %{state | conn: conn, websocket: websocket} + {:noreply, reply(state, {:ok, response})} + + {:error, conn, websocket, response} -> + state = %{state | conn: conn, websocket: websocket} + {:noreply, reply(state, {:error, response})} + + {:error, conn, response} -> + state = %{state | conn: conn} + {:noreply, reply(state, {:error, response})} + + {:error, _} -> + {:noreply, state} + end + end + + # Private + + defp reply(%{caller: nil} = state, response) do + Logger.warn("The caller is nil, so we can't reply the message: #{inspect(response)}") + state + end + + defp reply(state, response) do + GenServer.reply(state.caller, response) + + state + end +end diff --git a/lib/livebook_web/live/hub/edit_live.ex b/lib/livebook_web/live/hub/edit_live.ex index 0ad4d2585..681ccbf3d 100644 --- a/lib/livebook_web/live/hub/edit_live.ex +++ b/lib/livebook_web/live/hub/edit_live.ex @@ -18,13 +18,20 @@ defmodule LivebookWeb.Hub.EditLive do type = Provider.type(hub) if type == "local" do - {:ok, - socket |> redirect(to: "/") |> put_flash(:warning, "You can't edit the localhost Hub")} + {:noreply, + socket + |> redirect(to: "/") + |> put_flash(:warning, "You can't edit the localhost Hub")} else - {:ok, assign(socket, hub: hub, type: type, page_title: "Livebook - Hub", params: params)} + {:noreply, + assign(socket, + hub: hub, + type: type, + page_title: "Livebook - Hub", + params: params, + env_var_id: params["env_var_id"] + )} end - - {:noreply, assign(socket, hub: hub, type: type, env_var_id: params["env_var_id"])} end @impl true diff --git a/lib/livebook_web/live/hub/new/enterprise_component.ex b/lib/livebook_web/live/hub/new/enterprise_component.ex index 5bb16478f..457265ca4 100644 --- a/lib/livebook_web/live/hub/new/enterprise_component.ex +++ b/lib/livebook_web/live/hub/new/enterprise_component.ex @@ -4,7 +4,10 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do import Ecto.Changeset, only: [get_field: 2] alias Livebook.EctoTypes.HexColor - alias Livebook.Hubs.{Enterprise, EnterpriseClient} + alias Livebook.Hubs.Enterprise + alias Livebook.WebSocket + + @app_version Mix.Project.config()[:version] @impl true def update(assigns, socket) do @@ -104,25 +107,22 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do url = get_field(socket.assigns.changeset, :url) token = get_field(socket.assigns.changeset, :token) - with {:ok, _info} <- EnterpriseClient.fetch_info(url, token), - {:ok, %{"id" => id}} <- EnterpriseClient.fetch_me(url, token) do - base = %Enterprise{ - token: token, - url: url, - external_id: id, - hub_name: "Enterprise", - hub_color: HexColor.random() - } + case connect_with_enterprise(url, token) do + {:ok, {:session, session_response}} -> + base = %Enterprise{ + token: token, + url: url, + external_id: session_response.user.id, + hub_name: "Enterprise", + hub_color: HexColor.random() + } - changeset = Enterprise.change_hub(base) + changeset = Enterprise.change_hub(base) - {:noreply, assign(socket, changeset: changeset, base: base, connected: true)} - else - {:error, _message, reason} -> - {:noreply, - socket - |> put_flash(:error, message_from_reason(reason)) - |> push_patch(to: Routes.hub_path(socket, :new))} + {:noreply, assign(socket, changeset: changeset, base: base, connected: true)} + + {:error, reason} -> + handle_error(reason, socket) end end @@ -151,7 +151,34 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do {:noreply, assign(socket, changeset: Enterprise.change_hub(socket.assigns.base, attrs))} end - defp message_from_reason(:invalid_url), do: "Failed to connect with given URL" - defp message_from_reason(:unauthorized), do: "Failed to authorize with given token" - defp message_from_reason(:invalid_token), do: "Failed to authenticate with given token" + defp connect_with_enterprise(url, token) do + headers = [{"X-Auth-Token", token}] + session_request = LivebookProto.SessionRequest.new!(app_version: @app_version) + + with {:ok, connection, :connected} <- WebSocket.connect(url, headers), + {:ok, connection} <- WebSocket.send_request(connection, session_request), + {:ok, connection, session_response} <- WebSocket.receive_response(connection) do + WebSocket.disconnect(connection) + {:ok, session_response} + else + {:error, connection, response} -> + WebSocket.disconnect(connection) + {:error, response} + end + end + + def handle_error(%{reason: :econnrefused}, socket) do + show_connect_error("Failed to connect with given URL", socket) + end + + def handle_error(%{details: reason}, socket) do + show_connect_error(reason, socket) + end + + defp show_connect_error(message, socket) do + {:noreply, + socket + |> put_flash(:error, message) + |> push_patch(to: Routes.hub_path(socket, :new))} + end end diff --git a/mix.exs b/mix.exs index b2141f65f..3a2c0ef58 100644 --- a/mix.exs +++ b/mix.exs @@ -39,8 +39,8 @@ defmodule Livebook.MixProject do defp extra_applications(:app), do: [:wx] defp extra_applications(_), do: [] - defp elixirc_paths(:test), do: ["lib", "test/support"] - defp elixirc_paths(_), do: ["lib"] + defp elixirc_paths(:test), do: elixirc_paths(:dev) ++ ["test/support"] + defp elixirc_paths(_), do: ["lib", "proto/lib"] defp package do [ @@ -49,7 +49,7 @@ defmodule Livebook.MixProject do "GitHub" => "https://github.com/livebook-dev/livebook" }, files: - ~w(lib static config mix.exs mix.lock README.md LICENSE CHANGELOG.md iframe/priv/static/iframe) + ~w(lib static config mix.exs mix.lock README.md LICENSE CHANGELOG.md iframe/priv/static/iframe proto/lib) ] end @@ -57,7 +57,8 @@ defmodule Livebook.MixProject do [ "dev.setup": ["deps.get", "cmd npm install --prefix assets"], "dev.build": ["cmd npm run deploy --prefix ./assets"], - "format.all": ["format", "cmd npm run format --prefix ./assets"] + "format.all": ["format", "cmd npm run format --prefix ./assets"], + "protobuf.generate": ["cmd --cd proto mix protobuf.generate"] ] end @@ -100,6 +101,8 @@ defmodule Livebook.MixProject do {:aws_signature, "~> 0.3.0"}, {:ecto, "~> 3.9.0"}, {:phoenix_ecto, "~> 4.4.0"}, + {:mint_web_socket, "~> 1.0.0"}, + {:protobuf, "~> 0.8.0"}, {:phoenix_live_reload, "~> 1.2", only: :dev}, {:floki, ">= 0.27.0", only: :test}, {:bypass, "~> 2.1", only: :test} diff --git a/mix.lock b/mix.lock index 7ddf08d51..915c9a56e 100644 --- a/mix.lock +++ b/mix.lock @@ -8,14 +8,15 @@ "decimal": {:hex, :decimal, "2.0.0", "a78296e617b0f5dd4c6caf57c714431347912ffb1d0842e998e9792b5642d697", [:mix], [], "hexpm", "34666e9c55dea81013e77d9d87370fe6cb6291d1ef32f46a1600230b1d44f577"}, "earmark_parser": {:hex, :earmark_parser, "1.4.27", "755da957e2b980618ba3397d3f923004d85bac244818cf92544eaa38585cb3a8", [:mix], [], "hexpm", "8d02465c243ee96bdd655e7c9a91817a2a80223d63743545b2861023c4ff39ac"}, "ecto": {:hex, :ecto, "3.9.0", "7c74fc0d950a700eb7019057ff32d047ed7f19b57c1b2ca260cf0e565829101d", [:mix], [{:decimal, "~> 1.6 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "fed5ebc5831378b916afd0b5852a0c5bb3e7390665cc2b0ec8ab0c712495b73d"}, - "ex_machina": {:hex, :ex_machina, "2.7.0", "b792cc3127fd0680fecdb6299235b4727a4944a09ff0fa904cc639272cd92dc7", [:mix], [{:ecto, "~> 2.2 or ~> 3.0", [hex: :ecto, repo: "hexpm", optional: true]}, {:ecto_sql, "~> 3.0", [hex: :ecto_sql, repo: "hexpm", optional: true]}], "hexpm", "419aa7a39bde11894c87a615c4ecaa52d8f107bbdd81d810465186f783245bf8"}, "file_system": {:hex, :file_system, "0.2.10", "fb082005a9cd1711c05b5248710f8826b02d7d1784e7c3451f9c1231d4fc162d", [:mix], [], "hexpm", "41195edbfb562a593726eda3b3e8b103a309b733ad25f3d642ba49696bf715dc"}, "floki": {:hex, :floki, "0.33.1", "f20f1eb471e726342b45ccb68edb9486729e7df94da403936ea94a794f072781", [:mix], [{:html_entities, "~> 0.5.0", [hex: :html_entities, repo: "hexpm", optional: false]}], "hexpm", "461035fd125f13fdf30f243c85a0b1e50afbec876cbf1ceefe6fddd2e6d712c6"}, + "hpax": {:hex, :hpax, "0.1.1", "2396c313683ada39e98c20a75a82911592b47e5c24391363343bde74f82396ca", [:mix], [], "hexpm", "0ae7d5a0b04a8a60caf7a39fcf3ec476f35cc2cc16c05abea730d3ce6ac6c826"}, "html_entities": {:hex, :html_entities, "0.5.2", "9e47e70598da7de2a9ff6af8758399251db6dbb7eebe2b013f2bbd2515895c3c", [:mix], [], "hexpm", "c53ba390403485615623b9531e97696f076ed415e8d8058b1dbaa28181f4fdcc"}, "jason": {:hex, :jason, "1.4.0", "e855647bc964a44e2f67df589ccf49105ae039d4179db7f6271dfd3843dc27e6", [:mix], [{:decimal, "~> 1.0 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "79a3791085b2a0f743ca04cec0f7be26443738779d09302e01318f97bdb82121"}, - "libpe": {:hex, :libpe, "1.1.2", "16337b414c690e0ee9c49fe917b059622f001c399303102b98900c05c229cd9a", [:mix], [], "hexpm", "31df0639fafb603b20078c8db9596c8984f35a151c64ec2e483d9136ff9f428c"}, "mime": {:hex, :mime, "2.0.3", "3676436d3d1f7b81b5a2d2bd8405f412c677558c81b1c92be58c00562bb59095", [:mix], [], "hexpm", "27a30bf0db44d25eecba73755acf4068cbfe26a4372f9eb3e4ea3a45956bff6b"}, - "phoenix": {:hex, :phoenix, "1.6.13", "5b3152907afdb8d3a6cdafb4b149e8aa7aabbf1422fd9f7ef4c2a67ead57d24a", [:mix], [{:castore, ">= 0.0.0", [hex: :castore, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:phoenix_pubsub, "~> 2.0", [hex: :phoenix_pubsub, repo: "hexpm", optional: false]}, {:phoenix_view, "~> 1.0", [hex: :phoenix_view, repo: "hexpm", optional: false]}, {:plug, "~> 1.10", [hex: :plug, repo: "hexpm", optional: false]}, {:plug_cowboy, "~> 2.2", [hex: :plug_cowboy, repo: "hexpm", optional: true]}, {:plug_crypto, "~> 1.2", [hex: :plug_crypto, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "13d8806c31176e2066da4df2d7443c144211305c506ed110ad4044335b90171d"}, + "mint": {:hex, :mint, "1.4.2", "50330223429a6e1260b2ca5415f69b0ab086141bc76dc2fbf34d7c389a6675b2", [:mix], [{:castore, "~> 0.1.0", [hex: :castore, repo: "hexpm", optional: true]}, {:hpax, "~> 0.1.1", [hex: :hpax, repo: "hexpm", optional: false]}], "hexpm", "ce75a5bbcc59b4d7d8d70f8b2fc284b1751ffb35c7b6a6302b5192f8ab4ddd80"}, + "mint_web_socket": {:hex, :mint_web_socket, "1.0.0", "b33e534a938ec10736cef2b00cd485f6abd70aef68b9194f4d92fe2f7b8bba06", [:mix], [{:mint, "~> 1.4 and >= 1.4.1", [hex: :mint, repo: "hexpm", optional: false]}], "hexpm", "3d4fd81190fe60f16fef5ade89e008463d72e6a608a7f6af9041cd8b47458e30"}, + "phoenix": {:hex, :phoenix, "1.6.11", "29f3c0fd12fa1fc4d4b05e341578e55bc78d96ea83a022587a7e276884d397e4", [:mix], [{:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:phoenix_pubsub, "~> 2.0", [hex: :phoenix_pubsub, repo: "hexpm", optional: false]}, {:phoenix_view, "~> 1.0", [hex: :phoenix_view, repo: "hexpm", optional: false]}, {:plug, "~> 1.10", [hex: :plug, repo: "hexpm", optional: false]}, {:plug_cowboy, "~> 2.2", [hex: :plug_cowboy, repo: "hexpm", optional: true]}, {:plug_crypto, "~> 1.2", [hex: :plug_crypto, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "1664e34f80c25ea4918fbadd957f491225ef601c0e00b4e644b1a772864bfbc2"}, "phoenix_ecto": {:hex, :phoenix_ecto, "4.4.0", "0672ed4e4808b3fbed494dded89958e22fb882de47a97634c0b13e7b0b5f7720", [:mix], [{:ecto, "~> 3.3", [hex: :ecto, repo: "hexpm", optional: false]}, {:phoenix_html, "~> 2.14.2 or ~> 3.0", [hex: :phoenix_html, repo: "hexpm", optional: true]}, {:plug, "~> 1.9", [hex: :plug, repo: "hexpm", optional: false]}], "hexpm", "09864e558ed31ee00bd48fcc1d4fc58ae9678c9e81649075431e69dbabb43cc1"}, "phoenix_html": {:hex, :phoenix_html, "3.2.0", "1c1219d4b6cb22ac72f12f73dc5fad6c7563104d083f711c3fcd8551a1f4ae11", [:mix], [{:plug, "~> 1.5", [hex: :plug, repo: "hexpm", optional: true]}], "hexpm", "36ec97ba56d25c0136ef1992c37957e4246b649d620958a1f9fa86165f8bc54f"}, "phoenix_live_dashboard": {:hex, :phoenix_live_dashboard, "0.7.0", "9b5ab242e52c33596b132beaf97dccb9e59f7af941f41a22d0fa2465d0b63ab1", [:mix], [{:ecto, "~> 3.6.2 or ~> 3.7", [hex: :ecto, repo: "hexpm", optional: true]}, {:ecto_mysql_extras, "~> 0.5", [hex: :ecto_mysql_extras, repo: "hexpm", optional: true]}, {:ecto_psql_extras, "~> 0.7", [hex: :ecto_psql_extras, repo: "hexpm", optional: true]}, {:mime, "~> 1.6 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:phoenix_live_view, "~> 0.18.0", [hex: :phoenix_live_view, repo: "hexpm", optional: false]}, {:telemetry_metrics, "~> 0.6.0", [hex: :telemetry_metrics, repo: "hexpm", optional: false]}], "hexpm", "374d65e87e1e83528ea30852e34d4ad3022ddb92d642d43ec0b4e3c112046036"}, @@ -26,6 +27,7 @@ "plug": {:hex, :plug, "1.13.6", "187beb6b67c6cec50503e940f0434ea4692b19384d47e5fdfd701e93cadb4cc2", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:plug_crypto, "~> 1.1.1 or ~> 1.2", [hex: :plug_crypto, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.3 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "02b9c6b9955bce92c829f31d6284bf53c591ca63c4fb9ff81dfd0418667a34ff"}, "plug_cowboy": {:hex, :plug_cowboy, "2.5.2", "62894ccd601cf9597e2c23911ff12798a8a18d237e9739f58a6b04e4988899fe", [:mix], [{:cowboy, "~> 2.7", [hex: :cowboy, repo: "hexpm", optional: false]}, {:cowboy_telemetry, "~> 0.3", [hex: :cowboy_telemetry, repo: "hexpm", optional: false]}, {:plug, "~> 1.7", [hex: :plug, repo: "hexpm", optional: false]}], "hexpm", "ea6e87f774c8608d60c8d34022a7d073bd7680a0a013f049fc62bf35efea1044"}, "plug_crypto": {:hex, :plug_crypto, "1.2.3", "8f77d13aeb32bfd9e654cb68f0af517b371fb34c56c9f2b58fe3df1235c1251a", [:mix], [], "hexpm", "b5672099c6ad5c202c45f5a403f21a3411247f164e4a8fab056e5cd8a290f4a2"}, + "protobuf": {:hex, :protobuf, "0.8.0", "61b27d6fd50e7b1b2eb0ee17c1f639906121f4ef965ae0994644eb4c68d4647d", [:mix], [{:jason, "~> 1.2", [hex: :jason, repo: "hexpm", optional: true]}], "hexpm", "3644ed846fd6f5e3b5c2cd617aa8344641e230edf812a45365fee7622bccd25a"}, "ranch": {:hex, :ranch, "1.8.0", "8c7a100a139fd57f17327b6413e4167ac559fbc04ca7448e9be9057311597a1d", [:make, :rebar3], [], "hexpm", "49fbcfd3682fab1f5d109351b61257676da1a2fdbe295904176d5e521a2ddfe5"}, "telemetry": {:hex, :telemetry, "1.1.0", "a589817034a27eab11144ad24d5c0f9fab1f58173274b1e9bae7074af9cbee51", [:rebar3], [], "hexpm", "b727b2a1f75614774cff2d7565b64d0dfa5bd52ba517f16543e6fc7efcc0df48"}, "telemetry_metrics": {:hex, :telemetry_metrics, "0.6.1", "315d9163a1d4660aedc3fee73f33f1d355dcc76c5c3ab3d59e76e3edf80eef1f", [:mix], [{:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "7be9e0871c41732c233be71e4be11b96e56177bf15dde64a8ac9ce72ac9834c6"}, diff --git a/proto/.formatter.exs b/proto/.formatter.exs new file mode 100644 index 000000000..cea9b6d8d --- /dev/null +++ b/proto/.formatter.exs @@ -0,0 +1,5 @@ +# Used by "mix format" +[ + import_deps: [:protobuf], + inputs: ["{mix,.formatter}.exs", "{config,lib,test}/**/*.{ex,exs}"] +] diff --git a/proto/.gitignore b/proto/.gitignore new file mode 100644 index 000000000..17eae5ceb --- /dev/null +++ b/proto/.gitignore @@ -0,0 +1,26 @@ +# The directory Mix will write compiled artifacts to. +/_build/ + +# If you run "mix test --cover", coverage assets end up here. +/cover/ + +# The directory Mix downloads your dependencies sources to. +/deps/ + +# Where third-party dependencies like ExDoc output generated docs. +/doc/ + +# Ignore .fetch files in case you like to edit your project deps locally. +/.fetch + +# If the VM crashes, it generates a dump, let's ignore it too. +erl_crash.dump + +# Also ignore archive artifacts (built via "mix archive.build"). +*.ez + +# Ignore package tarball (built via "mix hex.build"). +livebook_proto-*.tar + +# Temporary files, for example, from tests. +/tmp/ diff --git a/proto/README.md b/proto/README.md new file mode 100644 index 000000000..c98086383 --- /dev/null +++ b/proto/README.md @@ -0,0 +1,11 @@ +# Livebook Protobuf + +Livebook uses Protobuf messages for cross-backend communication over WebSocket. + +## Generate message + +To generate the message modules, you'll need to run: + +```sh +mix protobuf.generate +``` diff --git a/proto/lib/livebook_proto.ex b/proto/lib/livebook_proto.ex new file mode 100644 index 000000000..054ef5eca --- /dev/null +++ b/proto/lib/livebook_proto.ex @@ -0,0 +1,10 @@ +defmodule LivebookProto do + @moduledoc false + + @mapping (for {_id, field_prop} <- LivebookProto.Request.__message_props__().field_props, + into: %{} do + {field_prop.type, field_prop.name_atom} + end) + + def request_type(module), do: Map.fetch!(@mapping, module) +end diff --git a/proto/lib/livebook_proto/error.pb.ex b/proto/lib/livebook_proto/error.pb.ex new file mode 100644 index 000000000..00558f6d5 --- /dev/null +++ b/proto/lib/livebook_proto/error.pb.ex @@ -0,0 +1,6 @@ +defmodule LivebookProto.Error do + @moduledoc false + use Protobuf, protoc_gen_elixir_version: "0.11.0", syntax: :proto3 + + field :details, 1, type: :string +end diff --git a/proto/lib/livebook_proto/request.pb.ex b/proto/lib/livebook_proto/request.pb.ex new file mode 100644 index 000000000..ddeb3a4e9 --- /dev/null +++ b/proto/lib/livebook_proto/request.pb.ex @@ -0,0 +1,8 @@ +defmodule LivebookProto.Request do + @moduledoc false + use Protobuf, protoc_gen_elixir_version: "0.11.0", syntax: :proto3 + + oneof :type, 0 + + field :session, 1, type: LivebookProto.SessionRequest, oneof: 0 +end diff --git a/proto/lib/livebook_proto/response.pb.ex b/proto/lib/livebook_proto/response.pb.ex new file mode 100644 index 000000000..e836076db --- /dev/null +++ b/proto/lib/livebook_proto/response.pb.ex @@ -0,0 +1,9 @@ +defmodule LivebookProto.Response do + @moduledoc false + use Protobuf, protoc_gen_elixir_version: "0.11.0", syntax: :proto3 + + oneof :type, 0 + + field :error, 1, type: LivebookProto.Error, oneof: 0 + field :session, 2, type: LivebookProto.SessionResponse, oneof: 0 +end diff --git a/proto/lib/livebook_proto/session_request.pb.ex b/proto/lib/livebook_proto/session_request.pb.ex new file mode 100644 index 000000000..bd5d0bc55 --- /dev/null +++ b/proto/lib/livebook_proto/session_request.pb.ex @@ -0,0 +1,6 @@ +defmodule LivebookProto.SessionRequest do + @moduledoc false + use Protobuf, protoc_gen_elixir_version: "0.11.0", syntax: :proto3 + + field :app_version, 1, type: :string, json_name: "appVersion" +end diff --git a/proto/lib/livebook_proto/session_response.pb.ex b/proto/lib/livebook_proto/session_response.pb.ex new file mode 100644 index 000000000..a87fd96e6 --- /dev/null +++ b/proto/lib/livebook_proto/session_response.pb.ex @@ -0,0 +1,7 @@ +defmodule LivebookProto.SessionResponse do + @moduledoc false + use Protobuf, protoc_gen_elixir_version: "0.11.0", syntax: :proto3 + + field :id, 1, type: :string + field :user, 2, type: LivebookProto.User +end diff --git a/proto/lib/livebook_proto/user.pb.ex b/proto/lib/livebook_proto/user.pb.ex new file mode 100644 index 000000000..78256da66 --- /dev/null +++ b/proto/lib/livebook_proto/user.pb.ex @@ -0,0 +1,7 @@ +defmodule LivebookProto.User do + @moduledoc false + use Protobuf, protoc_gen_elixir_version: "0.11.0", syntax: :proto3 + + field :id, 1, type: :string + field :email, 2, type: :string +end diff --git a/proto/messages.proto b/proto/messages.proto new file mode 100644 index 000000000..6dfd50787 --- /dev/null +++ b/proto/messages.proto @@ -0,0 +1,33 @@ +syntax = "proto3"; + +message User { + string id = 1; + string email = 2; +} + +message Error { + string details = 1; +} + +message SessionRequest { + string app_version = 1; +} + +message SessionResponse { + string id = 1; + User user = 2; +} + +message Request { + oneof type { + SessionRequest session = 1; + } +} + +message Response { + oneof type { + Error error = 1; + + SessionResponse session = 2; + } +} diff --git a/proto/mix.exs b/proto/mix.exs new file mode 100644 index 000000000..038074686 --- /dev/null +++ b/proto/mix.exs @@ -0,0 +1,33 @@ +defmodule LivebookProto.MixProject do + use Mix.Project + + def project do + [ + app: :livebook_proto, + version: "0.1.0", + elixir: "~> 1.14", + start_permanent: Mix.env() == :prod, + deps: deps(), + aliases: aliases() + ] + end + + def application do + [extra_applications: [:logger]] + end + + defp deps do + [{:protobuf, "~> 0.8.0"}] + end + + defp aliases do + [ + "protobuf.generate": [ + "cmd protoc --elixir_out=one_file_per_module=true:lib --elixir_opt=include_docs=true --elixir_opt=gen_struct=true --elixir_opt=package_prefix=livebook_proto messages.proto", + "format lib/livebook_proto/*.pb.ex", + "deps.get", + "compile" + ] + ] + end +end diff --git a/proto/mix.lock b/proto/mix.lock new file mode 100644 index 000000000..c392371d6 --- /dev/null +++ b/proto/mix.lock @@ -0,0 +1,3 @@ +%{ + "protobuf": {:hex, :protobuf, "0.8.0", "61b27d6fd50e7b1b2eb0ee17c1f639906121f4ef965ae0994644eb4c68d4647d", [:mix], [{:jason, "~> 1.2", [hex: :jason, repo: "hexpm", optional: true]}], "hexpm", "3644ed846fd6f5e3b5c2cd617aa8344641e230edf812a45365fee7622bccd25a"}, +} diff --git a/test/livebook/hubs/enterprise_client_test.exs b/test/livebook/hubs/enterprise_client_test.exs deleted file mode 100644 index 7a15a7ddb..000000000 --- a/test/livebook/hubs/enterprise_client_test.exs +++ /dev/null @@ -1,25 +0,0 @@ -defmodule Livebook.Hubs.EnterpriseClientTest do - use Livebook.EnterpriseIntegrationCase, async: true - - alias Livebook.Hubs.EnterpriseClient - - describe "fetch_info/1" do - test "fetches the token info", %{url: url, token: token} do - assert {:ok, %{"id" => _}} = EnterpriseClient.fetch_info(url, token) - end - - test "returns invalid_token when token is invalid", %{url: url} do - assert {:error, _, :invalid_token} = EnterpriseClient.fetch_info(url, "foo") - end - end - - describe "fetch_me/1" do - test "fetches the current user id", %{url: url, token: token} do - assert {:ok, %{"id" => _}} = EnterpriseClient.fetch_me(url, token) - end - - test "returns unauthorized when token is invalid", %{url: url} do - assert {:error, _, :unauthorized} = EnterpriseClient.fetch_me(url, "foo") - end - end -end diff --git a/test/livebook/web_socket/client_test.exs b/test/livebook/web_socket/client_test.exs new file mode 100644 index 000000000..b254b2d6d --- /dev/null +++ b/test/livebook/web_socket/client_test.exs @@ -0,0 +1,42 @@ +defmodule Livebook.WebSocket.ClientTest do + use Livebook.EnterpriseIntegrationCase, async: true + + alias Livebook.WebSocket.Client + + describe "connect/2" do + test "successfully authenticates the websocket connection", %{url: url, token: token} do + headers = [{"X-Auth-Token", token}] + + assert {:ok, conn, ref} = Client.connect(url, headers) + assert {:ok, conn, websocket, :connected} = Client.receive(conn, ref) + assert Client.disconnect(conn, websocket, ref) == :ok + end + + test "rejects the websocket with invalid address", %{token: token} do + headers = [{"X-Auth-Token", token}] + + assert {:error, %Mint.TransportError{reason: :econnrefused}} = + Client.connect("http://localhost:9999", headers) + end + + test "rejects the websocket connection with invalid credentials", %{url: url} do + headers = [{"X-Auth-Token", "foo"}] + + assert {:ok, conn, ref} = Client.connect(url, headers) + assert {:error, _conn, response} = Client.receive(conn, ref) + + assert response.status == 403 + + assert %{type: {:error, %{details: "Invalid Token"}}} = + LivebookProto.Response.decode(response.body) + + assert {:ok, conn, ref} = Client.connect(url) + assert {:error, _conn, response} = Client.receive(conn, ref) + + assert response.status == 401 + + assert %{type: {:error, %{details: "Token not found"}}} = + LivebookProto.Response.decode(response.body) + end + end +end diff --git a/test/livebook/web_socket/server_test.exs b/test/livebook/web_socket/server_test.exs new file mode 100644 index 000000000..13b69817d --- /dev/null +++ b/test/livebook/web_socket/server_test.exs @@ -0,0 +1,47 @@ +defmodule Livebook.WebSocket.ServerTest do + use Livebook.EnterpriseIntegrationCase, async: true + + alias Livebook.WebSocket.Server + + describe "connect/2" do + test "successfully authenticates the websocket connection", %{url: url, token: token} do + headers = [{"X-Auth-Token", token}] + + assert {:ok, pid} = Server.start_link() + assert {:ok, :connected} = Server.connect(pid, url, headers) + assert Server.disconnect(pid) == :ok + end + + test "rejects the websocket with invalid address", %{token: token} do + headers = [{"X-Auth-Token", token}] + + assert {:ok, pid} = Server.start_link() + + assert {:error, %Mint.TransportError{reason: :econnrefused}} = + Server.connect(pid, "http://localhost:9999", headers) + + assert Server.disconnect(pid) == :ok + end + + test "rejects the websocket connection with invalid credentials", %{url: url} do + headers = [{"X-Auth-Token", "foo"}] + + assert {:ok, pid} = Server.start_link() + assert {:error, response} = Server.connect(pid, url, headers) + + assert response.status == 403 + + assert %{type: {:error, %{details: "Invalid Token"}}} = + LivebookProto.Response.decode(response.body) + + assert {:error, response} = Server.connect(pid, url) + + assert response.status == 401 + + assert %{type: {:error, %{details: "Token not found"}}} = + LivebookProto.Response.decode(response.body) + + assert Server.disconnect(pid) == :ok + end + end +end diff --git a/test/livebook/web_socket_test.exs b/test/livebook/web_socket_test.exs new file mode 100644 index 000000000..b997b4be4 --- /dev/null +++ b/test/livebook/web_socket_test.exs @@ -0,0 +1,47 @@ +defmodule Livebook.WebSocketTest do + use Livebook.EnterpriseIntegrationCase, async: true + + @app_version Mix.Project.config()[:version] + + alias Livebook.WebSocket + + describe "authentication" do + test "successfully authenticates the web socket connection", %{url: url, token: token} do + headers = [{"X-Auth-Token", token}] + + assert {:ok, connection, :connected} = WebSocket.connect(url, headers) + assert WebSocket.disconnect(connection) == :ok + end + + test "rejects the web socket connection with invalid credentials", %{url: url} do + headers = [{"X-Auth-Token", "foo"}] + + assert {:error, connection, %{details: "Invalid Token"}} = WebSocket.connect(url, headers) + assert WebSocket.disconnect(connection) == :ok + + assert {:error, connection, %{details: "Token not found"}} = WebSocket.connect(url) + assert WebSocket.disconnect(connection) == :ok + end + end + + describe "send_request/2" do + test "receives the session response from server", %{url: url, token: token} do + headers = [{"X-Auth-Token", token}] + + assert {:ok, %WebSocket.Connection{} = connection, :connected} = + WebSocket.connect(url, headers) + + session_request = LivebookProto.SessionRequest.new!(app_version: @app_version) + + assert {:ok, %WebSocket.Connection{} = connection} = + WebSocket.send_request(connection, session_request) + + assert {:ok, connection, {:session, session_response}} = + WebSocket.receive_response(connection) + + assert WebSocket.disconnect(connection) == :ok + + assert session_response.user.email == "jake.peralta@mail.com" + end + end +end diff --git a/test/livebook_web/live/hub/new/enterprise_component_test.exs b/test/livebook_web/live/hub/new/enterprise_component_test.exs index e85e025bb..c305333ba 100644 --- a/test/livebook_web/live/hub/new/enterprise_component_test.exs +++ b/test/livebook_web/live/hub/new/enterprise_component_test.exs @@ -7,6 +7,8 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponentTest do describe "enterprise" do test "persists new hub", %{conn: conn, url: url, token: token} do + Livebook.Hubs.delete_hub("enterprise-bf1587a3-4501-4729-9f53-43679381e28b") + {:ok, view, _html} = live(conn, Routes.hub_path(conn, :new)) assert view @@ -22,9 +24,11 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponentTest do } }) - assert view - |> element("#connect") - |> render_click() =~ "Add Hub" + view + |> element("#connect") + |> render_click() + + assert render(view) =~ "bf1587a3-4501-4729-9f53-43679381e28b" attrs = %{ "url" => url, @@ -56,6 +60,31 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponentTest do assert hubs_html =~ "Enterprise" end + test "fails with invalid token", %{conn: conn, url: url} do + {:ok, view, _html} = live(conn, Routes.hub_path(conn, :new)) + token = "foo bar baz" + + assert view + |> element("#enterprise") + |> render_click() =~ "2. Configure your Hub" + + view + |> element("#enterprise-form") + |> render_change(%{ + "enterprise" => %{ + "url" => url, + "token" => token + } + }) + + view + |> element("#connect") + |> render_click() + + assert render(view) =~ "Invalid Token" + refute render(view) =~ "enterprise[hub_name]" + end + test "fails to create existing hub", %{conn: conn, url: url, token: token} do hub = insert_hub(:enterprise, @@ -80,9 +109,11 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponentTest do } }) - assert view - |> element("#connect") - |> render_click() =~ "Add Hub" + view + |> element("#connect") + |> render_click() + + assert render(view) =~ "bf1587a3-4501-4729-9f53-43679381e28b" attrs = %{ "url" => url, diff --git a/test/livebook_web/live/hub/new_live_test.exs b/test/livebook_web/live/hub/new_live_test.exs index cdbdbdbae..45adc3b58 100644 --- a/test/livebook_web/live/hub/new_live_test.exs +++ b/test/livebook_web/live/hub/new_live_test.exs @@ -118,124 +118,6 @@ defmodule LivebookWeb.Hub.NewLiveTest do end end - describe "enterprise" do - test "persists new hub", %{conn: conn} do - id = Livebook.Utils.random_short_id() - bypass = enterprise_bypass(id) - - {:ok, view, _html} = live(conn, Routes.hub_path(conn, :new)) - - assert view - |> element("#enterprise") - |> render_click() =~ "2. Configure your Hub" - - view - |> element("#enterprise-form") - |> render_change(%{ - "enterprise" => %{ - "url" => "http://localhost:#{bypass.port}", - "token" => "dummy access token" - } - }) - - assert view - |> element("#connect") - |> render_click() =~ "Add Hub" - - attrs = %{ - "url" => "http://localhost:#{bypass.port}", - "token" => "dummy access token", - "hub_name" => "Enterprise", - "hub_color" => "#FF00FF" - } - - view - |> element("#enterprise-form") - |> render_change(%{"enterprise" => attrs}) - - refute view - |> element("#enterprise-form .invalid-feedback") - |> has_element?() - - assert {:ok, view, _html} = - view - |> element("#enterprise-form") - |> render_submit(%{"enterprise" => attrs}) - |> follow_redirect(conn) - - assert render(view) =~ "Hub added successfully" - - assert view - |> element("#hubs") - |> render() =~ ~s/style="color: #FF00FF"/ - - assert view - |> element("#hubs") - |> render() =~ "/hub/enterprise-#{id}" - - assert view - |> element("#hubs") - |> render() =~ "Enterprise" - end - - test "fails to create existing hub", %{conn: conn} do - hub = insert_hub(:enterprise, id: "enterprise-foo", external_id: "foo") - bypass = enterprise_bypass(hub.external_id) - - {:ok, view, _html} = live(conn, Routes.hub_path(conn, :new)) - - assert view - |> element("#enterprise") - |> render_click() =~ "2. Configure your Hub" - - view - |> element("#enterprise-form") - |> render_change(%{ - "enterprise" => %{ - "url" => "http://localhost:#{bypass.port}", - "token" => "dummy access token" - } - }) - - assert view - |> element("#connect") - |> render_click() =~ "Add Hub" - - attrs = %{ - "url" => "http://localhost:#{bypass.port}", - "token" => "dummy access token", - "hub_name" => "Enterprise", - "hub_color" => "#FF00FF" - } - - view - |> element("#enterprise-form") - |> render_change(%{"enterprise" => attrs}) - - refute view - |> element("#enterprise-form .invalid-feedback") - |> has_element?() - - assert view - |> element("#enterprise-form") - |> render_submit(%{"enterprise" => attrs}) =~ "already exists" - - assert view - |> element("#hubs") - |> render() =~ ~s/style="color: #{hub.hub_color}"/ - - assert view - |> element("#hubs") - |> render() =~ Routes.hub_path(conn, :edit, hub.id) - - assert view - |> element("#hubs") - |> render() =~ hub.hub_name - - assert Hubs.fetch_hub!(hub.id) == hub - end - end - defp fly_bypass(app_id) do bypass = Bypass.open() Application.put_env(:livebook, :fly_graphql_endpoint, "http://localhost:#{bypass.port}") @@ -256,37 +138,6 @@ defmodule LivebookWeb.Hub.NewLiveTest do end) end - defp enterprise_bypass(id) do - bypass = Bypass.open() - - Bypass.expect(bypass, "POST", "/api/v1", fn conn -> - {:ok, body, conn} = Plug.Conn.read_body(conn) - body = Jason.decode!(body) - - response = - cond do - body["query"] =~ "info" -> - %{ - "data" => %{ - "info" => %{ - "id" => Livebook.Utils.random_short_id(), - "expire_at" => to_string(DateTime.utc_now()) - } - } - } - - body["query"] =~ "me" -> - %{"data" => %{"me" => %{"id" => id}}} - end - - conn - |> Plug.Conn.put_resp_content_type("application/json") - |> Plug.Conn.resp(200, Jason.encode!(response)) - end) - - bypass - end - defp fetch_apps_response(app_id) do app = %{ "id" => app_id, diff --git a/test/support/integration/enterprise_server.ex b/test/support/integration/enterprise_server.ex index 466fec8c6..52c242701 100644 --- a/test/support/integration/enterprise_server.ex +++ b/test/support/integration/enterprise_server.ex @@ -60,7 +60,8 @@ defmodule Livebook.EnterpriseServer do env = [ {~c"MIX_ENV", ~c"livebook"}, - {~c"LIVEBOOK_ENTERPRISE_PORT", String.to_charlist(app_port())} + {~c"LIVEBOOK_ENTERPRISE_PORT", String.to_charlist(app_port())}, + {~c"LIVEBOOK_ENTERPRISE_DEBUG", String.to_charlist(debug?())} ] args = [ @@ -117,6 +118,10 @@ defmodule Livebook.EnterpriseServer do System.get_env("ENTERPRISE_PORT", "4043") end + defp debug? do + System.get_env("ENTERPRISE_DEBUG", "false") + end + defp wait_on_start(port) do case :httpc.request(:get, {~c"#{url()}/public/health", []}, [], []) do {:ok, _} -> @@ -131,7 +136,8 @@ defmodule Livebook.EnterpriseServer do defp mix(args, opts \\ []) do env = [ {"MIX_ENV", "livebook"}, - {"LIVEBOOK_ENTERPRISE_PORT", app_port()} + {"LIVEBOOK_ENTERPRISE_PORT", app_port()}, + {"LIVEBOOK_ENTERPRISE_DEBUG", debug?()} ] cmd_opts = [stderr_to_stdout: true, env: env, cd: app_dir()]