diff --git a/test/mtp_test_cmd_rpc.erl b/test/mtp_test_cmd_rpc.erl new file mode 100644 index 0000000..ef3db61 --- /dev/null +++ b/test/mtp_test_cmd_rpc.erl @@ -0,0 +1,40 @@ +%% @doc Callback module for mtp_test_middle_server that supports some more tricky commands +-module(mtp_test_cmd_rpc). +-export([call/3, + packet_to_term/1]). +-export([init/1, + handle_rpc/2]). + +call(M, F, Opts) -> + true = erlang:function_exported(M, F, 3), + term_to_packet({M, F, Opts}). + +term_to_packet(Term) -> + RespBin = term_to_binary(Term), + RespSize = byte_size(RespBin), + PadSize = case (RespSize rem 16) of + 0 -> 0; + Rem -> 16 - Rem + end, + Pad = binary:copy(<<0>>, PadSize), + <>. + +packet_to_term(<>) -> + binary_to_term(Term). + +init(_) -> + #{}. + +handle_rpc({data, ConnId, Req}, St) -> + {M, F, Opts} = packet_to_term(Req), + case M:F(Opts, ConnId, St) of + {reply, Resp, St1} -> + {rpc, {proxy_ans, ConnId, term_to_packet(Resp)}, St1}; + {return, What} -> + What + end; +handle_rpc({remote_closed, ConnId}, St) -> + is_integer(maps:get(ConnId, St, undefined)) + orelse error({unexpected_closed, ConnId}), + {noreply, St#{ConnId := tombstone}}. + diff --git a/test/mtp_test_datacenter.erl b/test/mtp_test_datacenter.erl index 41754bf..1a89e3c 100644 --- a/test/mtp_test_datacenter.erl +++ b/test/mtp_test_datacenter.erl @@ -16,19 +16,28 @@ -define(SECRET_PATH, "/getProxySecret"). -define(CONFIG_PATH, "/getProxyConfig"). +-type dc_conf() :: [{DcId :: integer(), + Ip :: inet:ip4_address(), + Port :: inet:port_number()}]. start_dc() -> Secret = crypto:strong_rand_bytes(128), DcConf = [{1, {127, 0, 0, 1}, 8888}], {ok, _Cfg} = start_dc(Secret, DcConf, #{}). +-spec start_dc(binary(), dc_conf(), #{}) -> {ok, #{}}. start_dc(Secret, DcConf, Acc) -> Cfg = dc_list_to_config(DcConf), {ok, Acc1} = start_config_server({127, 0, 0, 1}, 3333, Secret, Cfg, Acc), + RpcHandler = maps:get(rpc_handler, Acc, mtp_test_echo_rpc), Ids = [begin Id = {?MODULE, DcId}, - {ok, _Pid} = mtp_test_middle_server:start(Id, #{port => Port, ip => Ip, secret => Secret}), + {ok, _Pid} = mtp_test_middle_server:start( + Id, #{port => Port, + ip => Ip, + secret => Secret, + rpc_handler => RpcHandler}), Id end || {DcId, Ip, Port} <- DcConf], {ok, Acc1#{srv_ids => Ids}}. diff --git a/test/mtp_test_echo_rpc.erl b/test/mtp_test_echo_rpc.erl new file mode 100644 index 0000000..4e623c7 --- /dev/null +++ b/test/mtp_test_echo_rpc.erl @@ -0,0 +1,15 @@ +%% @doc simple callback module for mtp_test_middle_server that echoes received packets back +-module(mtp_test_echo_rpc). +-export([init/1, + handle_rpc/2]). + +init(_) -> + #{}. + +handle_rpc({data, ConnId, Data}, St) -> + Cnt = maps:get(ConnId, St, 0), + {rpc, {proxy_ans, ConnId, Data}, St#{ConnId => Cnt + 1}}; +handle_rpc({remote_closed, ConnId}, St) -> + is_integer(maps:get(ConnId, St)) + orelse error({unexpected_closed, ConnId}), + {noreply, St#{ConnId := tombstone}}. diff --git a/test/mtp_test_middle_server.erl b/test/mtp_test_middle_server.erl index 40c2967..1c8770c 100644 --- a/test/mtp_test_middle_server.erl +++ b/test/mtp_test_middle_server.erl @@ -28,12 +28,14 @@ cli_ts, sender_pid, peer_pid, - srv_nonce}). + srv_nonce, + rpc_handler}). -record(t_state, {sock, transport, codec, - clients = #{} :: #{}}). + rpc_handler, + rpc_handler_state}). -define(RPC_NONCE, 170,135,203,122). -define(RPC_HANDSHAKE, 245,238,130,118). @@ -42,7 +44,7 @@ %% -type state_name() :: wait_nonce | wait_handshake | on_tunnel. %% Api -start(Id, Opts) -> +start(Id, #{port := _, secret := _} = Opts) -> {ok, _} = application:ensure_all_started(ranch), ranch:start_listener( Id, ranch_tcp, @@ -69,10 +71,12 @@ ranch_init({Ref, Transport, Opts}) -> init({Socket, Transport, Opts}) -> Codec = mtp_codec:new(mtp_noop_codec, mtp_noop_codec:new(), mtp_full, mtp_full:new(-2, -2)), - {ok, wait_nonce, #hs_state{sock = Socket, - transport = Transport, - secret = maps:get(secret, Opts), - codec = Codec}}. + State = #hs_state{sock = Socket, + transport = Transport, + secret = maps:get(secret, Opts), + codec = Codec, + rpc_handler = maps:get(rpc_handler, Opts, mtp_test_echo_rpc)}, + {ok, wait_nonce, State}. callback_mode() -> state_functions. @@ -131,12 +135,14 @@ wait_handshake(info, {tcp, _Sock, TcpData}, Answer = mtp_rpc:encode_handshake({handshake, SenderPID, PeerPID}), {ok, #hs_state{sock = Sock, transport = Transport, - codec = Codec2}} = hs_send(Answer, S#hs_state{codec = Codec1}), + codec = Codec2, + rpc_handler = Handler}} = hs_send(Answer, S#hs_state{codec = Codec1}), {next_state, on_tunnel, activate(#t_state{sock = Sock, transport = Transport, codec = Codec2, - clients = #{}})}; + rpc_handler = Handler, + rpc_handler_state = Handler:init([])})}; wait_handshake(Type, Event, S) -> handle_event(Type, Event, ?FUNCTION_NAME, S). @@ -178,15 +184,18 @@ activate(#t_state{transport = Transport, sock = Sock} = S) -> ok = Transport:setopts(Sock, [{active, once}]), S. -handle_rpc({data, ConnId, Data}, #t_state{clients = Clients} = S) -> - %% Echo data back - %% TODO: interptet Data to power some test scenarios, eg, client might - %% ask to close it's connection - {ok, S1} = t_send(mtp_rpc:srv_encode_packet({proxy_ans, ConnId, Data}), S), - Cnt = maps:get(ConnId, Clients, 0), - %% Increment can fail if there is a tombstone for this client - S1#t_state{clients = Clients#{ConnId => Cnt + 1}}; -handle_rpc({remote_closed, ConnId}, #t_state{clients = Clients} = S) -> - is_integer(maps:get(ConnId, Clients)) - orelse error({unexpected_closed, ConnId}), - S#t_state{clients = Clients#{ConnId := tombstone}}. +handle_rpc(RPC, #t_state{rpc_handler = Handler, rpc_handler_state = HSt} = S) -> + case Handler:handle_rpc(RPC, HSt) of + {rpc, Response, HSt1} -> + {ok, S1} = t_send(mtp_rpc:srv_encode_packet(Response), + S#t_state{rpc_handler_state = HSt1}), + S1; + {rpc_multi, Responses, HSt1} -> + lists:foldl( + fun(Response, S1) -> + {ok, S2} = t_send(mtp_rpc:srv_encode_packet(Response), S1), + S2 + end, S#t_state{rpc_handler_state = HSt1}, Responses); + {noreply, HSt1} -> + S#t_state{rpc_handler_state = HSt1} + end.