From 1183dbdd51124727fb6d54cfd4d632c7c36d5837 Mon Sep 17 00:00:00 2001
From: Sergey Prokhorov <sergey.prokhorov@klarna.com>
Date: Mon, 25 Jun 2018 15:13:49 +0200
Subject: [PATCH] Add support for mtp_intermediate upstream

---
 mtp_intermediate.erl   | 63 ++++++++++++++++++++++++++++++++++++++++++
 src/mtp_handler.erl    |  6 ++--
 src/mtp_obfuscated.erl |  8 +++---
 3 files changed, 70 insertions(+), 7 deletions(-)
 create mode 100644 mtp_intermediate.erl

diff --git a/mtp_intermediate.erl b/mtp_intermediate.erl
new file mode 100644
index 0000000..62f0095
--- /dev/null
+++ b/mtp_intermediate.erl
@@ -0,0 +1,63 @@
+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% MTProto intermediate protocol
+%%% @end
+%%% Created : 18 Jun 2018 by Sergey <me@seriyps.ru>
+
+-module(mtp_intermediate).
+
+-behaviour(mtp_layer).
+
+-export([new/0,
+         try_decode_packet/2,
+         encode_packet/2]).
+-export_type([codec/0]).
+
+-record(int_st,
+        {buffer = <<>> :: binary()}).
+-define(MAX_PACKET_SIZE, 1 * 1024 * 1024).      % 1mb
+-define(APP, mtproto_proxy).
+-define(MAX_SIZE, 16#80000000).
+-include("dbg.hrl").
+
+-opaque codec() :: #int_st{}.
+
+new() ->
+    #int_st{}.
+
+-spec try_decode_packet(binary(), codec()) -> {ok, binary(), codec()}
+                                                  | {incomplete, codec()}.
+try_decode_packet(<<Len:32/unsigned-little, _/binary>> = Data,
+                      #int_st{buffer = <<>>} = St) when Len < ?MAX_SIZE ->
+    Len1 = case Len < ?MAX_SIZE of
+               true -> Len;
+               false -> Len - ?MAX_SIZE
+           end,
+    (Len1 < ?MAX_PACKET_SIZE)
+        orelse
+        begin
+            metric:count_inc([?APP, protocol_error, total], 1, #{labels => [intermediate_max_size]}),
+            error({packet_too_large, Len1})
+        end,
+    try_decode_packet_len(Len1, Data, St);
+try_decode_packet(Bin, #int_st{buffer = Buf} = St) when byte_size(Buf) > 0 ->
+    try_decode_packet(<<Buf/binary, Bin/binary>>, St#int_st{buffer = <<>>});
+try_decode_packet(Bin, #int_st{buffer = <<>>} = St) ->
+    {incomplete, St#int_st{buffer = Bin}}.
+
+try_decode_packet_len(Len, Data, St) ->
+    case Data of
+        <<_:4/binary, Packet:Len/binary, Rest/binary>> ->
+            ?DBG("len=~p ~w -> ~w", [Len, Data, Packet]),
+            {ok, Packet, St#int_st{buffer = Rest}};
+        _ ->
+            {incomplete, St#int_st{buffer = Data}}
+    end.
+
+-spec encode_packet(iodata(), codec()) -> iodata().
+encode_packet(Data, St) ->
+    Size = iolist_size(Data),
+    Packet = [<<Size:32/little>> | Data],
+    ?DBG("~w -> ~w", [iolist_to_binary(Data), iolist_to_binary(Packet)]),
+    {Packet, St}.
diff --git a/src/mtp_handler.erl b/src/mtp_handler.erl
index c8dac1a..550454b 100644
--- a/src/mtp_handler.erl
+++ b/src/mtp_handler.erl
@@ -214,10 +214,10 @@ state_timeout(stop) ->
 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, DcId, ObfuscatedCodec} ->
+        {ok, DcId, PacketLayerMod, 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,
+            PacketLayer = mtp_layer:new(PacketLayerMod, PacketLayerMod:new()),
+            UpCodec = mtp_layer:new(mtp_wrap, mtp_wrap:new(PacketLayer,
                                                            ObfuscatedLayer)),
             handle_upstream_header(
               DcId,
diff --git a/src/mtp_obfuscated.erl b/src/mtp_obfuscated.erl
index 8a41dfb..a363377 100644
--- a/src/mtp_obfuscated.erl
+++ b/src/mtp_obfuscated.erl
@@ -60,7 +60,7 @@ init_down_encrypt(<<_:8/binary, Key:32/binary, IV:16/binary, _/binary>>) ->
 
 
 %% @doc creates new obfuscated stream (MTProto proxy format)
--spec from_header(binary(), binary()) -> {ok, inet:ip4_address(), codec()}.
+-spec from_header(binary(), binary()) -> {ok, integer(), mtp_layer:codec(), codec()}.
 from_header(Header, Secret) when byte_size(Header) == 64  ->
     {EncKey, EncIV} = init_up_encrypt(Header, Secret),
     {DecKey, DecIV} = init_up_decrypt(Header, Secret),
@@ -71,10 +71,10 @@ from_header(Header, Secret) when byte_size(Header) == 64  ->
     case NewHeader of
         <<_:56/binary, 16#ef, 16#ef, 16#ef, 16#ef, _/binary>> ->
             DcId = get_dc(NewHeader),
-            {ok, DcId, St1};
+            {ok, DcId, mtp_abridged, 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}};
+            DcId = get_dc(NewHeader),
+            {ok, DcId, mtp_intermediate, St1};
         _ ->
             metric:count_inc([?APP, protocol_error, total], 1, #{labels => [unknown]}),
             {error, unknown_protocol}