Add TestServer.SSH

Author: richard-ash

README.md

@@ -12,6 +12,7 @@ Features:
 - HTTP/1
 - HTTP/2
 - WebSocket
+- SSH
 - Built-in TLS with self-signed certificates
 - Plug route matching
 
@@ -140,6 +141,51 @@ end
 
 *Note: WebSocket is not supported by the `:httpd` adapter.*
 
+### SSH
+
+SSH exec and shell handlers can be registered with `TestServer.SSH.exec/1` and `TestServer.SSH.shell/1`. The server autostarts on first use and is torn down when the test exits.
+
+```elixir
+test "run remote command" do
+  TestServer.SSH.exec(to: fn _cmd, state ->
+    {:reply, {0, "file1\nfile2\n", ""}, state}
+  end)
+
+  {host, port} = TestServer.SSH.address()
+  {:ok, conn} = :ssh.connect(String.to_charlist(host), port,
+    user: ~c"test",
+    silently_accept_hosts: true,
+    user_interaction: false
+  )
+
+  {:ok, ch} = :ssh_connection.session_channel(conn, :infinity)
+  :success = :ssh_connection.exec(conn, ch, ~c"ls", :infinity)
+  # collect stdout, exit_status, closed messages...
+end
+```
+
+Password and public key credentials can be set with the `:credentials` option:
+
+```elixir
+# Password auth
+TestServer.SSH.start(credentials: [{"alice", "secret"}])
+
+# Public key auth
+TestServer.SSH.start(credentials: [{"bob", :public_key, pem_binary}])
+
+# No auth (default — accepts any connection)
+TestServer.SSH.start()
+```
+
+Handlers are matched FIFO and can be filtered with `:match`:
+
+```elixir
+TestServer.SSH.exec(
+  match: fn cmd, _state -> cmd == "ls" end,
+  to: fn _cmd, state -> {:reply, {0, "file1\n", ""}, state} end
+)
+```
+
 ### HTTP Server Adapter
 
 TestServer supports `Bandit`, `Plug.Cowboy`, and `:httpd` out of the box. The HTTP adapter will be selected in this order depending which is available in the dependencies. You can also explicitly set the http server in the configuration when calling `TestServer.HTTP.start/1`:

lib/test_server/ssh.ex

@@ -0,0 +1,200 @@
+defmodule TestServer.SSH do
+  @moduledoc false
+
+  alias TestServer.{InstanceManager, SSH}
+
+  @spec start(keyword()) :: {:ok, pid()}
+  def start(options \\ []) do
+    case ExUnit.fetch_test_supervisor() do
+      {:ok, _sup} -> start_with_ex_unit(options)
+      :error -> raise ArgumentError, "can only be called in a test process"
+    end
+  end
+
+  defp start_with_ex_unit(options) do
+    [_first_module_entry | stacktrace] = get_stacktrace()
+    caller = self()
+
+    options =
+      options
+      |> Keyword.put_new(:caller, caller)
+      |> Keyword.put_new(:stacktrace, stacktrace)
+
+    case InstanceManager.start_instance(caller, SSH.Instance.child_spec(options)) do
+      {:ok, instance} ->
+        put_ex_unit_on_exit_callback(instance)
+        {:ok, instance}
+
+      {:error, error} ->
+        raise_start_failure({:error, error})
+    end
+  end
+
+  defp put_ex_unit_on_exit_callback(instance) do
+    ExUnit.Callbacks.on_exit(fn ->
+      if Process.alive?(instance) do
+        verify_handlers!(:exec, instance)
+        verify_handlers!(:shell, instance)
+        stop(instance)
+      end
+    end)
+  end
+
+  defp verify_handlers!(type, instance) do
+    handlers_fn =
+      if type == :exec, do: &SSH.Instance.exec_handlers/1, else: &SSH.Instance.shell_handlers/1
+
+    instance
+    |> handlers_fn.()
+    |> Enum.reject(& &1.suspended)
+    |> case do
+      [] ->
+        :ok
+
+      active ->
+        raise """
+        #{SSH.Instance.format_instance(instance)} did not receive #{type} requests for these handlers before the test ended:
+
+        #{SSH.Instance.format_handlers(active)}
+        """
+    end
+  end
+
+  @spec stop(pid()) :: :ok | {:error, term()}
+  def stop(instance) do
+    instance_alive!(instance)
+    InstanceManager.stop_instance(instance)
+  end
+
+  @spec address() :: {binary(), :inet.port_number()}
+  def address, do: address(fetch_instance!())
+
+  @spec address(pid()) :: {binary(), :inet.port_number()}
+  def address(instance) do
+    instance_alive!(instance)
+    options = SSH.Instance.get_options(instance)
+    {"localhost", Keyword.fetch!(options, :port)}
+  end
+
+  @spec exec(keyword()) :: :ok
+  def exec(options) when is_list(options) do
+    {:ok, instance} = autostart()
+    exec(instance, options)
+  end
+
+  @spec exec(pid(), keyword()) :: :ok
+  def exec(instance, options) when is_pid(instance) and is_list(options) do
+    instance_alive!(instance)
+    [_first_module_entry | stacktrace] = get_stacktrace()
+    options = Keyword.put_new(options, :to, &default_exec_handler/2)
+    {:ok, _handler} = SSH.Instance.register(instance, {:exec, options, stacktrace})
+    :ok
+  end
+
+  defp default_exec_handler(_cmd, state), do: {:reply, {0, "", ""}, state}
+
+  @spec shell(keyword()) :: :ok
+  def shell(options) when is_list(options) do
+    {:ok, instance} = autostart()
+    shell(instance, options)
+  end
+
+  @spec shell(pid(), keyword()) :: :ok
+  def shell(instance, options) when is_pid(instance) and is_list(options) do
+    instance_alive!(instance)
+    [_first_module_entry | stacktrace] = get_stacktrace()
+    options = Keyword.put_new(options, :to, &default_shell_handler/2)
+    {:ok, _handler} = SSH.Instance.register(instance, {:shell, options, stacktrace})
+    :ok
+  end
+
+  defp default_shell_handler(data, state), do: {:reply, data, state}
+
+  defp autostart do
+    case fetch_instance() do
+      :error -> start()
+      {:ok, instance} -> {:ok, instance}
+    end
+  end
+
+  defp fetch_instance! do
+    case fetch_instance() do
+      :error -> raise "No current #{inspect(SSH.Instance)} running"
+      {:ok, instance} -> instance
+    end
+  end
+
+  defp fetch_instance do
+    instances = InstanceManager.get_by_caller(self()) || []
+    ssh_instances = Enum.filter(instances, &ssh_instance?/1)
+
+    case ssh_instances do
+      [] ->
+        :error
+
+      [instance] ->
+        {:ok, instance}
+
+      [_first | _rest] = multiple ->
+        [{m, f, a, _} | _] = get_stacktrace()
+
+        formatted =
+          multiple
+          |> Enum.map(&{&1, SSH.Instance.get_options(&1)})
+          |> Enum.with_index()
+          |> Enum.map_join("\n\n", fn {{instance, options}, index} ->
+            """
+            ##{index + 1}: #{SSH.Instance.format_instance(instance)}
+                #{Enum.map_join(options[:stacktrace], "\n    ", &Exception.format_stacktrace_entry/1)}
+            """
+          end)
+
+        raise """
+        Multiple #{inspect(SSH.Instance)}'s running, please pass instance to `#{inspect(m)}.#{f}/#{a}`.
+
+        #{formatted}
+        """
+    end
+  end
+
+  defp ssh_instance?(pid) do
+    SSH.Instance.get_options(pid)[:protocol] == :ssh
+  rescue
+    _ -> false
+  end
+
+  defp instance_alive!(instance) do
+    unless Process.alive?(instance),
+      do: raise("#{SSH.Instance.format_instance(instance)} is not running")
+  end
+
+  defp raise_start_failure({:error, {{:EXIT, reason}, _spec}}) do
+    raise_start_failure({:error, reason})
+  end
+
+  defp raise_start_failure({:error, error}) do
+    raise """
+    EXIT when starting #{inspect(SSH.Instance)}:
+
+    #{Exception.format_exit(error)}
+    """
+  end
+
+  defp get_stacktrace do
+    {:current_stacktrace, [{Process, :info, _, _} | stacktrace]} =
+      Process.info(self(), :current_stacktrace)
+
+    first_module_entry =
+      stacktrace
+      |> Enum.reverse()
+      |> Enum.find(fn {mod, _, _, _} -> mod == __MODULE__ end)
+
+    [first_module_entry] ++ prune_stacktrace(stacktrace)
+  end
+
+  defp prune_stacktrace([{__MODULE__, _, _, _} | t]), do: prune_stacktrace(t)
+  defp prune_stacktrace([{ExUnit.Assertions, _, _, _} | t]), do: prune_stacktrace(t)
+  defp prune_stacktrace([{ExUnit.Runner, _, _, _} | _]), do: []
+  defp prune_stacktrace([h | t]), do: [h | prune_stacktrace(t)]
+  defp prune_stacktrace([]), do: []
+end

lib/test_server/ssh/channel.ex

@@ -0,0 +1,120 @@
+defmodule TestServer.SSH.Channel do
+  @moduledoc false
+
+  @behaviour :ssh_server_channel
+
+  alias TestServer.SSH.Instance
+
+  defstruct [:instance, :channel_id, :connection, type: nil, handler_state: %{}]
+
+  @impl true
+  def init(instance: instance) do
+    {:ok, %__MODULE__{instance: instance}}
+  end
+
+  @impl true
+  def handle_msg({:ssh_channel_up, channel_id, connection}, state) do
+    {:ok, %{state | channel_id: channel_id, connection: connection}}
+  end
+
+  def handle_msg(_msg, state) do
+    {:ok, state}
+  end
+
+  @impl true
+  def handle_ssh_msg({:ssh_cm, conn, {:exec, ch_id, want_reply, command}}, state) do
+    command = to_string(command)
+    :ssh_connection.reply_request(conn, want_reply, :success, ch_id)
+
+    case GenServer.call(state.instance, {:dispatch, {:exec, command, state.handler_state}}) do
+      {:ok, {:reply, {exit_code, stdout, stderr}, new_handler_state}} ->
+        unless IO.iodata_length(stdout) == 0,
+          do: :ssh_connection.send(conn, ch_id, stdout)
+
+        unless IO.iodata_length(stderr) == 0,
+          do: :ssh_connection.send(conn, ch_id, 1, stderr)
+
+        :ssh_connection.exit_status(conn, ch_id, exit_code)
+        :ssh_connection.send_eof(conn, ch_id)
+        :ssh_connection.close(conn, ch_id)
+        {:stop, ch_id, %{state | handler_state: new_handler_state}}
+
+      {:ok, {:ok, new_handler_state}} ->
+        :ssh_connection.exit_status(conn, ch_id, 0)
+        :ssh_connection.send_eof(conn, ch_id)
+        :ssh_connection.close(conn, ch_id)
+        {:stop, ch_id, %{state | handler_state: new_handler_state}}
+
+      {:error, :not_found} ->
+        message =
+          "#{Instance.format_instance(state.instance)} received an unexpected SSH exec request: #{inspect(command)}"
+
+        report_error_and_close_exec(conn, ch_id, state, RuntimeError.exception(message), [])
+
+      {:error, {exception, stacktrace}} ->
+        report_error_and_close_exec(conn, ch_id, state, exception, stacktrace)
+    end
+  end
+
+  def handle_ssh_msg({:ssh_cm, conn, {:shell, ch_id, want_reply}}, state) do
+    :ssh_connection.reply_request(conn, want_reply, :success, ch_id)
+    {:ok, %{state | type: :shell, channel_id: ch_id, connection: conn}}
+  end
+
+  def handle_ssh_msg({:ssh_cm, conn, {:data, ch_id, 0, data}}, %{type: :shell} = state) do
+    :ssh_connection.adjust_window(conn, ch_id, byte_size(data))
+
+    case GenServer.call(state.instance, {:dispatch, {:shell, data, state.handler_state}}) do
+      {:ok, {:reply, output, new_handler_state}} ->
+        :ssh_connection.send(conn, ch_id, output)
+        {:ok, %{state | handler_state: new_handler_state}}
+
+      {:ok, {:ok, new_handler_state}} ->
+        {:ok, %{state | handler_state: new_handler_state}}
+
+      {:error, :not_found} ->
+        message =
+          "#{Instance.format_instance(state.instance)} received unexpected SSH shell data: #{inspect(data)}"
+
+        Instance.report_error(state.instance, {RuntimeError.exception(message), []})
+        {:ok, state}
+
+      {:error, {exception, stacktrace}} ->
+        Instance.report_error(state.instance, {exception, stacktrace})
+        {:ok, state}
+    end
+  end
+
+  def handle_ssh_msg({:ssh_cm, conn, {:pty, ch_id, want_reply, _pty_info}}, state) do
+    :ssh_connection.reply_request(conn, want_reply, :success, ch_id)
+    {:ok, state}
+  end
+
+  def handle_ssh_msg({:ssh_cm, conn, {:env, ch_id, want_reply, _name, _value}}, state) do
+    :ssh_connection.reply_request(conn, want_reply, :success, ch_id)
+    {:ok, state}
+  end
+
+  def handle_ssh_msg({:ssh_cm, _conn, {:eof, _ch_id}}, state) do
+    {:ok, state}
+  end
+
+  def handle_ssh_msg({:ssh_cm, _conn, {:closed, ch_id}}, state) do
+    {:stop, ch_id, state}
+  end
+
+  def handle_ssh_msg(_msg, state) do
+    {:ok, state}
+  end
+
+  @impl true
+  def terminate(_reason, _state), do: :ok
+
+  defp report_error_and_close_exec(conn, ch_id, state, exception, stacktrace) do
+    Instance.report_error(state.instance, {exception, stacktrace})
+    :ssh_connection.exit_status(conn, ch_id, 1)
+    :ssh_connection.send_eof(conn, ch_id)
+    :ssh_connection.close(conn, ch_id)
+    {:stop, ch_id, state}
+  end
+end