Simplify codecs. Use fixed/flat 2-layer codec instead of tree

src/mtp_abridged.erl

@@ -6,7 +6,7 @@
 %%% Created : 29 May 2018 by Sergey <me@seriyps.ru>
 
 -module(mtp_abridged).
--behaviour(mtp_layer).
+-behaviour(mtp_codec).
 
 -export([new/0,
          try_decode_packet/2,

src/mtp_aes_cbc.erl

@@ -6,7 +6,7 @@
 %%% Created :  6 Jun 2018 by Sergey <me@seriyps.ru>
 
 -module(mtp_aes_cbc).
--behaviour(mtp_layer).
+-behaviour(mtp_codec).
 
 -export([new/5,
          encrypt/2,

src/mtp_codec.erl

+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% This module provieds a comination of crypto and packet codecs.
+%%% Crypto is always outer layer and packet is inner:
+%%% ( --- packet --- )
+%%%   (-- crypto --)
+%%%      - tcp -
+%%% @end
+%%% Created :  6 Jun 2018 by Sergey <me@seriyps.ru>
+
+-module(mtp_codec).
+
+-export([new/4,
+         decompose/1,
+         try_decode_packet/2,
+         encode_packet/2,
+         fold_packets/4]).
+-export_type([codec/0]).
+
+-type state() :: any().
+-type crypto_codec() :: mtb_aes_cbc
+                      | mtp_obfuscated.
+-type packet_codec() :: mtp_abridged
+                      | mtp_full
+                      | mtp_intermediate
+                      | mtp_secure.
+
+-record(codec,
+        {crypto_mod :: crypto_codec(),
+         crypto_state :: any(),
+         packet_mod :: packet_codec(),
+         packet_state :: any()}).
+
+-define(APP, mtproto_proxy).
+
+-callback try_decode_packet(binary(), state()) ->
+    {ok, binary(), state()}
+        | {incomplete, state()}.
+
+-callback encode_packet(iodata(), state()) ->
+    {iodata(), state()}.
+
+-opaque codec() :: #codec{}.
+
+
+-spec new(crypto_codec(), state(), packet_codec(), state()) -> codec().
+new(CryptoMod, CryptoState, PacketMod, PacketState) ->
+    #codec{crypto_mod = CryptoMod,
+           crypto_state = CryptoState,
+           packet_mod = PacketMod,
+           packet_state = PacketState}.
+
+-spec decompose(codec()) -> {crypto_codec(), state(), packet_codec(), state()}.
+decompose(#codec{crypto_mod = CryptoMod, crypto_state = CryptoState,
+                 packet_mod = PacketMod, packet_state = PacketState}) ->
+    {CryptoMod, CryptoState, PacketMod, PacketState}.
+
+
+%% try_decode_packet(Inner) |> try_decode_packet(Outer)
+try_decode_packet(Bin, #codec{crypto_mod = CryptoMod,
+                              crypto_state = CryptoSt,
+                              packet_mod = PacketMod,
+                              packet_state = PacketSt} = S) ->
+    {Dec1, CryptoSt1} =
+        case CryptoMod:try_decode_packet(Bin, CryptoSt) of
+            {incomplete, PacketSt1_} ->
+                %% We have to check if something is left in packet's buffers
+                {<<>>, PacketSt1_};
+            {ok, Dec1_, PacketSt1_} ->
+                {Dec1_, PacketSt1_}
+        end,
+    case PacketMod:try_decode_packet(Dec1, PacketSt) of
+        {incomplete, PacketSt1} ->
+            {incomplete, S#codec{crypto_state = CryptoSt1,
+                                 packet_state = PacketSt1}};
+        {ok, Dec2, PacketSt1} ->
+            {ok, Dec2, S#codec{crypto_state = CryptoSt1,
+                               packet_state = PacketSt1}}
+    end.
+
+%% encode_packet(Outer) |> encode_packet(Inner)
+encode_packet(Bin, #codec{packet_mod = PacketMod,
+                          packet_state = PacketSt,
+                          crypto_mod = CryptoMod,
+                          crypto_state = CryptoSt} = S) ->
+    {Enc1, PacketSt1} = PacketMod:encode_packet(Bin, PacketSt),
+    {Enc2, CryptoSt1} = CryptoMod:encode_packet(Enc1, CryptoSt),
+    {Enc2, S#codec{crypto_state = CryptoSt1, packet_state = PacketSt1}}.
+
+
+-spec fold_packets(fun( (binary(), FoldSt) -> FoldSt ),
+                   FoldSt, binary(), codec()) ->
+                          {ok, FoldSt, codec()}
+                              when
+      FoldSt :: any().
+fold_packets(Fun, FoldSt, Data, Codec) ->
+    case try_decode_packet(Data, Codec) of
+        {ok, Decoded, Codec1} ->
+            FoldSt1 = Fun(Decoded, FoldSt),
+            fold_packets(Fun, FoldSt1, <<>>, Codec1);
+        {incomplete, Codec1} ->
+            {ok, FoldSt, Codec1}
+    end.

src/mtp_down_conn.erl

@@ -42,7 +42,7 @@
                 stage_state = [] :: any(),
                 sock :: gen_tcp:socket() | undefined,
                 addr_bin :: binary() | undefined,           % my external ip:port
-                codec :: mtp_layer:layer() | undefined,
+                codec :: mtp_codec:codec() | undefined,
                 upstreams = #{} :: #{mtp_handler:handle() => upstream()},
                 upstreams_rev = #{} :: #{mtp_rpc:conn_id() => mtp_handler:handle()},
                 pool :: pid(),
@@ -173,7 +173,7 @@ handle_upstream_closed(Upstream, #state{upstreams = Ups,
 handle_downstream_data(Bin, #state{stage = tunnel,
                                    codec = DownCodec} = S) ->
     {ok, S3, DownCodec1} =
-        mtp_layer:fold_packets(
+        mtp_codec:fold_packets(
           fun(Decoded, S1) ->
                   mtp_metric:histogram_observe(
                     [?APP, tg_packet_size, bytes],
@@ -184,7 +184,7 @@ handle_downstream_data(Bin, #state{stage = tunnel,
     {ok, S3#state{codec = DownCodec1}};
 handle_downstream_data(Bin, #state{stage = handshake_1,
                                    codec = DownCodec} = S) ->
-    case mtp_layer:try_decode_packet(Bin, DownCodec) of
+    case mtp_codec:try_decode_packet(Bin, DownCodec) of
         {ok, Packet, DownCodec1} ->
             down_handshake2(Packet, S#state{codec = DownCodec1});
         {incomplete, DownCodec1} ->
@@ -192,7 +192,7 @@ handle_downstream_data(Bin, #state{stage = handshake_1,
     end;
 handle_downstream_data(Bin, #state{stage = handshake_2,
                                    codec = DownCodec} = S) ->
-    case mtp_layer:try_decode_packet(Bin, DownCodec) of
+    case mtp_codec: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
@@ -223,7 +223,7 @@ handle_rpc({simple_ack, ConnId, Confirm}, S) ->
 -spec down_send(iodata(), #state{}) -> {ok, #state{}}.
 down_send(Packet, #state{sock = Sock, codec = Codec, dc_id = DcId} = St) ->
     %% lager:debug("Up>Down: ~w", [Packet]),
-    {Encoded, Codec1} = mtp_layer:encode_packet(Packet, Codec),
+    {Encoded, Codec1} = mtp_codec:encode_packet(Packet, Codec),
     mtp_metric:rt(
       [?APP, downstream_send_duration, seconds],
       fun() ->
@@ -292,16 +292,17 @@ down_handshake1(S) ->
             1:32/little,                        %AES
             CryptoTs:32/little,
             Nonce/binary>>,
-    Full = mtp_full:new(-2, -2),
-    S1 = S#state{codec = mtp_layer:new(mtp_full, Full),
-                 stage = handshake_1,
+    S1 = S#state{stage = handshake_1,
+                 %% Use fake encryption codec
+                 codec = mtp_codec:new(mtp_noop_codec, mtp_noop_codec:new(),
+                                       mtp_full, mtp_full:new(-2, -2)),
                  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},
-                                               sock = Sock,
-                                               codec = DownCodec} = S) ->
+                                               codec = Codec1,
+                                               sock = Sock} = 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}),
@@ -314,15 +315,17 @@ down_handshake2(<<Type:4/binary, KeySelector:4/binary, Schema:32/little, _Crypto
              clt_ip => MyIpBin, clt_port => MyPort, secret => Key},
     {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)),
+    {_, _, PacketMod, PacketState} = mtp_codec:decompose(Codec1),
+    CryptoState = mtp_aes_cbc:new(EncKey, EncIv, DecKey, DecIv, 16),
+    Codec = mtp_codec:new(mtp_aes_cbc, CryptoState,
+                          PacketMod, PacketState),
     SenderPID = PeerPID = <<"IPIPPRPDTIME">>,
     Handshake = [?RPC_HANDSHAKE,
                  ?RPC_FLAGS,
                  SenderPID,
                  PeerPID],
     down_send(Handshake,
-              S#state{codec = DownCodec1,
+              S#state{codec = Codec,
                       stage = handshake_2,
                       addr_bin = iolist_to_binary(mtp_rpc:encode_ip_port(MyIp, MyPort)),
                       stage_state = SenderPID}).

src/mtp_full.erl

@@ -9,7 +9,7 @@
 %%% Created :  6 Jun 2018 by Sergey <me@seriyps.ru>
 
 -module(mtp_full).
--behaviour(mtp_layer).
+-behaviour(mtp_codec).
 
 -export([new/0, new/2,
          try_decode_packet/2,

src/mtp_handler.erl

@@ -42,7 +42,7 @@
 
          sock :: gen_tcp:socket(),
          transport :: transport(),
-         codec = ident :: mtp_layer:layer(),
+         codec :: mtp_codec:codec() | undefined,
 
          down :: mtp_down_conn:handle(),
          dc_id :: integer(),
@@ -234,7 +234,7 @@ state_timeout(stop) ->
 handle_upstream_data(Bin, #state{stage = tunnel,
                                  codec = UpCodec} = S) ->
     {ok, S3, UpCodec1} =
-        mtp_layer:fold_packets(
+        mtp_codec:fold_packets(
           fun(Decoded, S1) ->
                   mtp_metric:histogram_observe(
                     [?APP, tg_packet_size, bytes],
@@ -247,16 +247,15 @@ handle_upstream_data(Bin, #state{stage = tunnel,
 handle_upstream_data(<<Header:64/binary, Rest/binary>>, #state{stage = init, stage_state = <<>>,
                                                                secret = Secret, listener = Listener} = S) ->
     case mtp_obfuscated:from_header(Header, Secret) of
-        {ok, DcId, PacketLayerMod, ObfuscatedCodec} ->
+        {ok, DcId, PacketLayerMod, CryptoCodecSt} ->
             mtp_metric:count_inc([?APP, protocol_ok, total],
                                  1, #{labels => [Listener, 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,
-                                                           ObfuscatedLayer)),
+            PacketCodec = PacketLayerMod:new(),
+            Codec = mtp_codec:new(mtp_obfuscated, CryptoCodecSt,
+                                  PacketLayerMod, PacketCodec),
             handle_upstream_header(
               DcId,
-              S#state{codec = UpCodec,
+              S#state{codec = Codec,
                       acc = Rest,
                       stage_state = undefined});
         {error, Reason} = Err ->
@@ -276,7 +275,7 @@ up_send(Packet, #state{stage = tunnel,
                        transport = Transport,
                        listener = Listener} = S) ->
     %% lager:debug(">Up: ~p", [Packet]),
-    {Encoded, UpCodec1} = mtp_layer:encode_packet(Packet, UpCodec),
+    {Encoded, UpCodec1} = mtp_codec:encode_packet(Packet, UpCodec),
     mtp_metric:rt([?APP, upstream_send_duration, seconds],
               fun() ->
                       case Transport:send(Sock, Encoded) of

src/mtp_intermediate.erl

@@ -7,7 +7,7 @@
 
 -module(mtp_intermediate).
 
--behaviour(mtp_layer).
+-behaviour(mtp_codec).
 
 -export([new/0,
          new/1,

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_intermediate
-               | mtp_secure
-               | 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_noop_codec.erl

+%%% @author Sergey <me@seriyps.ru>
+%%% @copyright (C) 2018, Sergey
+%%% @doc
+%%% Fake codec that returns it's input as output.
+%%% Used in downstream handshake flow
+%%% @end
+%%% Created : 31 Oct 2018 by Sergey <me@seriyps.ru>
+
+-module(mtp_noop_codec).
+-behaviour(mtp_codec).
+-export([new/0,
+         try_decode_packet/2,
+         encode_packet/2]).
+-export_type([codec/0]).
+
+-opaque codec() :: ?MODULE.
+
+-spec new() -> codec().
+new() ->
+    ?MODULE.
+
+-spec try_decode_packet(binary(), codec()) -> {ok, binary(), codec()}.
+try_decode_packet(Data, ?MODULE) ->
+    {ok, Data, ?MODULE}.
+
+-spec encode_packet(binary(), codec()) -> {binary(), codec()}.
+encode_packet(Data, ?MODULE) ->
+    {Data, ?MODULE}.

src/mtp_obfuscated.erl

@@ -6,7 +6,7 @@
 %%% Created : 29 May 2018 by Sergey <me@seriyps.ru>
 
 -module(mtp_obfuscated).
--behaviour(mtp_layer).
+-behaviour(mtp_codec).
 -export([create/0,
          create/1,
          from_header/2,

src/mtp_secure.erl

@@ -7,7 +7,7 @@
 
 -module(mtp_secure).
 
--behaviour(mtp_layer).
+-behaviour(mtp_codec).
 
 -export([new/0,
          try_decode_packet/2,

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}}.