Implement sponsored channels and work through official endpoints

rebar.config

@@ -3,7 +3,7 @@
             {parse_transform, lager_transform}]}.
 
 {deps, [{ranch, "1.5.0"},
-        {lager, "3.6.1"}
+        {lager, "3.6.3"}
         ]}.
 
 {relx, [{release, { mtp_proxy, "0.1.0" },

rebar.lock

 {"1.1.0",
 [{<<"goldrush">>,{pkg,<<"goldrush">>,<<"0.1.9">>},1},
- {<<"lager">>,{pkg,<<"lager">>,<<"3.6.1">>},0},
+ {<<"lager">>,{pkg,<<"lager">>,<<"3.6.3">>},0},
  {<<"ranch">>,{pkg,<<"ranch">>,<<"1.5.0">>},0}]}.
 [
 {pkg_hash,[
  {<<"goldrush">>, <<"F06E5D5F1277DA5C413E84D5A2924174182FB108DABB39D5EC548B27424CD106">>},
- {<<"lager">>, <<"9D29C5FF7F926D25ECD9899990867C9152DCF34EEE65BAC8EC0DFC0D16A26E0C">>},
+ {<<"lager">>, <<"FE78951D174616273F87F0DBC3374D1430B1952E5EFC4E1C995592D30A207294">>},
  {<<"ranch">>, <<"F04166F456790FEE2AC1AA05A02745CC75783C2BFB26D39FAF6AEFC9A3D3A58A">>}]}
 ].

src/metric.erl

+%%% @author sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, sergey
+%%% @doc
+%%% Interface for logging metrics (CODE WIPED)
+%%% @end
+%%% Created : 15 May 2018 by sergey <me@seriyps.ru>
+
+-module(metric).
+
+-export([count_inc/3,
+         gauge_set/3,
+         rt/2,
+         histogram_observe/3]).
+
+count_inc(_Name, _Value, _Extra) ->
+    noop.
+
+gauge_set(_Name, _Value, _Extra) ->
+    noop.
+
+histogram_observe(_Name, _Value, _Extra) ->
+    noop.
+
+rt(_Name, _Fun) ->
+    noop.

src/mtp_abridged.erl

+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% MTProto abridged packet format codec
+%%% @end
+%%% Created : 29 May 2018 by Sergey <me@seriyps.ru>
+
+-module(mtp_abridged).
+-behaviour(mtp_layer).
+
+-export([new/0,
+         try_decode_packet/2,
+         encode_packet/2]).
+-export_type([codec/0]).
+
+-record(st,
+        {buffer = <<>> :: binary()}).
+-define(MAX_PACKET_SIZE, 1 * 1024 * 1024).      % 1mb
+-define(APP, mtproto_proxy).
+
+-opaque codec() :: #st{}.
+
+new() ->
+    #st{}.
+
+-spec try_decode_packet(binary(), codec()) -> {ok, binary(), codec()}
+                                                  | {incomplete, codec()}.
+try_decode_packet(<<Flag, Len:24/unsigned-little-integer, Rest/binary>> = Data,
+                      #st{buffer = <<>>} = St) when Flag == 127; Flag == 255 ->
+    Len1 = Len * 4,
+    (Len1 < ?MAX_PACKET_SIZE)
+        orelse
+        begin
+            metric:count_inc([?APP, protocol_error, total], 1, #{labels => [abriged_max_size]}),
+            error({packet_too_large, Len1})
+        end,
+    try_decode_packet_len(Len1, Rest, Data, St);
+try_decode_packet(<<Len, Rest/binary>> = Data,
+                      #st{buffer = <<>>} = St) when Len >= 128 ->
+    Len1 = (Len - 128) * 4,
+    try_decode_packet_len(Len1, Rest, Data, St);
+try_decode_packet(<<Len, Rest/binary>> = Data,
+                      #st{buffer = <<>>} = St) when Len < 127 ->
+    Len1 = Len * 4,
+    try_decode_packet_len(Len1, Rest, Data, St);
+try_decode_packet(Bin, #st{buffer = Buf} = St) when byte_size(Buf) > 0 ->
+    try_decode_packet(<<Buf/binary, Bin/binary>>, St#st{buffer = <<>>});
+try_decode_packet(Bin, #st{buffer = <<>>} = St) ->
+    {incomplete, St#st{buffer = Bin}}.
+
+try_decode_packet_len(Len, LenStripped, Data, St) ->
+    case LenStripped of
+        <<Packet:Len/binary, Rest/binary>> ->
+            {ok, Packet, St#st{buffer = Rest}};
+        _ ->
+            {incomplete, St#st{buffer = Data}}
+    end.
+
+-spec encode_packet(binary(), codec()) -> iodata().
+encode_packet(Bin, St) ->
+    Size = byte_size(Bin),
+    Len = Size div 4,
+    Packet =
+        case Len < 127 of
+            true ->
+                [Len | Bin];
+            false ->
+                [<<127, Len:24/unsigned-little-integer>> | Bin]
+        end,
+    {Packet, St}.
+
+-ifdef(TEST).
+-include_lib("eunit/include/eunit.hrl").
+
+decode_none_test() ->
+    S = new(),
+    ?assertEqual(
+       {incomplete, S}, try_decode_packet(<<>>, S)).
+
+-endif.

src/mtp_aes_cbc.erl

+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% Block CBC AES codec with buffered decoding
+%%% @end
+%%% Created :  6 Jun 2018 by Sergey <me@seriyps.ru>
+
+-module(mtp_aes_cbc).
+-behaviour(mtp_layer).
+
+-export([new/5,
+         encrypt/2,
+         decrypt/2,
+         try_decode_packet/2,
+         encode_packet/2
+        ]).
+
+-export_type([codec/0]).
+
+-record(baes_st,
+        {decode_buf :: binary(),
+         block_size :: pos_integer(),
+         encrypt :: any(),                      % aes state
+         decrypt :: any()                       % aes state
+        }).
+
+-opaque codec() :: #baes_st{}.
+
+
+
+new(EncKey, EncIv, DecKey, DecIv, BlockSize) ->
+    #baes_st{
+       decode_buf = <<>>,
+       block_size = BlockSize,
+       encrypt = {EncKey, EncIv},
+       decrypt = {DecKey, DecIv}
+      }.
+
+-spec encrypt(iodata(), codec()) -> {binary(), codec()}.
+encrypt(Data, #baes_st{block_size = BSize,
+                       encrypt = {EncKey, EncIv}} = S) ->
+    ((iolist_size(Data) rem BSize) == 0)
+        orelse error({data_not_aligned, BSize, byte_size(Data)}),
+    Encrypted = crypto:block_encrypt(aes_cbc, EncKey, EncIv, Data),
+    {Encrypted, S#baes_st{encrypt = {EncKey, crypto:next_iv(aes_cbc, Encrypted)}}}.
+
+
+-spec decrypt(binary(), codec()) -> {binary(), codec()}.
+decrypt(Data, #baes_st{block_size = BSize,
+                       decode_buf = <<>>} = S) ->
+    Size = byte_size(Data),
+    Div = Size div BSize,
+    Rem = Size rem BSize,
+    case {Div, Rem} of
+        {0, _} ->
+            %% Not enough bytes
+            {<<>>, S#baes_st{decode_buf = Data}};
+        {_, 0} ->
+            %% Aligned
+            do_decrypt(Data, S);
+        {_, Tail} ->
+            %% N blocks + reminder
+            Head = Size - Tail,
+            <<ToDecode:Head/binary, Reminder/binary>> = Data,
+            do_decrypt(ToDecode, S#baes_st{decode_buf = Reminder})
+    end;
+decrypt(Data, #baes_st{decode_buf = Buf} = S) ->
+    decrypt(<<Buf/binary, Data/binary>>, S#baes_st{decode_buf = <<>>}).
+
+do_decrypt(Data, #baes_st{decrypt = {DecKey, DecIv}} = S) ->
+    Decrypted = crypto:block_decrypt(aes_cbc, DecKey, DecIv, Data),
+    NewDecIv = crypto:next_iv(aes_cbc, Data),
+    {Decrypted, S#baes_st{decrypt = {DecKey, NewDecIv}}}.
+
+%% To comply mtp_layer interface
+try_decode_packet(Bin, S) ->
+    case decrypt(Bin, S) of
+        {<<>>, S1} ->
+            {incomplete, S1};
+        {Dec, S1} ->
+            {ok, Dec, S1}
+    end.
+
+encode_packet(Bin, S) ->
+    encrypt(Bin, S).
+
+
+-ifdef(TEST).
+-include_lib("eunit/include/eunit.hrl").
+
+decode_none_test() ->
+    DecKey = <<21,211,191,127,143,222,184,152,232,213,25,173,243,163,243,224,133,131,199,13,206,156,146,141,67,172,85,114,190,203,176,215>>,
+    DecIV = <<9,156,175,247,37,161,219,155,52,115,93,76,122,195,158,194>>,
+    S = new(DecKey, DecIV, DecKey, DecIV, 16),
+    ?assertEqual(
+       {incomplete, S}, try_decode_packet(<<>>, S)).
+
+decrypt_test() ->
+    DecKey = <<21,211,191,127,143,222,184,152,232,213,25,173,243,163,243,224,133,131,199,13,206,156,146,141,67,172,85,114,190,203,176,215>>,
+    DecIV = <<9,156,175,247,37,161,219,155,52,115,93,76,122,195,158,194>>,
+    S = new(DecKey, DecIV, DecKey, DecIV, 16),
+    Samples =
+        [{<<36,170,147,95,53,27,44,255,252,105,70,8,90,40,77,226>>,
+          <<44,0,0,0,255,255,255,255,245,238,130,118,0,0,0,0>>},
+         {<<137,187,80,238,110,142,52,130,119,140,210,138,13,72,169,144,63,167,172,19,161,13,231,169,237,34,203,240,8,135,67,29>>,
+          <<134,153,66,10,95,9,134,49,221,133,21,91,73,80,73,80,80,82,80,68,84,73,77,69,133,250,76,84,4,0,0,0>>}
+        ],
+    lists:foldl(
+      fun({In, Out}, S1) ->
+              {Dec, S2} = decrypt(In, S1),
+              ?assertEqual(Out, Dec),
+              S2
+      end, S, Samples).
+
+encrypt_test() ->
+    EncKey = <<89,84,72,247,172,56,204,11,10,242,143,240,111,53,33,162,221,141,148,243,100,21,167,160,132,99,61,189,128,73,138,89>>,
+    EncIV = <<248,195,42,53,235,104,78,225,84,171,182,125,18,192,251,77>>,
+    S = new(EncKey, EncIV, EncKey, EncIV, 16),
+    Samples =
+        [{<<44,0,0,0,255,255,255,255,245,238,130,118,0,0,0,0,73,80,73,80,80,82,80,68,84,73,77,69,73,80,73,80,80,82,80,68,84,73,77,69,2,118,29,129,4,0,0,0>>,
+          <<161,206,198,191,175,240,48,162,245,192,234,210,104,195,161,214,55,147,145,157,174,33,243,198,84,188,29,201,116,128,185,149,73,241,149,122,244,193,59,112,153,188,141,134,90,24,75,216>>},
+         {<<136,0,0,0,0,0,0,0,238,241,206,54,8,16,2,64,195,43,106,127,211,218,156,102,65,65,65,65,65,65,65,65,65,65,65,65,65,65,65,65,65,65,65,65,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,66,24,0,0,0,174,38,30,219,16,220,190,143,20,147,250,76,217,171,48,8,145,192,181,179,38,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,56,220,23,91,20,0,0,0,120,151,70,96,26,49,96,74,221,243,104,13,173,13,132,192,238,22,97,126,247,233,151,22,4,0,0,0,4,0,0,0>>,
+          <<92,173,139,247,1,147,48,108,162,98,125,215,170,185,87,131,65,26,90,205,43,54,115,216,90,101,3,188,151,165,126,144,104,247,57,65,32,107,245,154,77,194,161,157,63,232,169,68,113,64,96,197,10,209,66,117,251,15,10,141,248,122,40,242,195,38,196,237,68,132,189,49,102,53,31,139,56,64,213,107,79,105,210,182,157,73,203,105,165,134,163,116,49,94,143,171,88,132,84,123,196,38,35,53,220,182,232,199,92,29,182,129,239,116,252,31,72,29,120,203,57,49,46,129,142,94,204,121,21,113,211,10,193,126,180,227,139,40,85,223,134,124,152,81>>}],
+    lists:foldl(
+      fun({In, Out}, S1) ->
+              {Enc, S2} = encrypt(In, S1),
+              ?assertEqual(Out, Enc),
+              S2
+      end, S, Samples).
+
+-endif.

src/mtp_full.erl

+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% MTProto "full" packet format with padding
+%%% ```
+%%% <<MsgLen:32/integer, SeqNo:32/integer, Body:MsgLen/binary, CRC:32/integer>>
+%%% ```
+%%% @end
+%%% Created :  6 Jun 2018 by Sergey <me@seriyps.ru>
+
+-module(mtp_full).
+-behaviour(mtp_layer).
+
+-export([new/0, new/2,
+         try_decode_packet/2,
+         encode_packet/2]).
+-export_type([codec/0]).
+
+-record(full_st,
+        {decode_buf = <<>> :: binary(),
+         enc_seq_no :: integer(),
+         dec_seq_no :: integer()}).
+-define(MIN_MSG_LEN, 12).
+-define(MAX_MSG_LEN,  16777216).                %2^24 - 16mb
+
+-define(BLOCK_SIZE, 16).
+-define(PAD, <<4:32/little>>).
+-define(APP, mtproto_proxy).
+
+-opaque codec() :: #full_st{}.
+
+
+new() ->
+    new(0, 0).
+
+new(EncSeqNo, DecSeqNo) ->
+    #full_st{enc_seq_no = EncSeqNo,
+             dec_seq_no = DecSeqNo}.
+
+try_decode_packet(<<4:32/little, Bin/binary>>, #full_st{decode_buf = <<>>} = S) ->
+    %% Skip padding
+    try_decode_packet(Bin, S);
+try_decode_packet(<<Len:32/little, PktSeqNo:32/signed-little, Tail/binary>> = Bin,
+                  #full_st{decode_buf = <<>>, dec_seq_no = SeqNo} = S) ->
+    ((Len rem byte_size(?PAD)) == 0)
+        orelse error({wrong_alignement, Len}),
+    ((?MIN_MSG_LEN =< Len) and (Len =< ?MAX_MSG_LEN))
+        orelse error({wrong_msg_len, Len}),
+    (SeqNo == PktSeqNo)
+        orelse error({wrong_seq_no, SeqNo, PktSeqNo}),
+    BodyLen = Len - 4 - 4 - 4,
+    case Tail of
+        <<Body:BodyLen/binary, CRC:32/little, Rest/binary>> ->
+            PacketCrc = erlang:crc32([<<Len:32/little, PktSeqNo:32/little>> | Body]),
+            (CRC == PacketCrc)
+                orelse error({wrong_checksum, CRC, PacketCrc}),
+            {ok, Body, S#full_st{decode_buf = Rest, dec_seq_no = SeqNo + 1}};
+        _ ->
+            {incomplete, S#full_st{decode_buf = Bin}}
+    end;
+try_decode_packet(Bin, #full_st{decode_buf = Buf} = S) when byte_size(Buf) > 0 ->
+    try_decode_packet(<<Buf/binary, Bin/binary>>, S#full_st{decode_buf = <<>>});
+try_decode_packet(Bin, #full_st{decode_buf = <<>>} = S) ->
+    {incomplete, S#full_st{decode_buf = Bin}}.
+
+encode_packet(Bin, #full_st{enc_seq_no = SeqNo} = S) ->
+    BodySize = iolist_size(Bin),
+    ((BodySize rem byte_size(?PAD)) == 0)
+        orelse error({wrong_alignment, BodySize}),
+    Len = BodySize + 4 + 4 + 4,
+    MsgNoChecksum =
+        [<<Len:32/unsigned-little-integer,
+           SeqNo:32/signed-little-integer>>
+             | Bin],
+    CheckSum = erlang:crc32(MsgNoChecksum),
+    FullMsg = [MsgNoChecksum | <<CheckSum:32/unsigned-little-integer>>],
+    Len = iolist_size(FullMsg),
+    %% XXX: is there a cleaner way?
+    PaddingSize = (?BLOCK_SIZE - (Len rem ?BLOCK_SIZE)) rem ?BLOCK_SIZE,
+    NPaddings = PaddingSize div byte_size(?PAD),
+    Padding = lists:duplicate(NPaddings, ?PAD),
+    {[FullMsg | Padding], S#full_st{enc_seq_no = SeqNo + 1}}.
+
+
+-ifdef(TEST).
+-include_lib("eunit/include/eunit.hrl").
+
+encode_nopadding_test() ->
+    S = new(),
+    {Enc, _S1} = encode_packet(<<1, 1, 1, 1>>, S),
+    ?assertEqual(
+       <<16,0,0,0,
+         0,0,0,0,
+         1,1,1,1,
+         22,39,175,160>>,
+       iolist_to_binary(Enc)).
+
+encode_padding_test() ->
+    S = new(),
+    {Enc, _S1} = encode_packet(<<1,1,1,1,1,1,1,1>>, S),
+    ?assertEqual(
+       <<20,0,0,0,0,0,0,0,                      %size, seq no
+         1,1,1,1,1,1,1,1,                       %data
+         246,196,46,149,                        %CRC
+         4,0,0,0,4,0,0,0,4,0,0,0>>,             %padding
+       iolist_to_binary(Enc)).
+
+encode_padding_seq_test() ->
+    S = new(),
+    {Enc1, S1} = encode_packet(binary:copy(<<9>>, 8), S),
+    ?assertEqual(
+       <<20,0,0,0,
+         0,0,0,0,
+         9,9,9,9,9,9,9,9,
+         229,35,162,164,
+         4, 0,0,0,4,0,0,0,4,0,0,0>>,
+       iolist_to_binary(Enc1)),
+    {Enc2, _S2} = encode_packet(binary:copy(<<8>>, 8), S1),
+    ?assertEqual(
+       <<20,0,0,0,
+         1,0,0,0,
+         8,8,8,8,8,8,8,8,
+         48,146,132,116,
+         4,0,0,0,4,0,0,0,4,0,0,0>>,
+       iolist_to_binary(Enc2)).
+
+decode_none_test() ->
+    S = new(),
+    ?assertEqual(
+       {incomplete, S}, try_decode_packet(<<>>, S)).
+
+codec_test() ->
+    S = new(),
+    Packets = [
+               binary:copy(<<0>>, 4),           %non-padded
+               binary:copy(<<1>>, 8),           %padded
+               binary:copy(<<2>>, 4),           %non-padded
+               binary:copy(<<2>>, 100)          %padded
+              ],
+    lists:foldl(
+      fun(B, S1) ->
+              {Encoded, S2} = encode_packet(B, S1),
+              BinEncoded = iolist_to_binary(Encoded),
+              {ok, Decoded, S3} = try_decode_packet(BinEncoded, S2),
+              ?assertEqual(B, Decoded, {BinEncoded, S2, S3}),
+              S3
+      end, S, Packets).
+
+codec_stream_test() ->
+    S = new(),
+    Packets = [
+               binary:copy(<<0>>, 4),           %non-padded
+               binary:copy(<<1>>, 8),           %padded
+               binary:copy(<<2>>, 4),           %non-padded
+               binary:copy(<<2>>, 100)          %padded
+              ],
+    {Encoded, SS} =
+        lists:foldl(
+          fun(B, {Enc1, S1}) ->
+                  {Enc2, S2} = encode_packet(B, S1),
+                  {[Enc1 | Enc2], S2}
+          end, {[], S}, Packets),
+    lists:foldl(
+      fun(B, {Enc, S1}) ->
+              {ok, Dec, S2} = try_decode_packet(Enc, S1),
+              ?assertEqual(B, Dec),
+              {<<>>, S2}
+      end, {iolist_to_binary(Encoded), SS}, Packets).
+
+-endif.

src/mtp_handler.erl

@@ -12,7 +12,7 @@
 %% API
 -export([start_link/4]).
 -export([hex/1]).
--export([key_str/0]).
+-export([keys_str/0]).
 
 %% Callbacks
 -export([ranch_init/1]).
@@ -20,16 +20,45 @@
          terminate/2, code_change/3]).
 
 -define(MAX_SOCK_BUF_SIZE, 1024 * 300).    % Decrease if CPU is cheaper than RAM
+-define(MAX_UP_INIT_BUF_SIZE, 1024 * 1024).     %1mb
+
+%% TODO: download from https://core.telegram.org/getProxyConfig
+-define(TG_MIDDLE_PROXIES_V4,
+        {
+          {{149, 154, 175, 50}, 8888},
+          {{149, 154, 162, 38}, 80},
+          {{149, 154, 175, 100}, 8888},
+          {{91, 108, 4, 136}, 8888},
+          {{91, 108, 56, 181}, 8888}
+        }).
+%% TODO: download from https://core.telegram.org/getProxySecret
+-define(PROXY_SECRET,
+        <<196,249,250,202,150,120,230,187,72,173,108,126,44,229,192,210,68,48,100,
+          93,85,74,221,235,85,65,158,3,77,166,39,33,208,70,234,171,110,82,171,20,
+          169,90,68,62,207,179,70,62,121,160,90,102,97,42,223,156,174,218,139,233,
+          168,13,166,152,111,176,166,255,56,122,248,77,136,239,58,100,19,113,62,92,
+          51,119,246,225,163,212,125,153,245,224,197,110,236,232,240,92,84,196,144,
+          176,121,227,27,239,130,255,14,232,242,176,163,39,86,210,73,197,242,18,105,
+          129,108,183,6,27,38,93,178,18>>).
+
 -define(APP, mtproto_proxy).
 
 -record(state,
         {stage = init :: stage(),
-         init_buf = <<>> :: binary(),
+         stage_state = <<>> :: any(),
+         up_acc = <<>> :: any(),
+
+         secret :: binary(),
+         proxy_tag :: binary(),
+
          up_sock :: gen_tcp:socket(),
          up_transport :: transport(),
-         up_codec :: mtp_obfuscated:codec(),
+         up_codec = ident :: mtp_layer:layer(),
+
          down_sock :: gen_tcp:socket(),
-         started :: pos_integer(),
+         down_codec = ident :: mtp_layer:layer(),
+
+         started_at :: pos_integer(),
          timer_state = init :: init | hibernate | stop,
          timer :: gen_timeout:tout()}).
 
@@ -40,11 +69,12 @@
 %% APIs
 
 start_link(Ref, Socket, Transport, Opts) ->
+    metric:count_inc([?APP, in_connection, total], 1, #{}),
     {ok, proc_lib:spawn_link(?MODULE, ranch_init, [{Ref, Socket, Transport, Opts}])}.
 
-key_str() ->
-    {ok, Secret} = application:get_env(?APP, secret),
-    hex(Secret).
+keys_str() ->
+    [{Name, Port, hex(Secret)}
+     || {Name, Port, Secret} <- application:get_env(?APP, ports, [])].
 
 %% Callbacks
 
@@ -55,14 +85,17 @@ ranch_init({Ref, Socket, Transport, _} = Opts) ->
             ok = ranch:accept_ack(Ref),
             ok = Transport:setopts(Socket,
                                    [{active, once},
+                                    %% {recbuf, ?MAX_SOCK_BUF_SIZE},
+                                    %% {sndbuf, ?MAX_SOCK_BUF_SIZE},
                                     {buffer, ?MAX_SOCK_BUF_SIZE}
                                    ]),
             gen_server:enter_loop(?MODULE, [], State);
         error ->
+            metric:count_inc([?APP, in_connection_closed, total], 1, #{}),
             exit(normal)
     end.
 
-init({_Ref, Socket, Transport, _}) ->
+init({_Ref, Socket, Transport, [Secret, Tag]}) ->
     case Transport:peername(Socket) of
         {ok, {Ip, Port}} ->
             lager:info("New connection ~s:~p", [inet:ntoa(Ip), Port]),
@@ -70,8 +103,10 @@ init({_Ref, Socket, Transport, _}) ->
             Timer = gen_timeout:new(
                       #{timeout => {env, ?APP, TimeoutKey, TimeoutDefault}}),
             State = #state{up_sock = Socket,
+                           secret = Secret,
+                           proxy_tag = Tag,
                            up_transport = Transport,
-                           started = erlang:system_time(second),
+                           started_at = erlang:system_time(second),
                            timer = Timer},
             {ok, State};
         {error, Reason} ->
@@ -89,6 +124,7 @@ handle_cast(_Msg, State) ->
 handle_info({tcp, Sock, Data}, #state{up_sock = Sock,
                                       up_transport = Transport} = S) ->
     %% client -> proxy
+    track(rx, Data),
     case handle_upstream_data(Data, S) of
         {ok, S1} ->
             ok = Transport:setopts(Sock, [{active, once}]),
@@ -99,30 +135,36 @@ handle_info({tcp, Sock, Data}, #state{up_sock = Sock,
     end;
 handle_info({tcp_closed, Sock}, #state{up_sock = Sock} = S) ->
     lager:debug("upstream sock closed"),
-    {stop, normal, maybe_close_out(S)};
+    {stop, normal, maybe_close_down(S)};
 handle_info({tcp_error, Sock, Reason}, #state{up_sock = Sock} = S) ->
     lager:info("upstream sock error: ~p", [Reason]),
-    {stop, Reason, maybe_close_out(S)};
+    {stop, Reason, maybe_close_down(S)};
 
 handle_info({tcp, Sock, Data}, #state{down_sock = Sock} = S) ->
     %% telegram server -> proxy
-    case handle_downstream_data(Data, S) of
+    track(tx, Data),
+    try handle_downstream_data(Data, S) of
         {ok, S1} ->
             ok = inet:setopts(Sock, [{active, once}]),
             {noreply, bump_timer(S1)};
         {error, Reason} ->
             lager:error("Error sending tunnelled data to in socket: ~p", [Reason]),
             {stop, normal, S}
+    catch throw:rpc_close ->
+            lager:info("downstream closed by RPC"),
+            #state{up_sock = USock, up_transport = UTrans} = S,
+            ok = UTrans:close(USock),
+            {stop, normal, maybe_close_down(S)}
     end;
 handle_info({tcp_closed, Sock}, #state{down_sock = Sock,
-                                       up_sock = ISock, up_transport = ITrans} = S) ->
+                                       up_sock = USock, up_transport = UTrans} = S) ->
     lager:debug("downstream sock closed"),
-    ok = ITrans:close(ISock),
+    ok = UTrans:close(USock),
     {stop, normal, S};
 handle_info({tcp_error, Sock, Reason}, #state{down_sock = Sock,
-                                              up_sock = ISock, up_transport = ITrans} = S) ->
+                                              up_sock = USock, up_transport = UTrans} = S) ->
     lager:info("downstream sock error: ~p", [Reason]),
-    ok = ITrans:close(ISock),
+    ok = UTrans:close(USock),
     {stop, Reason, S};
 
 
@@ -130,9 +172,11 @@ handle_info(timeout, #state{timer = Timer, timer_state = TState} = S) ->
     case gen_timeout:is_expired(Timer) of
         true when TState == stop;
                   TState == init ->
+            metric:count_inc([?APP, inactive_timeout, total], 1, #{}),
             lager:info("inactive timeout in state ~p", [TState]),
             {stop, normal, S};
         true when TState == hibernate ->
+            metric:count_inc([?APP, inactive_hibernate, total], 1, #{}),
             {noreply, switch_timer(S, stop), hibernate};
         false ->
             Timer1 = gen_timeout:reset(Timer),
@@ -143,14 +187,15 @@ handle_info(Other, S) ->
     {noreply, S}.
 
 terminate(_Reason, #state{}) ->
+    metric:count_inc([?APP, in_connection_closed, total], 1, #{}),
     lager:debug("terminate ~p", [_Reason]),
     ok.
 
 code_change(_OldVsn, State, _Extra) ->
     {ok, State}.
 
-maybe_close_out(#state{down_sock = undefined} = S) -> S;
-maybe_close_out(#state{down_sock = Out} = S) ->
+maybe_close_down(#state{down_sock = undefined} = S) -> S;
+maybe_close_down(#state{down_sock = Out} = S) ->
     gen_tcp:close(Out),
     S#state{down_sock = undefined}.
 
@@ -165,7 +210,9 @@ bump_timer(#state{timer = Timer, timer_state = TState} = S) ->
 
 switch_timer(#state{timer_state = TState} = S, TState) ->
     S;
-switch_timer(#state{timer_state = _FromState, timer = Timer} = S, ToState) ->
+switch_timer(#state{timer_state = FromState, timer = Timer} = S, ToState) ->
+    metric:count_inc([?APP, timer_switch, total], 1,
+                     #{labels => [FromState, ToState]}),
     {NewTimeKey, NewTimeDefault} = state_timeout(ToState),
     Timer1 = gen_timeout:set_timeout(
                {env, ?APP, NewTimeKey, NewTimeDefault}, Timer),
@@ -183,58 +230,129 @@ state_timeout(stop) ->
 %% Stream handlers
 
 %% Handle telegram client -> proxy stream
-handle_upstream_data(<<Header:64/binary, Rest/binary>>, #state{stage = init, init_buf = <<>>} = S) ->
-    {ok, Secret} = application:get_env(?APP, secret),
+handle_upstream_data(<<Header:64/binary, Rest/binary>>, #state{stage = init, stage_state = <<>>,
+                                                               secret = Secret} = S) ->
     case mtp_obfuscated:from_header(Header, Secret) of
-        {ok, Endpoint, Codec} ->
-            case handle_upstream_header(Endpoint, Codec, S) of
-                {ok, S1} ->
-                    handle_upstream_data(Rest, S1);
+        {ok, DcId, ObfuscatedCodec} ->
+            ObfuscatedLayer = mtp_layer:new(mtp_obfuscated, ObfuscatedCodec),
+            AbridgedLayer = mtp_layer:new(mtp_abridged, mtp_abridged:new()),
+            UpCodec = mtp_layer:new(mtp_wrap, mtp_wrap:new(AbridgedLayer,
+                                                           ObfuscatedLayer)),
+            handle_upstream_header(
+              DcId,
+              S#state{up_codec = UpCodec,
+                      up_acc = Rest,
+                      stage_state = undefined});
         Err ->
             Err
     end;
-        Err ->
-            Err
-    end;
-handle_upstream_data(Bin, #state{stage = init, init_buf = <<>>} = S) ->
-    {ok, S#state{init_buf = Bin}};
-handle_upstream_data(Bin, #state{stage = init, init_buf = Buf} = S) ->
-    handle_upstream_data(<<Buf/binary, Bin/binary>> , S#state{init_buf = <<>>});
+handle_upstream_data(Bin, #state{stage = init, stage_state = <<>>} = S) ->
+    {ok, S#state{stage_state = Bin}};
+handle_upstream_data(Bin, #state{stage = init, stage_state = Buf} = S) ->
+    handle_upstream_data(<<Buf/binary, Bin/binary>> , S#state{stage_state = <<>>});
 handle_upstream_data(Bin, #state{stage = tunnel,
-                                 up_codec = UpCodec,
-                                 down_sock = Sock} = S) ->
-    {Decoded, UpCodec1} = mtp_obfuscated:decrypt(Bin, UpCodec),
-    ok = gen_tcp:send(Sock, Decoded),
-    {ok, S#state{up_codec = UpCodec1}}.
+                                 up_codec = UpCodec} = S) ->
+    {ok, S3, UpCodec1} =
+        mtp_layer:fold_packets(
+          fun(Decoded, S1) ->
+                  metric:histogram_observe(
+                    [?APP, tg_packet_size, bytes],
+                    byte_size(Decoded),
+                    #{labels => [upstream_to_downstream]}),
+                  {ok, S2} = down_send(Decoded, S1),
+                  S2
+          end, S, Bin, UpCodec),
+    {ok, S3#state{up_codec = UpCodec1}};
+handle_upstream_data(Bin, #state{stage = Stage, up_acc = Acc} = S) when Stage =/= init,
+                                                                        Stage =/= tunnel ->
+    %% We are in downstream handshake; it would be better to leave socked in passive mode,
+    %% but let's do it in next iteration
+    ((byte_size(Bin) + byte_size(Acc)) < ?MAX_UP_INIT_BUF_SIZE)
+        orelse error(upstream_buffer_overflow),
+    {ok, S#state{up_acc = <<Acc/binary, Bin/binary>>}}.
 
 
 %% Handle telegram server -> proxy stream
+handle_downstream_data(Bin, #state{stage = down_handshake_1,
+                                   down_codec = DownCodec} = S) ->
+    case mtp_layer:try_decode_packet(Bin, DownCodec) of
+        {ok, Packet, DownCodec1} ->
+            down_handshake2(Packet, S#state{down_codec = DownCodec1});
+        {incomplete, DownCodec1} ->
+            {ok, S#state{down_codec = DownCodec1}}
+    end;
+handle_downstream_data(Bin, #state{stage = down_handshake_2,
+                                   proxy_tag = ProxyTag,
+                                   down_codec = DownCodec} = S) ->
+    case mtp_layer:try_decode_packet(Bin, DownCodec) of
+        {ok, Packet, DownCodec1} ->
+            %% TODO: There might be something in downstream buffers after stage3,
+            %% would be nice to run foldl
+            {ok, S1} = down_handshake3(Packet, ProxyTag, S#state{down_codec = DownCodec1}),
+            S2 = #state{up_acc = UpAcc} =  switch_timer(S1, hibernate),
+            %% Flush upstream accumulator
+            handle_upstream_data(UpAcc, S2#state{up_acc = []});
+        {incomplete, DownCodec1} ->
+            {ok, S#state{down_codec = DownCodec1}}
+    end;
 handle_downstream_data(Bin, #state{stage = tunnel,
+                                   down_codec = DownCodec} = S) ->
+    {ok, S3, DownCodec1} =
+        mtp_layer:fold_packets(
+          fun(Decoded, S1) ->
+                  metric:histogram_observe(
+                    [?APP, tg_packet_size, bytes],
+                    byte_size(Decoded),
+                    #{labels => [downstream_to_upstream]}),
+                  {ok, S2} = up_send(Decoded, S1),
+                  S2
+          end, S, Bin, DownCodec),
+    {ok, S3#state{down_codec = DownCodec1}}.
+
+
+up_send(Packet, #state{stage = tunnel,
                        up_codec = UpCodec,
                        up_sock = Sock,
                        up_transport = Transport} = S) ->
-    {Encoded, UpCodec1} = mtp_obfuscated:encrypt(Bin, UpCodec),
-    ok = Transport:send(Sock, Encoded),
+    {Encoded, UpCodec1} = mtp_layer:encode_packet(Packet, UpCodec),
+    metric:rt([?APP, upstream_send_duration, seconds],
+              fun() ->
+                      ok = Transport:send(Sock, Encoded)
+              end),
     {ok, S#state{up_codec = UpCodec1}}.
 
-
-%% Packet handlers
+down_send(Packet, #state{down_sock = Sock,
+                         down_codec = DownCodec} = S) ->
+    {Encoded, DownCodec1} = mtp_layer:encode_packet(Packet, DownCodec),
+    metric:rt([?APP, downstream_send_duration, seconds],
+              fun() ->
+                      ok = gen_tcp:send(Sock, Encoded)
+              end),
+    {ok, S#state{down_codec = DownCodec1}}.
 
 
 %% Internal
 
 
-handle_upstream_header(Endpoint, UpCodec, S) ->
-    case connect(Endpoint, 443) of
+handle_upstream_header(DcId, S) ->
+    {Addr, Port} =
+        try element(DcId, ?TG_MIDDLE_PROXIES_V4)
+        catch error:badarg ->
+                OtherDcId = (DcId rem tuple_size(?TG_MIDDLE_PROXIES_V4)) + 1,
+                lager:warning("Wrong DC id: ~p; will use ~p",
+                              [DcId, OtherDcId]),
+                element(OtherDcId, ?TG_MIDDLE_PROXIES_V4)
+        end,
+
+    case connect(Addr, Port) of
         {ok, Sock} ->
-            EndpointStr = inet:ntoa(Endpoint),
-            lager:info("Connected to ~s:~p", [EndpointStr, 443]),
-            ok = gen_tcp:send(Sock, <<239>>),
-            {ok, switch_timer(S#state{stage = tunnel,
-                                      down_sock = Sock,
-                                      up_codec = UpCodec},
-                              hibernate)};
-        {error, _Reason} = Err ->
+            AddrStr = inet:ntoa(Addr),
+            metric:count_inc([?APP, out_connect_ok, total], 1,
+                             #{labels => [AddrStr]}),
+            lager:info("Connected to ~s:~p", [AddrStr, Port]),
+            down_handshake1(S#state{down_sock = Sock});
+        {error, Reason} = Err ->
+            metric:count_inc([?APP, out_connect_error, total], 1, #{labels => [Reason]}),
             Err
     end.
 
@@ -248,22 +366,149 @@ connect(Host, Port) ->
                 {send_timeout, ?SEND_TIMEOUT},
                 %% {nodelay, true},
                 {keepalive, true}],
-    case gen_tcp:connect(Host, Port, SockOpts, ?CONN_TIMEOUT) of
+    case metric:rt([?APP, downstream_connect_duration, seconds],
+                   fun() ->
+                           gen_tcp:connect(Host, Port, SockOpts, ?CONN_TIMEOUT)
+                   end) of
         {ok, Sock} ->
-            ok = inet:setopts(Sock, [{buffer, ?MAX_SOCK_BUF_SIZE}]),
+            ok = inet:setopts(Sock, [%% {recbuf, ?MAX_SOCK_BUF_SIZE},
+                                     %% {sndbuf, ?MAX_SOCK_BUF_SIZE},
+                                     {buffer, ?MAX_SOCK_BUF_SIZE}]),
             {ok, Sock};
         {error, _} = Err ->
             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>> = ?PROXY_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>>,
+    Full = mtp_full:new(-2, -2),
+    S1 = S#state{down_codec = mtp_layer:new(mtp_full, Full),
+                 stage = down_handshake_1,
+                 stage_state = {KeySelector, Nonce, CryptoTs}},
+    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},
+                                               down_sock = Sock,
+                                               down_codec = DownCodec} = S) ->
+    (Type == ?RPC_NONCE) orelse error({wrong_rpc_type, Type}),
+    (Schema == 1) orelse error({wrong_schema, Schema}),
+    (KeySelector == MyKeySelector) orelse error({wrong_key_selector, KeySelector}),
+    {ok, {DownIp, DownPort}} = inet:peername(Sock),
+    {MyIp, MyPort} = get_external_ip(Sock),
+    DownIpBin = mtp_obfuscated:bin_rev(mtp_rpc:inet_pton(DownIp)),
+    MyIpBin = mtp_obfuscated:bin_rev(mtp_rpc:inet_pton(MyIp)),
+    Args = #{srv_n => SrvNonce, clt_n => CliNonce, clt_ts => MyTs,
+             srv_ip => DownIpBin, srv_port => DownPort,
+             clt_ip => MyIpBin, clt_port => MyPort, secret => ?PROXY_SECRET},
+    {EncKey, EncIv} = get_middle_key(Args#{purpose => <<"CLIENT">>}),
+    {DecKey, DecIv} = get_middle_key(Args#{purpose => <<"SERVER">>}),
+    CryptoCodec = mtp_layer:new(mtp_aes_cbc, mtp_aes_cbc:new(EncKey, EncIv, DecKey, DecIv, 16)),
+    DownCodec1 = mtp_layer:new(mtp_wrap, mtp_wrap:new(DownCodec, CryptoCodec)),
+    SenderPID = PeerPID = <<"IPIPPRPDTIME">>,
+    Handshake = [?RPC_HANDSHAKE,
+                 ?RPC_FLAGS,
+                 SenderPID,
+                 PeerPID],
+    down_send(Handshake, S#state{down_codec = DownCodec1,
+                                 stage = down_handshake_2,
+                                 stage_state = {MyIp, MyPort, SenderPID}}).
+
+get_middle_key(#{srv_n := Nonce, clt_n := MyNonce, clt_ts := MyTs, srv_ip := SrvIpBinBig, srv_port := SrvPort,
+                 clt_ip := CltIpBinBig, clt_port := CltPort, secret := Secret, purpose := Purpose} = _Args) ->
+    Msg =
+        <<Nonce/binary,
+          MyNonce/binary,
+          MyTs:32/little,
+          SrvIpBinBig/binary,
+          CltPort:16/little,
+          Purpose/binary,
+          CltIpBinBig/binary,
+          SrvPort:16/little,
+          Secret/binary,
+          Nonce/binary,
+          %% IPv6
+          MyNonce/binary
+        >>,
+    <<_, ForMd51/binary>> = Msg,
+    <<_, _, ForMd52/binary>> = Msg,
+    <<Key1:12/binary, _/binary>> = crypto:hash(md5, ForMd51),
+    ShaSum = crypto:hash(sha, Msg),
+    Key = <<Key1/binary, ShaSum/binary>>,
+    IV = crypto:hash(md5, ForMd52),
+    {Key, IV}.
+
+
+down_handshake3(<<Type:4/binary, _Flags:4/binary, _SenderPid:12/binary, PeerPid:12/binary>>,
+                ProxyTag,
+                #state{stage_state = {MyIp, MyPort, PrevSenderPid},
+                       down_codec = DownCodec,
+                       up_sock = Sock,
+                       up_transport = Transport} = S) ->
+    (Type == ?RPC_HANDSHAKE) orelse error({wrong_rpc_type, Type}),
+    (PeerPid == PrevSenderPid) orelse error({wrong_sender_pid, PeerPid}),
+    {ok, {ClientIp, ClientPort}} = Transport:peername(Sock),
+    RpcCodec = mtp_layer:new(mtp_rpc, mtp_rpc:new(ClientIp, ClientPort, MyIp, MyPort, ProxyTag)),
+    DownCodec1 = mtp_layer:new(mtp_wrap, mtp_wrap:new(RpcCodec, DownCodec)),
+    {ok, S#state{down_codec = DownCodec1,
+                 stage = tunnel,
+                 stage_state = undefined}}.
 %% Internal
 
+get_external_ip(Sock) ->
+    {ok, {MyIp, MyPort}} = inet:sockname(Sock),
+    case application:get_env(?APP, external_ip) of
+        {ok, IpStr} ->
+            {ok, IP} = inet:parse_ipv4strict_address(IpStr),
+            {IP, MyPort};
+        undefined ->
+            {MyIp, MyPort}
+    end.
+
 hex(Bin) ->
-    [begin
+    <<begin
          if N < 10 ->
-                 48 + N;
+                 <<($0 + N)>>;
             true ->
-                 87 + N
+                 <<($W + N)>>
          end
-     end || <<N:4>> <= Bin].
+     end || <<N:4>> <= Bin>>.
+
+track(Direction, Data) ->
+    Size = byte_size(Data),
+    metric:count_inc([?APP, tracker, bytes], Size, #{labels => [Direction]}),
+    metric:histogram_observe([?APP, tracker_packet_size, bytes], Size, #{labels => [Direction]}).
+
+-ifdef(TEST).
+-include_lib("eunit/include/eunit.hrl").
+
+middle_key_test() ->
+    Args = #{srv_port => 80,
+             srv_ip => mtp_obfuscated:bin_rev(mtp_rpc:inet_pton({149, 154, 162, 38})),
+             srv_n => <<247,40,210,56,65,12,101,170,216,155,14,253,250,238,219,226>>,
+             clt_n => <<24,49,53,111,198,10,235,180,230,112,92,78,1,201,106,105>>,
+             clt_ip => mtp_obfuscated:bin_rev(mtp_rpc:inet_pton({80, 211, 29, 34})),
+             clt_ts => 1528396015,
+             clt_port => 54208,
+             purpose => <<"CLIENT">>,
+             secret => ?PROXY_SECRET
+            },
+    Key = <<165,158,127,49,41,232,187,69,38,29,163,226,183,146,28,67,225,224,134,191,207,152,255,166,152,66,169,196,54,135,50,188>>,
+    IV = <<33,110,125,221,183,121,160,116,130,180,156,249,52,111,37,178>>,
+    ?assertEqual(
+       {Key, IV},
+       get_middle_key(Args)).
+
+-endif.

src/mtp_layer.erl

+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% Behaviour for MTProto layer codec
+%%% @end
+%%% Created :  6 Jun 2018 by Sergey <me@seriyps.ru>
+
+-module(mtp_layer).
+
+-export([new/2,
+         try_decode_packet/2,
+         encode_packet/2]).
+
+-export([fold_packets/4]).
+-export_type([codec/0,
+              layer/0]).
+
+-type state() :: any().
+-type codec() :: mtb_aes_cbc
+               | mtp_abridged
+               | mtp_full
+               | mtp_obfuscated
+               | mtp_rpc
+               | mtp_wrap.
+-type layer() :: {codec(), state()} | ident.
+
+-callback try_decode_packet(binary(), state()) ->
+    {ok, binary(), state()}
+        | {incomplete, state()}.
+
+-callback encode_packet(binary(), state()) ->
+    {binary(), state()}.
+
+new(Mod, S) ->
+    {Mod, S}.
+
+encode_packet(Msg, ident) ->
+    {Msg, ident};
+encode_packet(Msg, {Mod, St}) ->
+    {Enc, St1} = Mod:encode_packet(Msg, St),
+    {Enc, {Mod, St1}}.
+
+try_decode_packet(Msg, ident) ->
+    {ok, Msg, ident};
+try_decode_packet(Msg, {Mod, St}) ->
+    case Mod:try_decode_packet(Msg, St) of
+        {ok, Dec, St1} ->
+            {ok, Dec, {Mod, St1}};
+        {incomplete, St1} ->
+            {incomplete, {Mod, St1}}
+    end.
+
+-spec fold_packets(fun( (binary(), FoldSt) -> FoldSt ),
+                   FoldSt, binary(), layer()) ->
+                          {ok, FoldSt, layer()}
+                              when
+      FoldSt :: any().
+fold_packets(Fun, FoldSt, Data, ident) ->
+    FoldSt1 = Fun(Data, FoldSt),
+    {ok, FoldSt1, ident};
+fold_packets(Fun, FoldSt, Data, Layer) ->
+    case try_decode_packet(Data, Layer) of
+        {ok, Decoded, L1} ->
+            FoldSt1 = Fun(Decoded, FoldSt),
+            fold_packets(Fun, FoldSt1, <<>>, L1);
+        {incomplete, L1} ->
+            {ok, FoldSt, L1}
+    end.

src/mtp_obfuscated.erl

@@ -6,11 +6,17 @@
 %%% Created : 29 May 2018 by Sergey <me@seriyps.ru>
 
 -module(mtp_obfuscated).
--export([new/0,
-         new/1,
+-behaviour(mtp_layer).
+-export([create/0,
+         create/1,
          from_header/2,
+         new/4,
          encrypt/2,
-         decrypt/2]).
+         decrypt/2,
+         try_decode_packet/2,
+         encode_packet/2
+        ]).
+-export([bin_rev/1]).
 
 -export_type([codec/0]).
 
@@ -19,36 +25,26 @@
          decrypt :: any()                       % aes state
         }).
 
--define(ENDPOINTS, {
-          {149, 154, 175, 50},
-          {149, 154, 167, 51},
-          {149, 154, 175, 100},
-          {149, 154, 167, 91},
-          {149, 154, 171, 5}
-         }).
 -define(APP, mtproto_proxy).
-%% -define(DBG(Fmt, Args), io:format(user, Fmt, Args)).
--define(DBG(_F, _A), ok).
 
 -opaque codec() :: #st{}.
 
 %% @doc Creates new obfuscated stream (usual format)
--spec new() -> {ok, Header :: binary(), codec()}.
-new() ->
-    new(crypto:strong_rand_bytes(60)).
+-spec create() -> {ok, Header :: binary(), codec()}.
+create() ->
+    create(crypto:strong_rand_bytes(60)).
 
--spec new(binary()) -> {ok, Header :: binary(), codec()}.
-new(<<Left:56/binary, Right:4/binary>>) ->
+-spec create(binary()) -> {ok, Header :: binary(), codec()}.
+create(<<Left:56/binary, Right:4/binary>>) ->
     DownHeader = <<Left/binary,
                    16#ef, 16#ef, 16#ef, 16#ef,
                    Right/binary>>,
     new2(DownHeader).
 
 new2(<<Left:56/binary, _/binary>> = DownHeader) ->
-    Encrypt = init_down_encrypt(DownHeader),
-    Decrypt = init_down_decrypt(DownHeader),
-    St = #st{decrypt = Decrypt,
-             encrypt = Encrypt},
+    {EncKey, EncIV} = init_down_encrypt(DownHeader),
+    {DecKey, DecIV} = init_down_decrypt(DownHeader),
+    St = new(EncKey, EncIV, DecKey, DecIV),
     {<<_:56/binary, Rep:8/binary, _/binary>>, St1} = encrypt(DownHeader, St),
     {ok,
      <<Left/binary, Rep/binary>>,
@@ -57,30 +53,30 @@ new2(<<Left:56/binary, _/binary>> = DownHeader) ->
 init_down_decrypt(<<_:8/binary, ToRev:48/binary, _/binary>>) ->
     Reversed = bin_rev(ToRev),
     <<KeyRev:32/binary, RevIV:16/binary>> = Reversed,
-    ?DBG("down-DEC Key: ~w;~nIV: ~w~n", [KeyRev, RevIV]),
-    crypto:stream_init('aes_ctr', KeyRev, RevIV).
+    {KeyRev, RevIV}.
 
 init_down_encrypt(<<_:8/binary, Key:32/binary, IV:16/binary, _/binary>>) ->
-    ?DBG("down-ENC Key: ~w;~nIV: ~w~n", [Key, IV]),
-    crypto:stream_init('aes_ctr', Key, IV).
+    {Key, IV}.
 
 
 %% @doc creates new obfuscated stream (MTProto proxy format)
 -spec from_header(binary(), binary()) -> {ok, inet:ip4_address(), codec()}.
 from_header(Header, Secret) when byte_size(Header) == 64  ->
-    Encrypt = init_up_encrypt(Header, Secret),
-    Decrypt = init_up_decrypt(Header, Secret),
-    {Decrypt1, <<_:56/binary, Bin1:8/binary, _/binary>>} = crypto:stream_encrypt(Decrypt, Header),
+    {EncKey, EncIV} = init_up_encrypt(Header, Secret),
+    {DecKey, DecIV} = init_up_decrypt(Header, Secret),
+    St = new(EncKey, EncIV, DecKey, DecIV),
+    {<<_:56/binary, Bin1:8/binary, _/binary>>, St1} = decrypt(Header, St),
     <<HeaderPart:56/binary, _/binary>> = Header,
     NewHeader = <<HeaderPart/binary, Bin1/binary>>,
     case NewHeader of
         <<_:56/binary, 16#ef, 16#ef, 16#ef, 16#ef, _/binary>> ->
-            Endpoint = get_endpoint(NewHeader),
-            {ok, Endpoint, #st{decrypt = Decrypt1,
-                               encrypt = Encrypt}};
+            DcId = get_dc(NewHeader),
+            {ok, DcId, St1};
         <<_:56/binary, 16#ee, 16#ee, 16#ee, 16#ee, _/binary>> ->
+            metric:count_inc([?APP, protocol_error, total], 1, #{labels => [intermediate]}),
             {error, {protocol_not_supported, intermediate}};
         _ ->
+            metric:count_inc([?APP, protocol_error, total], 1, #{labels => [unknown]}),
             {error, unknown_protocol}
     end.
 
@@ -90,30 +86,44 @@ init_up_encrypt(Bin, Secret) ->
     <<KeyRev:32/binary, RevIV:16/binary, _/binary>> = Rev,
     %% <<_:32/binary, RevIV:16/binary, _/binary>> = Bin,
     KeyRevHash = crypto:hash('sha256', <<KeyRev/binary, Secret/binary>>),
-    ?DBG("up-ENC Key: ~p;~nIV: ~p~n", [KeyRevHash, RevIV]),
-    crypto:stream_init('aes_ctr', KeyRevHash, RevIV).
+    {KeyRevHash, RevIV}.
 
 init_up_decrypt(Bin, Secret) ->
-    <<_:8/binary, Key:32/binary, _/binary>> = Bin,
-    <<_:40/binary, IV:16/binary, _/binary>> = Bin,
+    <<_:8/binary, Key:32/binary, IV:16/binary, _/binary>> = Bin,
     KeyHash = crypto:hash('sha256', <<Key/binary, Secret/binary>>),
-    ?DBG("up-DEC Key: ~p;~nIV: ~p~n", [KeyHash, IV]),
-    crypto:stream_init('aes_ctr', KeyHash, IV).
+    {KeyHash, IV}.
+
+get_dc(<<_:60/binary, DcId:16/signed-little-integer, _/binary>>) ->
+    abs(DcId).
+
 
-get_endpoint(<<_:60/binary, DcId:16/signed-little-integer, _/binary>>) ->
-    element(abs(DcId), ?ENDPOINTS).
+new(EncKey, EncIV, DecKey, DecIV) ->
+    #st{decrypt = crypto:stream_init('aes_ctr', DecKey, DecIV),
+        encrypt = crypto:stream_init('aes_ctr', EncKey, EncIV)}.
 
+-spec encrypt(iodata(), codec()) -> {binary(), codec()}.
 encrypt(Data, #st{encrypt = Enc} = St) ->
     {Enc1, Encrypted} = crypto:stream_encrypt(Enc, Data),
-    ?DBG("encrypt: IN:~p~nOUT:~p~n", [Data, Encrypted]),
     {Encrypted, St#st{encrypt = Enc1}}.
 
+-spec decrypt(iodata(), codec()) -> {binary(), codec()}.
 decrypt(Encrypted, #st{decrypt = Dec} = St) ->
     {Dec1, Data} = crypto:stream_encrypt(Dec, Encrypted),
-    ?DBG("decrypt: IN:~p~nOUT:~p~n", [Encrypted, Data]),
     {Data, St#st{decrypt = Dec1}}.
 
+%% To comply with mtp_layer interface
+-spec try_decode_packet(iodata(), codec()) -> {ok, Decoded :: binary(), codec()}
+                                                  | {incomplete, codec()}.
+try_decode_packet(Encrypted, St) ->
+    {Decrypted, St1} = decrypt(Encrypted, St),
+    {ok, Decrypted, St1}.
+
+-spec encode_packet(iodata(), codec()) -> {iodata(), codec()}.
+encode_packet(Msg, S) ->
+    encrypt(Msg, S).
+
 
 %% Helpers
 bin_rev(Bin) ->
+    %% binary:encode_unsigned(binary:decode_unsigned(Bin, little)).
     list_to_binary(lists:reverse(binary_to_list(Bin))).

src/mtp_rpc.erl

+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% MTProto RPC codec
+%%% @end
+%%% Created :  6 Jun 2018 by Sergey <me@seriyps.ru>
+
+-module(mtp_rpc).
+-behaviour(mtp_layer).
+
+-export([new/5,
+         try_decode_packet/2,
+         encode_packet/2]).
+-export([inet_pton/1]).
+-export_type([codec/0]).
+
+-record(rpc_st,
+        {client_addr :: binary(),
+         proxy_addr :: binary(),
+         proxy_tag :: binary(),
+         req_id = 1 :: non_neg_integer()}).
+
+-define(APP, mtproto_proxy).
+-define(RPC_PROXY_ANS, 13,218,3,68).
+-define(RPC_CLOSE_EXT, 162,52,182,94).
+
+-opaque codec() :: #rpc_st{}.
+
+new(ClientIp, ClientPort, ProxyIp, ProxyPort, ProxyTag) ->
+    #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}.
+
+%% It expects that packet segmentation was done on previous layer
+try_decode_packet(<<?RPC_PROXY_ANS, _AnsFlags:4/binary, _ReqId:8/binary, Data/binary>> = _Msg, S) ->
+    {ok, Data, S};
+try_decode_packet(<<?RPC_CLOSE_EXT, _/binary>> = _Msg, _S) ->
+    %% Use throw as short-circuit
+    throw(rpc_close);
+try_decode_packet(<<>>, S) ->
+    {incomplete, S}.
+
+
+encode_packet(Msg, #rpc_st{client_addr = ClientAddr, proxy_addr = ProxyAddr,
+                           req_id = ReqId, proxy_tag = ProxyTag} = S) ->
+    ((iolist_size(Msg) rem 4) == 0)
+        orelse error(not_aligned),
+    Req =
+        [<<238,241,206,54,                          %RPC_PROXY_REQ
+           8,16,2,64,                               %Flags
+           ReqId:64/little                          %ReqId
+         >>,
+         ClientAddr, ProxyAddr,
+         <<24:32/little,                            %ExtraSize
+           174,38,30,219,                           %ProxyTag
+           (byte_size(ProxyTag)),
+           ProxyTag/binary,
+           0, 0, 0                                  %Padding
+         >>
+             | Msg
+        ],
+    {Req, S#rpc_st{req_id = ReqId + 1}}.
+
+encode_ip_port(IPv4, Port) when tuple_size(IPv4) == 4 ->
+    IpBin = inet_pton(IPv4),
+    [lists:duplicate(10, <<0>>)
+     | <<255,255,
+         IpBin/binary,
+         Port:32/little>>];
+encode_ip_port(IPv6, Port) when tuple_size(IPv6) == 8 ->
+    IpBin = inet_pton(IPv6),
+    [IpBin, <<Port:32/little>>].
+
+inet_pton({X1, X2, X3, X4}) ->
+    <<X1, X2, X3, X4>>;
+inet_pton(IPv6) when tuple_size(IPv6) == 8 ->
+    << <<I:16/big-integer>> || I <- tuple_to_list(IPv6)>>.
+    
+-ifdef(TEST).
+-include_lib("eunit/include/eunit.hrl").
+
+tst_new() ->
+    ClientIp = {109, 238, 131, 159},
+    ClientPort = 1128,
+    ProxyIp = {80, 211, 29, 34},
+    ProxyPort = 53634,
+    ProxyTag = <<220,190,143,20,147,250,76,217,171,48,8,145,192,181,179,38>>,
+    new(ClientIp, ClientPort, ProxyIp, ProxyPort, ProxyTag).
+
+decode_none_test() ->
+    S = tst_new(),
+    ?assertEqual(
+       {incomplete, S}, try_decode_packet(<<>>, S)).
+
+encode_test() ->
+    S = tst_new(),
+    Samples =
+        [{<<0,0,0,0,0,0,0,0,0,0,0,0,61,2,24,91,20,0,0,0,120,151,70,96,153,197,142,238,245,139,85,208,160,241,68,89,106,7,118,167>>,
+          <<238,241,206,54,8,16,2,64,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,255,255,109,238,131,159,104,4,0,0,0,0,0,0,0,0,0,0,0,0,255,255,80,211,29,34,130,209,0,0,24,0,0,0,174,38,30,219,16,220,190,143,20,147,250,76,217,171,48,8,145,192,181,179,38,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,61,2,24,91,20,0,0,0,120,151,70,96,153,197,142,238,245,139,85,208,160,241,68,89,106,7,118,167>>},
+         {<<14,146,6,159,99,150,29,221,115,87,68,198,122,39,38,249,153,87,37,105,4,111,147,70,54,179,134,12,90,4,223,155,206,220,167,201,203,176,123,181,103,176,49,216,163,106,54,148,133,51,206,212,81,90,47,26,3,161,149,251,182,90,190,51,213,7,107,176,112,220,25,144,183,249,149,182,172,194,218,146,161,191,247,4,250,123,230,251,41,181,139,177,55,171,253,198,153,183,61,53,119,115,46,174,172,245,90,166,215,99,181,58,236,129,103,80,218,244,81,45,142,128,177,146,26,131,184,155,22,217,218,187,209,155,156,64,219,235,175,40,249,235,77,82,212,73,11,133,52,4,222,157,67,176,251,46,254,241,15,192,215,192,186,82,233,68,147,234,88,250,96,14,172,179,7,159,28,11,237,48,44,33,137,185,166,166,173,103,136,174,31,35,77,151,76,55,176,211,230,176,118,144,139,77,0,213,68,179,73,58,58,80,238,120,197,67,241,210,210,156,72,105,60,125,239,98,7,19,234,249,222,194,166,37,46,100,1,65,225,224,244,57,147,119,49,20,1,160,4,51,247,161,142,11,131,11,27,166,159,110,145,78,55,205,126,246,126,68,44,114,91,191,213,241,242,9,33,16,30,228>>,
+          <<238,241,206,54,8,16,2,64,2,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,255,255,109,238,131,159,104,4,0,0,0,0,0,0,0,0,0,0,0,0,255,255,80,211,29,34,130,209,0,0,24,0,0,0,174,38,30,219,16,220,190,143,20,147,250,76,217,171,48,8,145,192,181,179,38,0,0,0,14,146,6,159,99,150,29,221,115,87,68,198,122,39,38,249,153,87,37,105,4,111,147,70,54,179,134,12,90,4,223,155,206,220,167,201,203,176,123,181,103,176,49,216,163,106,54,148,133,51,206,212,81,90,47,26,3,161,149,251,182,90,190,51,213,7,107,176,112,220,25,144,183,249,149,182,172,194,218,146,161,191,247,4,250,123,230,251,41,181,139,177,55,171,253,198,153,183,61,53,119,115,46,174,172,245,90,166,215,99,181,58,236,129,103,80,218,244,81,45,142,128,177,146,26,131,184,155,22,217,218,187,209,155,156,64,219,235,175,40,249,235,77,82,212,73,11,133,52,4,222,157,67,176,251,46,254,241,15,192,215,192,186,82,233,68,147,234,88,250,96,14,172,179,7,159,28,11,237,48,44,33,137,185,166,166,173,103,136,174,31,35,77,151,76,55,176,211,230,176,118,144,139,77,0,213,68,179,73,58,58,80,238,120,197,67,241,210,210,156,72,105,60,125,239,98,7,19,234,249,222,194,166,37,46,100,1,65,225,224,244,57,147,119,49,20,1,160,4,51,247,161,142,11,131,11,27,166,159,110,145,78,55,205,126,246,126,68,44,114,91,191,213,241,242,9,33,16,30,228>>}],
+    lists:foldl(
+      fun({In, Out}, S1) ->
+              {Enc, S2} = encode_packet(In, S1),
+              ?assertEqual(Out, iolist_to_binary(Enc)),
+              S2
+      end, S, Samples).
+
+decode_test() ->
+    S = tst_new(),
+    Samples =
+        [{<<13,218,3,68,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,52,62,238,60,2,24,91,64,0,0,0,99,36,22,5,153,197,142,238,245,139,85,208,160,241,68,89,106,7,118,167,146,202,163,241,63,158,32,27,246,203,226,70,177,46,106,225,8,34,202,206,241,19,38,121,245,0,0,0,21,196,181,28,1,0,0,0,33,107,232,108,2,43,180,195>>,
+          <<0,0,0,0,0,0,0,0,1,52,62,238,60,2,24,91,64,0,0,0,99,36,22,5,153,197,142,238,245,139,85,208,160,241,68,89,106,7,118,167,146,202,163,241,63,158,32,27,246,203,226,70,177,46,106,225,8,34,202,206,241,19,38,121,245,0,0,0,21,196,181,28,1,0,0,0,33,107,232,108,2,43,180,195>>},
+         {<<13,218,3,68,0,0,0,0,2,0,0,0,0,0,0,0,14,146,6,159,99,150,29,221,85,233,237,52,236,18,11,0,174,214,89,213,69,89,250,18,116,192,128,240,217,221,210,144,123,9,182,152,60,206,88,187,101,178,53,107,44,98,190,195,149,114,0,19,90,218,101,133,183,249,183,170,90,21,86,24,42,81,224,152,13,58,90,84,41,158,177,99,57,83,123,99,138,127,29,238,162,49,71,65,165,168,218,220,245,202,24,135,152,1,28,38,85,197,8,232,201,163,65,118,202,89,204,67,48,21,51,106,188,7,167,61,185,82,39,210,164,21,97,99,63,167,2,143,69,126,214,75,95,142,69,68,243,49,11,121,28,177,159,0,154,134,206,34>>,
+          <<14,146,6,159,99,150,29,221,85,233,237,52,236,18,11,0,174,214,89,213,69,89,250,18,116,192,128,240,217,221,210,144,123,9,182,152,60,206,88,187,101,178,53,107,44,98,190,195,149,114,0,19,90,218,101,133,183,249,183,170,90,21,86,24,42,81,224,152,13,58,90,84,41,158,177,99,57,83,123,99,138,127,29,238,162,49,71,65,165,168,218,220,245,202,24,135,152,1,28,38,85,197,8,232,201,163,65,118,202,89,204,67,48,21,51,106,188,7,167,61,185,82,39,210,164,21,97,99,63,167,2,143,69,126,214,75,95,142,69,68,243,49,11,121,28,177,159,0,154,134,206,34>>}],
+    lists:foldl(
+      fun({In, Out}, S1) ->
+              {ok, Dec, S2} = try_decode_packet(In, S1),
+              ?assertEqual(Out, iolist_to_binary(Dec)),
+              S2
+      end, S, Samples).
+
+%% decode_close_test() ->
+%%     S = tst_new(),
+%%     In = <<>>,
+%%     ?assertError(rpc_close, try_decode_packet(In, S)).
+
+-endif.

src/mtp_wrap.erl

+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% Abstraction that allows to wrap one mtp_layer into another mtp_layer
+%%% @end
+%%% Created :  6 Jun 2018 by Sergey <me@seriyps.ru>
+
+-module(mtp_wrap).
+-behaviour(mtp_layer).
+
+-export([new/2,
+         try_decode_packet/2,
+         encode_packet/2]).
+-export_type([codec/0]).
+
+-record(wrap_st,
+        {outer :: mtp_layer:layer(),
+         inner :: mtp_layer:layer()}).
+
+-define(APP, mtproto_proxy).
+
+-opaque codec() :: #wrap_st{}.
+
+new(Outer, Inner) ->
+    #wrap_st{outer = Outer,
+             inner = Inner}.
+
+%% try_decode_packet(Inner) |> try_decode_packet(Outer)
+try_decode_packet(Bin, #wrap_st{outer = Outer,
+                                inner = Inner} = S) ->
+    {Dec1, Inner1} =
+        case mtp_layer:try_decode_packet(Bin, Inner) of
+            {incomplete, Inner1_} ->
+                %% We have to check if something is left in inner's buffers
+                {<<>>, Inner1_};
+            {ok, Dec1_, Inner1_} ->
+                {Dec1_, Inner1_}
+        end,
+    case mtp_layer:try_decode_packet(Dec1, Outer) of
+        {incomplete, Outer1} ->
+            {incomplete, S#wrap_st{inner = Inner1,
+                                   outer = Outer1}};
+        {ok, Dec2, Outer1} ->
+            {ok, Dec2, S#wrap_st{inner = Inner1,
+                                 outer = Outer1}}
+    end.
+
+%% encode_packet(Outer) |> encode_packet(Inner)
+encode_packet(Bin, #wrap_st{outer = Outer,
+                            inner = Inner} = S) ->
+    {Enc1, Outer1} = mtp_layer:encode_packet(Bin, Outer),
+    {Enc2, Inner1} = mtp_layer:encode_packet(Enc1, Inner),
+    {Enc2, S#wrap_st{outer = Outer1, inner = Inner1}}.

src/mtproto_proxy.app.src

@@ -16,10 +16,14 @@
         {ready_timeout_sec, 1200},
 
         {ip, {0, 0, 0, 0}},
-        {ports, [{mtp_handler, 4430}]},
+        %% {external_ip, "YOUR.SERVER.EXTERNAL.IP"},
+        {ports, [#{name => mtp_handler,
+                   port => 1443,
+                   secret => <<208,214,225,17,186,218,85,17,252,206,149,132,222,173,190,239>>,
+                   tag => <<220,190,143,20,147,250,76,217,171,48,8,145,192,181,179,38>>}]},
         {num_acceptors, 60},
-        {max_connections, 1024},
-        {secret, <<15,61,100,92,77,206,249,22,29,186,62,205,196,61,65,194>>}]},
+        {max_connections, 20480}
+       ]},
   {modules, []},
 
   {maintainers, []},

src/mtproto_proxy_app.erl

@@ -16,36 +16,35 @@
 %%====================================================================
 start(_StartType, _StartArgs) ->
     Res = mtproto_proxy_sup:start_link(),
+    io:format("+++++++++++++++++++++++++++++++++++++++~n"
+              "Erlang MTProto proxy by @seriyps https://github.com/seriyps/mtproto_proxy~n"
+              "Sponsored by and powers @socksy_bot~n"),
     [start_proxy(Where) || Where <- application:get_env(?APP, ports, [])],
     Res.
 
 %%--------------------------------------------------------------------
 stop(_State) ->
+    [stop_proxy(Where) || Where <- application:get_env(?APP, ports, [])],
     ok.
 
 %%====================================================================
 %% Internal functions
 %%====================================================================
-
-start_proxy({Name, Port}) ->
+start_proxy(#{name := Name, port := Port, secret := Secret, tag := Tag}) ->
     ListenIp = application:get_env(?APP, ip, {0, 0, 0, 0}),
     NumAcceptors = application:get_env(?APP, num_acceptors, 60),
-    MaxConnections = application:get_env(?APP, max_connections, 1024),
-    Res = {ok, Pid} =
+    MaxConnections = application:get_env(?APP, max_connections, 10240),
+    Res =
         ranch:start_listener(
           Name, ranch_tcp,
           [{ip, ListenIp},
            {port, Port},
            {num_acceptors, NumAcceptors},
            {max_connections, MaxConnections}],
-          mtp_handler, []),
-    KeyStr = mtp_handler:key_str(),
-    io:format("+++++++++++++++++++++++++++++++++++++++~n"
-              "Erlang MTProto proxy by @seriyps https://github.com/seriyps/mtproto_proxy~n"
-              "Sponsored by and powers @socksy_bot~n"
-              "Proxy started on ~s:~p with key: ~s~n~n"
-              "+++++++++++++++++++++++++++++++++++++++~n",
-              [inet:ntoa(ListenIp), Port, KeyStr]),
-    lager:info("mtproto=~p listening on addr=~s:~p with key ~s",
-               [Pid, inet:ntoa(ListenIp), Port, KeyStr]),
+          mtp_handler, [Secret, Tag]),
+    io:format("Proxy started on ~s:~p with key: ~s~n",
+              [inet:ntoa(ListenIp), Port, mtp_handler:hex(Secret)]),
     Res.
+
+stop_proxy(#{name := Name}) ->
+    ranch:stop_listener(Name).