fix handler

README.md

@@ -13,10 +13,18 @@ Run
 
     $ rebar3 shell
 
-The UDP listener is configured in `config/sys.config` and listens on port `7000`
+The UDP listener is configured in `config/sys.config` and listens on port `16380`
-by default. It echoes received datagrams.
+by default. It accepts RelayKit UDP relay frames.
 
-Test
+Protocol
-----
+--------
 
-    $ echo "ping" | nc -u -w1 127.0.0.1 7000
+Each UDP datagram is one `RKP1` frame:
+
+    magic(4) = "RKP1", version(1) = 1, type(1), reserved(2),
+    stream_id(8), payload_len(4), payload(payload_len)
+
+Frame types are `open = 1`, `data = 2`, `close = 3`, and `error = 4`.
+`open` payload is RelayKit JSON: `host`, `port`, `useTLS`, and `method`.
+The server is a transparent TCP relay: `useTLS` is parsed for client
+compatibility, but the upstream connection is always raw TCP.

apps/relay_server/src/relay_server_udp_handler.erl

@@ -7,25 +7,40 @@
 
 -behaviour(gen_server).
 
--export([start_link/3, handle_packet/2]).
+-export([start_link/4]).
 
 -export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]).
 
+-define(TCP_OPTIONS, [binary, {packet, raw}, {active, once}, {nodelay, true}]).
+
+-record(stream, {
+    id :: non_neg_integer(),
+    socket :: inet:socket(),
+    host :: binary(),
+    port :: inet:port_number()
+}).
+
+-type stream() :: #stream{}.
+
 -record(state, {
     server :: pid(),
     peer :: {inet:ip_address(), inet:port_number()},
-    idle_timeout :: timeout()
+    idle_timeout :: timeout(),
+    connect_timeout :: timeout(),
+    streams = #{} :: #{non_neg_integer() => stream()},
+    sockets = #{} :: #{inet:socket() => non_neg_integer()}
 }).
 
-start_link(Transport, Peer, IdleTimeout) ->
+start_link(Transport, Peer, IdleTimeout, ConnectTimeout) ->
-    gen_server:start_link(?MODULE, [Transport, Peer, IdleTimeout], []).
+    gen_server:start_link(?MODULE, [Transport, Peer, IdleTimeout, ConnectTimeout], []).
 
-init([{udp, Server, _Sock}, Peer, IdleTimeout]) ->
+init([{udp, Server, _Sock}, Peer, IdleTimeout, ConnectTimeout]) ->
     logger:debug("UDP peer connected: ~s", [esockd:format(Peer)]),
     {ok, #state{
         server = Server,
         peer = Peer,
-        idle_timeout = IdleTimeout
+        idle_timeout = IdleTimeout,
+        connect_timeout = ConnectTimeout
     }, IdleTimeout}.
 
 handle_call(_Request, _From, State = #state{idle_timeout = IdleTimeout}) ->
@@ -37,21 +52,176 @@ handle_cast(_Request, State = #state{idle_timeout = IdleTimeout}) ->
 handle_info({datagram, Server, <<"stop">>}, State = #state{server = Server, peer = Peer}) ->
     logger:debug("UDP peer stopped: ~s", [esockd:format(Peer)]),
     {stop, normal, State};
-handle_info({datagram, Server, Packet}, State = #state{server = Server, peer = Peer, idle_timeout = IdleTimeout}) ->
+handle_info({datagram, Server, Packet}, State = #state{server = Server}) ->
-    Reply = handle_packet(Peer, Packet),
+    handle_datagram(Packet, State);
-    Server ! {datagram, Peer, Reply},
+handle_info({tcp, Socket, Data}, State) ->
-    {noreply, State, IdleTimeout};
+    handle_remote_data(Socket, Data, State);
+handle_info({tcp_closed, Socket}, State) ->
+    handle_remote_closed(Socket, State);
+handle_info({tcp_error, Socket, Reason}, State) ->
+    handle_remote_error(Socket, Reason, State);
 handle_info(timeout, State = #state{peer = Peer}) ->
     logger:debug("UDP peer idle timeout: ~s", [esockd:format(Peer)]),
     {stop, normal, State};
 handle_info(_Info, State = #state{idle_timeout = IdleTimeout}) ->
     {noreply, State, IdleTimeout}.
 
-terminate(_Reason, _State) ->
+terminate(_Reason, #state{streams = Streams}) ->
+    maps:foreach(fun(_StreamId, Stream) -> close_remote(Stream) end, Streams),
     ok.
 
 code_change(_OldVsn, State, _Extra) ->
     {ok, State}.
 
-handle_packet(_Peer, Packet) ->
+handle_datagram(Packet, State = #state{idle_timeout = IdleTimeout}) ->
-    Packet.
+    case relay_server_udp_protocol:decode(Packet) of
+        {ok, {open, StreamId, Payload}} ->
+            {noreply, handle_open(StreamId, Payload, State), IdleTimeout};
+        {ok, {data, StreamId, Payload}} ->
+            {noreply, handle_data(StreamId, Payload, State), IdleTimeout};
+        {ok, {close, StreamId, _Payload}} ->
+            {noreply, close_stream(StreamId, State), IdleTimeout};
+        {ok, {error, StreamId, Payload}} ->
+            logger:warning("UDP relay stream ~p client error: ~ts", [StreamId, Payload]),
+            {noreply, close_stream(StreamId, State), IdleTimeout};
+        {error, Reason} ->
+            logger:debug("UDP relay ignored invalid frame: ~p", [Reason]),
+            {noreply, State, IdleTimeout}
+    end.
+
+handle_open(StreamId, Payload, State = #state{streams = Streams}) ->
+    case maps:is_key(StreamId, Streams) of
+        true ->
+            send_error(StreamId, <<"stream already open">>, State),
+            State;
+        false ->
+            case relay_server_udp_protocol:decode_open_request(Payload) of
+                {ok, Request} ->
+                    open_stream(StreamId, Request, State);
+                {error, Reason} ->
+                    send_error(StreamId, format_error(Reason), State),
+                    State
+            end
+    end.
+
+open_stream(StreamId, #{host := Host, port := Port}, State) ->
+    case connect_remote(Host, Port, State#state.connect_timeout) of
+        {ok, Socket} ->
+            Stream = #stream{
+                id = StreamId,
+                socket = Socket,
+                host = Host,
+                port = Port
+            },
+            logger:debug("UDP relay stream ~p opened to ~ts:~p", [StreamId, Host, Port]),
+            put_stream(Stream, State);
+        {error, Reason} ->
+            Message = iolist_to_binary(io_lib:format("connect failed: ~p", [Reason])),
+            send_error(StreamId, Message, State),
+            State
+    end.
+
+handle_data(StreamId, Payload, State = #state{streams = Streams}) ->
+    case maps:get(StreamId, Streams, undefined) of
+        undefined ->
+            send_error(StreamId, <<"unknown stream">>, State),
+            State;
+        Stream ->
+            case send_remote(Stream, Payload) of
+                ok ->
+                    State;
+                {error, Reason} ->
+                    send_error(StreamId, format_error(Reason), State),
+                    close_stream(StreamId, State)
+            end
+    end.
+
+handle_remote_data(Socket, Data, State = #state{sockets = Sockets}) ->
+    case maps:get(Socket, Sockets, undefined) of
+        undefined ->
+            {noreply, State, State#state.idle_timeout};
+        StreamId ->
+            send_frame(data, StreamId, Data, State),
+            NState = case maps:get(StreamId, State#state.streams, undefined) of
+                undefined ->
+                    State;
+                Stream ->
+                    case set_active_once(Stream) of
+                        ok -> State;
+                        {error, _Reason} -> close_stream(StreamId, State)
+                    end
+            end,
+            {noreply, NState, NState#state.idle_timeout}
+    end.
+
+handle_remote_closed(Socket, State = #state{sockets = Sockets}) ->
+    case maps:get(Socket, Sockets, undefined) of
+        undefined ->
+            {noreply, State, State#state.idle_timeout};
+        StreamId ->
+            send_frame(close, StreamId, <<>>, State),
+            NState = remove_stream(StreamId, State),
+            {noreply, NState, NState#state.idle_timeout}
+    end.
+
+handle_remote_error(Socket, Reason, State = #state{sockets = Sockets}) ->
+    case maps:get(Socket, Sockets, undefined) of
+        undefined ->
+            {noreply, State, State#state.idle_timeout};
+        StreamId ->
+            send_error(StreamId, format_error(Reason), State),
+            NState = close_stream(StreamId, State),
+            {noreply, NState, NState#state.idle_timeout}
+    end.
+
+connect_remote(Host, Port, Timeout) ->
+    case gen_tcp:connect(binary_to_list(Host), Port, ?TCP_OPTIONS, Timeout) of
+        {ok, Socket} -> {ok, Socket};
+        {error, Reason} -> {error, Reason}
+    end.
+
+send_remote(#stream{socket = Socket}, Payload) ->
+    gen_tcp:send(Socket, Payload).
+
+set_active_once(#stream{socket = Socket}) ->
+    inet:setopts(Socket, [{active, once}]).
+
+put_stream(Stream = #stream{id = StreamId, socket = Socket}, State) ->
+    State#state{
+        streams = maps:put(StreamId, Stream, State#state.streams),
+        sockets = maps:put(Socket, StreamId, State#state.sockets)
+    }.
+
+close_stream(StreamId, State = #state{streams = Streams}) ->
+    case maps:get(StreamId, Streams, undefined) of
+        undefined ->
+            State;
+        Stream ->
+            close_remote(Stream),
+            remove_stream(StreamId, State)
+    end.
+
+remove_stream(StreamId, State = #state{streams = Streams, sockets = Sockets}) ->
+    case maps:get(StreamId, Streams, undefined) of
+        undefined ->
+            State;
+        #stream{socket = Socket} ->
+            State#state{
+                streams = maps:remove(StreamId, Streams),
+                sockets = maps:remove(Socket, Sockets)
+            }
+    end.
+
+close_remote(#stream{socket = Socket}) ->
+    gen_tcp:close(Socket).
+
+send_frame(Type, StreamId, Payload, #state{server = Server, peer = Peer}) ->
+    Server ! {datagram, Peer, relay_server_udp_protocol:encode(Type, StreamId, Payload)}.
+
+send_error(StreamId, Message, State) ->
+    send_frame(error, StreamId, Message, State).
+
+format_error(Reason) when is_binary(Reason) ->
+    Reason;
+format_error(Reason) ->
+    iolist_to_binary(io_lib:format("~p", [Reason])).

apps/relay_server/src/relay_server_udp_listener.erl

@@ -13,8 +13,9 @@
 
 -define(SERVER, ?MODULE).
 -define(DEFAULT_LISTENER, 'relay_server/udp').
--define(DEFAULT_LISTEN_ON, 7000).
+-define(DEFAULT_LISTEN_ON, 16380).
 -define(DEFAULT_IDLE_TIMEOUT, 30000).
+-define(DEFAULT_CONNECT_TIMEOUT, 5000).
 -define(DEFAULT_MAX_CONNECTIONS, 1024).
 -define(DEFAULT_UDP_OPTIONS, [binary, {reuseaddr, true}]).
 
@@ -34,9 +35,14 @@ init([]) ->
     Props = application:get_env(relay_server, udp_server, []),
 
     Listener = proplists:get_value(listener, Props, ?DEFAULT_LISTENER),
-    ListenOn = proplists:get_value(port, Props, ?DEFAULT_LISTEN_ON),
+    ListenOn = proplists:get_value(
+        listen_on,
+        Props,
+        proplists:get_value(port, Props, ?DEFAULT_LISTEN_ON)
+    ),
 
     IdleTimeout = proplists:get_value(idle_timeout, Props, ?DEFAULT_IDLE_TIMEOUT),
+    ConnectTimeout = proplists:get_value(connect_timeout, Props, ?DEFAULT_CONNECT_TIMEOUT),
     UdpOptions = proplists:get_value(udp_options, Props, ?DEFAULT_UDP_OPTIONS),
     MaxConnections = proplists:get_value(max_connections, Props, ?DEFAULT_MAX_CONNECTIONS),
     AccessRules = proplists:get_value(access_rules, Props, [{allow, all}]),
@@ -48,7 +54,7 @@ init([]) ->
     ],
     Opts = maybe_add(max_conn_rate, Props, BaseOpts),
 
-    MFA = {relay_server_udp_handler, start_link, [IdleTimeout]},
+    MFA = {relay_server_udp_handler, start_link, [IdleTimeout, ConnectTimeout]},
     case esockd:open_udp(Listener, ListenOn, Opts, MFA) of
         {ok, Pid} ->
             logger:info("UDP listener ~p started on ~s",
@@ -94,4 +100,4 @@ maybe_add(Key, Props, Opts) ->
     case proplists:get_value(Key, Props) of
         undefined -> Opts;
         Value -> [{Key, Value} | Opts]
-    end.
+    end.

apps/relay_server/src/relay_server_udp_protocol.erl

@@ -0,0 +1,72 @@
+%%%-------------------------------------------------------------------
+%% @doc RelayKit UDP relay frame codec.
+%% @end
+%%%-------------------------------------------------------------------
+
+-module(relay_server_udp_protocol).
+
+-export([decode/1, encode/3, decode_open_request/1]).
+
+-define(MAGIC, <<"RKP1">>).
+-define(VERSION, 1).
+-define(HEADER_LENGTH, 20).
+
+decode(<<"RKP1", ?VERSION:8, TypeNo:8, _Reserved:16, StreamId:64, PayloadLen:32, Rest/binary>>) ->
+    case {type(TypeNo), byte_size(Rest) >= PayloadLen} of
+        {undefined, _} ->
+            {error, unknown_frame_type};
+        {_, false} ->
+            {error, truncated_frame};
+        {Type, true} ->
+            <<Payload:PayloadLen/binary, _Trailing/binary>> = Rest,
+            {ok, {Type, StreamId, Payload}}
+    end;
+decode(Packet) when byte_size(Packet) < ?HEADER_LENGTH ->
+    {error, frame_too_short};
+decode(_Packet) ->
+    {error, bad_frame_header}.
+
+encode(Type, StreamId, Payload) when is_binary(Payload) ->
+    TypeNo = type_no(Type),
+    PayloadLen = byte_size(Payload),
+    <<"RKP1", ?VERSION:8, TypeNo:8, 0:16, StreamId:64, PayloadLen:32, Payload/binary>>.
+
+decode_open_request(Payload) ->
+    try json:decode(Payload) of
+        #{
+            <<"host">> := Host,
+            <<"port">> := Port,
+            <<"useTLS">> := UseTLS,
+            <<"method">> := Method
+        } when
+            is_binary(Host),
+            byte_size(Host) > 0,
+            is_integer(Port),
+            Port > 0,
+            Port =< 65535,
+            is_boolean(UseTLS),
+            is_binary(Method)
+        ->
+            {ok, #{
+                host => Host,
+                port => Port,
+                use_tls => UseTLS,
+                method => Method
+            }};
+        _Other ->
+            {error, invalid_open_request}
+    catch
+        _Class:_Reason ->
+            {error, invalid_open_request}
+    end.
+
+type(1) -> open;
+type(2) -> data;
+type(3) -> close;
+type(4) -> error;
+type(_) -> undefined.
+
+type_no(open) -> 1;
+type_no(data) -> 2;
+type_no(close) -> 3;
+type_no(error) -> 4.

config/sys.config

@@ -3,8 +3,9 @@
     {udp_server, [
       {enabled, true},
       {listener, 'relay_server/udp'},
-      {port, 7000},
+      {listen_on, 16380},
       {idle_timeout, 30000},
+      {connect_timeout, 5000},
       {max_connections, 1024},
       {udp_options, [binary, {reuseaddr, true}]}
     ]}
@@ -37,4 +38,4 @@
     ]}
   ]}
 
-].
+].