%% This Source Code Form is subject to the terms of the Mozilla Public
%% License, v. 2.0. If a copy of the MPL was not distributed with this
%% file, You can obtain one at https://mozilla.org/MPL/2.0/.
%%
%% Copyright (c) 2007-2026 Broadcom. All Rights Reserved. The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries. All rights reserved.
%%

-module(rabbit_web_stomp_handler).
-behaviour(cowboy_websocket).
-behaviour(cowboy_sub_protocol).

-include_lib("kernel/include/logger.hrl").
-include_lib("rabbitmq_stomp/include/rabbit_stomp.hrl").
-include_lib("rabbitmq_stomp/include/rabbit_stomp_frame.hrl").
-include_lib("amqp_client/include/amqp_client.hrl").
-include_lib("rabbit_common/include/logging.hrl").

%% Websocket.
-export([
    init/2,
    websocket_init/1,
    websocket_handle/2,
    websocket_info/2,
    terminate/3
]).
-export([close_connection/2]).

%% cowboy_sub_protocol
-export([upgrade/4,
         upgrade/5,
         takeover/7]).

-record(state, {
    frame_type,
    heartbeat_mode,
    heartbeat,
    heartbeat_sup,
    parse_state,
    proc_state,
    state,
    conserve_resources,
    socket,
    peername,
    auth_hd,
    stats_timer,
    connection
}).

-define(APP, rabbitmq_web_stomp).

%% cowboy_sub_protcol
upgrade(Req, Env, Handler, HandlerState) ->
    upgrade(Req, Env, Handler, HandlerState, #{}).

%% We create a proxy socket for HTTP/2 even if no proxy was used,
%% and add a special field 'http_version' to indicate this is HTTP/2.
upgrade(Req=#{version := 'HTTP/2', pid := Parent, peer := Peer, sock := Sock},
        Env, Handler, HandlerState, Opts) ->
    %% Cowboy doesn't expose the socket when HTTP/2 is used.
    %% We take it directly from the connection's state.
    %%
    %% @todo Ideally we would not need the real socket for
    %%       normal operations. But we currently need it for
    %%       the heartbeat processes to do their job. In the
    %%       future we should not rely on those processes
    %%       and instead do the heartbeating directly in the
    %%       Websocket handler.
    RealSocket = element(4,element(1,sys:get_state(Parent))),
    ProxyInfo = case Req of
        #{proxy_header := ProxyHeader} ->
            ProxyHeader#{http_version => 'HTTP/2'};
        _ ->
            {SrcAddr, SrcPort} = Peer,
            {DestAddr, DestPort} = Sock,
            #{
                http_version => 'HTTP/2',
                src_address => SrcAddr,
                src_port => SrcPort,
                dest_address => DestAddr,
                dest_port => DestPort
            }
    end,
    ProxySocket = {rabbit_proxy_socket, RealSocket, ProxyInfo},
    cowboy_websocket:upgrade(Req, Env, Handler, HandlerState#state{socket = ProxySocket}, Opts);
upgrade(Req, Env, Handler, HandlerState, Opts) ->
    cowboy_websocket:upgrade(Req, Env, Handler, HandlerState, Opts).

%% This is only called for HTTP/1.1.
takeover(Parent, Ref, Socket, Transport, Opts, Buffer, {Handler, HandlerState}) ->
    Sock = case HandlerState#state.socket of
               undefined ->
                   Socket;
               ProxyInfo ->
                   {rabbit_proxy_socket, Socket, ProxyInfo}
           end,
    Env0 = maps:get(env, Opts, #{}),
    Env = Env0#{socket => Sock},
    cowboy_websocket:takeover(Parent, Ref, Socket, Transport, maps:put(env, Env, Opts), Buffer,
                              {Handler, HandlerState#state{socket = Sock}}).

%% Websocket.
init(Req0, Opts) ->
    case check_origin(Req0) of
        ok ->
            logger:set_process_metadata(#{domain => ?RMQLOG_DOMAIN_CONN}),
            {PeerAddr, _PeerPort} = maps:get(peer, Req0),
            {_, KeepaliveSup} = lists:keyfind(keepalive_sup, 1, Opts),
            SockInfo = maps:get(proxy_header, Req0, undefined),
            Req = case cowboy_req:parse_header(<<"sec-websocket-protocol">>, Req0) of
                undefined  -> Req0;
                Protocols ->
                    case filter_stomp_protocols(Protocols) of
                        [] -> Req0;
                        [StompProtocol|_] ->
                            cowboy_req:set_resp_header(<<"sec-websocket-protocol">>,
                                StompProtocol, Req0)
                    end
            end,
            WsOpts0 = proplists:get_value(ws_opts, Opts, #{}),
            MaxFrameSize = application:get_env(
                rabbitmq_stomp, max_frame_size_unauthenticated,
                ?DEFAULT_MAX_FRAME_SIZE_UNAUTHENTICATED) + 4096,
            WsOpts = maps:merge(#{compress => true,
                                   max_frame_size => MaxFrameSize}, WsOpts0),
            {?MODULE, Req, #state{
                frame_type         = proplists:get_value(type, Opts, text),
                heartbeat_sup      = KeepaliveSup,
                heartbeat          = {none, none},
                heartbeat_mode     = heartbeat,
                state              = running,
                conserve_resources = false,
                socket             = SockInfo,
                peername           = PeerAddr,
                auth_hd            = cowboy_req:header(<<"authorization">>, Req)
            }, WsOpts#{data_delivery => relay}};
        {error, origin_not_allowed} ->
            ?LOG_WARNING("Web STOMP: WebSocket connection rejected, "
                         "origin not in allow_origins: ~tp",
                         [cowboy_req:header(<<"origin">>, Req0)]),
            {ok,
             cowboy_req:reply(403, #{<<"connection">> => <<"close">>}, Req0),
             #state{}}
    end.

websocket_init(State) ->
    process_flag(trap_exit, true),
    rabbit_access_control:set_max_heap_size_unauthenticated(rabbitmq_web_stomp),
    {ok, ProcessorState} = init_processor_state(State),
    LoginTimeout = application:get_env(rabbitmq_stomp, login_timeout, 10_000),
    erlang:send_after(LoginTimeout, self(), login_timeout),
    {ok, rabbit_event:init_stats_timer(
           State#state{proc_state     = ProcessorState,
                       parse_state    = rabbit_stomp_frame:initial_state()},
           #state.stats_timer)}.

-spec close_connection(pid(), string()) -> 'ok'.
close_connection(Pid, Reason) ->
    ?LOG_INFO("Web STOMP: will terminate connection process ~tp, reason: ~ts",
                               [Pid, Reason]),
    try
        sys:terminate(Pid, Reason)
    catch
        exit:{noproc, _} -> ok
    end.

init_processor_state(#state{socket=Sock, peername=PeerAddr, auth_hd=AuthHd}) ->
    Self = self(),
    SendFun = fun(Data) ->
                      Self ! {send, Data},
                      ok
              end,

    SSLLogin = application:get_env(rabbitmq_stomp, ssl_cert_login, false),
    StompConfig0 = #stomp_configuration{ssl_cert_login = SSLLogin, implicit_connect = false},
    UseHTTPAuth = application:get_env(rabbitmq_web_stomp, use_http_auth, false),
    UserConfig = application:get_env(rabbitmq_stomp, default_user, undefined),
    StompConfig1 = rabbit_stomp:parse_default_user(UserConfig, StompConfig0),
    StompConfig2 = case UseHTTPAuth of
        true ->
            case AuthHd of
                undefined ->
                    %% We fall back to the default STOMP credentials.
                    StompConfig1#stomp_configuration{force_default_creds = true};
                _ ->
                    {basic, HTTPLogin, HTTPPassCode}
                        = cow_http_hd:parse_authorization(AuthHd),
                    StompConfig0#stomp_configuration{
                      default_login = HTTPLogin,
                      default_passcode = HTTPPassCode,
                      force_default_creds = true}
            end;
        false ->
            StompConfig1
    end,

    AdapterInfo = amqp_connection:socket_adapter_info(Sock, {'Web STOMP', 0}),
    RealSocket = rabbit_net:unwrap_socket(Sock),
    LoginNameFromCertificate = rabbit_stomp_reader:ssl_login_name(RealSocket, StompConfig2),
    ProcessorState = rabbit_stomp_processor:initial_state(
        StompConfig2,
        {SendFun, AdapterInfo, LoginNameFromCertificate, PeerAddr}),
    {ok, ProcessorState}.

websocket_handle({text, Data}, State) ->
    handle_data(Data, State);
websocket_handle({binary, Data}, State) ->
    handle_data(Data, State);
websocket_handle(_Frame, State) ->
    {ok, State}.

websocket_info({send, Msg}, State=#state{frame_type=FrameType}) ->
    {reply, {FrameType, Msg}, State};

websocket_info({conserve_resources, Conserve}, State) ->
    NewState = State#state{conserve_resources = Conserve},
    handle_credits(control_throttle(NewState));
websocket_info({bump_credit, Msg}, State) ->
    credit_flow:handle_bump_msg(Msg),
    handle_credits(control_throttle(State));

websocket_info(#'basic.consume_ok'{}, State) ->
    {ok, State};
websocket_info(#'basic.cancel_ok'{}, State) ->
    {ok, State};
websocket_info(#'basic.ack'{delivery_tag = Tag, multiple = IsMulti},
               State=#state{ proc_state = ProcState0 }) ->
    ProcState = rabbit_stomp_processor:flush_pending_receipts(Tag,
                                                              IsMulti,
                                                              ProcState0),
    {ok, State#state{ proc_state = ProcState }};
websocket_info({Delivery = #'basic.deliver'{},
               #amqp_msg{props = Props, payload = Payload},
               DeliveryCtx},
               State=#state{ proc_state = ProcState0 }) ->
    ProcState = rabbit_stomp_processor:send_delivery(Delivery,
                                                     Props,
                                                     Payload,
                                                     DeliveryCtx,
                                                     ProcState0),
    {ok, State#state{ proc_state = ProcState }};
websocket_info({Delivery = #'basic.deliver'{},
               #amqp_msg{props = Props, payload = Payload}},
               State=#state{ proc_state = ProcState0 }) ->
    ProcState = rabbit_stomp_processor:send_delivery(Delivery,
                                                     Props,
                                                     Payload,
                                                     undefined,
                                                     ProcState0),
    {ok, State#state{ proc_state = ProcState }};
websocket_info(#'basic.cancel'{consumer_tag = Ctag},
               State=#state{ proc_state = ProcState0 }) ->
    case rabbit_stomp_processor:cancel_consumer(Ctag, ProcState0) of
      {ok, ProcState, _Connection} ->
        {ok, State#state{ proc_state = ProcState }};
      {stop, _Reason, ProcState} ->
        stop(State#state{ proc_state = ProcState })
    end;

websocket_info({start_heartbeats, _},
               State = #state{heartbeat_mode = no_heartbeat}) ->
    {ok, State};

websocket_info({start_heartbeats, {0, 0}}, State) ->
    {ok, State};
websocket_info({start_heartbeats, {SendTimeout, ReceiveTimeout}},
               State = #state{socket         = Sock,
                              heartbeat_sup  = SupPid,
                              heartbeat_mode = heartbeat}) ->
    Self = self(),
    SendFun = fun () -> Self ! {send, <<$\n>>}, ok end,
    ReceiveFun = fun() -> Self ! client_timeout end,
    Heartbeat = case Sock of
        {rabbit_proxy_socket, RealSocket, #{http_version := 'HTTP/2'}} ->
            %% HTTP/2 Websocket may have multiple Websocket sessions
            %% on a single connection (this can happen for example
            %% when refreshing a page). As a result we need to attach
            %% the heartbeat processes to the session. We do this via
            %% a link for now. @todo In the future we will have a
            %% mechanism in Cowboy to attach them to the stream.
            {ok, Sender} = rabbit_heartbeat:start_heartbeat_sender(RealSocket,
                SendTimeout, SendFun, {heartbeat_sender, unknown}),
            {ok, Receiver} = rabbit_heartbeat:start_heartbeat_receiver(RealSocket,
                ReceiveTimeout, ReceiveFun, {heartbeat_receiver, unknown}),
            {Sender, Receiver};
        _ ->
            rabbit_heartbeat:start(SupPid, Sock, SendTimeout,
                                   SendFun, ReceiveTimeout, ReceiveFun)
    end,
    {ok, State#state{heartbeat = Heartbeat}};
websocket_info(client_timeout, State) ->
    stop(State);

%%----------------------------------------------------------------------------
websocket_info({'EXIT', From, Reason},
               State=#state{ proc_state = ProcState0 }) ->
  case rabbit_stomp_processor:handle_exit(From, Reason, ProcState0) of
    {stop, _Reason, ProcState} ->
        stop(State#state{ proc_state = ProcState });
    unknown_exit ->
        %% Allow the server to send remaining error messages
        self() ! close_websocket,
        {ok, State}
  end;
websocket_info(close_websocket, State) ->
    stop(State);

%%----------------------------------------------------------------------------

websocket_info(emit_stats, State) ->
    {ok, emit_stats(State)};

websocket_info(increase_max_frame_size, State) ->
    MaxFrameSize = application:get_env(
        rabbitmq_stomp, max_frame_size, ?DEFAULT_MAX_FRAME_SIZE) + 4096,
    {[{set_options, #{max_frame_size => MaxFrameSize}}], State};

websocket_info(login_timeout, State = #state{connection = C})
  when C =:= none; C =:= undefined ->
    ?LOG_ERROR("Web STOMP: closing connection (login timeout)"),
    stop(State);
websocket_info(login_timeout, State) ->
    {ok, State};

websocket_info(Msg, State) ->
    ?LOG_INFO("Web STOMP: unexpected message ~tp",
                    [Msg]),
    {ok, State}.

terminate(_Reason, _Req, State = #state{proc_state = undefined}) ->
    terminate_heartbeaters(State);
terminate(_Reason, _Req, State = #state{proc_state = ProcState}) ->
    _ = rabbit_stomp_processor:flush_and_die(ProcState),
    terminate_heartbeaters(State).

terminate_heartbeaters(#state{heartbeat = {none, none}}) ->
    ok;
terminate_heartbeaters(#state{heartbeat = {SPid, RPid}}) ->
    %% Don't assume that these processes are still alive.
    safe_terminate(SPid),
    safe_terminate(RPid),
    ok.

safe_terminate(none) ->
    ok;
safe_terminate(Pid) ->
    try
        sys:terminate(Pid, shutdown)
    catch
        exit:{noproc, _} -> ok
    end.

%%----------------------------------------------------------------------------

check_origin(Req) ->
    case application:get_env(?APP, allow_origins, []) of
        [] ->
            ok;
        AllowedOrigins ->
            case cowboy_req:header(<<"origin">>, Req) of
                undefined ->
                    ok;
                Origin ->
                    case lists:member(binary_to_list(Origin), AllowedOrigins) of
                        true -> ok;
                        false -> {error, origin_not_allowed}
                    end
            end
    end.

%% The protocols v10.stomp, v11.stomp and v12.stomp are registered
%% at IANA: https://www.iana.org/assignments/websocket/websocket.xhtml

filter_stomp_protocols(Protocols) ->
    lists:reverse(lists:sort(lists:filter(
        fun(<< "v1", C, ".stomp">>)
            when C =:= $2; C =:= $1; C =:= $0 -> true;
           (_) ->
            false
        end,
        Protocols))).

%%----------------------------------------------------------------------------

handle_data(Data, State0) ->
    case handle_data1(Data, State0) of
        {ok, State1 = #state{state = blocked}} ->
            {[{active, false}], State1};
        {error, Error0} ->
            Error1 = rabbit_misc:format("~tp", [Error0]),
            ?LOG_ERROR("STOMP detected framing error '~ts'", [Error1]),
            stop(State0, 1007, Error1);
        Other ->
            Other
    end.

handle_data1(<<>>, State) ->
    {ok, ensure_stats_timer(State)};
handle_data1(Bytes, State = #state{proc_state  = ProcState,
                                   parse_state = ParseState,
                                   connection  = OldConn}) ->
    case rabbit_stomp_frame:parse(Bytes, ParseState) of
        {more, ParseState1} ->
            {ok, ensure_stats_timer(State#state{ parse_state = ParseState1 })};
        {ok, Frame, Rest} ->
            case rabbit_stomp_processor:process_frame(Frame, ProcState) of
                {ok, ProcState1, ConnPid} ->
                    maybe_increase_max_frame_size(OldConn, ConnPid),
                    ParseState1 = rabbit_stomp_frame:initial_state(),
                    State1 = maybe_block(State, Frame),
                    handle_data1(
                      Rest,
                      State1 #state{ parse_state = ParseState1,
                                     proc_state  = ProcState1,
                                     connection  = ConnPid });
                {stop, _Reason, ProcState1} ->
                    %% do not exit here immediately, because we need to wait for messages eventually enqueued by process_request
                    self() ! close_websocket,
                    {ok, State#state{ proc_state = ProcState1 }}
            end;
        Other ->
            Other
    end.

maybe_increase_max_frame_size(OldConn, ConnPid)
  when (OldConn =:= none orelse OldConn =:= undefined) andalso
       is_pid(ConnPid) ->
    self() ! increase_max_frame_size;
maybe_increase_max_frame_size(_, _) ->
    ok.

maybe_block(State = #state{state = blocking, heartbeat = Heartbeat},
            #stomp_frame{command = "SEND"}) ->
    rabbit_heartbeat:pause_monitor(Heartbeat),
    State#state{state = blocked};
maybe_block(State, _) ->
    State.

stop(State) ->
    stop(State, 1000, "STOMP died").

stop(State = #state{proc_state = ProcState}, CloseCode, Error0) ->
    maybe_emit_stats(State),
    _ = rabbit_stomp_processor:flush_and_die(ProcState),
    Error1 = rabbit_data_coercion:to_binary(Error0),
    {[{close, CloseCode, Error1}], State}.

%%----------------------------------------------------------------------------

handle_credits(State0) ->
    case control_throttle(State0) of
        State = #state{state = running} ->
            {[{active, true}], State};
        State ->
            {ok, State}
    end.

control_throttle(State = #state{state              = CS,
                                conserve_resources = Mem}) ->
    case {CS, Mem orelse credit_flow:blocked()} of
        {running,   true} -> blocking(State);
        {blocking, false} -> running(State);
        {blocked,  false} -> running(State);
        {_,            _} -> State
    end.

blocking(State) ->
    State#state{state = blocking}.

running(State = #state{heartbeat=Heartbeat}) ->
    rabbit_heartbeat:resume_monitor(Heartbeat),
    State#state{state = running}.

%%----------------------------------------------------------------------------

ensure_stats_timer(State) ->
    rabbit_event:ensure_stats_timer(State, #state.stats_timer, emit_stats).

maybe_emit_stats(State) ->
    rabbit_event:if_enabled(State, #state.stats_timer,
                                fun() -> emit_stats(State) end).

emit_stats(State=#state{connection = C}) when C == none; C == undefined ->
    %% Avoid emitting stats on terminate when the connection has not yet been
    %% established, as this causes orphan entries on the stats database
    State1 = rabbit_event:reset_stats_timer(State, #state.stats_timer),
    State1;
emit_stats(State=#state{socket=Sock, state=RunningState, connection=Conn}) ->
    SockInfos = case rabbit_net:getstat(Sock,
            [recv_oct, recv_cnt, send_oct, send_cnt, send_pend]) of
        {ok,    SI} -> SI;
        {error,  _} -> []
    end,
    Infos = [{pid, Conn}, {state, RunningState}|SockInfos],
    rabbit_core_metrics:connection_stats(Conn, Infos),
    rabbit_event:notify(connection_stats, Infos),
    State1 = rabbit_event:reset_stats_timer(State, #state.stats_timer),
    State1.
