Move handshake packet encoding/decoding from mtp_down_con to mtp_rpc

src/mtp_down_conn.erl

@@ -17,6 +17,9 @@
          shutdown/1,
          send/2,
          ack/3]).
+-ifdef(TEST).
+-export([get_middle_key/1]).
+-endif.
 
 %% gen_server callbacks
 -export([init/1, handle_call/3, handle_cast/2, handle_info/2,
@@ -388,20 +391,12 @@ tcp_connect(Host, Port) ->
             Err
     end.
 
--define(RPC_NONCE, <<170,135,203,122>>).
--define(RPC_HANDSHAKE, <<245,238,130,118>>).
--define(RPC_FLAGS, <<0, 0, 0, 0>>).
-
 down_handshake1(S) ->
-    RpcNonce = ?RPC_NONCE,
     <<KeySelector:4/binary, _/binary>> = Key = mtp_config:get_secret(),
     CryptoTs = os:system_time(seconds),
     Nonce = crypto:strong_rand_bytes(16),
-    Msg = <<RpcNonce/binary,
-            KeySelector/binary,
-            1:32/little,                        %AES
-            CryptoTs:32/little,
-            Nonce/binary>>,
+    Schema = 1,                                 %AES
+    Msg = mtp_rpc:encode_nonce({nonce, KeySelector, Schema, CryptoTs, Nonce}),
     S1 = S#state{stage = handshake_1,
                  %% Use fake encryption codec
                  codec = mtp_codec:new(mtp_noop_codec, mtp_noop_codec:new(),
@@ -409,11 +404,10 @@ down_handshake1(S) ->
                  stage_state = {KeySelector, Nonce, CryptoTs, Key}},
     down_send(Msg, S1).
 
-down_handshake2(<<Type:4/binary, KeySelector:4/binary, Schema:32/little, _CryptoTs:4/binary,
-                  SrvNonce:16/binary>>, #state{stage_state = {MyKeySelector, CliNonce, MyTs, Key},
-                                               codec = Codec1,
-                                               sock = Sock} = S) ->
-    (Type == ?RPC_NONCE) orelse error({wrong_rpc_type, Type}),
+down_handshake2(Pkt, #state{stage_state = {MyKeySelector, CliNonce, MyTs, Key},
+                            codec = Codec1,
+                            sock = Sock} = S) ->
+    {nonce, KeySelector, Schema, _CryptoTs, SrvNonce} = mtp_rpc:decode_nonce(Pkt),
     (Schema == 1) orelse error({wrong_schema, Schema}),
     (KeySelector == MyKeySelector) orelse error({wrong_key_selector, KeySelector}),
     {ok, {DownIp, DownPort}} = inet:peername(Sock),
@@ -430,10 +424,7 @@ down_handshake2(<<Type:4/binary, KeySelector:4/binary, Schema:32/little, _Crypto
     Codec = mtp_codec:new(mtp_aes_cbc, CryptoState,
                           PacketMod, PacketState),
     SenderPID = PeerPID = <<"IPIPPRPDTIME">>,
-    Handshake = [?RPC_HANDSHAKE,
-                 ?RPC_FLAGS,
-                 SenderPID,
-                 PeerPID],
+    Handshake = mtp_rpc:encode_handshake({handshake, SenderPID, PeerPID}),
     down_send(Handshake,
               S#state{codec = Codec,
                       stage = handshake_2,
@@ -465,10 +456,9 @@ get_middle_key(#{srv_n := Nonce, clt_n := MyNonce, clt_ts := MyTs, srv_ip := Srv
     {Key, IV}.
 
 
-down_handshake3(<<Type:4/binary, _Flags:4/binary, _SenderPid:12/binary, PeerPid:12/binary>>,
-                #state{stage_state = PrevSenderPid, pool = Pool,
-                       netloc = {Addr, Port}} = S) ->
-    (Type == ?RPC_HANDSHAKE) orelse error({wrong_rpc_type, Type}),
+down_handshake3(Pkt, #state{stage_state = PrevSenderPid, pool = Pool,
+                            netloc = {Addr, Port}} = S) ->
+    {handshake, _SenderPid, PeerPid} = mtp_rpc:decode_handshake(Pkt),
     (PeerPid == PrevSenderPid) orelse error({wrong_sender_pid, PeerPid}),
     ok = mtp_dc_pool:ack_connected(Pool, self()),
     lager:info("~s:~w: handshake complete", [inet:ntoa(Addr), Port]),

src/mtp_rpc.erl

@@ -8,7 +8,15 @@
 -module(mtp_rpc).
 
 -export([decode_packet/1,
+         decode_nonce/1,
+         decode_handshake/1,
+         encode_nonce/1,
+         encode_handshake/1,
          encode_packet/2]).
+%% For tests
+-export([srv_decode_packet/1,
+         srv_encode_packet/1]).
+%% Helpers
 -export([inet_pton/1,
          encode_ip_port/2]).
 -export_type([codec/0]).
@@ -29,6 +37,11 @@
 -define(RPC_SIMPLE_ACK, 155,64,172,59).         %0x3bac409b
 -define(TL_PROXY_TAG, 174,38,30,219).
 
+-define(RPC_NONCE, 170,135,203,122).
+-define(RPC_HANDSHAKE, 245,238,130,118).
+-define(RPC_FLAGS, 0, 0, 0, 0).
+
+
 -define(FLAG_NOT_ENCRYPTED , 16#2).
 -define(FLAG_HAS_AD_TAG    , 16#8).
 -define(FLAG_MAGIC         , 16#1000).
@@ -45,15 +58,31 @@
                 | {close_ext, conn_id()}
                 | {simple_ack, conn_id(), binary()}.
 
-%% new(ClientIp, ClientPort, ProxyIp, ProxyPort, ProxyTag) ->
-%%     new(ClientIp, ClientPort, ProxyIp, ProxyPort, ProxyTag,
-%%         erlang:unique_integer()).
+decode_nonce(<<?RPC_NONCE,
+               KeySelector:4/binary,
+               Schema:32/little,
+               CryptoTs:32/little,
+               CliNonce:16/binary>>) ->
+    {nonce, KeySelector, Schema, CryptoTs, CliNonce}.
 
-%% new(ClientIp, ClientPort, ProxyIp, ProxyPort, ProxyTag, ConnId) ->
-%%     #rpc_st{client_addr = iolist_to_binary(encode_ip_port(ClientIp, ClientPort)),
-%%             proxy_addr = iolist_to_binary(encode_ip_port(ProxyIp, ProxyPort)),
-%%             proxy_tag = ProxyTag,
-%%             conn_id = ConnId}.
+decode_handshake(<<?RPC_HANDSHAKE,
+                       ?RPC_FLAGS,
+                       SenderPID:12/binary,
+                       PeerPID:12/binary>>) ->
+    {handshake, SenderPID, PeerPID}.
+
+encode_nonce({nonce, KeySelector, Schema, CryptoTs, SrvNonce}) ->
+    <<?RPC_NONCE,
+      KeySelector:4/binary,
+      Schema:32/little,
+      CryptoTs:32/little,
+      SrvNonce:16/binary>>.
+
+encode_handshake({handshake, SenderPID, PeerPID}) ->
+    <<?RPC_HANDSHAKE,
+      0, 0, 0, 0,                    %Some flags
+      SenderPID:12/binary,
+      PeerPID:12/binary>>.
 
 %% It expects that packet segmentation was done on previous layer
 %% See mtproto/mtproto-proxy.c:process_client_packet
@@ -99,6 +128,29 @@ encode_packet({data, Msg}, {{ConnId, ClientAddr, ProxyTag}, ProxyAddr}) ->
 encode_packet(remote_closed, ConnId) ->
     <<?RPC_CLOSE_CONN, ConnId:64/little-signed>>.
 
+%%
+%% Middle-proxy side encoding and decodong (FOR TESTS ONLY!)
+%%
+
+%% opposite of encode_packet
+srv_decode_packet(<<?RPC_PROXY_REQ, _Flags:32/little, ConnId:64/little-signed,
+                    _ClientAddrBin:20/binary, _ProxyAddrBin:20/binary,
+                    24:32/little, ?TL_PROXY_TAG, 16, _ProxyTag:16/binary, 0, 0, 0,
+                    Data/binary>>) ->
+    {data, ConnId, Data};
+srv_decode_packet(<<?RPC_CLOSE_CONN, ConId:64/little-signed>>) ->
+    {remote_closed, ConId}.
+
+%% Opposite of decode_packet
+srv_encode_packet({proxy_ans, ConnId, Data}) ->
+    <<?RPC_PROXY_ANS,
+      0, 0, 0, 0,                               %some flags
+      ConnId:64/signed-little,
+      Data/binary>>;
+srv_encode_packet({close_ext, ConnId}) ->
+    <<?RPC_CLOSE_EXT, ConnId:64/little-signed>>.
+
+%% IP and port as 10 + 2 + 4 + 4 = 20b
 -spec encode_ip_port(inet:ip_address(), inet:port_number()) -> iodata().
 encode_ip_port(IPv4, Port) when tuple_size(IPv4) == 4 ->
     IpBin = inet_pton(IPv4),