tcp relay in erlang

relay.app

{application, relay,
 [{description, "relay"},
  {vsn, "0.01"},
  {modules, [
    relay_app,
    relay_listener,
    relay_sup,
    relay_worker
  ]},
  {registered, []},
  {mod, {relay_app, []}},
  {env, []},
  {applications, [kernel, stdlib, sasl]}]}.
 

relay_app.erl

-module(relay_app).
-behaviour(application).
-export([start/2, stop/1]).

start(_Type, _StartArgs) ->
    relay_sup:start_link().

stop(_State) ->
    ok.

relay_listen.erl

-module(relay_listen).

-export([start_link/1, start/0, init/1]).

start_link(Port) ->
    proc_lib:start_link(?MODULE, init, [{self(), Port}]).

init({Parent, Port}) ->
    register(?MODULE, self()),
    proc_lib:init_ack(Parent, {ok, self()}),
    {ok, ListenSock} = gen_tcp:listen(Port, [{packet, 0}, {active, false}, {reuseaddr, true}]),
    accept(ListenSock).

accept(ListenSock) ->
    case gen_tcp:accept(ListenSock) of
        {ok, Client} ->
            {ok, _} = relay_worker:start(Client),
            accept(ListenSock);
        _Error ->
            ok
    end.

start() ->
    start_link(9000).

relay_sup.erl

-module(relay_sup).

-behaviour(supervisor).

-export([start_link/0, init/1]).

-define(SERVER, ?MODULE).

start_link() ->
    supervisor:start_link({local, ?SERVER}, ?MODULE, []).

init([]) ->
    RestartStrategy = one_for_one,
    MaxRestarts = 1000,
    MaxSecondsBetweenRestarts = 3600,

    SupFlags = {RestartStrategy, MaxRestarts, MaxSecondsBetweenRestarts},

    Restart = permanent,
    Shutdown = 2000,
    Type = worker,

    AChild = {relay_listen, {relay_listen, start_link, [9000]},
          Restart, Shutdown, Type, [relay_listen]},

    {ok, {SupFlags, [AChild]}}.

relay_worker.erl

-module(relay_worker).

-behaviour(gen_server).

-export([start/1]).

%% gen_server callbacks
-export([init/1, handle_cast/2, handle_info/2,
         terminate/2, code_change/3]).

-define(SERVER, ?MODULE).
-define(HOST, "127.0.0.1").
-define(PORT, 4567).

-record(state, {client, server}).

% spawn a worker for connection, assign controll to it, then tell the server it can go!

start(Client) ->
    {ok, Pid} = gen_server:start(?MODULE, Client, []),
    gen_tcp:controlling_process(Client, Pid),
    gen_server:cast(Pid, setup_socket),
    {ok, Pid}.

init(Client) ->
    {ok, #state{client=Client}}.

% setup_socket connects to the backend and tells both sockets to be active

handle_cast(setup_socket, #state{client=Client}=State) ->
    inet:setopts(Client, [{active, once}]),
    case gen_tcp:connect(?HOST, ?PORT, [{active, once}, {packet, 0}]) of
        {ok, Server} ->
            {noreply, #state{client=Client, server=Server}};
        Error ->
            {stop, io_lib:format("Relay exception: ~p~n", [Error]), State}
    end.

% handle connection termination

handle_info({tcp_closed, Socket}, #state{client=Client, server=Server}=State) ->
    case Socket of
        Client ->
            gen_tcp:close(Server);
        Server ->
            gen_tcp:close(Client)
    end,
    {stop, shutdown, State};


% shuttle data between sockets

% FIXME - optimize this by adding a buffer for the client (or have another process that spoon feeds)

handle_info({tcp, Client, Data}, #state{client=Client, server=Server}=State) ->
    gen_tcp:send(Server, Data),
    inet:setopts(Client, [{active, once}]),
    {noreply, State};

handle_info({tcp, Server, Data}, #state{client=Client, server=Server}=State) ->
    gen_tcp:send(Client, Data),
    inet:setopts(Server, [{active, once}]),
    {noreply, State}.

terminate(_Reason, _State) ->
    ok.

code_change(_OldVsn, State, _Extra) ->
    {ok, State}.