%%   The contents of this file are subject to the Mozilla Public License

%%   Version 1.1 (the "License"); you may not use this file except in

%%   compliance with the License. You may obtain a copy of the License at

%%   http://www.mozilla.org/MPL/

%%

%%   Software distributed under the License is distributed on an "AS IS"

%%   basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See the

%%   License for the specific language governing rights and limitations

%%   under the License.

%%

%%   The Original Code is RabbitMQ.

%%

%%   The Initial Developers of the Original Code are LShift Ltd.,

%%   Cohesive Financial Technologies LLC., and Rabbit Technologies Ltd.

%%

%%   Portions created by LShift Ltd., Cohesive Financial

%%   Technologies LLC., and Rabbit Technologies Ltd. are Copyright (C) 

%%   2007 LShift Ltd., Cohesive Financial Technologies LLC., and Rabbit 

%%   Technologies Ltd.; 

%%

%%   All Rights Reserved.

%%

%%   Contributor(s): ______________________________________.

%%

%% rabbit_stomp implements STOMP messaging semantics, as per protocol

%% "version 1.0", at http://stomp.codehaus.org/Protocol

-module(rabbit_stomp).

-export([kickstart/0,

	 start/1,

	 listener_started/2, listener_stopped/2, start_client/1,

	 start_link/0, init/1, mainloop/1]).

-include("rabbit.hrl").

-include("rabbit_framing.hrl").

-include("stomp_frame.hrl").

-record(state, {socket, session_id, channel, parse_state, ticket}).

kickstart() ->

    {ok, StompListeners} = application:get_env(stomp_listeners),

    ok = start(StompListeners).

start(Listeners) ->

    {ok,_} = supervisor:start_child(

               rabbit_sup,

               {rabbit_stomp_client_sup,

                {tcp_client_sup, start_link,

                 [{local, rabbit_stomp_client_sup},

                  {rabbit_stomp,start_link,[]}]},

                transient, infinity, supervisor, [tcp_client_sup]}),

    ok = start_listeners(Listeners),

    ok.

start_listeners([]) ->

    ok;

start_listeners([{Host, Port} | More]) ->

    {IPAddress, Name} = rabbit_networking:check_tcp_listener_address(rabbit_stomp_listener_sup,

								     Host,

								     Port),

    {ok,_} = supervisor:start_child(

               rabbit_sup,

               {Name,

                {tcp_listener_sup, start_link,

		 [IPAddress, Port,

		  [{packet, raw},

		   {reuseaddr, true}],

		  {?MODULE, listener_started, []},

		  {?MODULE, listener_stopped, []},

		  {?MODULE, start_client, []}]},

                transient, infinity, supervisor, [tcp_listener_sup]}),

    start_listeners(More).

listener_started(_IPAddress, _Port) ->

    ok.

listener_stopped(_IPAddress, _Port) ->

    ok.

start_client(Sock) ->

    {ok, Child} = supervisor:start_child(rabbit_stomp_client_sup, []),

    ok = gen_tcp:controlling_process(Sock, Child),

    Child ! {go, Sock},

    Child.

start_link() ->

    {ok, proc_lib:spawn_link(?MODULE, init, [self()])}.

init(_Parent) ->

    receive

        {go, Sock} ->

	    ok = inet:setopts(Sock, [{active, true}]),

	    process_flag(trap_exit, true),

	    ?MODULE:mainloop(#state{socket = Sock,

				    channel = none,

				    parse_state = stomp_frame:initial_state()})

    end.

mainloop(State) ->

    receive

	E = {'EXIT', _Pid, _Reason} ->

	    handle_exit(E, State);

	{tcp, _Sock, Bytes} ->

	    process_received_bytes(Bytes, State);

	{tcp_closed, _Sock} ->

	    case State#state.channel of

		none ->

		    done;

		ChPid ->

		    rabbit_channel:shutdown(ChPid),

		    ?MODULE:mainloop(State)

	    end;

	{send_command, Command} ->

	    ?MODULE:mainloop(send_reply(Command, State));

	{send_command_and_notify, QPid, TxPid, Method, Content} ->

	    State1 = send_reply(Method, Content, State),

	    rabbit_amqqueue:notify_sent(QPid, TxPid),

	    ?MODULE:mainloop(State1);

	{send_command_and_shutdown, Command} ->

	    send_reply(Command, State),

	    done;

	shutdown ->

	    done;

	Data ->

	    error_logger:error_msg("Internal error: unknown STOMP Data: ~p~n", [Data]),

	    ?MODULE:mainloop(State)

    end.

simple_method_sync_rpc(Method, State0) ->

    State = send_method(Method, State0),

    receive

	E = {'EXIT', _Pid, _Reason} ->

	    handle_exit(E, State);

	{send_command, Reply} ->

	    {ok, Reply, State}

    end.

handle_exit({'EXIT', _Pid, {amqp, Code, Method}}, State) ->

    explain_amqp_death(Code, Method, State),

    done;

handle_exit({'EXIT', Pid, Reason}, State) ->

    send_error("Error", "Process ~p exited with reason:~n~p~n", [Pid, Reason], State),

    done.

process_received_bytes([], State) ->

    ?MODULE:mainloop(State);

process_received_bytes(Bytes, State = #state{parse_state = ParseState}) ->

    case stomp_frame:parse(Bytes, ParseState) of

	{more, ParseState1} ->

	    ?MODULE:mainloop(State#state{parse_state = ParseState1});

	{ok, Frame = #stomp_frame{command = Command}, Rest} ->

	    %% io:format("Frame: ~p~n", [Frame]),

	    case catch process_frame(Command, Frame,

				     State#state{parse_state = stomp_frame:initial_state()}) of

		{'EXIT', {amqp, Code, Method}} ->

		    explain_amqp_death(Code, Method, State),

		    done;

		{'EXIT', Reason} ->

		    send_error("Processing error", "~p~n", [Reason], State),

		    done;

		{ok, NewState} ->

		    process_received_bytes(Rest, NewState);

		stop ->

		    done

	    end;

	{error, Reason} ->

	    send_error("Invalid frame", "Could not parse frame: ~p~n", [Reason], State),

	    done

    end.

explain_amqp_death(Code, Method, State) ->

    send_error(atom_to_list(Code), "Method was ~p~n", [Method], State).

send_reply(#'channel.close_ok'{}, State) ->

    State;

send_reply(Command, State) ->

    error_logger:error_msg("STOMP Reply command unhandled: ~p~n", [Command]),

    State.

send_reply(#'basic.deliver'{consumer_tag = ConsumerTag,

			    delivery_tag = DeliveryTag,

			    exchange = Exchange,

			    routing_key = RoutingKey},

	   #content{properties = #'P_basic'{headers = Headers},

		    payload_fragments_rev = BodyFragmentsRev},

	   State = #state{session_id = SessionId}) ->

    send_frame("MESSAGE",

	       [{"destination", binary_to_list(RoutingKey)},

		{"exchange", binary_to_list(Exchange)},

		{"message-id", SessionId ++ "_" ++ integer_to_list(DeliveryTag)}]

	       ++ case ConsumerTag of

		      <<"Q_", _/binary>> ->

			  [];

		      <<"T_", Id/binary>> ->

			  [{"subscription", binary_to_list(Id)}]

		  end

	       ++ adhoc_convert_headers(case Headers of

					    undefined -> [];

					    _ -> Headers

					end),

	       lists:concat(lists:reverse(lists:map(fun erlang:binary_to_list/1,

						    BodyFragmentsRev))),

	       State);

send_reply(Command, Content, State) ->

    error_logger:error_msg("STOMP Reply command unhandled: ~p~n~p~n", [Command, Content]),

    State.

adhoc_convert_headers([]) ->

    [];

adhoc_convert_headers([{K, longstr, V} | Rest]) ->

    [{"X-" ++ binary_to_list(K), binary_to_list(V)} | adhoc_convert_headers(Rest)];

adhoc_convert_headers([{K, signedint, V} | Rest]) ->

    [{"X-" ++ binary_to_list(K), integer_to_list(V)} | adhoc_convert_headers(Rest)];

adhoc_convert_headers([_ | Rest]) ->

    adhoc_convert_headers(Rest).

send_frame(Frame, State = #state{socket = Sock}) ->

    %% io:format("Sending ~p~n", [Frame]),

    ok = gen_tcp:send(Sock, stomp_frame:serialize(Frame)),

    State.

send_frame(Command, Headers, Body, State) ->

    send_frame(#stomp_frame{command = Command,

			    headers = Headers,

			    body = Body},

	       State).

send_error(Message, Detail, State) ->

    error_logger:error_msg("STOMP error frame sent:~nMessage: ~p~nDetail: ~p~n",

			   [Message, Detail]),

    send_frame("ERROR", [{"message", Message}], Detail, State).

send_error(Message, Format, Args, State) ->

    send_error(Message, lists:flatten(io_lib:format(Format, Args)), State).

process_frame("CONNECT", Frame, State = #state{channel = none}) ->

    {ok, DefaultVHost} = application:get_env(default_vhost),

    do_login(stomp_frame:header(Frame, "login"),

	     stomp_frame:header(Frame, "passcode"),

	     stomp_frame:header(Frame, "virtual-host", binary_to_list(DefaultVHost)),

	     stomp_frame:header(Frame, "realm", "/data"),

	     State);

process_frame("DISCONNECT", _Frame, _State = #state{channel = none}) ->

    stop;

process_frame(_Command, _Frame, State = #state{channel = none}) ->

    {ok, send_error("Illegal command",

		    "You must log in using CONNECT first\n",

		    State)};

process_frame(Command, Frame, State) ->

    case process_command(Command, Frame, State) of

	{ok, State1} ->

	    {ok, case stomp_frame:header(Frame, "receipt") of

		     {ok, Id} ->

			 send_frame("RECEIPT", [{"receipt-id", Id}], "", State1);

		     not_found ->

			 State1

		 end};

	stop ->

	    stop

    end.

send_method(Method, State = #state{channel = ChPid}) ->

    ok = rabbit_channel:do(ChPid, Method),

    State.

send_method(Method, Properties, Body, State = #state{channel = ChPid}) ->

    ok = rabbit_channel:do(ChPid,

			   Method,

			   #content{class_id = 60, %% basic

				    properties = Properties,

				    properties_bin = none,

				    payload_fragments_rev = [list_to_binary(Body)]}),

    State.

do_login({ok, Login}, {ok, Passcode}, VirtualHost, Realm, State) ->

    U = rabbit_access_control:user_pass_login(list_to_binary(Login),

					      list_to_binary(Passcode)),

    ok = rabbit_access_control:check_vhost_access(U, list_to_binary(VirtualHost)),

    ChPid = 

	rabbit_channel:start_link(self(), self(), U#user.username, list_to_binary(VirtualHost)),

    {ok, #'channel.open_ok'{}, State1} =

	simple_method_sync_rpc(#'channel.open'{out_of_band = <<"">>},

			       State#state{channel = ChPid}),

    SessionId = rabbit_misc:string_guid("session"),

    {ok, #'access.request_ok'{ticket = Ticket}, State2} =

	simple_method_sync_rpc(#'access.request'{realm = list_to_binary(Realm),

						 exclusive = false,

						 passive = true,

						 active = true,

						 write = true,

						 read = true},

			       send_frame("CONNECTED",

					  [{"session", SessionId}],

					  "",

					  State1#state{session_id = SessionId})),

    {ok, State2#state{ticket = Ticket}};

do_login(_, _, _, _, State) ->

    {ok, send_error("Bad CONNECT", "Missing login or passcode header(s)\n", State)}.

user_header_key("X-" ++ UserKey) -> UserKey;

user_header_key(_) -> false.

make_string_table(_KeyFilter, []) -> [];

make_string_table(KeyFilter, [{K, V} | Rest]) ->

    case KeyFilter(K) of

	false ->

	    make_string_table(KeyFilter, Rest);

	NewK ->

	    [{list_to_binary(NewK), longstr, list_to_binary(V)}

	     | make_string_table(KeyFilter, Rest)]

    end.

transactional(Frame) ->

    case stomp_frame:header(Frame, "transaction") of

	{ok, Transaction} ->

	    {yes, Transaction};

	not_found ->

	    no

    end.

transactional_action(Frame, Name, Fun, State) ->

    case transactional(Frame) of

	{yes, Transaction} ->

	    Fun(Transaction, State);

	no ->

	    {ok, send_error("Missing transaction",

			    Name ++ " must include a 'transaction' header\n",

			    State)}

    end.

with_transaction(Transaction, State, Fun) ->

    case get({transaction, Transaction}) of

	undefined ->

	    {ok, send_error("Bad transaction",

			    "Invalid transaction identifier: ~p~n", [Transaction],

			    State)};

	Actions ->

	    Fun(Actions, State)

    end.

begin_transaction(Transaction, State) ->

    put({transaction, Transaction}, []),

    {ok, State}.

extend_transaction(Transaction, Action, State0) ->

    with_transaction(Transaction, State0,

		     fun (Actions, State) ->

			     put({transaction, Transaction}, [Action | Actions]),

			     {ok, State}

		     end).

commit_transaction(Transaction, State0) ->

    with_transaction(Transaction, State0,

		     fun (Actions, State) ->

			     FinalState = lists:foldr(fun perform_transaction_action/2,

						      State,

						      Actions),

			     erase({transaction, Transaction}),

			     {ok, FinalState}

		     end).

abort_transaction(Transaction, State0) ->

    with_transaction(Transaction, State0,

		     fun (_Actions, State) ->

			     erase({transaction, Transaction}),

			     {ok, State}

		     end).

perform_transaction_action({Method}, State) ->

    send_method(Method, State);

perform_transaction_action({Method, Props, Body}, State) ->

    send_method(Method, Props, Body, State).

process_command("BEGIN", Frame, State) ->

    transactional_action(Frame, "BEGIN", fun begin_transaction/2, State);

process_command("SEND",

		Frame = #stomp_frame{headers = Headers, body = Body},

		State = #state{ticket = Ticket}) ->

    case stomp_frame:header(Frame, "destination") of

	{ok, RoutingKeyStr} ->

	    ExchangeStr = stomp_frame:header(Frame, "exchange", ""),

	    Props = #'P_basic'{

	      content_type = stomp_frame:binary_header(Frame, "content-type", <<"text/plain">>),

	      headers = make_string_table(fun user_header_key/1, Headers),

	      delivery_mode = stomp_frame:integer_header(Frame, "delivery-mode", undefined),

	      priority = stomp_frame:integer_header(Frame, "priority", undefined),

	      correlation_id = stomp_frame:binary_header(Frame, "correlation-id", undefined),

	      reply_to = stomp_frame:binary_header(Frame, "reply-to", undefined),

	      message_id = stomp_frame:binary_header(Frame, "message-id", undefined)

	     },

	    Method = #'basic.publish'{ticket = Ticket,

				      exchange = list_to_binary(ExchangeStr),

				      routing_key = list_to_binary(RoutingKeyStr),

				      mandatory = false,

				      immediate = false},

	    case transactional(Frame) of

		{yes, Transaction} ->

		    extend_transaction(Transaction, {Method, Props, Body}, State);

		no ->

		    {ok, send_method(Method, Props, Body, State)}

	    end;

	not_found ->

	    {ok, send_error("Missing destination",

			    "SEND must include a 'destination', and optional 'exchange' header\n",

			    State)}

    end;

process_command("ACK", Frame, State = #state{session_id = SessionId}) ->

    case stomp_frame:header(Frame, "message-id") of

	{ok, IdStr} ->

	    IdPrefix = SessionId ++ "_",

	    case string:substr(IdStr, 1, length(IdPrefix)) of

		IdPrefix ->

		    DeliveryTag = list_to_integer(string:substr(IdStr, length(IdPrefix) + 1)),

		    Method = #'basic.ack'{delivery_tag = DeliveryTag,

					  multiple = false},

		    case transactional(Frame) of

			{yes, Transaction} ->

			    extend_transaction(Transaction, {Method}, State);

			no ->

			    {ok, send_method(Method, State)}

		    end;

		_ ->

		    rabbit_misc:die(command_invalid, 'basic.ack')

	    end;

	not_found ->

	    {ok, send_error("Missing message-id",

			    "ACK must include a 'message-id' header\n",

			    State)}

    end;

process_command("COMMIT", Frame, State) ->

    transactional_action(Frame, "COMMIT", fun commit_transaction/2, State);

process_command("ABORT", Frame, State) ->

    transactional_action(Frame, "ABORT", fun abort_transaction/2, State);

process_command("SUBSCRIBE",

		Frame = #stomp_frame{headers = Headers},

		State = #state{ticket = Ticket}) ->

    AckMode = case stomp_frame:header(Frame, "ack", "auto") of

		  "auto" -> auto;

		  "client" -> client

	      end,

    case stomp_frame:header(Frame, "destination") of

	{ok, QueueStr} ->

	    ConsumerTag = case stomp_frame:header(Frame, "id") of

			      {ok, Str} ->

				  list_to_binary("T_" ++ Str);

			      not_found ->

				  list_to_binary("Q_" ++ QueueStr)

			  end,

	    Queue = list_to_binary(QueueStr),

	    {ok, send_method(#'basic.consume'{ticket = Ticket,

					      queue = Queue,

					      consumer_tag = ConsumerTag,

					      no_local = false,

					      no_ack = (AckMode == auto),

					      exclusive = false,

					      nowait = true},

			     send_method(#'queue.declare'{ticket = Ticket,

							  queue = Queue,

							  passive = false,

							  durable = false,

							  exclusive = falxse,

							  auto_delete = true,

							  nowait = true,

							  arguments =

							    make_string_table(fun user_header_key/1,

									      Headers)},

					 State))};

	not_found ->

	    {ok, send_error("Missing destination",

			    "SUBSCRIBE must include a 'destination' header\n",

			    State)}

    end;

process_command("UNSUBSCRIBE", Frame, State) ->

    ConsumerTag = case stomp_frame:header(Frame, "id") of

		      {ok, IdStr} ->

			  list_to_binary("T_" ++ IdStr);

		      not_found ->

			  case stomp_frame:header(Frame, "destination") of

			      {ok, QueueStr} ->

				  list_to_binary("Q_" ++ QueueStr);

			      not_found ->

				  missing

			  end

		  end,

    if

	ConsumerTag == missing ->

	    {ok, send_error("Missing destination or id",

			    "UNSUBSCRIBE must include a 'destination' or 'id' header\n",

			    State)};

	true ->

	    {ok, send_method(#'basic.cancel'{consumer_tag = ConsumerTag,

					     nowait = true},

			    State)}

    end;

process_command("DISCONNECT", _Frame, State) ->

    {ok, send_method(#'channel.close'{reply_code = 200, reply_text = <<"">>,

				      class_id = 0, method_id = 0}, State)};

process_command(Command, _Frame, State) ->

    {ok, send_error("Bad command",

		    "Could not interpret command " ++ Command ++ "\n",

		    State)}.