Fix bug in mtp_codec

Buffer was not saved if crypto codec returns `incomplete`

src/mtp_codec.erl

@@ -74,7 +74,7 @@ decode_crypto(Bin, #codec{crypto_mod = CryptoMod,
                           crypto_buf = <<>>} = S) ->
     case CryptoMod:try_decode_packet(Bin, CryptoSt) of
         {incomplete, CryptoSt1} ->
-            decode_packet(<<>>, CryptoSt1, <<>>, S);
+            decode_packet(<<>>, CryptoSt1, Bin, S);
         {ok, Dec1, Tail1, CryptoSt1} ->
             decode_packet(Dec1, CryptoSt1, Tail1, S)
     end;

test/prop_mtp_codec.erl

@@ -5,7 +5,8 @@
 
 -export([prop_obfuscated_secure_stream/1,
          prop_obfuscated_secure_duplex/1,
-         prop_obfuscated_secure_duplex_multi/1]).
+         prop_obfuscated_secure_duplex_multi/1,
+         prop_fullcbc_stream/1]).
 
 
 prop_obfuscated_secure_stream(doc) ->
@@ -23,7 +24,7 @@ stream_arg_set() ->
       ]).
 
 obfuscated_secure_stream(Key, Iv, Stream) ->
-    Codec0 = mk_codec(Key, Iv, Key, Iv),
+    Codec0 = mk_secure_codec(Key, Iv, Key, Iv),
     {BinStream, Codec2} =
         lists:foldl(
           fun(Bin, {Acc, Codec1}) ->
@@ -31,12 +32,9 @@ obfuscated_secure_stream(Key, Iv, Stream) ->
                   {<<Acc/binary, (iolist_to_binary(Data))/binary>>,
                    Codec2}
           end, {<<>>, Codec0}, Stream),
-    {ok, RevStream, _Codec3} =
-        mtp_codec:fold_packets(
-          fun(Decoded, Acc, Codec3) ->
-                  {[Decoded | Acc], Codec3}
-          end, [], BinStream, Codec2),
-    Stream == lists:reverse(RevStream).
+    {ResStream, _Codec3} = parse_stream(BinStream, Codec2),
+    ?assertEqual(Stream, ResStream),
+    true.
 
 
 prop_obfuscated_secure_duplex(doc) ->
@@ -55,8 +53,8 @@ duplex_arg_set() ->
      mtp_prop_gen:stream_4b()}.
 
 obfuscated_secure_duplex(TxKey, TxIv, RxKey, RxIv, Stream) ->
-    CliCodec0 = mk_codec(TxKey, TxIv, RxKey, RxIv),
-    SrvCodec0 = mk_codec(RxKey, RxIv, TxKey, TxIv),
+    CliCodec0 = mk_secure_codec(TxKey, TxIv, RxKey, RxIv),
+    SrvCodec0 = mk_secure_codec(RxKey, RxIv, TxKey, TxIv),
     {_, _, BackStream} = roundtrip(CliCodec0, SrvCodec0, Stream),
     Stream == BackStream.
 
@@ -71,8 +69,8 @@ prop_obfuscated_secure_duplex_multi() ->
             obfuscated_secure_duplex_multi(N, TxKey, TxIv, RxKey, RxIv, Stream)).
 
 obfuscated_secure_duplex_multi(N, TxKey, TxIv, RxKey, RxIv, Stream0) ->
-    CliCodec0 = mk_codec(TxKey, TxIv, RxKey, RxIv),
-    SrvCodec0 = mk_codec(RxKey, RxIv, TxKey, TxIv),
+    CliCodec0 = mk_secure_codec(TxKey, TxIv, RxKey, RxIv),
+    SrvCodec0 = mk_secure_codec(RxKey, RxIv, TxKey, TxIv),
     {_, _, Stream} =
         lists:foldl(
           fun(I, {CliCodec1, SrvCodec1, Stream1}) ->
@@ -109,8 +107,65 @@ roundtrip(CliCodec0, SrvCodec0, Stream) ->
           end, [], SrvBinStream, CliCodec1),
     {CliCodec2, SrvCodec1, lists:reverse(RevCliDecodedStream)}.
 
-mk_codec(EncKey, EncIv, DecKey, DecIv) ->
+mk_secure_codec(EncKey, EncIv, DecKey, DecIv) ->
     Crypto = mtp_obfuscated:new(EncKey, EncIv, DecKey, DecIv),
     Packet = mtp_secure:new(),
     mtp_codec:new(mtp_obfuscated, Crypto,
                   mtp_secure, Packet).
+
+
+prop_fullcbc_stream(doc) ->
+    "Tests that any number of packets can be encrypted and decrypted with mtp_full + mtp_aes_cbc"
+        " It emulates downstream codec set".
+
+prop_fullcbc_stream() ->
+    ?FORALL({Key, Iv, Stream}, fullcbc_arg_set(), fullcbc_stream(Key, Iv, Stream)).
+
+fullcbc_arg_set() ->
+    proper_types:tuple(
+      [mtp_prop_gen:key(),
+       mtp_prop_gen:iv(),
+       mtp_prop_gen:stream_16b()
+      ]).
+
+fullcbc_stream(Key, Iv, Stream) ->
+    Codec0 = mk_fullcbc_codec(Key, Iv, Key, Iv),
+    {BinStream, Codec2} =
+        lists:foldl(
+          fun(Bin, {Acc, Codec1}) ->
+                  {Data, Codec2} = mtp_codec:encode_packet(Bin, Codec1),
+                  {<<Acc/binary, (iolist_to_binary(Data))/binary>>,
+                   Codec2}
+          end, {<<>>, Codec0}, Stream),
+    {ResStream, Codec3} = parse_stream(BinStream, Codec2),
+    ?assertEqual(Stream, ResStream, #{codec => Codec3}),
+    true.
+
+mk_fullcbc_codec(EncKey, EncIv, DecKey, DecIv) ->
+    Crypto = mtp_aes_cbc:new(EncKey, EncIv, DecKey, DecIv, 16),
+    Packet = mtp_full:new(1, 1),
+    mtp_codec:new(mtp_aes_cbc, Crypto,
+                  mtp_full, Packet).
+
+parse_stream(Bin, Codec0) ->
+    %% We want to split solid stream to smaller chunks to emulate network packet fragmentation
+    Chunks = split_stream(Bin),
+    {DecodedRev, Codec} =
+        lists:foldl(
+          fun(Chunk, {Acc1, Codec1}) ->
+                  {ok, Acc3, Codec3} =
+                      mtp_codec:fold_packets(
+                        fun(Decoded, Acc2, Codec2) ->
+                                {[Decoded | Acc2], Codec2}
+                        end, Acc1, Chunk, Codec1),
+                  {Acc3, Codec3}
+          end, {[], Codec0}, Chunks),
+    {lists:reverse(DecodedRev), Codec}.
+
+split_stream(<<>>) -> [];
+split_stream(Bin) when byte_size(Bin) < 4 -> [Bin];
+split_stream(Bin) ->
+    %% TODO: should have deterministic seed for rand!
+    Size = rand:uniform(byte_size(Bin) div 2),
+    <<Chunk:Size/binary, Tail/binary>> = Bin,
+    [Chunk | split_stream(Tail)].