fix: improve receiving proxy protocol information

savonarola · savonarola · commit 73361401c06d · 2023-04-26T17:50:18.000+03:00
* handle proxy protocol timeout correctly
* handle recv errors while receiving proxy protocol info
diff --git a/src/esockd.app.src b/src/esockd.app.src
@@ -1,7 +1,7 @@
 {application, esockd,
  [{description, "General Non-blocking TCP/SSL and UDP/DTLS Server"},
   {id, "esockd"},
-  {vsn, "5.9.6"},
+  {vsn, "5.9.7"},
   {modules, []},
   {registered, []},
   {applications, [kernel, stdlib, sasl, ssl]},
diff --git a/src/esockd.appup.src b/src/esockd.appup.src
@@ -1,7 +1,11 @@
 %%-*- mode: erlang; -*-
 
-{"5.9.6",
- [{"5.9.5",
+{"5.9.7",
+ [{"5.9.6",
+   [ {load_module,esockd_proxy_protocol,brutal_purge,soft_purge,[]}
+   ]
+  },
+  {"5.9.5",
    [ {load_module,esockd_transport,brutal_purge,soft_purge,[]}
    , {load_module,esockd_proxy_protocol,brutal_purge,soft_purge,[]}
    , {load_module,esockd_acceptor,brutal_purge,soft_purge,[]}
@@ -46,7 +50,17 @@
   },
   {<<".*">>, []}
  ],
- [{"5.9.4",
+ [{"5.9.6",
+   [ {load_module,esockd_proxy_protocol,brutal_purge,soft_purge,[]}
+   ]
+  },
+  {"5.9.5",
+   [ {load_module,esockd_transport,brutal_purge,soft_purge,[]}
+   , {load_module,esockd_proxy_protocol,brutal_purge,soft_purge,[]}
+   , {load_module,esockd_acceptor,brutal_purge,soft_purge,[]}
+   ]
+  },
+  {"5.9.4",
    [ {load_module,esockd_transport,brutal_purge,soft_purge,[]}
    , {load_module,esockd_proxy_protocol,brutal_purge,soft_purge,[]}
    , {load_module,esockd_acceptor,brutal_purge,soft_purge,[]}
diff --git a/src/esockd_proxy_protocol.erl b/src/esockd_proxy_protocol.erl
@@ -40,7 +40,7 @@
 
 -define(SPACE, 16#20).
 
--define(TIMEOUT, 5000).
+-define(MAX_TIMEOUT, 17000).
 
 %% Proxy Protocol Additional Fields
 -define(PP2_TYPE_ALPN,           16#01).
@@ -62,6 +62,7 @@
 -spec(recv(module(), inet:socket() | #ssl_socket{}, timeout()) ->
       {ok, #proxy_socket{}} | {error, term()}).
 recv(Transport, Sock, Timeout) ->
+    Deadline = deadline(Timeout),
     {ok, OriginOpts} = Transport:getopts(Sock, [mode, active, packet]),
     ok = Transport:setopts(Sock, [binary, {active, once}, {packet, line}]),
     receive
@@ -75,16 +76,7 @@ recv(Transport, Sock, Timeout) ->
             {ok, Sock};
         %% V2 TCP
         {_, _Sock, <<"\r\n">>} ->
-            Transport:setopts(Sock, [{active, false}, {packet, raw}]),
-            {ok, Header} = Transport:recv(Sock, 14, 1000),
-            <<?SIG, 2:4, Cmd:4, AF:4, Trans:4, Len:16>> = Header,
-            case Transport:recv(Sock, Len, 1000) of
-                {ok, ProxyInfo} ->
-                    Transport:setopts(Sock, OriginOpts),
-                    parse_v2(Cmd, Trans, ProxyInfo, #proxy_socket{inet = inet_family(AF), socket = Sock});
-                {error, Reason} ->
-                    {error, {recv_proxy_info_error, Reason}}
-            end;
+            recv_v2(Transport, Sock, OriginOpts, Deadline);
         {tcp_error, _Sock, Reason} ->
             {error, {recv_proxy_info_error, Reason}};
         {tcp_closed, _Sock} ->
@@ -98,6 +90,31 @@ recv(Transport, Sock, Timeout) ->
             {error, proxy_proto_timeout}
     end.
 
+recv_v2(Transport, Sock, OriginOpts, Deadline) ->
+    Transport:setopts(Sock, [{active, false}, {packet, raw}]),
+    with_remaining_timeout(Deadline, fun(HeaderTimeout) ->
+        case Transport:recv(Sock, 14, HeaderTimeout) of
+            {ok, <<?SIG, 2:4, Cmd:4, AF:4, Trans:4, Len:16>>} ->
+                with_remaining_timeout(Deadline, fun(ProxyInfoTimeout) ->
+                    case Transport:recv(Sock, Len, ProxyInfoTimeout) of
+                        {ok, ProxyInfo} ->
+                            Transport:setopts(Sock, OriginOpts),
+                            parse_v2(Cmd, Trans, ProxyInfo, #proxy_socket{inet = inet_family(AF), socket = Sock});
+                        {error, closed} ->
+                            {error, proxy_proto_close};
+                        {error, Reason} ->
+                            {error, {recv_proxy_info_error, Reason}}
+                    end
+                end);
+            {ok, UnknownHeader} ->
+                {error, {invalid_proxy_info, UnknownHeader}};
+            {error, closed} ->
+                {error, proxy_proto_close};
+            {error, Reason} ->
+                {error, {recv_proxy_info_error, Reason}}
+        end
+    end).
+
 parse_v1(ProxyInfo, ProxySock) ->
     [SrcAddrBin, DstAddrBin, SrcPortBin, DstPortBin]
         = binary:split(ProxyInfo, [<<" ">>, <<"\r\n">>], [global, trim]),
@@ -210,3 +227,22 @@ inet_family(?UNIX)   -> unix.
 bool(1) -> true;
 bool(_) -> false.
 
+maybe_limit_timeout(infinity) -> ?MAX_TIMEOUT;
+maybe_limit_timeout(Timeout) when is_integer(Timeout) andalso Timeout > ?MAX_TIMEOUT ->
+    ?MAX_TIMEOUT;
+maybe_limit_timeout(Timeout) when is_integer(Timeout) andalso Timeout > 0 ->
+    Timeout.
+
+deadline(Timeout) ->
+    erlang:monotonic_time(millisecond) + maybe_limit_timeout(Timeout).
+
+timeout_left(Deadline) ->
+    Deadline - erlang:monotonic_time(millisecond).
+
+with_remaining_timeout(Timer, Fun) ->
+    case timeout_left(Timer) of
+        Timeout when Timeout > 0 ->
+            Fun(Timeout);
+        _ ->
+            {error, proxy_proto_timeout}
+    end.
diff --git a/test/esockd_proxy_protocol_SUITE.erl b/test/esockd_proxy_protocol_SUITE.erl
@@ -182,7 +182,79 @@ t_ppv2_connect(_) ->
                                                      12,127,50,210,1,210,21,16,142,250,32,1,181>>),
                            ok = gen_tcp:send(Sock, <<"Hello">>),
                            {ok, <<"Hello">>} = gen_tcp:recv(Sock, 0)
-                   end).
+                   end),
+    with_pp_server(fun(Sock) ->
+                            ok = gen_tcp:send(Sock, <<13,10,13,10,0,13,10,81,85,73,84,10,33,17,0,
+                                                    12,127,50,210,1,210,21,16,142,250,32,1,181>>),
+                            ok = gen_tcp:send(Sock, <<"Hello">>),
+                            {ok, <<"Hello">>} = gen_tcp:recv(Sock, 0)
+                   end,
+                   [{proxy_protocol_timeout, infinity}]).
+
+
+t_ppv2_connect_first_marker_timeout(_) ->
+    with_pp_server(
+        fun(Sock) ->
+            timer:sleep(200),
+            {error, closed} = gen_tcp:recv(Sock, 0)
+        end,
+        [{proxy_protocol_timeout, 100}]),
+    with_pp_server(
+        fun(Sock) ->
+            timer:sleep(400),
+            ok = gen_tcp:send(Sock, <<13,10,13,10,0,13,10,81,85,73,84,10,33,17,0,
+                                      12,127,50,210,1,210,21,16,142,250,32,1,181>>),
+            ok = gen_tcp:send(Sock, <<"Hello">>),
+            {ok, <<"Hello">>} = gen_tcp:recv(Sock, 0)
+        end,
+        [{proxy_protocol_timeout, 500}]).
+
+t_ppv2_connect_header_timeout(_) ->
+    with_pp_server(
+        fun(Sock) ->
+            timer:sleep(200),
+            ok = gen_tcp:send(Sock, <<13,10>>),
+            timer:sleep(200),
+            ok = gen_tcp:send(Sock, <<13,10,0,13,10,81,85,73,84,10,33,17,0,
+                                     12,127,50,210,1,210,21,16,142,250,32,1,181>>),
+            ok = gen_tcp:send(Sock, <<"Hello">>),
+            {ok, <<"Hello">>} = gen_tcp:recv(Sock, 0)
+        end,
+        [{proxy_protocol_timeout, 500}]),
+    with_pp_server(
+        fun(Sock) ->
+            timer:sleep(200),
+            ok = gen_tcp:send(Sock, <<13,10>>),
+            timer:sleep(400),
+            {error, closed} = gen_tcp:recv(Sock, 0)
+        end,
+        [{proxy_protocol_timeout, 500}]).
+
+t_ppv2_connect_proxy_info_timeout(_) ->
+    with_pp_server(
+        fun(Sock) ->
+            timer:sleep(100),
+            ok = gen_tcp:send(Sock, <<13,10>>),
+            timer:sleep(100),
+            ok = gen_tcp:send(Sock, <<13,10,0,13,10,81,85,73,84,10,33,17,0,
+                                        12,127,50,210,1,210,21,16,142>>),
+            timer:sleep(200),
+            ok = gen_tcp:send(Sock, <<250,32,1,181>>),
+            ok = gen_tcp:send(Sock, <<"Hello">>),
+            {ok, <<"Hello">>} = gen_tcp:recv(Sock, 0)
+        end,
+        [{proxy_protocol_timeout, 500}]),
+    with_pp_server(
+        fun(Sock) ->
+            timer:sleep(100),
+            ok = gen_tcp:send(Sock, <<13,10>>),
+            timer:sleep(100),
+            ok = gen_tcp:send(Sock, <<13,10,0,13,10,81,85,73,84,10,33,17,0,
+                                        12,127,50,210,1,210,21,16,142>>),
+            timer:sleep(400),
+            {error, closed} = gen_tcp:recv(Sock, 0)
+        end,
+        [{proxy_protocol_timeout, 500}]).
 
 t_ppv1_unknown(_) ->
     with_pp_server(fun(Sock) ->
@@ -198,11 +270,15 @@ t_ppv1_garbage_data(_) ->
                    end).
 
 with_pp_server(TestFun) ->
-    {ok, _} = esockd:open(echo, 5000, [{tcp_options, [binary]},
-                                       proxy_protocol,
-                                       {proxy_protocol_timeout, 3000}
-                                      ],
-                          {echo_server, start_link, []}),
+    PPOpts = [{proxy_protocol_timeout, 3000}],
+    with_pp_server(TestFun, PPOpts).
+
+with_pp_server(TestFun, PPOpts) ->
+    Opts = PPOpts ++
+    [
+        {tcp_options, [binary]},
+        proxy_protocol
+    ],
+    {ok, _} = esockd:open(echo, 5000, Opts, {echo_server, start_link, []}),
     {ok, Sock} = gen_tcp:connect({127,0,0,1}, 5000, [binary, {active, false}]),
     try TestFun(Sock) after ok = esockd:close(echo, 5000) end.
-
