Decode messages like a bitch

include/wsecli.hrl

@@ -23,3 +23,9 @@
     extended_payload_len_cont :: integer(),
     masking_key :: binary(),
     payload :: binary()}).
+
+-record(message, {
+    frames = [] :: list(#frame{}),
+    payload :: string() | binary(), % FALSE!!! what about control message with code + message
+    type :: {text, binary, control, fragmented}
+  }).

src/wsecli_framing.erl

@@ -40,15 +40,8 @@ from_binary(<<Head:9, PayloadLen:7, Payload:PayloadLen/binary, Rest/binary>>, Ac
 from_binary(<<>>, Acc) ->
   Acc.
 
-decode_frame(Data) ->
+decode_frame(Data = <<Fin:1, Rsv1:1, Rsv2:1, Rsv3:1, Opcode:4, Mask:1, _/bits>> ) ->
   % TODO: ensure that Mask is not set
-  <<
-    Fin:1,
-    Rsv1:1, Rsv2:1, Rsv3:1,
-    Opcode:4,
-    Mask:1,
-    _/bits
-  >> = Data,
 
   Frame = #frame{
     fin = Fin,
@@ -86,9 +79,7 @@ binary_payload(Data, Frame) ->
   end,
 
   case Frame#frame.opcode of
-    ?OP_CODE_TEXT ->
-      Frame#frame{ payload = Payload };
-    ?OP_CODE_CONT ->
+    _ ->
       Frame#frame{ payload = Payload }
   end.
 

src/wsecli_message.erl

@@ -1,9 +1,10 @@
 -module(wsecli_message).
 -include("wsecli.hrl").
 
--export([encode/2]).
+-export([encode/2, decode/1, decode/2]).
 
 -define(FRAGMENT_SIZE, 4096).
+-type message_type() :: begin_message | continue_message.
 
 -spec encode(Data::string() | binary(), Type::atom()) -> binary().
 encode(Data, Type) when is_list(Data)->
@@ -12,6 +13,123 @@ encode(Data, Type) when is_list(Data)->
 encode(Data, Type)->
   lists:reverse(encode(Data, Type, [])).
 
+-spec decode(Data::binary()) -> list(#message{}).
+decode(Data) ->
+  decode(Data, begin_message, #message{}).
+
+-spec decode(Data::binary(), Message::#message{}) -> list(#message{}).
+decode(Data, Message) ->
+  decode(Data, continue_message, Message).
+
+-spec decode(Data::binary(), Type :: message_type(), Message::#message{}) -> list(#message{}).
+decode(Data, begin_message, _Message) ->
+  Frames = wsecli_framing:from_binary(Data),
+  lists:reverse(process_frames(begin_message, Frames, []));
+
+decode(Data, continue_message, Message) ->
+  Frames = wsecli_framing:from_binary(Data),
+  lists:reverse(process_frames(continue_message, Frames, [Message | []])).
+
+-spec process_frames(Type:: message_type(), Frames :: list(#frame{}), Messages :: list(#message{})) -> list(#message{}).
+process_frames(begin_message, [Frame | Frames], Acc) ->
+  case process_frame(Frame, begin_message, #message{}) of
+    {fragmented, Message} ->
+      process_frames(continue_message, Frames, [Message#message{type = fragmented} | Acc]);
+    {completed, Message} ->
+      process_frames(begin_message, Frames, [Message | Acc])
+  end;
+
+process_frames(continue_message, [Frame | Frames], [FramgmentedMessage | Acc]) ->
+  case process_frame(Frame, continue_message, FramgmentedMessage) of
+    {fragmented, Message} ->
+      process_frames(continue_message, Frames, [Message#message{type = fragmented} | Acc]);
+    {completed, Message} ->
+      process_frames(begin_message, Frames, [Message | Acc])
+  end;
+
+process_frames(_, [], Acc) ->
+  Acc.
+
+-spec process_frame(Frame :: #frame{}, MessageType :: message_type(), Message :: #message{})-> {fragmented | completed, #message{}}.
+process_frame(Frame, begin_message, Message) ->
+  case contextualize_frame(Frame) of
+    open_close ->
+      BuiltMessage = build_message(Message, [Frame]),
+      {completed, BuiltMessage};
+    open_continue ->
+      Frames = Message#message.frames,
+      {fragmented, Message#message{frames = [Frame | Frames]}}
+  end;
+
+process_frame(Frame, continue_message, Message) ->
+  case contextualize_frame(Frame) of
+    continue ->
+      Frames = Message#message.frames,
+      {fragmented, Message#message{frames = [Frame | Frames]}};
+    continue_close ->
+      BuiltMessage = build_message(Message, lists:reverse([Frame | Message#message.frames])),
+      {completed, BuiltMessage}
+  end.
+
+-spec contextualize_frame(Frame :: #frame{}) -> continue_close | open_continue | continue | open_close.
+contextualize_frame(Frame) ->
+  case {Frame#frame.fin, Frame#frame.opcode} of
+    {1, 0} -> continue_close;
+    {0, 0} -> continue;
+    {1, _} -> open_close;
+    {0, _} -> open_continue
+  end.
+
+build_message(Message, Frames) ->
+  [HeadFrame | _] = Frames,
+
+  case HeadFrame#frame.opcode of
+    1 ->
+      Payload = build_payload_from_frames(text, Frames),
+      Message#message{type = text, payload = Payload};
+    2 ->
+      Payload = build_payload_from_frames(binary, Frames),
+      Message#message{type = binary, payload = Payload}
+  end.
+
+build_payload_from_frames(binary, Frames) ->
+  contatenate_payload_from_frames(Frames);
+
+build_payload_from_frames(text, Frames) ->
+  Payload = contatenate_payload_from_frames(Frames),
+  binary_to_list(Payload).
+
+contatenate_payload_from_frames(Frames) ->
+  contatenate_payload_from_frames(Frames, <<>>).
+
+contatenate_payload_from_frames([Frame | Rest], Acc) ->
+  contatenate_payload_from_frames(Rest, <<Acc/binary, (Frame#frame.payload)/binary>>);
+
+contatenate_payload_from_frames([], Acc) ->
+  Acc.
+
+-spec process(Frames ::list(#frame{}), Messages::list(#message{})) -> list(#message{}).
+process([Frame | Tail], Acc) ->
+  Message = case Frame#frame.opcode of
+    1 ->
+      #message{type = text, payload = binary_to_list(Frame#frame.payload)};
+    2 ->
+      #message{type = binary, payload = Frame#frame.payload}
+  end,
+
+  process(Tail, [Message | Acc]);
+
+process([], Acc) ->
+  Acc.
+
+
+%-spec process(Frames::list(#frame{}), IncompleteMessage::#message{}) -> {text | binary | control | fragmented, #message{}}.
+%process(Frames, IncompleteMessage) ->
+%  [Head | Tail] = Frames,
+
+%
+% Internal
+%
 -spec encode(Data::binary(), Type :: atom(), Acc ::list()) -> list().
 encode(<<Data:?FRAGMENT_SIZE/binary>>, Type, Acc) ->
   [frame(Data, [fin, {opcode, Type}]) | Acc];

test/spec/wsecli_message_spec.erl

@@ -113,22 +113,207 @@ spec() ->
     end),
   it("wsecli_message:control"),
   describe("decode", fun()->
+        describe("fragmented messages", fun() ->
+              it("should complain when control messages are fragmented"),
+              it("should return a fragmented message with undefined payload when message is not complete", fun() ->
+                    Payload = crypto:rand_bytes(20),
+                    <<
+                    Payload1:10/binary,
+                    Payload2:5/binary,
+                    _Payload3/binary
+                    >> = Payload,
+
+                    FakeFragment1 = get_binary_frame(0, 0, 0, 0, 2, 0, 10, 0, Payload1),
+                    FakeFragment2 = get_binary_frame(0, 0, 0, 0, 0, 0, 5, 0, Payload2),
+
+                    Data = <<FakeFragment1/binary, FakeFragment2/binary>>,
+
+                    [Message] = wsecli_message:decode(Data),
+
+                    assert_that(Message#message.type, is(fragmented)),
+                    assert_that(length(Message#message.frames), is(2))
+                end),
+              it("should decode data containing a complete fragmented binary message", fun() ->
+                    Payload = crypto:rand_bytes(40),
+                    <<
+                    Payload1:10/binary,
+                    Payload2:10/binary,
+                    Payload3:10/binary,
+                    Payload4:10/binary
+                    >> = Payload,
+
+                    FakeFragment1 = get_binary_frame(0, 0, 0, 0, 2, 0, 10, 0, Payload1),
+                    FakeFragment2 = get_binary_frame(0, 0, 0, 0, 0, 0, 10, 0, Payload2),
+                    FakeFragment3 = get_binary_frame(0, 0, 0, 0, 0, 0, 10, 0, Payload3),
+                    FakeFragment4 = get_binary_frame(1, 0, 0, 0, 0, 0, 10, 0, Payload4),
+
+                    Data = << FakeFragment1/binary, FakeFragment2/binary, FakeFragment3/binary, FakeFragment4/binary>>,
+
+                    [Message] = wsecli_message:decode(Data),
+
+                    assert_that(Message#message.type, is(binary)),
+                    assert_that(Message#message.payload, is(Payload))
+                end),
+              it("should decode data containing a complete fragmented text message", fun() ->
+                    Text = "asasdasdasdasdasdasdasdasdasdasdasdasdasdasdasd",
+                    Payload = list_to_binary(Text),
+                    <<
+                    Payload1:5/binary,
+                    Payload2:2/binary,
+                    Payload3/binary
+                    >> = Payload,
+
+                    FakeFragment1 = get_binary_frame(0, 0, 0, 0, 1, 0, byte_size(Payload1), 0, Payload1),
+                    FakeFragment2 = get_binary_frame(0, 0, 0, 0, 0, 0, byte_size(Payload2), 0, Payload2),
+                    FakeFragment3 = get_binary_frame(1, 0, 0, 0, 0, 0, byte_size(Payload3), 0, Payload3),
+
+                    Data = << FakeFragment1/binary, FakeFragment2/binary, FakeFragment3/binary>>,
+
+                    [Message] = wsecli_message:decode(Data),
+
+                    assert_that(Message#message.type, is(text)),
+                    assert_that(Message#message.payload, is(Text))
+                end),
+              it("should complete a fragmented message", fun() ->
+                    Payload = crypto:rand_bytes(20),
+                    <<
+                    Payload1:10/binary,
+                    Payload2:5/binary,
+                    Payload3/binary
+                    >> = Payload,
+
+                    FakeFragment1 = get_binary_frame(0, 0, 0, 0, 2, 0, 10, 0, Payload1),
+                    FakeFragment2 = get_binary_frame(0, 0, 0, 0, 0, 0, 5, 0, Payload2),
+                    FakeFragment3 = get_binary_frame(1, 0, 0, 0, 0, 0, 5, 0, Payload3),
+
+
+                    Data1 = <<FakeFragment1/binary, FakeFragment2/binary>>,
+                    Data2 = <<FakeFragment3/binary>>,
+
+                    [Message1] = wsecli_message:decode(Data1),
+                    [Message2] = wsecli_message:decode(Data2, Message1),
+
+                    assert_that(Message1#message.type, is(fragmented)),
+                    assert_that(Message2#message.type, is(binary)),
+                    assert_that(Message2#message.payload, is(Payload))
+                end),
+              it("should decode data with complete fragmented messages and part of fragmented one", fun() ->
+                    BinPayload1 = crypto:rand_bytes(30),
+                    <<
+                    Payload1:10/binary,
+                    Payload2:10/binary,
+                    Payload3/binary
+                    >> = BinPayload1,
+
+                    FakeFragment1 = get_binary_frame(0, 0, 0, 0, 2, 0, 10, 0, Payload1),
+                    FakeFragment2 = get_binary_frame(0, 0, 0, 0, 0, 0, 10, 0, Payload2),
+                    FakeFragment3 = get_binary_frame(1, 0, 0, 0, 0, 0, 10, 0, Payload3),
+
+                    BinPayload2 = crypto:rand_bytes(10),
+                    <<
+                    Payload4:10/binary,
+                    _/binary
+                    >> = BinPayload2,
+                    FakeFragment4 = get_binary_frame(0, 0, 0, 0, 2, 0, 10, 0, Payload4),
+
+                    Data = << FakeFragment1/binary, FakeFragment2/binary, FakeFragment3/binary, FakeFragment4/binary>>,
+
+                    [Message1, Message2] = wsecli_message:decode(Data),
+
+                    assert_that(Message1#message.type, is(binary)),
+                    assert_that(Message1#message.payload, is(BinPayload1)),
+                    assert_that(Message2#message.type, is(fragmented)),
+                    assert_that(length(Message2#message.frames), is(1))
+                end)
+
+          end),
         describe("unfragmented messages", fun()->
-              it("shit")
-              %it("decodes a text message", fun() ->
-              %      Payload = "Iepa yei!",
-
-              %      Fin = 1,
-              %      Rsv = 0,
-              %      Opcode = 1, %Text
-              %      Mask = 0,
-              %      PayloadLength = length(Payload),
-              %      PayloadData = list_to_binary(Payload),
-
-              %      FakeMessage =
-              %        <<Fin:1, Rsv:3, Opcode:4, Mask:1, PayloadLength:7, PayloadData/bits>>,
-
-              %      assert_that(wsecli_message:decode(FakeMessage), is({text, Payload}))
-              %  end)
+              it("control messages"),
+              it("should decode data containing various text messages", fun()->
+                  Text1 = "Churras churras",
+                  Payload1 = list_to_binary(Text1),
+                  PayloadLength1 = byte_size(Payload1),
+
+                  Text2 = "Pitas pitas",
+                  Payload2 = list_to_binary(Text2),
+                  PayloadLength2 = byte_size(Payload2),
+
+                  Text3 = "Pero que jallo eh",
+                  Payload3 = list_to_binary(Text3),
+                  PayloadLength3 = byte_size(Payload3),
+
+                  FakeMessage1 = get_binary_frame(1, 0, 0, 0, 1, 0, PayloadLength1, 0, Payload1),
+                  FakeMessage2 = get_binary_frame(1, 0, 0, 0, 1, 0, PayloadLength2, 0, Payload2),
+                  FakeMessage3 = get_binary_frame(1, 0, 0, 0, 1, 0, PayloadLength3, 0, Payload3),
+
+                  Data = << FakeMessage1/binary, FakeMessage2/binary, FakeMessage3/binary>>,
+
+                  [Message1, Message2, Message3] = wsecli_message:decode(Data),
+
+                  assert_that(Message1#message.type, is(text)),
+                  assert_that(Message1#message.payload, is(Text1)),
+                  assert_that(Message2#message.type, is(text)),
+                  assert_that(Message2#message.payload, is(Text2)),
+                  assert_that(Message3#message.type, is(text)),
+                  assert_that(Message3#message.payload, is(Text3))
+                end),
+              it("should decode data containing text and binary messages", fun()->
+                  Text1 = "Churras churras",
+                  Payload1 = list_to_binary(Text1),
+                  PayloadLength1 = byte_size(Payload1),
+
+                  Payload2 = crypto:rand_bytes(20),
+                  PayloadLength2 = 20,
+
+                  Text3 = "Pero que jallo eh",
+                  Payload3 = list_to_binary(Text3),
+                  PayloadLength3 = byte_size(Payload3),
+
+                  FakeMessage1 = get_binary_frame(1, 0, 0, 0, 1, 0, PayloadLength1, 0, Payload1),
+                  FakeMessage2 = get_binary_frame(1, 0, 0, 0, 2, 0, PayloadLength2, 0, Payload2),
+                  FakeMessage3 = get_binary_frame(1, 0, 0, 0, 1, 0, PayloadLength3, 0, Payload3),
+
+                  Data = << FakeMessage1/binary, FakeMessage2/binary, FakeMessage3/binary>>,
+
+                  [Message1, Message2, Message3] = wsecli_message:decode(Data),
+
+                  assert_that(Message1#message.type, is(text)),
+                  assert_that(Message1#message.payload, is(Text1)),
+                  assert_that(Message2#message.type, is(binary)),
+                  assert_that(Message2#message.payload, is(Payload2)),
+                  assert_that(Message3#message.type, is(text)),
+                  assert_that(Message3#message.payload, is(Text3))
+                end),
+              it("should decode data containing all message types"),
+              it("should decode data containing a binary message", fun() ->
+                    Payload = crypto:rand_bytes(45),
+                    %")
+                    FakeMessage = get_binary_frame(1, 0, 0, 0, 2, 0, 45, 0, Payload),
+                    [Message] = wsecli_message:decode(FakeMessage),
+
+                    assert_that( Message#message.payload, is(Payload))
+                end),
+              it("should decode data containing a text message", fun() ->
+                    Payload = "Iepa yei!",
+                    PayloadLength = length(Payload),
+                    PayloadData = list_to_binary(Payload),
+
+                    FakeMessage = get_binary_frame(1, 0, 0, 0, 1, 0, PayloadLength, 0, PayloadData),
+                    [Message] = wsecli_message:decode(FakeMessage),
+
+                    assert_that( Message#message.payload, is(Payload))
+                end)
           end)
       end).
+
+get_binary_frame(Fin, Rsv1, Rsv2, Rsv3, Opcode, Mask, Length, ExtendedPayloadLength, Payload) ->
+  Head = <<Fin:1, Rsv1:1, Rsv2:1, Rsv3:1, Opcode:4, Mask:1, Length:7>>,
+
+  case Length of
+    126 ->
+      <<Head/binary, ExtendedPayloadLength:16, Payload/binary>>;
+    127 ->
+      <<Head/binary, ExtendedPayloadLength:64, Payload/binary>>;
+    _ ->
+      <<Head/binary, Payload/binary>>
+  end.