%% 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-2020 VMware, Inc. or its affiliates.  All rights reserved.
%%

-module(rabbit_connection_tracking).

-behaviour(gen_server).

%% Abstracts away how tracked connection records are stored
%% and queried.
%%
%% See also:
%%
%%  * rabbit_connection_tracking_handler
%%  * rabbit_reader
%%  * rabbit_event

-export([boot/0,
         ensure_tracked_connections_table_for_node/1,
         ensure_per_vhost_tracked_connections_table_for_node/1,
         ensure_tracked_connections_table_for_this_node/0,
         ensure_per_vhost_tracked_connections_table_for_this_node/0,
         tracked_connection_table_name_for/1, tracked_connection_per_vhost_table_name_for/1,
         delete_tracked_connections_table_for_node/1, delete_per_vhost_tracked_connections_table_for_node/1,
         clear_tracked_connection_tables_for_this_node/0,
         register_connection/1, unregister_connection/1,
         list/0, list/1, list_on_node/1, list_on_node/2, list_of_user/1,
         tracked_connection_from_connection_created/1,
         tracked_connection_from_connection_state/1,
         count_connections_in/1,
         lookup/1,
         count/0]).

-include_lib("rabbit.hrl").

-import(rabbit_misc, [pget/2]).

-export([start_link/0]).

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

-export([close_connections/3]).

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

%%
%% GenServer API
%%

init([]) ->
    {ok, nostate}.

handle_call(_Msg, _From, nostate) ->
    {reply, ok, nostate}.


handle_cast({connection_created, Details}, nostate) ->
    ThisNode = node(),
    case pget(node, Details) of
        ThisNode ->
            TConn = tracked_connection_from_connection_created(Details),
            ConnId = TConn#tracked_connection.id,
            try
                register_connection(TConn)
            catch
                error:{no_exists, _} ->
                    Msg = "Could not register connection ~p for tracking, "
                          "its table is not ready yet or the connection terminated prematurely",
                    rabbit_log_connection:warning(Msg, [ConnId]),
                    ok;
                error:Err ->
                    Msg = "Could not register connection ~p for tracking: ~p",
                    rabbit_log_connection:warning(Msg, [ConnId, Err]),
                    ok
            end;
        _OtherNode ->
            %% ignore
            ok
    end,
    {noreply, nostate};
handle_cast({connection_closed, Details}, nostate) ->
    ThisNode = node(),
    case pget(node, Details) of
      ThisNode ->
          %% [{name,<<"127.0.0.1:64078 -> 127.0.0.1:5672">>},
          %%  {pid,<0.1774.0>},
          %%  {node, rabbit@hostname}]
          unregister_connection(
            {pget(node, Details),
             pget(name, Details)});
      _OtherNode ->
        %% ignore
        ok
    end,
    {noreply, nostate};
handle_cast({vhost_deleted, Details}, nostate) ->
    VHost = pget(name, Details),
    rabbit_log_connection:info("Closing all connections in vhost '~s' because it's being deleted", [VHost]),
    close_connections(rabbit_connection_tracking:list(VHost),
                      rabbit_misc:format("vhost '~s' is deleted", [VHost])),
    {noreply, nostate};
%% Note: under normal circumstances this will be called immediately
%% after the vhost_deleted above. Therefore we should be careful about
%% what we log and be more defensive.
handle_cast({vhost_down, Details}, nostate) ->
    VHost = pget(name, Details),
    Node = pget(node, Details),
    rabbit_log_connection:info("Closing all connections in vhost '~s' on node '~s'"
                               " because the vhost is stopping",
                               [VHost, Node]),
    close_connections(rabbit_connection_tracking:list_on_node(Node, VHost),
                      rabbit_misc:format("vhost '~s' is down", [VHost])),
    {noreply, nostate};
handle_cast({user_deleted, Details}, nostate) ->
    Username = pget(name, Details),
    rabbit_log_connection:info("Closing all connections from user '~s' because it's being deleted", [Username]),
    close_connections(rabbit_connection_tracking:list_of_user(Username),
                      rabbit_misc:format("user '~s' is deleted", [Username])),
    {noreply, nostate};
%% A node had been deleted from the cluster.
handle_cast({node_deleted, Details}, nostate) ->
    Node = pget(node, Details),
    rabbit_log_connection:info("Node '~s' was removed from the cluster, deleting its connection tracking tables...", [Node]),
    rabbit_connection_tracking:delete_tracked_connections_table_for_node(Node),
    rabbit_connection_tracking:delete_per_vhost_tracked_connections_table_for_node(Node),
    {noreply, nostate};
handle_cast(_Msg, nostate) ->
    {noreply, nostate}.

handle_info(_Info, nostate) ->
    {noreply, nostate}.

terminate(_Reason, nostate) ->
    ok.

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

%%
%% API
%%

-spec boot() -> ok.

%% Sets up and resets connection tracking tables for this
%% node.
boot() ->
    ensure_tracked_connections_table_for_this_node(),
    rabbit_log:info("Setting up a table for connection tracking on this node: ~p",
                    [tracked_connection_table_name_for(node())]),
    ensure_per_vhost_tracked_connections_table_for_this_node(),
    rabbit_log:info("Setting up a table for per-vhost connection counting on this node: ~p",
                    [tracked_connection_per_vhost_table_name_for(node())]),
    clear_tracked_connection_tables_for_this_node(),
    ok.


-spec ensure_tracked_connections_table_for_this_node() -> ok.

ensure_tracked_connections_table_for_this_node() ->
    ensure_tracked_connections_table_for_node(node()).


-spec ensure_per_vhost_tracked_connections_table_for_this_node() -> ok.

ensure_per_vhost_tracked_connections_table_for_this_node() ->
    ensure_per_vhost_tracked_connections_table_for_node(node()).


-spec ensure_tracked_connections_table_for_node(node()) -> ok.

ensure_tracked_connections_table_for_node(Node) ->
    TableName = tracked_connection_table_name_for(Node),
    case mnesia:create_table(TableName, [{record_name, tracked_connection},
                                         {attributes, record_info(fields, tracked_connection)}]) of
        {atomic, ok}                   -> ok;
        {aborted, {already_exists, _}} -> ok;
        {aborted, Error}               ->
            rabbit_log:error("Failed to create a tracked connection table for node ~p: ~p", [Node, Error]),
            ok
    end.


-spec ensure_per_vhost_tracked_connections_table_for_node(node()) -> ok.

ensure_per_vhost_tracked_connections_table_for_node(Node) ->
    TableName = tracked_connection_per_vhost_table_name_for(Node),
    case mnesia:create_table(TableName, [{record_name, tracked_connection_per_vhost},
                                         {attributes, record_info(fields, tracked_connection_per_vhost)}]) of
        {atomic, ok}                   -> ok;
        {aborted, {already_exists, _}} -> ok;
        {aborted, Error}               ->
            rabbit_log:error("Failed to create a per-vhost tracked connection table for node ~p: ~p", [Node, Error]),
            ok
    end.


-spec clear_tracked_connection_tables_for_this_node() -> ok.

clear_tracked_connection_tables_for_this_node() ->
    case mnesia:clear_table(tracked_connection_table_name_for(node())) of
        {atomic, ok} -> ok;
        {aborted, _} -> ok
    end,
    case mnesia:clear_table(tracked_connection_per_vhost_table_name_for(node())) of
        {atomic, ok} -> ok;
        {aborted, _} -> ok
    end.


-spec delete_tracked_connections_table_for_node(node()) -> ok.

delete_tracked_connections_table_for_node(Node) ->
    TableName = tracked_connection_table_name_for(Node),
    case mnesia:delete_table(TableName) of
        {atomic, ok}              -> ok;
        {aborted, {no_exists, _}} -> ok;
        {aborted, Error} ->
            rabbit_log:error("Failed to delete a tracked connection table for node ~p: ~p", [Node, Error]),
            ok
    end.


-spec delete_per_vhost_tracked_connections_table_for_node(node()) -> ok.

delete_per_vhost_tracked_connections_table_for_node(Node) ->
    TableName = tracked_connection_per_vhost_table_name_for(Node),
    case mnesia:delete_table(TableName) of
        {atomic, ok}              -> ok;
        {aborted, {no_exists, _}} -> ok;
        {aborted, Error} ->
            rabbit_log:error("Failed to delete a per-vhost tracked connection table for node ~p: ~p", [Node, Error]),
            ok
    end.


-spec tracked_connection_table_name_for(node()) -> atom().

tracked_connection_table_name_for(Node) ->
    list_to_atom(rabbit_misc:format("tracked_connection_on_node_~s", [Node])).

-spec tracked_connection_per_vhost_table_name_for(node()) -> atom().

tracked_connection_per_vhost_table_name_for(Node) ->
    list_to_atom(rabbit_misc:format("tracked_connection_per_vhost_on_node_~s", [Node])).


-spec register_connection(rabbit_types:tracked_connection()) -> ok.
-dialyzer([{nowarn_function, [register_connection/1]}, race_conditions]).

register_connection(#tracked_connection{vhost = VHost, id = ConnId, node = Node} = Conn) when Node =:= node() ->
    TableName = tracked_connection_table_name_for(Node),
    PerVhostTableName = tracked_connection_per_vhost_table_name_for(Node),
    %% upsert
    case mnesia:dirty_read(TableName, ConnId) of
      []    ->
          mnesia:dirty_write(TableName, Conn),
          mnesia:dirty_update_counter(
            PerVhostTableName, VHost, 1);
      [_Row] ->
          ok
    end,
    ok.

-spec unregister_connection(rabbit_types:connection_name()) -> ok.

unregister_connection(ConnId = {Node, _Name}) when Node =:= node() ->
    TableName = tracked_connection_table_name_for(Node),
    PerVhostTableName = tracked_connection_per_vhost_table_name_for(Node),
    case mnesia:dirty_read(TableName, ConnId) of
        []     -> ok;
        [Row] ->
            mnesia:dirty_update_counter(
                PerVhostTableName, Row#tracked_connection.vhost, -1),
            mnesia:dirty_delete(TableName, ConnId)
    end.

-spec lookup(rabbit_types:connection_name()) -> rabbit_types:tracked_connection() | 'not_found'.

lookup(Name) ->
    Nodes = rabbit_mnesia:cluster_nodes(running),
    lookup(Name, Nodes).

lookup(_, []) ->
    not_found;
lookup(Name, [Node | Nodes]) ->
    TableName = tracked_connection_table_name_for(Node),
    case mnesia:dirty_read(TableName, {Node, Name}) of
        [] -> lookup(Name, Nodes);
        [Row] -> Row
    end.

-spec list() -> [rabbit_types:tracked_connection()].

list() ->
    lists:foldl(
      fun (Node, Acc) ->
              Tab = tracked_connection_table_name_for(Node),
              Acc ++ mnesia:dirty_match_object(Tab, #tracked_connection{_ = '_'})
      end, [], rabbit_mnesia:cluster_nodes(running)).

-spec count() -> non_neg_integer().

count() ->
    lists:foldl(
      fun (Node, Acc) ->
              Tab = tracked_connection_table_name_for(Node),
              Acc + mnesia:table_info(Tab, size)
      end, 0, rabbit_mnesia:cluster_nodes(running)).

-spec list(rabbit_types:vhost()) -> [rabbit_types:tracked_connection()].

list(VHost) ->
    lists:foldl(
      fun (Node, Acc) ->
              Tab = tracked_connection_table_name_for(Node),
              Acc ++ mnesia:dirty_match_object(Tab, #tracked_connection{vhost = VHost, _ = '_'})
      end, [], rabbit_mnesia:cluster_nodes(running)).


-spec list_on_node(node()) -> [rabbit_types:tracked_connection()].

list_on_node(Node) ->
    try mnesia:dirty_match_object(
          tracked_connection_table_name_for(Node),
          #tracked_connection{_ = '_'})
    catch exit:{aborted, {no_exists, _}} -> []
    end.

-spec list_on_node(node(), rabbit_types:vhost()) -> [rabbit_types:tracked_connection()].

list_on_node(Node, VHost) ->
    try mnesia:dirty_match_object(
          tracked_connection_table_name_for(Node),
          #tracked_connection{vhost = VHost, _ = '_'})
    catch exit:{aborted, {no_exists, _}} -> []
    end.


-spec list_of_user(rabbit_types:username()) -> [rabbit_types:tracked_connection()].

list_of_user(Username) ->
    lists:foldl(
      fun (Node, Acc) ->
              Tab = tracked_connection_table_name_for(Node),
              Acc ++ mnesia:dirty_match_object(
                       Tab,
                       #tracked_connection{username = Username, _ = '_'})
      end, [], rabbit_mnesia:cluster_nodes(running)).

-spec count_connections_in(rabbit_types:vhost()) -> non_neg_integer().

count_connections_in(VirtualHost) ->
    lists:foldl(fun (Node, Acc) ->
                        Tab = tracked_connection_per_vhost_table_name_for(Node),
                        try
                            N = case mnesia:dirty_read(Tab, VirtualHost) of
                                    []    -> 0;
                                    [Val] -> Val#tracked_connection_per_vhost.connection_count
                                end,
                            Acc + N
                        catch _:Err  ->
                                rabbit_log:error(
                                  "Failed to fetch number of connections in vhost ~p on node ~p:~n~p~n",
                                  [VirtualHost, Err, Node]),
                                Acc
                        end
                end, 0, rabbit_mnesia:cluster_nodes(running)).

%% Returns a #tracked_connection from connection_created
%% event details.
%%
%% @see rabbit_connection_tracking_handler.
tracked_connection_from_connection_created(EventDetails) ->
    %% Example event:
    %%
    %% [{type,network},
    %%  {pid,<0.329.0>},
    %%  {name,<<"127.0.0.1:60998 -> 127.0.0.1:5672">>},
    %%  {port,5672},
    %%  {peer_port,60998},
    %%  {host,{0,0,0,0,0,65535,32512,1}},
    %%  {peer_host,{0,0,0,0,0,65535,32512,1}},
    %%  {ssl,false},
    %%  {peer_cert_subject,''},
    %%  {peer_cert_issuer,''},
    %%  {peer_cert_validity,''},
    %%  {auth_mechanism,<<"PLAIN">>},
    %%  {ssl_protocol,''},
    %%  {ssl_key_exchange,''},
    %%  {ssl_cipher,''},
    %%  {ssl_hash,''},
    %%  {protocol,{0,9,1}},
    %%  {user,<<"guest">>},
    %%  {vhost,<<"/">>},
    %%  {timeout,14},
    %%  {frame_max,131072},
    %%  {channel_max,65535},
    %%  {client_properties,
    %%      [{<<"capabilities">>,table,
    %%        [{<<"publisher_confirms">>,bool,true},
    %%         {<<"consumer_cancel_notify">>,bool,true},
    %%         {<<"exchange_exchange_bindings">>,bool,true},
    %%         {<<"basic.nack">>,bool,true},
    %%         {<<"connection.blocked">>,bool,true},
    %%         {<<"authentication_failure_close">>,bool,true}]},
    %%       {<<"product">>,longstr,<<"Bunny">>},
    %%       {<<"platform">>,longstr,
    %%        <<"ruby 2.3.0p0 (2015-12-25 revision 53290) [x86_64-darwin15]">>},
    %%       {<<"version">>,longstr,<<"2.3.0.pre">>},
    %%       {<<"information">>,longstr,
    %%        <<"http://rubybunny.info">>}]},
    %%  {connected_at,1453214290847}]
    Name = pget(name, EventDetails),
    Node = pget(node, EventDetails),
    #tracked_connection{id           = {Node, Name},
                        name         = Name,
                        node         = Node,
                        vhost        = pget(vhost, EventDetails),
                        username     = pget(user, EventDetails),
                        connected_at = pget(connected_at, EventDetails),
                        pid          = pget(pid, EventDetails),
                        type         = pget(type, EventDetails),
                        peer_host    = pget(peer_host, EventDetails),
                        peer_port    = pget(peer_port, EventDetails)}.

tracked_connection_from_connection_state(#connection{
                                            vhost = VHost,
                                            connected_at = Ts,
                                            peer_host = PeerHost,
                                            peer_port = PeerPort,
                                            user = Username,
                                            name = Name
                                           }) ->
    tracked_connection_from_connection_created(
      [{name, Name},
       {node, node()},
       {vhost, VHost},
       {user, Username},
       {user_who_performed_action, Username},
       {connected_at, Ts},
       {pid, self()},
       {type, network},
       {peer_port, PeerPort},
       {peer_host, PeerHost}]).

close_connections(Tracked, Message) ->
    close_connections(Tracked, Message, 0).

close_connections(Tracked, Message, Delay) ->
    [begin
         close_connection(Conn, Message),
         timer:sleep(Delay)
     end || Conn <- Tracked],
    ok.

close_connection(#tracked_connection{pid = Pid, type = network}, Message) ->
    try
        rabbit_networking:close_connection(Pid, Message)
    catch error:{not_a_connection, _} ->
            %% could has been closed concurrently, or the input
            %% is bogus. In any case, we should not terminate
            ok;
          _:Err ->
            %% ignore, don't terminate
            rabbit_log:warning("Could not close connection ~p: ~p", [Pid, Err]),
            ok
    end;
close_connection(#tracked_connection{pid = Pid, type = direct}, Message) ->
    %% Do an RPC call to the node running the direct client.
    Node = node(Pid),
    rpc:call(Node, amqp_direct_connection, server_close, [Pid, 320, Message]).

