From 2eb3fedd7071ead3fa90bab641713b9ac5a3cdf0 Mon Sep 17 00:00:00 2001
From: Sergey Prokhorov <seriy.pr@gmail.com>
Date: Wed, 16 Oct 2019 21:15:19 +0200
Subject: [PATCH] Early detection of non-tls connect when tls_only

This is a minor optimization for a common case when outdated
clients are trying to use tls-only proxy and storm it with
incorrect non-tls connect attempts.
---
 src/mtp_handler.erl | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/mtp_handler.erl b/src/mtp_handler.erl
index 35e900c..01cd888 100644
--- a/src/mtp_handler.erl
+++ b/src/mtp_handler.erl
@@ -327,6 +327,9 @@ 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,
                             sock = Sock, transport = Transport} = S) ->
+    AllowedProtocols = allowed_protocols(),
+    (not is_tls_only(AllowedProtocols)) orelse
+        error({protocol_error, tls_client_hello_expected, Header}),
     case mtp_obfuscated:from_header(Header, Secret) of
         {ok, DcId, PacketLayerMod, CryptoCodecSt} ->
             maybe_check_replay(Header),
@@ -335,7 +338,7 @@ parse_upstream_data(<<Header:64/binary, Rest/binary>>,
                     {true, _} when PacketLayerMod == mtp_secure ->
                         {mtp_secure_fake_tls, PState0};
                     {false, _} ->
-                        assert_protocol(PacketLayerMod),
+                        assert_protocol(PacketLayerMod, AllowedProtocols),
                         check_policy(Listener, Ip, undefined),
                         %FIXME: if any codebelow fail, we will get counter policy leak
                         {PacketLayerMod, {ok, undefined}}
@@ -372,8 +375,18 @@ parse_upstream_data(Bin, #state{stage = Stage, codec = Codec0} = S) when Stage =
     Codec = mtp_codec:push_back(first, Bin, Codec0),
     {incomplete, S#state{codec = Codec}}.
 
-assert_protocol(Protocol) ->
+
+allowed_protocols() ->
     {ok, AllowedProtocols} = application:get_env(?APP, allowed_protocols),
+    AllowedProtocols.
+
+is_tls_only([mtp_fake_tls]) -> true;
+is_tls_only(_) -> false.
+
+assert_protocol(Protocol) ->
+    assert_protocol(Protocol, allowed_protocols()).
+
+assert_protocol(Protocol, AllowedProtocols) ->
     lists:member(Protocol, AllowedProtocols)
         orelse error({protocol_error, disabled_protocol, Protocol}).