websocket_client

-module(wsc_lib).
%% @author Jeremy Ong
%% @author Michael Coles
%% Purely functional aspects of websocket client comms.
%%
%% Herein live all the functions for pure data processing.
-compile([export_all]).
-include("websocket_req.hrl").

-export([create_auth_header/3]).
-export([create_handshake/2]).
-export([encode_frame/1]).
-export([generate_ws_key/0]).

-export([validate_handshake/2]).
-export([decode_frame/2]).
-export([decode_frame/5]).

-spec create_auth_header(Type :: basic, User :: binary(), Pass :: binary()) ->
    {binary(), binary()}.
create_auth_header(basic, User, Pass) ->
    PlainAuth = << User/binary, ":", Pass/binary >>,
    B64Auth = base64:encode(PlainAuth),
    {<<"Authorization">>, <<"Basic ", B64Auth/binary>>}.

-spec create_handshake(websocket_req:req(), [{string(), string()}]) ->
    iolist().
create_handshake(WSReq, ExtraHeaders) ->
    [Path, Host, Key] = websocket_req:get([path, host, key], WSReq),
    ["GET ", Path, " HTTP/1.1\r\n"
     "Host: ", Host, "\r\n"
     "Connection: Upgrade\r\n"
     "Sec-WebSocket-Version: 13\r\n"
     "Sec-WebSocket-Key: ", Key, "\r\n"
     "Upgrade: websocket\r\n",
     [ [Header, ": ", Value, "\r\n"] || {Header, Value} <- ExtraHeaders],
     "\r\n"].

%% @doc Validate handshake response challenge
-spec validate_handshake(HandshakeResponse :: binary(), Key :: binary()) ->
    {ok, binary()}
    | {notfound, binary()}
    | {error, term()}.
validate_handshake(HandshakeResponse, Key) ->
    case re:run(HandshakeResponse, "\\r\\n\\r\\n") of
        {match, _} ->
            Challenge = base64:encode(
                          crypto:hash(sha, << Key/binary, "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" >>)),
            %% Consume the response...
            {ok, Status, Header, Buffer} = consume_response(HandshakeResponse),
            {_Version, Code, Message} = Status,
            case Code of
                % 101 means Switching Protocol
                101 ->
                    %% ...and make sure the challenge is valid.
                    case proplists:get_value(<<"Sec-Websocket-Accept">>, Header) of
                        Challenge -> {ok, Buffer};
                        _Invalid   -> {error, invalid_handshake}
                    end;
                _ -> {error, {Code, Message}}
            end;
        _ ->
            {notfound, HandshakeResponse}
    end.

%% @doc Consumes the HTTP response and extracts status, header and the body.
consume_response(Response) ->
    {ok, {http_response, Version, Code, Message}, Header} =
        erlang:decode_packet(http_bin, Response, []),
    consume_response({Version, Code, Message}, Header, []).
consume_response(Status, Response, HeaderAcc) ->
    case erlang:decode_packet(httph_bin, Response, []) of
        {ok, {http_header, _Length, Field, _Reserved, Value}, Rest} ->
            consume_response(Status, Rest, [{Field, Value} | HeaderAcc]);
        {ok, http_eoh, Body} ->
            {ok, Status, HeaderAcc, Body}
    end.

unpack_frame(<< Fin:1, RSV:3, OpCode:4, Mask:1, Len:7, Payload/bits >>)
  when Len < 126 ->
    {ok, Fin, RSV, OpCode, Len, unmask_frame(Mask, Len, Payload)};
unpack_frame(<< Fin:1, RSV:3, OpCode:4, Mask:1, 126:7, Len:16, Payload/bits>>)
  when Len > 125, OpCode < 8 ->
    {ok, Fin, RSV, OpCode, Len, unmask_frame(Mask, Len, Payload)};
unpack_frame(<< Fin:1, RSV:3, OpCode:4, Mask:1, 127:7, 0:1, Len:63, Payload/bits>>)
  when Len > 16#ffff, OpCode < 8 ->
    {ok, Fin, RSV, OpCode, Len, unmask_frame(Mask, Len, Payload)};
unpack_frame(Data) ->
    {incomplete, Data}.

%% @doc Start or continue continuation payload with length less than 126 bytes
decode_frame(WSReq, Frame) when is_binary(Frame) ->
    case unpack_frame(Frame) of
        {incomplete, Data} -> {recv, WSReq, Data};
        {ok, 0, 0, OpCode, Len, Payload} ->
            WSReq1 = set_continuation_if_empty(WSReq, OpCode),
            WSReq2 = websocket_req:fin(0, WSReq1),
            decode_frame(WSReq2, OpCode, Len, Payload, <<>>);
        {ok, 1, 0, OpCode, Len, Payload} ->
            WSReq1 = websocket_req:fin(1, WSReq),
            decode_frame(WSReq1, OpCode, Len, Payload, <<>>)
    end.

unmask_frame(0, _, Payload) -> Payload;
unmask_frame(1, Len, << Mask:32, Rest/bits >>) ->
    << Payload:Len/binary, NextFrame/bits >> = Rest,
    << (mask_payload(Mask, Payload))/bits, NextFrame/binary >>.

-spec decode_frame(websocket_req:req(),
               Opcode :: websocket_req:opcode(),
               Len :: non_neg_integer(),
               Data :: binary(),
               Buffer :: binary()) ->
    {recv, websocket_req:req(), IncompleteFrame :: binary()}
    | {frame, {OpcodeName :: atom(), Payload :: binary()},
                   websocket_req:req(), Rest :: binary()}
    | {close, Reason :: term(), websocket_req:req()}.
%% @doc Length known and still missing data
decode_frame(WSReq, Opcode, Len, Data, Buffer)
  when byte_size(Data) < Len ->
    Remaining = Len - byte_size(Data),
    WSReq1 = websocket_req:remaining(Remaining, WSReq),
    WSReq2 = websocket_req:opcode(Opcode, WSReq1),
    {recv, WSReq2, << Buffer/bits, Data/bits >>};
%% @doc Length known and remaining data is appended to the buffer
decode_frame(WSReq, Opcode, Len, Data, Buffer) ->
    [Continuation, ContinuationOpcode] =
        websocket_req:get([continuation, continuation_opcode], WSReq),
    Fin = websocket_req:fin(WSReq),
    << Payload:Len/binary, Rest/bits >> = Data,
    FullPayload = << Buffer/binary, Payload/binary >>,
    OpcodeName = websocket_req:opcode_to_name(Opcode),
    case OpcodeName of
        close when byte_size(FullPayload) >= 2 ->
            << CodeBin:2/binary, ClosePayload/binary >> = FullPayload,
            Code = binary:decode_unsigned(CodeBin),
            Reason = case Code of
                         1000 -> {normal, ClosePayload};
                         1002 -> {error, badframe, ClosePayload};
                         1007 -> {error, badencoding, ClosePayload};
                         1011 -> {error, handler, ClosePayload};
                         _ -> {remote, Code, ClosePayload}
                     end,
            {close, Reason, WSReq};
        close ->
            {close, {remote, <<>>}, WSReq};
        %% Non-control continuation frame
        _ when Opcode < 8, Continuation =/= undefined, Fin == 0 ->
            %% Append to previously existing continuation payloads and continue
            Continuation1 = << Continuation/binary, FullPayload/binary >>,
            WSReq1 = websocket_req:continuation(Continuation1, WSReq),
            decode_frame(WSReq1, Rest);
        %% Terminate continuation frame sequence with non-control frame
        _ when Opcode < 8, Continuation =/= undefined, Fin == 1 ->
            DefragPayload = << Continuation/binary, FullPayload/binary >>,
            WSReq1 = websocket_req:continuation(undefined, WSReq),
            WSReq2 = websocket_req:continuation_opcode(undefined, WSReq1),
            ContinuationOpcodeName = websocket_req:opcode_to_name(ContinuationOpcode),
            {frame, {ContinuationOpcodeName, DefragPayload}, WSReq2, Rest};
        _ ->
            {frame, {OpcodeName, FullPayload}, WSReq, Rest}
    end.

%% @doc Encodes the data with a header (including a masking key) and
%% masks the data
-spec encode_frame(websocket_req:frame()) -> binary().
encode_frame({Type, Payload}) ->
    Opcode = websocket_req:name_to_opcode(Type),
    Len = iolist_size(Payload),
    BinLen = payload_length_to_binary(Len),
    MaskingKeyBin = crypto:strong_rand_bytes(4),
    << MaskingKey:32 >> = MaskingKeyBin,
    Header = << 1:1, 0:3, Opcode:4, 1:1, BinLen/bits, MaskingKeyBin/bits >>,
    MaskedPayload = mask_payload(MaskingKey, Payload),
    << Header/binary, MaskedPayload/binary >>;
encode_frame(Type) when is_atom(Type) ->
    encode_frame({Type, <<>>}).

%% @doc The payload is masked using a masking key byte by byte.
%% Can do it in 4 byte chunks to save time until there is left than 4 bytes left
mask_payload(MaskingKey, Payload) when is_integer(MaskingKey), is_binary(Payload) ->
    mask_payload(MaskingKey, Payload, <<>>).
mask_payload(_, <<>>, Acc) ->
    Acc;
mask_payload(MaskingKey, << D:32, Rest/bits >>, Acc) ->
    T = D bxor MaskingKey,
    mask_payload(MaskingKey, Rest, << Acc/binary, T:32 >>);
mask_payload(MaskingKey, << D:24 >>, Acc) ->
    << MaskingKeyPart:24, _:8 >> = << MaskingKey:32 >>,
    T = D bxor MaskingKeyPart,
    << Acc/binary, T:24 >>;
mask_payload(MaskingKey, << D:16 >>, Acc) ->
    << MaskingKeyPart:16, _:16 >> = << MaskingKey:32 >>,
    T = D bxor MaskingKeyPart,
    << Acc/binary, T:16 >>;
mask_payload(MaskingKey, << D:8 >>, Acc) ->
    << MaskingKeyPart:8, _:24 >> = << MaskingKey:32 >>,
    T = D bxor MaskingKeyPart,
    << Acc/binary, T:8 >>.

%% @doc Encode the payload length as binary in a variable number of bits.
%% See RFC Doc for more details
payload_length_to_binary(Len) when Len =<125 ->
    << Len:7 >>;
payload_length_to_binary(Len) when Len =< 16#ffff ->
    << 126:7, Len:16 >>;
payload_length_to_binary(Len) when Len =< 16#7fffffffffffffff ->
    << 127:7, Len:64 >>.

%% @doc If this is the first continuation frame, set the opcode and initialize
%% continuation to an empty binary. Otherwise, return the request object untouched.
-spec set_continuation_if_empty(WSReq :: websocket_req:req(),
                                Opcode :: websocket_req:opcode()) ->
                                       websocket_req:req().
set_continuation_if_empty(WSReq, Opcode) ->
    case websocket_req:continuation(WSReq) of
        undefined ->
            WSReq1 = websocket_req:continuation_opcode(Opcode, WSReq),
            websocket_req:continuation(<<>>, WSReq1);
        _ ->
            WSReq
    end.

%% @doc Key sent in initial handshake
-spec generate_ws_key() -> binary().
generate_ws_key() ->
    base64:encode(crypto:strong_rand_bytes(16)).