websockex

defmodule WebSockex do
  @handshake_guid "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

  @moduledoc ~S"""
  A client handles negotiating the connection, then sending frames, receiving
  frames, closing, and reconnecting that connection.

  A simple client implementation would be:

  ```
  defmodule WsClient do
    use WebSockex

    def start_link(url, state) do
      WebSockex.start_link(url, __MODULE__, state)
    end

    def handle_frame({:text, msg}, state) do
      IO.puts "Received a message: #{msg}"
      {:ok, state}
    end
  end
  ```
  """

  @type frame :: {:ping | :ping, nil | message :: binary}
                 | {:text | :binary, message :: binary}

  @typedoc """
  The frame sent when the negotiating a connection closure.
  """
  @type close_frame :: {close_code, message :: binary}

  @typedoc """
  An integer between 1000 and 4999 that specifies the reason for closing the connection.
  """
  @type close_code :: integer

  @type options :: [option]

  @typedoc """
  Options values for `start_link`.

  - `:async` - Replies with `{:ok, pid}` before establishing the connection.
    This is useful for when attempting to connect indefinitely, this way the
    process doesn't block trying to establish a connection.
  - `:handle_initial_conn_failure` - When set to `true` a connection failure
    while establishing the initial connection won't immediately return an error
    and instead will invoke the `c:handle_disconnect/2` callback. This option
    only matters during process initialization. The `handle_disconnect`
    callback is always invoked if an established connection is lost.

  Other possible option values include: `t:WebSockex.connection_option/0`
  """
  @type option :: WebSockex.Conn.connection_option
                  | {:async, boolean}
                  | {:handle_initial_conn_failure, boolean}

  @typedoc """
  The reason a connection was closed.

  A `:normal` reason is the same as a `1000` reason with no payload.

  If the peer closes the connection abruptly without a close frame then the
  close reason is `{:remote, :closed}`.
  """
  @type close_reason :: {:remote | :local, :normal}
                        | {:remote | :local, close_code, message :: binary}
                        | {:remote, :closed}
                        | {:error, term}

  @typedoc """
  The error returned when a connection fails to be established.
  """
  @type connection_error :: %WebSockex.RequestError{} | %WebSockex.ConnError{}

  @typedoc """
  A map that contains information about the failure to connect.

  This map contains the error, attempt number, and the `t:WebSockex.Conn.t/0`
  that was used to attempt the connection.
  """
  @type connection_status_map :: %{reason: close_reason | connection_error,
                                   attempt_number: integer,
                                   conn: WebSockex.Conn.t}

  @doc """
  Invoked after a connection is established.

  This is invoked after both the initial connection and a reconnect.
  """
  @callback handle_connect(state :: term, conn :: WebSockex.Conn.t) ::
    {:ok, new_state :: term}

  @doc """
  Invoked on the reception of a frame on the socket.

  The control frames have possible payloads, when they don't have a payload
  then the frame will have `nil` as the payload. e.g. `{:ping, nil}`
  """
  @callback handle_frame(frame, state :: term) ::
    {:ok, new_state}
    | {:reply, frame, new_state}
    | {:close, new_state}
    | {:close, close_frame, new_state} when new_state: term

  @doc """
  Invoked to handle asynchronous `cast/2` messages.
  """
  @callback handle_cast(msg :: term, state :: term) ::
    {:ok, new_state}
    | {:reply, frame, new_state}
    | {:close, new_state}
    | {:close, close_frame, new_state} when new_state: term

  @doc """
  Invoked to handle all other non-WebSocket messages.
  """
  @callback handle_info(msg :: term, state :: term) ::
    {:ok, new_state}
    | {:reply, frame, new_state}
    | {:close, new_state}
    | {:close, close_frame, new_state} when new_state: term

  @doc """
  Invoked when the WebSocket disconnects from the server.

  This callback is only invoked in the even of a connection failure. In cases
  of crashes or other errors then the process will terminate immediately
  skipping this callback.

  If the `handle_initial_conn_failure: true` option is provided during process
  startup, then this callback will be invoked if the process fails to establish
  an initial connection.

  If a connection is established by reconnecting, the `c:handle_connect/2`
  callback will be invoked.

  The possible returns for this callback are:

  - `{:ok, state}` will continue the process termination.
  - `{:reconnect, state}` will attempt to reconnect instead of terminating.
  - `{:reconnect, conn, state}` will attempt to reconnect with the connection
    data in `conn`. `conn` is expected to be a `t:WebSockex.Conn.t/0`.
  """
  @callback handle_disconnect(connection_status_map, state :: term) ::
    {:ok, new_state}
    | {:reconnect, new_state}
    | {:reconnect, new_conn :: WebSockex.Conn.t, new_state} when new_state: term

  @doc """
  Invoked when the Websocket receives a ping frame
  """
  @callback handle_ping(ping_frame :: :ping | {:ping, binary}, state :: term) ::
    {:ok, new_state}
    | {:reply, frame, new_state}
    | {:close, new_state}
    | {:close, close_frame, new_state} when new_state: term

  @doc """
  Invoked when the Websocket receives a pong frame.
  """
  @callback handle_pong(pong_frame :: :pong | {:pong, binary}, state :: term) ::
    {:ok, new_state}
    | {:reply, frame, new_state}
    | {:close, new_state}
    | {:close, close_frame, new_state} when new_state: term

  @doc """
  Invoked when the process is terminating.
  """
  @callback terminate(close_reason, state :: term) :: any

  @doc """
  Invoked when a new version the module is loaded during runtime.
  """
  @callback code_change(old_vsn :: term | {:down, term},
                        state :: term, extra :: term) ::
    {:ok, new_state :: term}
    | {:error, reason :: term}

  @optional_callbacks [handle_disconnect: 2, handle_ping: 2, handle_pong: 2, terminate: 2,
                       code_change: 3]

  defmacro __using__(_) do
    quote location: :keep do
      @behaviour WebSockex

      @doc false
      def handle_connect(state, _conn) do
        {:ok, state}
      end

      @doc false
      def handle_frame(frame, _state) do
        raise "No handle_frame/2 clause in #{__MODULE__} provided for #{inspect frame}"
      end

      @doc false
      def handle_cast(message, _state) do
        raise "No handle_cast/2 clause in #{__MODULE__} provided for #{inspect message}"
      end

      @doc false
      def handle_info(message, state) do
        require Logger
        Logger.error "No handle_info/2 clause in #{__MODULE__} provided for #{inspect message}"
        {:ok, state}
      end

      @doc false
      def handle_disconnect(_connection_status_map, state) do
        {:ok, state}
      end

      @doc false
      def handle_ping(:ping, state) do
        {:reply, :pong, state}
      end
      def handle_ping({:ping, msg}, state) do
        {:reply, {:pong, msg}, state}
      end

      @doc false
      def handle_pong(:pong, state), do: {:ok, state}
      def handle_pong({:pong, _}, state), do: {:ok, state}

      @doc false
      def terminate(_close_reason, _state), do: :ok

      @doc false
      def code_change(_old_vsn, state, _extra), do: {:ok, state}

      defoverridable [handle_connect: 2, handle_frame: 2, handle_cast: 2, handle_info: 2, handle_ping: 2,
                      handle_pong: 2, handle_disconnect: 2, terminate: 2, code_change: 3]
    end
  end

  @doc """
  Starts a `WebSockex` process.

  For available option values see `t:option/0`.
  """
  @spec start(String.t, module, term, options) :: {:ok, pid} | {:error, term}
  def start(url, module, state, opts \\ []) do
    case parse_uri(url) do
      {:ok, uri} ->
        :proc_lib.start(__MODULE__, :init, [self(), uri, module, state, opts])
      {:error, error} ->
        {:error, error}
    end
  end

  @doc """
  Starts a `WebSockex` process linked to the current process.

  For available option values see `t:option/0`.

  The callback `c:handle_connect/2` is invoked after the connection is
  established.
  """
  @spec start_link(String.t, module, term, options) :: {:ok, pid} | {:error, term}
  def start_link(url, module, state, opts \\ []) do
    case parse_uri(url) do
      {:ok, uri} ->
        :proc_lib.start_link(__MODULE__, :init, [self(), uri, module, state, opts])
      {:error, error} ->
        {:error, error}
    end
  end

  @doc """
  Asynchronously sends a message to a client that is handled by `c:handle_cast/2`.
  """
  @spec cast(pid, term) :: :ok
  def cast(client, message) do
    send(client, {:"$websockex_cast", message})
    :ok
  end

  @doc """
  Queue a frame to be sent asynchronously.
  """
  @spec send_frame(pid, frame) :: :ok | {:error, WebSockex.FrameEncodeError.t}
  def send_frame(pid, frame) do
    with {:ok, binary_frame} <- WebSockex.Frame.encode_frame(frame) do
      send(pid, {:"$websockex_send", binary_frame})
      :ok
    end
  end

  @doc false
  @spec init(pid, URI.t, module, term, options) :: {:ok, pid} | {:error, term}
  def init(parent, uri, module, module_state, opts) do
    # OTP stuffs
    debug = :sys.debug_options([])
    conn = WebSockex.Conn.new(uri, opts)

    reply_fun = case Keyword.get(opts, :async, false) do
                  true ->
                    :proc_lib.init_ack(parent, {:ok, self()})
                    &async_init_fun/1
                  false ->
                    &:proc_lib.init_ack(parent, &1)
                end

    state = %{conn: conn,
              module: module,
              module_state: module_state,
              reply_fun: reply_fun}

    handle_conn_failure = Keyword.get(opts, :handle_initial_conn_failure, false)

    case open_connection(state.conn) do
      {:ok, new_conn} ->
        module_init(parent, debug, %{state | conn: new_conn})
      {:error, error} when handle_conn_failure == true ->
        on_disconnect(error, parent, debug, state, success: &module_init/3, failure: &init_failure/4)
      {:error, error} ->
        state.reply_fun.({:error, error})
    end
  end

  ## OTP Stuffs

  @doc false
  def system_continue(parent, debug, state) do
    websocket_loop(parent, debug, state)
  end

  @doc false
  def system_terminate(reason, parent, debug, state) do
    terminate(reason, parent, debug, state)
  end

  @doc false
  def system_get_state(state) do
    {:ok, state, state}
  end

  @doc false
  def system_replace_state(fun, state) do
    new_state = fun.(state)
    {:ok, new_state, new_state}
  end

  # Internals! Yay

  defp on_disconnect(reason, parent, debug, state, callbacks \\ [], attempt \\ 1) do
    case handle_disconnect(reason, state, attempt) do
      {:ok, new_module_state} ->
        callback = Keyword.get(callbacks, :failure, &terminate/4)
        callback.(reason, parent, debug, %{state | module_state: new_module_state})
      {:reconnect, new_conn, new_module_state} ->
        state = %{state | conn: new_conn, module_state: new_module_state}
        case open_connection(new_conn) do
          {:ok, success_conn} ->
            callback = Keyword.get(callbacks, :success, &reconnect/3)
            callback.(parent, debug, %{state | conn: success_conn})
          {:error, new_reason} ->
            on_disconnect(new_reason, parent, debug, state, callbacks, attempt+1)
        end
    end
  end

  defp reconnect(parent, debug, state) do
    case apply(state.module, :handle_connect, [state.module_state, state.conn]) do
      {:ok, new_module_state} ->
         state = Map.merge(state, %{buffer: <<>>,
                                    fragment: nil,
                                    module_state: new_module_state})
          websocket_loop(parent, debug, state)
      badreply ->
        raise %WebSockex.BadResponseError{module: state.module, function: :handle_connect,
          args: [state.module_state, state.conn], response: badreply}
    end
  end

  defp open_connection(conn) do
    with {:ok, conn} <- WebSockex.Conn.open_socket(conn),
         key <- :crypto.strong_rand_bytes(16) |> Base.encode64,
         {:ok, request} <- WebSockex.Conn.build_request(conn, key),
         :ok <- WebSockex.Conn.socket_send(conn, request),
         {:ok, headers} <- WebSockex.Conn.handle_response(conn),
         :ok <- validate_handshake(headers, key),
         :ok <- WebSockex.Conn.set_active(conn),
    do: {:ok, conn}
  end

  defp validate_handshake(headers, key) do
    challenge = :crypto.hash(:sha, key <> @handshake_guid) |> Base.encode64

    {_, res} = List.keyfind(headers, "Sec-Websocket-Accept", 0)

    if challenge == res do
      :ok
    else
      {:error, %WebSockex.HandshakeError{response: res, challenge: challenge}}
    end
  end

  defp websocket_loop(parent, debug, state) do
    case WebSockex.Frame.parse_frame(state.buffer) do
      {:ok, frame, buffer} ->
        handle_frame(frame, parent, debug, %{state | buffer: buffer})
      :incomplete ->
        transport = state.conn.transport
        socket = state.conn.socket
        receive do
          {:system, from, req} ->
            :sys.handle_system_msg(req, from, parent, __MODULE__, debug, state)
          {:"$websockex_cast", msg} ->
            common_handle({:handle_cast, msg}, parent, debug, state)
          {:"$websockex_send", binary_frame} ->
            handle_send(binary_frame, parent, debug, state)
          {^transport, ^socket, message} ->
            buffer = <<state.buffer::bitstring, message::bitstring>>
            websocket_loop(parent, debug, %{state | buffer: buffer})
          {:tcp_closed, ^socket} ->
            handle_close({:remote, :closed}, parent, debug, state)
          :"websockex_close_timeout" ->
            websocket_loop(parent, debug, state)
          msg ->
            common_handle({:handle_info, msg}, parent, debug, state)
        end
    end
  end

  defp handle_frame(:ping, parent, debug, state) do
    common_handle({:handle_ping, :ping}, parent, debug, state)
  end
  defp handle_frame({:ping, msg}, parent, debug, state) do
    common_handle({:handle_ping, {:ping, msg}}, parent, debug, state)
  end
  defp handle_frame(:pong, parent, debug, state) do
    common_handle({:handle_pong, :pong}, parent, debug, state)
  end
  defp handle_frame({:pong, msg}, parent, debug, state) do
    common_handle({:handle_pong, {:pong, msg}}, parent, debug, state)
  end
  defp handle_frame(:close, parent, debug, state) do
    handle_close({:remote, :normal}, parent, debug, state)
  end
  defp handle_frame({:close, code, reason}, parent, debug, state) do
    handle_close({:remote, code, reason}, parent, debug, state)
  end
  defp handle_frame({:fragment, _, _} = fragment, parent, debug, state) do
    handle_fragment(fragment, parent, debug, state)
  end
  defp handle_frame({:continuation, _} = fragment, parent, debug, state) do
    handle_fragment(fragment, parent, debug, state)
  end
  defp handle_frame({:finish, _} = fragment, parent, debug, state) do
    handle_fragment(fragment, parent, debug, state)
  end
  defp handle_frame(frame, parent, debug, state) do
    common_handle({:handle_frame, frame}, parent, debug, state)
  end

  defp common_handle({function, msg}, parent, debug, state) do
    case apply(state.module, function, [msg, state.module_state]) do
      {:ok, new_state} ->
        websocket_loop(parent, debug, %{state | module_state: new_state})
      {:reply, frame, new_state} ->
        with {:ok, binary_frame} <- WebSockex.Frame.encode_frame(frame),
             :ok <- WebSockex.Conn.socket_send(state.conn, binary_frame) do
          websocket_loop(parent, debug, %{state | module_state: new_state})
        else
          {:error, error} ->
            raise error
        end
      {:close, new_state} ->
        handle_close({:local, :normal}, parent, debug, %{state | module_state: new_state})
      {:close, {close_code, message}, new_state} ->
        handle_close({:local, close_code, message}, parent, debug, %{state | module_state: new_state})
      badreply ->
        raise %WebSockex.BadResponseError{module: state.module,
          function: function, args: [msg, state.module_state],
          response: badreply}
    end
  rescue
    exception ->
      terminate({exception, System.stacktrace}, parent, debug, state)
  end

  defp handle_close({:remote, :closed} = reason, parent, debug, state) do
    new_conn = %{state.conn | socket: nil}
    on_disconnect(reason, parent, debug, %{state | conn: new_conn})
  end
  defp handle_close({:remote, _} = reason, parent, debug, state) do
    handle_remote_close(reason, parent, debug, state)
  end
  defp handle_close({:remote, _, _} = reason, parent, debug, state) do
    handle_remote_close(reason, parent, debug, state)
  end
  defp handle_close({:local, _} = reason, parent, debug, state) do
    handle_local_close(reason, parent, debug, state)
  end
  defp handle_close({:local, _, _} = reason, parent, debug, state) do
    handle_local_close(reason, parent, debug, state)
  end

  defp module_init(parent, debug, state) do
    case apply(state.module, :handle_connect, [state.module_state, state.conn]) do
      {:ok, new_module_state} ->
         state.reply_fun.({:ok, self()})
         state = Map.merge(state, %{buffer: <<>>,
                                    fragment: nil,
                                    module_state: new_module_state})
                 |> Map.delete(:reply_fun)

          websocket_loop(parent, debug, state)
      badreply ->
        raise %WebSockex.BadResponseError{module: state.module, function: :handle_connect,
          args: [state.module_state, state.conn], response: badreply}
    end
  end

  defp terminate(reason, _parent, _debug, %{module: mod, module_state: mod_state}) do
    mod.terminate(reason, mod_state)
    case reason do
      {_, :normal} ->
        exit(:normal)
      {_, 1000, _} ->
        exit(:normal)
      _ ->
        exit(reason)
    end
  end

  defp handle_send(binary_frame, parent, debug, %{conn: conn} = state) do
    case WebSockex.Conn.socket_send(conn, binary_frame) do
      :ok ->
        websocket_loop(parent, debug, state)
      {:error, error} ->
        terminate(error, parent, debug, state)
    end
  end

  defp handle_fragment({:fragment, type, part}, parent, debug, %{fragment: nil} = state) do
    websocket_loop(parent, debug, %{state | fragment: {type, part}})
  end
  defp handle_fragment({:fragment, _, _}, parent, debug, state) do
    handle_close({:local, 1002, "Endpoint tried to start a fragment without finishing another"}, parent, debug, state)
  end
  defp handle_fragment({:continuation, _}, parent, debug, %{fragment: nil} = state) do
    handle_close({:local, 1002, "Endpoint sent a continuation frame without starting a fragment"}, parent, debug, state)
  end
  defp handle_fragment({:continuation, next}, parent, debug, %{fragment: {type, part}} = state) do
    websocket_loop(parent, debug, %{state | fragment: {type, <<part::binary, next::binary>>}})
  end
  defp handle_fragment({:finish, next}, parent, debug, %{fragment: {type, part}} = state) do
    handle_frame({type, <<part::binary, next::binary>>}, parent, debug, %{state | fragment: nil})
  end

  defp handle_remote_close(reason, parent, debug, state) do
    # If the socket is already closed then that's ok, but the spec says to send
    # the close frame back in response to receiving it.
    send_close_frame(reason, state.conn)

    Process.send_after(self(), :"$websockex_close_timeout", 5000)
    close_loop(reason, parent, debug, state)
  end

  defp handle_local_close(reason, parent, debug, state) do
    case send_close_frame(reason, state.conn) do
      :ok ->
        Process.send_after(self(), :"$websockex_close_timeout", 5000)
        close_loop(reason, parent, debug, state)
      {:error, %WebSockex.ConnError{original: :closed}} ->
        close_loop({:remote, :closed}, parent, debug, state)
    end
  end

  defp handle_disconnect(reason, state, attempt) do
    status_map = %{conn: state.conn,
                   reason: reason,
                   attempt_number: attempt}

    case apply(state.module, :handle_disconnect, [status_map, state.module_state]) do
      {:ok, new_state} ->
        {:ok, new_state}
      {:reconnect, new_state} ->
        {:reconnect, state.conn, new_state}
      {:reconnect, new_conn, new_state} ->
        {:reconnect, new_conn, new_state}
      badreply ->
        raise %WebSockex.BadResponseError{module: state.module,
                                          function: :handle_disconnect,
                                          args: [status_map, state.module_state],
                                          response: badreply}
    end
  end

  defp send_close_frame(reason, conn) do
    with {:ok, binary_frame} <- build_close_frame(reason),
    do: WebSockex.Conn.socket_send(conn, binary_frame)
  end

  defp build_close_frame({_, :normal}) do
    WebSockex.Frame.encode_frame(:close)
  end
  defp build_close_frame({_, code, msg}) do
    WebSockex.Frame.encode_frame({:close, code, msg})
  end

  defp init_failure(reason, _parent, _debug, state) do
    state.reply_fun.({:error, reason})
  end

  defp async_init_fun({:ok, _}), do: :noop
  defp async_init_fun(exit_reason), do: exit(exit_reason)

  defp close_loop(reason, parent, debug, %{conn: conn} = state) do
    transport = state.conn.transport
    socket = state.conn.socket
    receive do
      {^transport, ^socket, _} ->
        close_loop(reason, parent, debug, state)
      {:tcp_closed, ^socket} ->
        new_conn = %{conn | socket: nil}
        on_disconnect(reason, parent, debug, %{state | conn: new_conn})
      :"$websockex_close_timeout" ->
        new_conn = WebSockex.Conn.close_socket(conn)
        on_disconnect(reason, parent, debug, %{state | conn: new_conn})
    end
  end

  defp parse_uri(url) do
    case URI.parse(url) do
      # This is confusing to look at. But it's just a match with multiple guards
      %URI{host: host, port: port, scheme: protocol}
      when is_nil(host)
      when is_nil(port)
      when not protocol in ["ws", "wss"] ->
        {:error, %WebSockex.URLError{url: url}}
      {:error, error} ->
        {:error, error}
      %URI{} = uri ->
        {:ok, uri}
    end
  end
end