ns

%% Copyright (c) 2012-2020, DNSimple Corporation
%%
%% Permission to use, copy, modify, and/or distribute this software for any
%% purpose with or without fee is hereby granted, provided that the above
%% copyright notice and this permission notice appear in all copies.
%%
%% THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
%% WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
%% MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
%% ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
%% WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
%% ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
%% OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

%% @doc Worker module that asynchronously accepts a single DNS packet and
%% hands it to a worker process that has a set timeout.
-module(erldns_worker).

-include_lib("dnssec/include/dns.hrl").
-include_lib("opentelemetry_api/include/otel_tracer.hrl").
-include_lib("opentelemetry_api/include/opentelemetry.hrl").

-behaviour(gen_server).

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

-record(state, {worker_process_sup, worker_process}).

erlang_proc_message_queue_len() ->
    case process_info(self(), message_queue_len) of
        undefined -> <<"undefined">>;
        {message_queue_len, Count} -> Count
    end.

start_link(Args) ->
    gen_server:start_link(?MODULE, Args, []).

init([WorkerId]) ->
    {ok, WorkerProcessSup} = erldns_worker_process_sup:start_link([WorkerId]),
    WorkerProcess = lists:last(supervisor:which_children(WorkerProcessSup)),
    {ok, #state{worker_process_sup = WorkerProcessSup, worker_process = WorkerProcess}}.

handle_call(_Request, From, State) ->
    lager:debug("Received unexpected call (from: ~p)", [From]),
    {reply, ok, State}.

handle_cast({tcp_query, Socket, Bin}, State) ->
    ?with_span(<<"erldns_tcp_worker">>, #{},
        fun(SpanCtx) ->
            case handle_tcp_dns_query(Socket, Bin, SpanCtx, {State#state.worker_process_sup, State#state.worker_process}) of
                ok ->
                    {noreply, State};
                {error, timeout, NewWorkerPid} ->
                    {Id, _, Type, Modules} = State#state.worker_process,
                    {noreply, State#state{worker_process = {Id, NewWorkerPid, Type, Modules}}};
                Error ->
                    lager:error("Error handling TCP query (module: ~p, event: ~p, error: ~p)", [?MODULE, handle_tcp_query_error, Error]),
                    {noreply, State}
            end
        end
    );
handle_cast({udp_query, Socket, Host, Port, Bin}, State) ->
    ?with_span(<<"erldns_udp_worker">>, #{},
           fun(_SpanCtx) ->
                ?set_attributes([{host, Host}, {port, Port}, {worker_process, State#state.worker_process}]),
                ?set_attributes([{erlang_port_count, erlang:system_info(port_count)},
                                 {erlang_proc_count, erlang:system_info(process_count)},
                                 {erlang_run_queue, erlang:statistics(run_queue)},
                                 {erlang_proc_message_queue_len, erlang_proc_message_queue_len()}
                            ]),

                case handle_udp_dns_query(Socket, Host, Port, Bin, ?current_span_ctx, {State#state.worker_process_sup, State#state.worker_process}) of
                    ok ->
                        ?set_attributes([{status, <<"ok">>}]),
                        {noreply, State};
                    {error, timeout, NewWorkerPid} ->
                        ?set_attributes([{status, <<"timeout">>}]),
                        {Id, _, Type, Modules} = State#state.worker_process,
                        {noreply, State#state{worker_process = {Id, NewWorkerPid, Type, Modules}}};
                    Error ->
                        ?set_attributes([{status, <<"error">>}]),
                        lager:error("Error handling UDP query (module: ~p, event: ~p, error: ~p)", [?MODULE, handle_udp_query_error, Error]),
                        {noreply, State}
                end
            end
    );
handle_cast(_Msg, State) ->
    {noreply, State}.

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

terminate(_Reason, _State) ->
    ok.

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

%% @doc Handle DNS query that comes in over TCP
-spec handle_tcp_dns_query(gen_tcp:socket(), iodata(), otel_tracer:tracer_ctx(), {pid(), term()}) -> ok | {error, timeout} | {error, timeout, pid()}.
handle_tcp_dns_query(Socket, <<_Len:16, Bin/binary>>, SpanCtx, {WorkerProcessSup, WorkerProcess}) ->
    ?set_current_span(SpanCtx),
    ?with_span(<<"handle_tcp_dns_query">>, #{},
                fun(_SpanCtx) ->
                    case inet:peername(Socket) of
                        {ok, {Address, _Port}} ->
                            try
                                erldns_events:notify({?MODULE, start_tcp, [{host, Address}]}),
                                case Bin of
                                    <<>> ->
                                        ok;
                                    _ ->
                                        case erldns_decoder:decode_message(Bin) of
                                            {trailing_garbage, DecodedMessage, TrailingGarbage} ->
                                                Query = lists:last(DecodedMessage#dns_message.questions),
                                                ?set_attributes([
                                                    {status, <<"trailing_garbage">>},
                                                    {qr, DecodedMessage#dns_message.qr},
                                                    {rd, DecodedMessage#dns_message.rd},
                                                    {ad, DecodedMessage#dns_message.ad},
                                                    {qname, Query#dns_query.name},
                                                    {qtype, dns:type_name(Query#dns_query.type)}
                                                ]),
                                                lager:info("Decoded message included trailing garbage (module: ~p, event: ~p, message: ~p, garbage: ~p)",
                                                        [?MODULE, decode_message_trailing_garbage, DecodedMessage, TrailingGarbage]),
                                                % erldns_events:notify({?MODULE, decode_message_trailing_garbage, {DecodedMessage, TrailingGarbage}}),
                                                handle_decoded_tcp_message(DecodedMessage, Socket, Address, SpanCtx, {WorkerProcessSup, WorkerProcess});
                                            {Error, Message, _} ->
                                                otel_span:set_status(SpanCtx, #status{code = error, message = <<"Error decoding message">>}),
                                                lager:error("Error decoding message (module: ~p, event: ~p, error: ~p, message: ~p)",
                                                            [?MODULE, decode_message_error, Error, Message]),
                                                % erldns_events:notify({?MODULE, decode_message_error, {Error, Message}}),
                                                ok;
                                            DecodedMessage ->
                                                Query = lists:last(DecodedMessage#dns_message.questions),
                                                ?set_attributes([
                                                    {status, <<"ok">>},
                                                    {qr, DecodedMessage#dns_message.qr},
                                                    {rd, DecodedMessage#dns_message.rd},
                                                    {ad, DecodedMessage#dns_message.ad},
                                                    {qname, Query#dns_query.name},
                                                    {qtype, dns:type_name(Query#dns_query.type)}
                                                ]),
                                                handle_decoded_tcp_message(DecodedMessage, Socket, Address, SpanCtx, {WorkerProcessSup, WorkerProcess})
                                        end
                                end
                            of
                                Result ->
                                    erldns_events:notify({?MODULE, end_tcp, [{host, Address}]}),
                                    Result
                            catch
                                Exception:Reason:Stacktrace ->
                                    otel_span:set_status(SpanCtx, #status{code = error, message = <<"">>}),
                                    otel_span:record_exception(SpanCtx, Exception, Reason, Stacktrace, []),
                                    {error, Exception, Reason}
                            after
                                gen_tcp:close(Socket)
                            end;
                        {error, Reason} ->
                            lager:debug("Notifying error reason: ~p", [Reason]),
                            erldns_events:notify({?MODULE, tcp_error, Reason})
                    end
                end
    );
handle_tcp_dns_query(Socket, BadPacket, SpanCtx, _) ->
    ?set_current_span(SpanCtx),
    ?with_span(<<"handle_tcp_dns_query">>, #{},
                fun(_SpanCtx) ->
                    ?set_attributes([{status, <<"bad_packet">>}]),
                    lager:error("Received bad packet (module: ~p, event: ~p, protocol: ~p, packet: ~p)", [?MODULE, bad_packet, tcp, BadPacket]),
                    % erldns_events:notify({?MODULE, bad_packet, {tcp, BadPacket}}),
                    gen_tcp:close(Socket)
                end
    ).

handle_decoded_tcp_message(DecodedMessage, Socket, Address, SpanCtx, {WorkerProcessSup, {WorkerProcessId, WorkerProcessPid, _, _}}) ->
    ?set_current_span(SpanCtx),
    ?with_span(<<"handle_decoded_tcp_message">>, #{},
        fun(_SpanCtx) ->
            case DecodedMessage#dns_message.qr of
                false ->
                    try gen_server:call(WorkerProcessPid, {process, DecodedMessage, Socket, {tcp, Address}, SpanCtx}, _Timeout = erldns_config:ingress_tcp_request_timeout()) of
                        _ ->
                            ok
                    catch
                        exit:{timeout, _} ->
                            ?set_attributes([{status, <<"timeout">>}]),
                            erldns_events:notify({?MODULE, timeout}),
                            handle_timeout(WorkerProcessSup, WorkerProcessId);
                        Error:Reason ->
                            ?set_attributes([{status, <<"error">>}]),
                            lager:error("Worker process crashed (module: ~p, event: ~p, protocol: ~p, error: ~p, reason: ~p, message: ~p)",
                                        [?MODULE, process_crashed, tcp, Error, Reason, DecodedMessage]),
                            {error, {Error, Reason}}
                    end;
                true ->
                    {error, not_a_question}
            end
        end
    ).

%% @doc Handle DNS query that comes in over UDP
-spec handle_udp_dns_query(gen_udp:socket(), gen_udp:ip(), inet:port_number(), binary(), otel_tracer:tracer_ctx(), {pid(), term()}) ->
                              ok | {error, not_owner | timeout | inet:posix() | atom()} | {error, timeout, pid()}.
handle_udp_dns_query(Socket, Host, Port, Bin, SpanCtx, {WorkerProcessSup, WorkerProcess}) ->
    erldns_events:notify({?MODULE, start_udp, [{host, Host}]}),
    ?set_current_span(SpanCtx),
    Result = ?with_span(<<"handle_udp_dns_query">>, #{},
                fun(_SpanCtx) ->
                    case erldns_decoder:decode_message(Bin) of
                        {trailing_garbage, DecodedMessage, TrailingGarbage} ->
                            ?set_attributes([{status, <<"trailing_garbage">>}]),
                            lager:info("Decoded message included trailing garbage (module: ~p, event: ~p, message: ~p, garbage: ~p)",
                                    [?MODULE, decode_message_trailing_garbage, DecodedMessage, TrailingGarbage]),
                            %erldns_events:notify({?MODULE, decode_message_trailing_garbage, {DecodedMessage, TrailingGarbage}}),
                            handle_decoded_udp_message(DecodedMessage, Socket, Host, Port, SpanCtx, {WorkerProcessSup, WorkerProcess});
                        {Error, Message, _} ->
                            ?set_attributes([{status, <<"error">>}]),
                            lager:error("Error decoding message (module: ~p, event: ~p, error: ~p, message: ~p)", [?MODULE, decode_message_error, Error, Message]),
                            % erldns_events:notify({?MODULE, decode_message_error, {Error, Message}}),
                            ok;
                        DecodedMessage ->
                            Query = lists:last(DecodedMessage#dns_message.questions),
                            ?set_attributes([
                                {status, <<"ok">>},
                                {qr, DecodedMessage#dns_message.qr},
                                {rd, DecodedMessage#dns_message.rd},
                                {ad, DecodedMessage#dns_message.ad},
                                {qname, Query#dns_query.name},
                                {qtype, dns:type_name(Query#dns_query.type)}
                            ]),
                            handle_decoded_udp_message(DecodedMessage, Socket, Host, Port, SpanCtx, {WorkerProcessSup, WorkerProcess})
                    end
                end
    ),
    erldns_events:notify({?MODULE, end_udp, [{host, Host}]}),
    Result.

-spec handle_decoded_udp_message(dns:message(), gen_udp:socket(), gen_udp:ip(), inet:port_number(), otel_tracer:tracer_ctx(), {pid(), term()}) ->
                                    ok | {error, not_owner | timeout | inet:posix() | atom()} | {error, timeout, term()}.
handle_decoded_udp_message(DecodedMessage, Socket, Host, Port, SpanCtx, {WorkerProcessSup, {WorkerProcessId, WorkerProcessPid, _, _}}) ->
    ?set_current_span(SpanCtx),
    ?with_span(<<"handle_decoded_udp_message">>, #{},
        fun(_SpanCtx) ->
            case DecodedMessage#dns_message.qr of
                false ->
                    try gen_server:call(WorkerProcessPid, {process, DecodedMessage, Socket, Port, {udp, Host}, SpanCtx}, _Timeout = erldns_config:ingress_udp_request_timeout()) of
                        _ ->
                            ok
                    catch
                        exit:{timeout, _} ->
                            ?set_attributes([{status, <<"timeout">>}]),
                            lager:info("Worker timeout (module: ~p, event: ~p, protocol: ~p, message: ~p)", [?MODULE, timeout, udp, DecodedMessage]),
                            erldns_events:notify({?MODULE, timeout}),
                            handle_timeout(WorkerProcessSup, WorkerProcessId);
                        Error:Reason ->
                            ?set_attributes([{status, <<"error">>}]),
                            lager:error("Worker process crashed (module: ~p, event: ~p, protocol: ~p, error: ~p, reason: ~p, message: ~p)",
                                        [?MODULE, process_crashed, udp, Error, Reason, DecodedMessage]),
                            % erldns_events:notify({?MODULE, process_crashed, {udp, Error, Reason, DecodedMessage}}),
                            {error, {Error, Reason}}
                    end;
                true ->
                    {error, not_a_question}
            end
    end).

-spec handle_timeout(pid(), term()) -> {error, timeout, term()} | {error, timeout}.
handle_timeout(WorkerProcessSup, WorkerProcessId) ->
    TerminateResult = supervisor:terminate_child(WorkerProcessSup, WorkerProcessId),
    lager:debug("Terminate result: ~p", [TerminateResult]),

    case supervisor:restart_child(WorkerProcessSup, WorkerProcessId) of
        {ok, NewChild} ->
            {error, timeout, NewChild};
        {ok, NewChild, _} ->
            {error, timeout, NewChild};
        {error, Error} ->
            erldns_events:notify({?MODULE, restart_failed, {Error}}),
            {error, timeout}
    end.