CERBERUS/mtproto_proxy
1
0
Fork 0
mirror of https://github.com/seriyps/mtproto_proxy.git synced 2026-08-04 15:19:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

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

Browse source
This commit is contained in:
Сергей Прохоров 2018-10-31 02:49:01 +01:00
parent 590a633ab7
commit 447639224f
No known key found for this signature in database
GPG key ID: 1C570244E4EF3337
12 changed files with 162 additions and 151 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/mtp_abridged.erl
View file

@ -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,

2
src/mtp_aes_cbc.erl
View file

@ -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,

104
src/mtp_codec.erl Normal file
View file

@ -0,0 +1,104 @@
%%% @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.

29
src/mtp_down_conn.erl
View file

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

2
src/mtp_full.erl
View file

@ -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,

17
src/mtp_handler.erl
View file

@ -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

2
src/mtp_intermediate.erl
View file

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

70
src/mtp_layer.erl
View file

@ -1,70 +0,0 @@
%%% @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.

28
src/mtp_noop_codec.erl Normal file
View file

@ -0,0 +1,28 @@
%%% @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}.

2
src/mtp_obfuscated.erl
View file

@ -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,

2
src/mtp_secure.erl
View file

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

53
src/mtp_wrap.erl
View file

@ -1,53 +0,0 @@
%%% @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}}.