aboutsummaryrefslogtreecommitdiffstats
path: root/lib/ssh/test/ssh_relay.erl
diff options
context:
space:
mode:
Diffstat (limited to 'lib/ssh/test/ssh_relay.erl')
-rw-r--r--lib/ssh/test/ssh_relay.erl407
1 files changed, 407 insertions, 0 deletions
diff --git a/lib/ssh/test/ssh_relay.erl b/lib/ssh/test/ssh_relay.erl
new file mode 100644
index 0000000000..a4f2bad2e2
--- /dev/null
+++ b/lib/ssh/test/ssh_relay.erl
@@ -0,0 +1,407 @@
+%%%-------------------------------------------------------------------
+%%% @author Simon Cornish <[email protected]>
+%%% @copyright (C) 2015, Simon Cornish
+%%% @doc
+%%% Provide manipulatable TCP-level relaying for testing SSH
+%%% @end
+%%% Created : 7 May 2015 by Simon Cornish <[email protected]>
+%%%-------------------------------------------------------------------
+-module(ssh_relay).
+
+-behaviour(gen_server).
+
+%% API
+-export([start_link/4]).
+-export([stop/1]).
+-export([hold/4, release/2, release_next/3]).
+
+%% gen_server callbacks
+-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
+ terminate/2, code_change/3]).
+
+-record(hold, {
+ port,
+ n,
+ tmo,
+ tref,
+ q = []
+ }).
+
+-record(state, {
+ local_addr,
+ local_port,
+ peer_addr,
+ peer_port,
+ lpid,
+ local,
+ peer,
+ tx_hold,
+ rx_hold
+ }).
+
+-define(ACCEPT_TMO, 200).
+%%%===================================================================
+%%% API
+%%%===================================================================
+%%--------------------------------------------------------------------
+%% @doc
+%% Hold N (or 'all') messages in given direction.
+%% Messages will be released after the N+1th message or
+%% Tmo ms or 'infinity'
+%%
+%% Dir is 'tx' for direction local -> peer
+%% and 'rx' for direction peer -> local
+%%
+%% An Error, ealready, is returned if there is already a hold
+%% in the given direction
+%%
+%% @spec hold(Srv, Dir, N, Tmo) -> ok | {error, Error}
+%% @end
+%%--------------------------------------------------------------------
+hold(Srv, Dir, N, Tmo) ->
+ gen_server:call(Srv, {hold, Dir, N, Tmo}).
+
+%%--------------------------------------------------------------------
+%% @doc
+%% Release all held messages in given direction.
+%%
+%% An Error, enoent, is returned if there is no hold
+%% in the given direction
+%%
+%% @spec release(Srv, Dir) -> ok | {error, Error}
+%% @end
+%%--------------------------------------------------------------------
+release(Srv, Dir) ->
+ gen_server:call(Srv, {release, Dir}).
+
+%%--------------------------------------------------------------------
+%% @doc
+%% Release all held messages in given direction after the
+%% next message in the trigger direction
+%%
+%% An Error, enoent, is returned if there is no hold
+%% in the given direction
+%%
+%% @spec release_next(Srv, Dir, TriggerDir) -> ok | {error, Error}
+%% @end
+%%--------------------------------------------------------------------
+release_next(Srv, Dir, TriggerDir) ->
+ gen_server:call(Srv, {release_next, Dir, TriggerDir}).
+
+%%--------------------------------------------------------------------
+%% @doc
+%% Starts the server
+%%
+%% @spec start_link() -> {ok, Pid} | ignore | {error, Error}
+%% @end
+%%--------------------------------------------------------------------
+start_link(ListenAddr, ListenPort, PeerAddr, PeerPort) ->
+ gen_server:start_link(?MODULE, [ListenAddr, ListenPort, PeerAddr, PeerPort], []).
+
+stop(Srv) ->
+ unlink(Srv),
+ Srv ! stop.
+
+%%%===================================================================
+%%% gen_server callbacks
+%%%===================================================================
+
+%%--------------------------------------------------------------------
+%% @private
+%% @doc
+%% Initializes the server
+%%
+%% @spec init(Args) -> {ok, State} |
+%% {ok, State, Timeout} |
+%% ignore |
+%% {stop, Reason}
+%% @end
+%%--------------------------------------------------------------------
+init([ListenAddr, ListenPort, PeerAddr, PeerPort | Options]) ->
+ IfAddr = case ListenAddr of
+ {0,0,0,0} ->
+ [];
+ _ ->
+ [{ifaddr, ListenAddr}]
+ end,
+ case gen_tcp:listen(ListenPort, [{reuseaddr, true}, {backlog, 1}, {active, false}, binary | IfAddr]) of
+ {ok, LSock} ->
+ Parent = self(),
+ {LPid, _LMod} = spawn_monitor(fun() -> listen(Parent, LSock) end),
+ S = #state{local_addr = ListenAddr,
+ local_port = ListenPort,
+ lpid = LPid,
+ peer_addr = PeerAddr,
+ peer_port = PeerPort
+ },
+ {ok, S};
+ Error ->
+ {stop, Error}
+ end.
+
+%%--------------------------------------------------------------------
+%% @private
+%% @doc
+%% Handling call messages
+%%
+%% @spec handle_call(Request, From, State) ->
+%% {reply, Reply, State} |
+%% {reply, Reply, State, Timeout} |
+%% {noreply, State} |
+%% {noreply, State, Timeout} |
+%% {stop, Reason, Reply, State} |
+%% {stop, Reason, State}
+%% @end
+%%--------------------------------------------------------------------
+handle_call({hold, Dir, N, Tmo}, _From, State) ->
+ case Dir of
+ tx ->
+ do_hold(#state.tx_hold, State#state.peer, N, Tmo, State);
+ rx ->
+ do_hold(#state.rx_hold, State#state.local, N, Tmo, State);
+ _ ->
+ {reply, {error, einval}, State}
+ end;
+handle_call({release, Dir}, _From, State) ->
+ case Dir of
+ tx ->
+ do_release(#state.tx_hold, State);
+ rx ->
+ do_release(#state.rx_hold, State);
+ _ ->
+ {reply, {error, einval}, State}
+ end;
+handle_call({release_next, _Dir, _TriggerDir}, _From, State) ->
+ {reply, {error, nyi}, State};
+
+handle_call(Request, _From, State) ->
+ Reply = {unhandled, Request},
+ {reply, Reply, State}.
+
+%%--------------------------------------------------------------------
+%% @private
+%% @doc
+%% Handling cast messages
+%%
+%% @spec handle_cast(Msg, State) -> {noreply, State} |
+%% {noreply, State, Timeout} |
+%% {stop, Reason, State}
+%% @end
+%%--------------------------------------------------------------------
+handle_cast(_Msg, State) ->
+ {noreply, State}.
+
+%%--------------------------------------------------------------------
+%% @private
+%% @doc
+%% Handling all non call/cast messages
+%%
+%% @spec handle_info(Info, State) -> {noreply, State} |
+%% {noreply, State, Timeout} |
+%% {stop, Reason, State}
+%% @end
+%%--------------------------------------------------------------------
+handle_info({tcp, Local, Data}, S) when S#state.local == Local ->
+ S1 = do_local(Data, S),
+ {noreply, S1};
+
+handle_info({tcp_error, Local, Error}, S) when S#state.local == Local ->
+ S1 = do_local({error, Error}, S),
+ {noreply, S1};
+
+handle_info({tcp_closed, Local}, S) when S#state.local == Local ->
+ S1 = do_local(closed, S),
+ {noreply, S1};
+
+handle_info({tcp, Peer, Data}, S) when S#state.peer == Peer ->
+ S1 = do_peer(Data, S),
+ {noreply, S1};
+
+handle_info({tcp_error, Peer, Error}, S) when S#state.peer == Peer ->
+ S1 = do_peer({error, Error}, S),
+ {noreply, S1};
+
+handle_info({tcp_closed, Peer}, S) when S#state.peer == Peer ->
+ S1 = do_peer(closed, S),
+ {noreply, S1};
+
+handle_info({accept, Local}, S) ->
+ S1 = do_accept(Local, S),
+ {noreply, S1};
+
+handle_info({activate, Local}, State) ->
+ inet:setopts(Local, [{active, true}]),
+ {noreply, State};
+
+handle_info({release, Pos}, S) ->
+ {reply, _, S1} = do_release(Pos,S),
+ {noreply, S1};
+
+handle_info(stop, State) ->
+ {stop, normal, State};
+
+handle_info({'DOWN', _Ref, _process, LPid, Reason}, S) when S#state.lpid == LPid ->
+ io:format("Acceptor has finished: ~p~n", [Reason]),
+ {noreply, S};
+
+handle_info(_Info, State) ->
+ io:format("Unhandled info: ~p~n", [_Info]),
+ {noreply, State}.
+
+%%--------------------------------------------------------------------
+%% @private
+%% @doc
+%% This function is called by a gen_server when it is about to
+%% terminate. It should be the opposite of Module:init/1 and do any
+%% necessary cleaning up. When it returns, the gen_server terminates
+%% with Reason. The return value is ignored.
+%%
+%% @spec terminate(Reason, State) -> void()
+%% @end
+%%--------------------------------------------------------------------
+terminate(_Reason, _State) ->
+ ok.
+
+%%--------------------------------------------------------------------
+%% @private
+%% @doc
+%% Convert process state when code is changed
+%%
+%% @spec code_change(OldVsn, State, Extra) -> {ok, NewState}
+%% @end
+%%--------------------------------------------------------------------
+code_change(_OldVsn, State, _Extra) ->
+ {ok, State}.
+
+%%%===================================================================
+%%% Internal functions
+%%%===================================================================
+do_hold(Pos, _Port, _N, _Tmo, S) when element(Pos, S) /= undefined ->
+ {reply, {error, ealready}, S};
+do_hold(Pos, Port, N, Tmo, S) ->
+ TRef = if is_integer(Tmo) andalso Tmo > 0 ->
+ erlang:send_after(Tmo, self(), {release, Pos});
+ true ->
+ undefined
+ end,
+ Hold = #hold{port = Port, n = N, tmo = Tmo, tref = TRef},
+ {reply, ok, setelement(Pos, S, Hold)}.
+
+do_release(HPos, S) when element(HPos, S) == undefined ->
+ {reply, {error, enoent}, S};
+do_release(HPos, S) ->
+ #hold{port = Port, tref = TRef, q = Q} = element(HPos, S),
+ lists:foreach(fun(M) -> gen_tcp:send(Port, M), erlang:yield() end, Q),
+ catch erlang:cancel_timer(TRef),
+ receive
+ {release, HPos} -> ok
+ after 0 ->
+ ok
+ end,
+ {reply, ok, setelement(HPos, S, undefined)}.
+
+listen(Parent, LSock) ->
+ monitor(process, Parent),
+ do_listen(Parent, LSock).
+
+do_listen(Parent, LSock) ->
+ %% So annoying there is no select-like sematic for this
+ case gen_tcp:accept(LSock, ?ACCEPT_TMO) of
+ {ok, Sock} ->
+ Parent ! {accept, Sock},
+ gen_tcp:controlling_process(Sock, Parent),
+ Parent ! {activate, Sock},
+ do_flush(Parent, Sock),
+ gen_tcp:close(LSock);
+ {error, timeout} ->
+ receive
+ DOWN when element(1, DOWN) == 'DOWN' ->
+ ok;
+ stop ->
+ ok
+ after 1 ->
+ do_listen(Parent, LSock)
+ end;
+ Error ->
+ gen_tcp:close(LSock),
+ exit({accept,Error})
+ end.
+
+do_flush(Parent, Sock) ->
+ receive
+ {Tcp, Sock, _} = Msg when Tcp == tcp; Tcp == tcp_error ->
+ Parent ! Msg,
+ do_flush(Parent, Sock);
+ {tcp_closed, Sock} = Msg ->
+ Parent ! Msg,
+ do_flush(Parent, Sock)
+ after 1 ->
+ ok
+ end.
+
+do_accept(Local, S) ->
+ case gen_tcp:connect(S#state.peer_addr, S#state.peer_port, [{active, true}, binary]) of
+ {ok, Peer} ->
+ S#state{local = Local, peer = Peer};
+ Error ->
+ exit({connect, Error})
+ end.
+
+do_local(Data, S) when is_binary(Data) ->
+ TxH = S#state.tx_hold,
+ if TxH == undefined ->
+ gen_tcp:send(S#state.peer, Data),
+ S;
+ TxH#hold.n == 0 ->
+ lists:foreach(fun(M) -> gen_tcp:send(S#state.peer, M) end, TxH#hold.q),
+ gen_tcp:send(S#state.peer, Data),
+ catch erlang:cancel_timer(TxH#hold.tref),
+ TxP = #state.tx_hold,
+ receive
+ {release, TxP} ->
+ ok
+ after 0 ->
+ ok
+ end,
+ S#state{tx_hold = undefined};
+ true ->
+ Q = TxH#hold.q ++ [Data],
+ N = if is_integer(TxH#hold.n) ->
+ TxH#hold.n -1;
+ true ->
+ TxH#hold.n
+ end,
+ S#state{tx_hold = TxH#hold{q = Q, n = N}}
+ end;
+do_local(Error, _S) ->
+ exit({local, Error}).
+
+do_peer(Data, S) when is_binary(Data) ->
+ RxH = S#state.rx_hold,
+ if RxH == undefined ->
+ gen_tcp:send(S#state.local, Data),
+ S;
+ RxH#hold.n == 0 ->
+ lists:foreach(fun(M) -> gen_tcp:send(S#state.local, M) end, RxH#hold.q),
+ gen_tcp:send(S#state.local, Data),
+ catch erlang:cancel_timer(RxH#hold.tref),
+ RxP = #state.rx_hold,
+ receive
+ {release, RxP} ->
+ ok
+ after 0 ->
+ ok
+ end,
+ S#state{rx_hold = undefined};
+ true ->
+ Q = RxH#hold.q ++ [Data],
+ N = if is_integer(RxH#hold.n) ->
+ RxH#hold.n -1;
+ true ->
+ RxH#hold.n
+ end,
+ S#state{rx_hold = RxH#hold{q = Q, n = N}}
+ end;
+do_peer(Error, _S) ->
+ exit({peer, Error}).
+