defmodule ExWs.Reader do import Bitwise, only: [band: 2] alias __MODULE__ alias ExWs.{Writer, Mask} @compile {:inline, atomize_op: 1} @error_invalid_op "invalid op" |> Writer.close(1002) |> Writer.to_binary() @error_control_len "large control" |> Writer.close(1002) |> Writer.to_binary() @error_control_fin "non-fin control" |> Writer.close(1002) |> Writer.to_binary() @error_preamble_unmasked "unmasked frame" |> Writer.close(1002) |> Writer.to_binary() @error_invalid_continuation "invalid continuation" |> Writer.close(1002) |> Writer.to_binary() @enforce_keys [:op, :fin, :len, :mask, :data, :chain] defstruct @enforce_keys defguard is_control(op) when op in [:close, :ping, :pong] def new() do %Reader{op: nil, fin: nil, len: nil, mask: {4, <<>>}, data: nil, chain: nil} end def received(data, frame), do: parse(data, frame, []) defp parse(<<>>, frame, []), do: {:ok, frame} defp parse(<<>>, frame, acc), do: {:ok, acc, frame} defp parse(<>, %{op: nil} = frame, acc) do {op, fin} = case band(b1, 128) == 128 do true -> {b1 - 128, true} false -> {b1, false} end res = with {:ok, op} <- atomize_op(op) do has_chain? = frame.chain != nil is_control? = is_control(op) cond do op == :cont && !has_chain? -> {:close, @error_invalid_continuation} op != :cont && has_chain? && not is_control? -> {:close, @error_invalid_continuation} is_control? && !fin -> {:close, @error_control_fin} true -> parse(data, %Reader{frame | op: op, fin: fin}, acc) end end case {res, acc} do {{:close, _} = close, []} -> {:ok, close, frame} {{:close, _} = close, acc} -> {:ok, Enum.reverse([close | acc]), frame} {ok, _} -> ok end end defp parse(<>, %{len: nil} = frame, acc) do cond do masked == 0 -> {:close, @error_preamble_unmasked} len > 125 && is_control(frame.op) -> {:close, @error_control_len} len == 127 -> parse(data, %Reader{frame | len: {8, <<>>}}, acc) len == 126 -> parse(data, %Reader{frame | len: {2, <<>>}}, acc) len == 0 -> parse(data, %Reader{frame | len: 0, data: <<>>}, acc) true -> parse(data, %Reader{frame | len: len, data: {len, []}}, acc) end end defp parse(data, %{len: {missing, known}} = frame, acc) do case data do <> -> len = known <> len len = case byte_size(len) do 2 -> <> = len; len 8 -> <> = len; len end parse(data, %Reader{frame | len: len, data: {len, []}}, acc) _ -> {:ok, %Reader{frame | len: {missing - byte_size(data), known <> data}}} end end defp parse(data, %{mask: {missing, known}} = frame, acc) do case data do <> -> frame = %Reader{frame | mask: known <> mask} case frame.len == 0 do true -> finalize_frame(data, frame, acc) false -> parse(data, frame, acc) end _ -> {:ok, %Reader{frame | mask: {missing - byte_size(data), known <> data}}} end end defp parse(data, %{data: {missing, known}} = frame, acc) do case data do <> -> finalize_frame(extra, %Reader{frame | data: [known, data]}, acc) _ -> {:ok, %Reader{frame | data: {missing - byte_size(data), [known, data]}}} end end defp finalize_frame(extra, frame, acc) do data = Mask.apply(frame.mask, :erlang.iolist_to_binary(frame.data)) {next_frame, acc} = continue(frame, data, acc) case {extra, acc} do {<<>>, []} -> {:ok, next_frame} {<<>>, acc} -> {:ok, Enum.reverse(acc), next_frame} {_, _} -> parse(extra, next_frame, acc) end end defp continue(%{fin: true, chain: nil} = frame, unmasked, acc) do acc = [{frame.op, unmasked} | acc] {new(), acc} end defp continue(%{fin: true, op: :cont, chain: {op, data}}, unmasked, acc) do acc = [{op, [data, unmasked]} | acc] {new(), acc} end defp continue(%{fin: true} = frame, unmasked, acc) do acc = [{frame.op, unmasked} | acc] {%Reader{new() | chain: frame.chain}, acc} end defp continue(%{fin: false} = frame, unmasked, acc) do chain = case frame.chain do nil -> {frame.op, [unmasked]} {op, data} -> {op, [data, unmasked]} end {%Reader{new() | chain: chain}, acc} end for {value, op} <- [cont: 0, txt: 1, bin: 2, close: 8, ping: 9, pong: 10] do defp atomize_op(unquote(op)), do: {:ok, unquote(value)} end defp atomize_op(_op), do: {:close, @error_invalid_op} end