From d6fb2a480fbc91b46d1edc769aef64dbcaffed60 Mon Sep 17 00:00:00 2001 From: anlicheng <244108715@qq.com> Date: Sat, 27 Jun 2026 15:13:50 +0800 Subject: [PATCH] fix handler --- README.md | 18 +- .../src/relay_server_udp_handler.erl | 196 ++++++++++++++++-- .../src/relay_server_udp_listener.erl | 14 +- .../src/relay_server_udp_protocol.erl | 72 +++++++ config/sys.config | 5 +- 5 files changed, 281 insertions(+), 24 deletions(-) create mode 100644 apps/relay_server/src/relay_server_udp_protocol.erl diff --git a/README.md b/README.md index 4f0447b..12f4816 100644 --- a/README.md +++ b/README.md @@ -13,10 +13,18 @@ Run $ rebar3 shell -The UDP listener is configured in `config/sys.config` and listens on port `7000` -by default. It echoes received datagrams. +The UDP listener is configured in `config/sys.config` and listens on port `16380` +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. diff --git a/apps/relay_server/src/relay_server_udp_handler.erl b/apps/relay_server/src/relay_server_udp_handler.erl index ef921c3..fb9a431 100644 --- a/apps/relay_server/src/relay_server_udp_handler.erl +++ b/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) -> - gen_server:start_link(?MODULE, [Transport, Peer, IdleTimeout], []). +start_link(Transport, Peer, IdleTimeout, ConnectTimeout) -> + 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}) -> - Reply = handle_packet(Peer, Packet), - Server ! {datagram, Peer, Reply}, - {noreply, State, IdleTimeout}; +handle_info({datagram, Server, Packet}, State = #state{server = Server}) -> + handle_datagram(Packet, State); +handle_info({tcp, Socket, Data}, State) -> + 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) -> - Packet. +handle_datagram(Packet, State = #state{idle_timeout = IdleTimeout}) -> + 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])). diff --git a/apps/relay_server/src/relay_server_udp_listener.erl b/apps/relay_server/src/relay_server_udp_listener.erl index 2b62a1a..8e8b088 100644 --- a/apps/relay_server/src/relay_server_udp_listener.erl +++ b/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. \ No newline at end of file + end. diff --git a/apps/relay_server/src/relay_server_udp_protocol.erl b/apps/relay_server/src/relay_server_udp_protocol.erl new file mode 100644 index 0000000..826c981 --- /dev/null +++ b/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} -> + <> = 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. diff --git a/config/sys.config b/config/sys.config index 048d9b8..af7d0d2 100644 --- a/config/sys.config +++ b/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 @@ ]} ]} -]. \ No newline at end of file +].