From 9403f11e9e73040455abaa2d8d8eec379b550c9d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=A1=D0=B5=D1=80=D0=B3=D0=B5=D0=B9=20=D0=9F=D1=80=D0=BE?=
 =?UTF-8?q?=D1=85=D0=BE=D1=80=D0=BE=D0=B2?= <seriy.pr@gmail.com>
Date: Fri, 29 Jun 2018 17:09:23 +0200
Subject: [PATCH] Add "secure" protocol with padding

---
 src/mtp_handler.erl      |  6 +++++-
 src/mtp_intermediate.erl | 18 ++++++++++++++----
 src/mtp_layer.erl        |  2 ++
 src/mtp_obfuscated.erl   |  3 +++
 src/mtp_secure.erl       | 29 +++++++++++++++++++++++++++++
 5 files changed, 53 insertions(+), 5 deletions(-)
 create mode 100644 src/mtp_secure.erl

diff --git a/src/mtp_handler.erl b/src/mtp_handler.erl
index 550454b..80110ee 100644
--- a/src/mtp_handler.erl
+++ b/src/mtp_handler.erl
@@ -215,6 +215,8 @@ handle_upstream_data(<<Header:64/binary, Rest/binary>>, #state{stage = init, sta
                                                                secret = Secret} = S) ->
     case mtp_obfuscated:from_header(Header, Secret) of
         {ok, DcId, PacketLayerMod, ObfuscatedCodec} ->
+            metric:count_inc([?APP, protocol_ok, total],
+                             1, #{labels => [PacketLayerMod]}),
             ObfuscatedLayer = mtp_layer:new(mtp_obfuscated, ObfuscatedCodec),
             PacketLayer = mtp_layer:new(PacketLayerMod, PacketLayerMod:new()),
             UpCodec = mtp_layer:new(mtp_wrap, mtp_wrap:new(PacketLayer,
@@ -224,7 +226,9 @@ handle_upstream_data(<<Header:64/binary, Rest/binary>>, #state{stage = init, sta
               S#state{up_codec = UpCodec,
                       up_acc = Rest,
                       stage_state = undefined});
-        Err ->
+        {error, unknown_protocol} = Err ->
+            metric:count_inc([?APP, protocol_error, total],
+                             1, #{labels => [unknown]}),
             Err
     end;
 handle_upstream_data(Bin, #state{stage = init, stage_state = <<>>} = S) ->
diff --git a/src/mtp_intermediate.erl b/src/mtp_intermediate.erl
index 8868d71..6676768 100644
--- a/src/mtp_intermediate.erl
+++ b/src/mtp_intermediate.erl
@@ -10,12 +10,14 @@
 -behaviour(mtp_layer).
 
 -export([new/0,
+         new/1,
          try_decode_packet/2,
          encode_packet/2]).
 -export_type([codec/0]).
 
 -record(int_st,
-        {buffer = <<>> :: binary()}).
+        {padding = false :: boolean(),
+         buffer = <<>> :: binary()}).
 -define(MAX_PACKET_SIZE, 1 * 1024 * 1024).      % 1mb
 -define(APP, mtproto_proxy).
 -define(MAX_SIZE, 16#80000000).
@@ -23,7 +25,10 @@
 -opaque codec() :: #int_st{}.
 
 new() ->
-    #int_st{}.
+    new(#{}).
+
+new(Opts) ->
+    #int_st{padding = maps:get(padding, Opts, false)}.
 
 -spec try_decode_packet(binary(), codec()) -> {ok, binary(), codec()}
                                                   | {incomplete, codec()}.
@@ -45,9 +50,14 @@ try_decode_packet(Bin, #int_st{buffer = Buf} = St) when byte_size(Buf) > 0 ->
 try_decode_packet(Bin, #int_st{buffer = <<>>} = St) ->
     {incomplete, St#int_st{buffer = Bin}}.
 
-try_decode_packet_len(Len, Data, St) ->
+try_decode_packet_len(Len, Data, #int_st{padding = Pad} = St) ->
+    Padding = case Pad of
+                  true -> Len rem 4;
+                  false -> 0
+              end,
+    NopadLen = Len - Padding,
     case Data of
-        <<_:4/binary, Packet:Len/binary, Rest/binary>> ->
+        <<_:4/binary, Packet:NopadLen/binary, _Padding:Padding/binary, Rest/binary>> ->
             {ok, Packet, St#int_st{buffer = Rest}};
         _ ->
             {incomplete, St#int_st{buffer = Data}}
diff --git a/src/mtp_layer.erl b/src/mtp_layer.erl
index 455b266..a975e3e 100644
--- a/src/mtp_layer.erl
+++ b/src/mtp_layer.erl
@@ -20,6 +20,8 @@
                | mtp_abridged
                | mtp_full
                | mtp_obfuscated
+               | mtp_intermediate
+               | mtp_secure
                | mtp_rpc
                | mtp_wrap.
 -type layer() :: {codec(), state()} | ident.
diff --git a/src/mtp_obfuscated.erl b/src/mtp_obfuscated.erl
index a363377..740af94 100644
--- a/src/mtp_obfuscated.erl
+++ b/src/mtp_obfuscated.erl
@@ -75,6 +75,9 @@ from_header(Header, Secret) when byte_size(Header) == 64  ->
         <<_:56/binary, 16#ee, 16#ee, 16#ee, 16#ee, _/binary>> ->
             DcId = get_dc(NewHeader),
             {ok, DcId, mtp_intermediate, St1};
+        <<_:56/binary, 16#dd, 16#dd, 16#dd, 16#dd, _/binary>> ->
+            DcId = get_dc(NewHeader),
+            {ok, DcId, mtp_secure, St1};
         _ ->
             metric:count_inc([?APP, protocol_error, total], 1, #{labels => [unknown]}),
             {error, unknown_protocol}
diff --git a/src/mtp_secure.erl b/src/mtp_secure.erl
new file mode 100644
index 0000000..0828c0f
--- /dev/null
+++ b/src/mtp_secure.erl
@@ -0,0 +1,29 @@
+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% MTProto intermediate protocol with random padding ("secure")
+%%% @end
+%%% Created : 29 Jun 2018 by Sergey <me@seriyps.ru>
+
+-module(mtp_secure).
+
+-behaviour(mtp_layer).
+
+-export([new/0,
+         try_decode_packet/2,
+         encode_packet/2]).
+-export_type([codec/0]).
+
+-opaque codec() :: mtp_intermediate:codec().
+
+new() ->
+    mtp_intermediate:new(#{padding => true}).
+
+-spec try_decode_packet(binary(), codec()) -> {ok, binary(), codec()}
+                                                  | {incomplete, codec()}.
+try_decode_packet(Data, St) ->
+    mtp_intermediate:try_decode_packet(Data, St).
+
+-spec encode_packet(iodata(), codec()) -> iodata().
+encode_packet(Data, St) ->
+    mtp_intermediate:encode_packet(Data, St).