emqx/apps/emqx_ctl/src/emqx_ctl.erl

%%--------------------------------------------------------------------
%% Copyright (c) 2017-2023 EMQ Technologies Co., Ltd. All Rights Reserved.
%%
%% Licensed under the Apache License, Version 2.0 (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.apache.org/licenses/LICENSE-2.0
%%
%% Unless required by applicable law or agreed to in writing, software
%% distributed under the License is distributed on an "AS IS" BASIS,
%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
%% See the License for the specific language governing permissions and
%% limitations under the License.
%%--------------------------------------------------------------------

-module(emqx_ctl).

-behaviour(gen_server).

-include_lib("kernel/include/logger.hrl").

-export([start_link/0, stop/0]).

-export([
    register_command/2,
    register_command/3,
    unregister_command/1
]).

-export([
    run_command/1,
    run_command/2,
    lookup_command/1,
    get_commands/0
]).

-export([
    print/1,
    print/2,
    usage/1,
    usage/2
]).

%% Exports mainly for test cases
-export([
    format/2,
    format_usage/1,
    format_usage/2
]).

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

-record(state, {seq = 0}).

-type cmd() :: atom().
-type cmd_params() :: string().
-type cmd_descr() :: string().
-type cmd_usage() :: {cmd_params(), cmd_descr()}.

-define(SERVER, ?MODULE).
-define(CMD_TAB, emqx_command).

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

-spec stop() -> ok.
stop() ->
    try
        gen_server:stop(?SERVER)
    catch
        exit:R when R =:= noproc orelse R =:= timeout ->
            ok
    end.

-spec register_command(cmd(), {module(), atom()}) -> ok.
register_command(Cmd, MF) when is_atom(Cmd) ->
    register_command(Cmd, MF, []).

-spec register_command(cmd(), {module(), atom()}, list()) -> ok.
register_command(Cmd, MF, Opts) when is_atom(Cmd) ->
    call({register_command, Cmd, MF, Opts}).

-spec unregister_command(cmd()) -> ok.
unregister_command(Cmd) when is_atom(Cmd) ->
    cast({unregister_command, Cmd}).

call(Req) -> gen_server:call(?SERVER, Req).

cast(Msg) -> gen_server:cast(?SERVER, Msg).

-spec run_command(list(string())) -> ok | {error, term()}.
run_command([]) ->
    run_command(help, []);
run_command([Cmd | Args]) ->
    case safe_to_existing_atom(Cmd) of
        {ok, Cmd1} ->
            run_command(Cmd1, Args);
        _ ->
            help(),
            {error, cmd_not_found}
    end.

-spec run_command(cmd(), list(string())) -> ok | {error, term()}.
run_command(help, []) ->
    help();
run_command(Cmd, Args) when is_atom(Cmd) ->
    case lookup_command(Cmd) of
        [{Mod, Fun}] ->
            try
                _ = apply(Mod, Fun, [Args]),
                ok
            catch
                _:Reason:Stacktrace ->
                    ?LOG_ERROR(#{
                        msg => "ctl_command_crashed",
                        stacktrace => Stacktrace,
                        reason => Reason
                    }),
                    {error, Reason}
            end;
        [] ->
            help(),
            {error, cmd_not_found}
    end.

-spec lookup_command(cmd()) -> [{module(), atom()}].
lookup_command(Cmd) when is_atom(Cmd) ->
    case ets:match(?CMD_TAB, {{'_', Cmd}, '$1', '_'}) of
        [El] -> El;
        [] -> []
    end.

-spec get_commands() -> list({cmd(), module(), atom()}).
get_commands() ->
    [{Cmd, M, F} || {{_Seq, Cmd}, {M, F}, _Opts} <- ets:tab2list(?CMD_TAB)].

help() ->
    case ets:tab2list(?CMD_TAB) of
        [] ->
            print("No commands available.~n");
        Cmds ->
            print("Usage: ~ts~n", ["emqx ctl"]),
            lists:foreach(
                fun({_, {Mod, Cmd}, _}) ->
                    print("~110..-s~n", [""]),
                    apply(Mod, Cmd, [usage])
                end,
                Cmds
            )
    end.

-spec print(io:format()) -> ok.
print(Msg) ->
    io:format("~ts", [format(Msg, [])]).

-spec print(io:format(), [term()]) -> ok.
print(Format, Args) ->
    io:format("~ts", [format(Format, Args)]).

-spec usage([cmd_usage()]) -> ok.
usage(UsageList) ->
    io:format(format_usage(UsageList)).

-spec usage(cmd_params(), cmd_descr()) -> ok.
usage(CmdParams, Desc) ->
    io:format(format_usage(CmdParams, Desc)).

-spec format(io:format(), [term()]) -> string().
format(Format, Args) ->
    lists:flatten(io_lib:format(Format, Args)).

-spec format_usage([cmd_usage()]) -> [string()].
format_usage(UsageList) ->
    Width = lists:foldl(
        fun({CmdStr, _}, W) ->
            max(iolist_size(CmdStr), W)
        end,
        0,
        UsageList
    ),
    lists:map(
        fun({CmdParams, Desc}) ->
            format_usage(CmdParams, Desc, Width)
        end,
        UsageList
    ).

-spec format_usage(cmd_params(), cmd_descr()) -> string().
format_usage(CmdParams, Desc) ->
    format_usage(CmdParams, Desc, 0).

format_usage(CmdParams, Desc, 0) ->
    format_usage(CmdParams, Desc, iolist_size(CmdParams));
format_usage(CmdParams, Desc, Width) ->
    CmdLines = split_cmd(CmdParams),
    DescLines = split_cmd(Desc),
    Zipped = zip_cmd(CmdLines, DescLines),
    Fmt = "~-" ++ integer_to_list(Width + 1) ++ "s# ~ts~n",
    lists:foldl(
        fun({CmdStr, DescStr}, Usage) ->
            Usage ++ format(Fmt, [CmdStr, DescStr])
        end,
        "",
        Zipped
    ).

%%--------------------------------------------------------------------
%% gen_server callbacks
%%--------------------------------------------------------------------

init([]) ->
    _ = ets:new(?CMD_TAB, [named_table, protected, ordered_set]),
    {ok, #state{seq = 0}}.

handle_call({register_command, Cmd, MF, Opts}, _From, State = #state{seq = Seq}) ->
    case ets:match(?CMD_TAB, {{'$1', Cmd}, '_', '_'}) of
        [] ->
            ets:insert(?CMD_TAB, {{Seq, Cmd}, MF, Opts}),
            {reply, ok, next_seq(State)};
        [[OriginSeq] | _] ->
            ?LOG_WARNING(#{msg => "CMD_overidden", cmd => Cmd, mf => MF}),
            true = ets:insert(?CMD_TAB, {{OriginSeq, Cmd}, MF, Opts}),
            {reply, ok, State}
    end;
handle_call(Req, _From, State) ->
    ?LOG_ERROR(#{msg => "unexpected_call", call => Req}),
    {reply, ignored, State}.

handle_cast({unregister_command, Cmd}, State) ->
    ets:match_delete(?CMD_TAB, {{'_', Cmd}, '_', '_'}),
    noreply(State);
handle_cast(Msg, State) ->
    ?LOG_ERROR(#{msg => "unexpected_cast", cast => Msg}),
    noreply(State).

handle_info(Info, State) ->
    ?LOG_ERROR(#{msg => "unexpected_info", info => Info}),
    noreply(State).

terminate(_Reason, _State) ->
    ok.

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

%%------------------------------------------------------------------------------
%% Internal Function
%%------------------------------------------------------------------------------

noreply(State) ->
    {noreply, State, hibernate}.

next_seq(State = #state{seq = Seq}) ->
    State#state{seq = Seq + 1}.

split_cmd(CmdStr) ->
    Lines = string:split(CmdStr, "\n", all),
    [L || L <- Lines, L =/= []].

zip_cmd([X | Xs], [Y | Ys]) -> [{X, Y} | zip_cmd(Xs, Ys)];
zip_cmd([X | Xs], []) -> [{X, ""} | zip_cmd(Xs, [])];
zip_cmd([], [Y | Ys]) -> [{"", Y} | zip_cmd([], Ys)];
zip_cmd([], []) -> [].

safe_to_existing_atom(Str) ->
    try
        {ok, list_to_existing_atom(Str)}
    catch
        _:badarg ->
            undefined
    end.