From c18cd95f5109635e712660beb4fa726f3f5ba468 Mon Sep 17 00:00:00 2001
From: Sergey Prokhorov <seriy.pr@gmail.com>
Date: Fri, 30 Aug 2019 00:36:30 +0200
Subject: [PATCH] Option to close client sockets with RST

---
 src/mtp_handler.erl       | 23 ++++++++++++++++++-----
 src/mtproto_proxy.app.src | 15 +++++++++++++++
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/mtp_handler.erl b/src/mtp_handler.erl
index c4345c7..35e900c 100644
--- a/src/mtp_handler.erl
+++ b/src/mtp_handler.erl
@@ -81,15 +81,19 @@ ranch_init({Ref, Transport, Opts}) ->
     {ok, Socket} = ranch:handshake(Ref),
     case init({Socket, Transport, Opts}) of
         {ok, State} ->
-            BufSize = application:get_env(?APP, upstream_socket_buffer_size,
-                                          ?MAX_SOCK_BUF_SIZE),
+            BufSize = application:get_env(?APP, upstream_socket_buffer_size, ?MAX_SOCK_BUF_SIZE),
+            Linger = case application:get_env(?APP, reset_close_socket, off) of
+                         off -> [];
+                         _ ->
+                             [{linger, {true, 0}}]
+                     end,
             ok = Transport:setopts(
                    Socket,
                    [{active, once},
                     %% {recbuf, ?MAX_SOCK_BUF_SIZE},
                     %% {sndbuf, ?MAX_SOCK_BUF_SIZE},
                     {buffer, BufSize}
-                   ]),
+                    | Linger]),
             gen_server:enter_loop(?MODULE, [], State);
         {stop, error} ->
             exit(normal)
@@ -214,7 +218,8 @@ handle_info(Other, S) ->
     {noreply, S}.
 
 terminate(_Reason, #state{started_at = Started, listener = Listener,
-                          addr = {Ip, _}, policy_state = PolicyState} = S) ->
+                          addr = {Ip, _}, policy_state = PolicyState,
+                          sock = Sock, transport = Trans} = S) ->
     case PolicyState of
         {ok, TlsDomain} ->
             try mtp_policy:dec(
@@ -228,6 +233,7 @@ terminate(_Reason, #state{started_at = Started, listener = Listener,
             ok
     end,
     maybe_close_down(S),
+    ok = Trans:close(Sock),
     mtp_metric:count_inc([?APP, in_connection_closed, total], 1, #{labels => [Listener]}),
     Lifetime = erlang:system_time(millisecond) - Started,
     mtp_metric:histogram_observe(
@@ -319,7 +325,8 @@ parse_upstream_data(<<?TLS_START, _/binary>> = Data, #state{stage = init} = S) -
     parse_upstream_data(Data, S#state{stage = tls_hello});
 parse_upstream_data(<<Header:64/binary, Rest/binary>>,
                      #state{stage = init, secret = Secret, listener = Listener, codec = Codec0,
-                            ad_tag = Tag, addr = {Ip, _} = Addr, policy_state = PState0} = S) ->
+                            ad_tag = Tag, addr = {Ip, _} = Addr, policy_state = PState0,
+                            sock = Sock, transport = Transport} = S) ->
     case mtp_obfuscated:from_header(Header, Secret) of
         {ok, DcId, PacketLayerMod, CryptoCodecSt} ->
             maybe_check_replay(Header),
@@ -335,6 +342,12 @@ parse_upstream_data(<<Header:64/binary, Rest/binary>>,
                 end,
             mtp_metric:count_inc([?APP, protocol_ok, total],
                                  1, #{labels => [Listener, ProtoToReport]}),
+            case application:get_env(?APP, reset_close_socket, off) of
+                handshake_error ->
+                    ok = Transport:setopts(Sock, [{linger, {false, 0}}]);
+                _ ->
+                    ok
+            end,
             Codec1 = mtp_codec:replace(crypto, mtp_obfuscated, CryptoCodecSt, Codec0),
             PacketCodec = PacketLayerMod:new(),
             Codec2 = mtp_codec:replace(packet, PacketLayerMod, PacketCodec, Codec1),
diff --git a/src/mtproto_proxy.app.src b/src/mtproto_proxy.app.src
index 35d6fe4..0166f03 100644
--- a/src/mtproto_proxy.app.src
+++ b/src/mtproto_proxy.app.src
@@ -85,6 +85,21 @@
         %% Default: off
         {replay_check_server_error_filter, first},
 
+        %% This option controls how proxy closes client sockets
+        %% (SO_LINGER timeout=0 socket option)
+        %% Can be useful if you have too many connection attempts with wrong
+        %% secret and protocol, which creates lots of sockets in
+        %% TIME_WAIT, ORPHANED and CLOSED state
+        %% Allowed values:
+        %% - off - never send RST
+        %% - handshake_error - (recommended) only send if client handshake failed because of:
+        %%   - wrong secret
+        %%   - disabled protocol
+        %%   - replay attack detected
+        %%   - policy check failed
+        %% - always - always close socket with RST
+        {reset_close_socket, off},
+
         %% Options for `mtp_session_storage` replay attack check
         %% Those settings are not precise! They are checked not in realtime, but
         %% once per minute.