From d43152e9be477377dedc12d457c21388a95ac0bf Mon Sep 17 00:00:00 2001
From: Sergey Prokhorov <seriy.pr@gmail.com>
Date: Sat, 6 Jun 2020 15:25:01 +0200
Subject: [PATCH] tests: handle situation when ServerHello is split to multiple
 packets

---
 src/mtp_fake_tls.erl     |  9 +++++++--
 test/mtp_test_client.erl | 28 +++++++++++++++++++---------
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/src/mtp_fake_tls.erl b/src/mtp_fake_tls.erl
index 66064d8..a0a20a4 100644
--- a/src/mtp_fake_tls.erl
+++ b/src/mtp_fake_tls.erl
@@ -298,9 +298,14 @@ add_padding_ext(RealExtensions, ExtLen) ->
 %% Parses "ServerHello" (the one produced by from_client_hello/2). Used for tests only.
 parse_server_hello(<<?TLS_REC_HANDSHAKE, ?TLS_12_VERSION, HSLen:?u16, Handshake:HSLen/binary,
                      ?TLS_REC_CHANGE_CIPHER, ?TLS_12_VERSION, CCLen:?u16, ChangeCipher:CCLen/binary,
-                     ?TLS_REC_DATA, ?TLS_12_VERSION, DLen:?u16, Data:DLen/binary,
+                     ?TLS_REC_DATA, ?TLS_12_VERSION, DLen:?u16,
                      Tail/binary>>) ->
-    {Handshake, ChangeCipher, Data, Tail};
+    case Tail of
+        <<Data:DLen/binary, Tail2/binary>> ->
+            {Handshake, ChangeCipher, Data, Tail2};
+        _ ->
+            incomplete
+    end;
 parse_server_hello(B) when byte_size(B) < (512 + 5) ->
     incomplete.
 
diff --git a/test/mtp_test_client.erl b/test/mtp_test_client.erl
index 9c07a38..9db6eb8 100644
--- a/test/mtp_test_client.erl
+++ b/test/mtp_test_client.erl
@@ -35,31 +35,31 @@ connect(Host, Port, Secret, DcId, Protocol) ->
               binary(), binary(), integer(),
               mtp_codec:packet_codec() | {mtp_fake_tls, binary()}) -> client().
 connect(Host, Port, Seed, Secret, DcId, Protocol0) ->
+    Timeout = 5000,
     Opts = [{packet, raw},
             {mode, binary},
             {active, false},
             {buffer, 1024},
-            {send_timeout, 5000}],
-    {ok, Sock} = gen_tcp:connect(Host, Port, Opts, 1000),
-    {Protocol, TlsEnabled, TlsSt} =
+            {send_timeout, Timeout}],
+    {ok, Sock} = gen_tcp:connect(Host, Port, Opts, Timeout),
+    {Protocol, TlsEnabled, TlsSt, Tail} =
         case Protocol0 of
             {mtp_fake_tls, Domain} ->
                 ClientHello = mtp_fake_tls:make_client_hello(Secret, Domain),
                 ok = gen_tcp:send(Sock, ClientHello),
                 %% Let's hope whole server hello will arrive in a single chunk
-                {ok, ServerHello} = gen_tcp:recv(Sock, 0, 5000),
-                %% TODO: if Tail is not empty, use codec:push_back(first, ..)
-                {_HS, _CC, _D, <<>>} = mtp_fake_tls:parse_server_hello(ServerHello),
-                {mtp_secure, true, mtp_fake_tls:new()};
-            _ -> {Protocol0, false, undefined}
+                Tail_ = recv_server_hello(Sock, Timeout, <<>>),
+                {mtp_secure, true, mtp_fake_tls:new(), Tail_};
+            _ -> {Protocol0, false, undefined, <<>>}
         end,
     {Header0, _, _, CryptoLayer} = mtp_obfuscated:client_create(Seed, Secret, Protocol, DcId),
     NoopSt = mtp_noop_codec:new(),
     %% First, create codec with just TLS (which might be noop as well) to encode "obfuscated" header
-    Codec0 = mtp_codec:new(mtp_noop_codec, NoopSt,
+    Codec00 = mtp_codec:new(mtp_noop_codec, NoopSt,
                            mtp_noop_codec, NoopSt,
                            TlsEnabled, TlsSt,
                            25 * 1024 * 1024),
+    Codec0 = mtp_codec:push_back(first, Tail, Codec00),
     {Header, Codec1} = mtp_codec:encode_packet(Header0, Codec0),
     ok = gen_tcp:send(Sock, Header),
     PacketLayer = Protocol:new(),
@@ -68,6 +68,16 @@ connect(Host, Port, Seed, Secret, DcId, Protocol0) ->
     #client{sock = Sock,
             codec = Codec3}.
 
+recv_server_hello(Sock, Timeout, Acc) ->
+    {ok, ServerHelloPart} = gen_tcp:recv(Sock, 0, Timeout),
+    ServerHello = <<Acc/binary, ServerHelloPart/binary>>,
+    case mtp_fake_tls:parse_server_hello(ServerHello) of
+        {_HS, _CC, _D, Tail} ->
+            Tail;
+        incomplete ->
+            recv_server_hello(Sock, Timeout, ServerHello)
+    end.
+
 send(Data, #client{sock = Sock, codec = Codec} = Client) ->
     {Enc, Codec1} = mtp_codec:encode_packet(Data, Codec),
     ok = gen_tcp:send(Sock, Enc),