-module(bouncer).
-export([bouncer/1]).

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% A simple TCP bouncer. bouncer_task opens a listening socket, and
%% accept_loop accepts connections. When a connection arrives, it is handled
%% by handle_client, that relays data between endpoints.
%%
%% This code is released under GPL and is (C) Matteo Cicuttin 2009
%% As always, this code is the result of 10 minutes of work, so expect lots of
%% bugs :)
%%

%% Some options
-define(TCP_OPTIONS,[binary, {packet, 0}, {active, false}, {reuseaddr, true}]).
-define(TCP6_OPTIONS,[binary, {packet, 0}, {active, false}, {reuseaddr, true}, 
inet6]).

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Relay data between two sockets
relay_data(FromSocket, ToSocket) ->
    case gen_tcp:recv(FromSocket, 0) of
        {ok, Packet} ->
            gen_tcp:send(ToSocket, Packet),
            relay_data(FromSocket, ToSocket);
        {error, Reason} ->
            exit(Reason)
    end.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Handle a client: two processes are spawned to move data between endpoints,
%% one for each direction.
handle_client(LocalSocket, {RemoteHost, RemotePort}) ->
    case gen_tcp:connect(RemoteHost, RemotePort, ?TCP_OPTIONS) of
        {ok, RemoteSocket} ->
            %% Trap exit signals, so we can clean up things
            process_flag(trap_exit,true),
            %% Spawn processes
            PidLR = spawn( fun() -> relay_data(LocalSocket, RemoteSocket) end ),
            link(PidLR),
            PidRL = spawn( fun() -> relay_data(RemoteSocket, LocalSocket) end ),
            link(PidRL),
            %% Wait their termination
            receive
                {'EXIT', _, Why} ->
                    %% we got an exit signal from our spawned processes,
                    %% cleanup and shutdown
                    io:format("handle_client: relay_data died: ~s~n", [Why]),
                    inet:close(LocalSocket),
                    inet:close(RemoteSocket),
                    exit(Why)
            end;
        {error, Reason} ->
            %% For some reason we cannot connect to the remote host
            inet:close(LocalSocket),
            io:format("handle_client: ~s~n", [Reason]),
            exit(Reason)
    end.
    
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Accept clients. When a connection is received, handle_client is spawned
%% and data is bounced between endpoints.
accept_loop(ListenSocket, RemoteParams) ->
    case gen_tcp:accept(ListenSocket) of
        {ok, Socket} ->
            {ok, {Address, _}} = inet:peername(Socket),
            io:format("accept_loop: accepting a new client from ~p~n", [Address]),
            spawn(fun() -> handle_client(Socket, RemoteParams) end),
            accept_loop(ListenSocket, RemoteParams);
        {error, Reason} ->
            io:format("accept_loop: accept failed: ~s~n", [Reason]),
            accept_loop(ListenSocket, RemoteParams)
    end.
            

bouncer_task({LocalPort, RemoteHost, RemotePort}) ->
    io:format("bouncer_task: Local port: ~w Remote: ~s:~w~n", [LocalPort, RemoteHost, RemotePort]),
    case gen_tcp:listen(LocalPort,?TCP_OPTIONS) of
        {ok, ListenSocket} ->
            accept_loop(ListenSocket, {RemoteHost, RemotePort});
        {error, Reason} ->
            io:format("bouncer_task exiting: ~w Remote: ~s:~w~n", [LocalPort, RemoteHost, RemotePort]),
            exit(Reason)
    end.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Spawn BNC threads
launchbouncers([]) -> 
    receive
        {'EXIT', _, Reason} ->
            io:format("launchbouncers: a bouncer is died: ~s~n", [Reason])
    end;
launchbouncers([BncParam|BncParams]) ->
    Pid = spawn(fun() -> bouncer_task(BncParam) end),
    link(Pid),
    launchbouncers(BncParams).
    
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Main entry point. Takes a list of bounces and spawns bouncer threads
%%
%% Example: bouncer:bouncer([{1234,"www.google.com",80}]).
bouncer([]) ->
    io:format("No bounces, exiting~n");
bouncer([BP|BPs]) ->
    launchbouncers([BP|BPs]).