<commit_before>%% ===================================================================
%% @author <NAME>
%% @copyright (C): 2014 ACK CYFRONET AGH
%% This software is released under the MIT license
%% cited in 'LICENSE.txt'.
%% @end
%% ===================================================================
%% @doc: This module handles DNS protocol requests (TCP and UDP) and allows using
%% custom query handlers.
%% @end
%% ===================================================================
-module(dns_server).

-include("dns/dns.hrl").
-include("logging.hrl").

%% Local name of the process waiting for dns udp messages
-define(DNS_UDP_LISTENER, dns_udp).

%% Module name of dns tcp ranch listener
-define(DNS_TCP_LISTENER, dns_tcp).

%% Timeout to wait for DNS listeners to start. After it, they are assumed to have failed to start.
-define(LISTENERS_START_TIMEOUT, 10000).


%% ====================================================================
%% API
%% ====================================================================
-export([start/8, stop/1, handle_query/2, start_listening/5]).

% Functions useful in qury handler modules
-export([answer_record/3, authority_record/3, additional_record/3, authoritative_answer_flag/1]).

% Server configuration (in runtime)
-export([set_handler_module/1, get_handler_module/0, set_dns_response_ttl/1, get_dns_response_ttl/0, set_max_edns_udp_size/1, get_max_edns_udp_size/0]).

-export([test/1, test/2]).

test(Domain) ->
    test(Domain, {8, 8, 8, 8}).

test(Domain, NS) ->
    Query = inet_dns:encode(
        #dns_rec{
            header = #dns_header{
                id = crypto:rand_uniform(1, 16#FFFF),
                opcode = 'query',
                rd = true
            },
            qdlist = [#dns_query{
                domain = Domain,
                type = soa,
                class = in
            }],
            arlist = [{dns_rr_opt, ".", opt, 1280, 0, 0, 0, <<>>}]
        }),
    {ok, Socket} = gen_udp:open(0, [binary, {active, false}]),
    gen_udp:send(Socket, NS, 53, Query),
    {ok, {NS, 53, Reply}} = gen_udp:recv(Socket, 65535),
    inet_dns:decode(Reply).


%% start/8
%% ====================================================================
%% @doc Starts a DNS server. The server will listen on chosen port (UDP and TCP).
%% QueryHandlerModule must conform to dns_query_handler_behaviour.
%% The server is started in a new process.
%% OnFailureFun is evalueated by the process if the server fails to start.
%% @end
%% ====================================================================
-spec start(SupervisorName :: atom(), DNSPort :: integer(), QueryHandlerModule :: atom(), DNSResponseTTL :: integer(), EdnsMaxUdpSize :: integer(),
    TCPNumAcceptors :: integer(), TCPTimeout :: integer(), OnFailureFun :: function()) -> ok | {error, Reason :: term()}.
%% ====================================================================
start(SupervisorName, DNSPort, QueryHandlerModule, DNSResponseTTL, EdnsMaxUdpSize, TCPNumAcceptors, TCPTimeout, OnFailureFun) ->
    set_handler_module(QueryHandlerModule),
    set_dns_response_ttl(DNSResponseTTL),
    set_max_edns_udp_size(EdnsMaxUdpSize),
    % Listeners start is done in another process.
    % This is because this function is often called during the init of the supervisor process
    % that supervises the listeners - which causes a deadlock.
    case proc_lib:start(?MODULE, start_listening, [SupervisorName, DNSPort, TCPNumAcceptors, TCPTimeout, OnFailureFun], ?LISTENERS_START_TIMEOUT) of
        ok -> ok;
        {error, Reason} -> {error, Reason}
    end.


%% stop/1
%% ====================================================================
%% @doc Stops the DNS server and cleans up.
%% @end
%% ====================================================================
-spec stop(SupervisorName :: atom()) -> ok.
%% ====================================================================
stop(SupervisorName) ->
    % Cleaning up must be done in another process. This is because this function is often called during the init of the supervisor process
    % that supervises the listeners - which causes a deadlock.
    spawn(fun() -> clear_children_and_listeners(SupervisorName, true) end),
    ok.


%% handle_query/2
%% ====================================================================
%% @doc Handles a DNS request. This function is called from dns_upd_handler or dns_tcp_handler after
%% they have received a request. After evaluation, the response is sent back to the client.
%% @end
%% ====================================================================
-spec handle_query(Packet :: binary(), Transport :: udp | tcp) -> {ok, Packet :: binary()} | {error, term()}.
%% ====================================================================
handle_query(Packet, Transport) ->
    case inet_dns:decode(Packet) of
        {ok, #dns_rec{header = Header, qdlist = QDList, anlist = _AnList, nslist = _NSList, arlist = ARList} = DNSRecWithAdditionalSection} ->
            % Detach OPT RR from the DNS query record an proceed with processing it - the OPT RR will be added during answer generation
            DNSRec = DNSRecWithAdditionalSection#dns_rec{arlist = []},
            OPTRR = case ARList of
                        [] -> undefined;
                        [#dns_rr_opt{} = OptRR] -> OptRR
                    end,
            case validate_query(DNSRec) of
                form_error ->
                    generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?FORMERR)}, OPTRR, Transport);
                bad_version ->
                    generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?BADVERS)}, OPTRR, Transport);
                ok ->
                    HandlerModule = get_handler_module(),
                    [#dns_query{domain = Domain, type = Type, class = Class}] = QDList,
                    case call_handler_module(HandlerModule, string:to_lower(Domain), Type) of
                        serv_fail ->
                            generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?SERVFAIL)}, OPTRR, Transport);
                        nx_domain ->
                            generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?NXDOMAIN)}, OPTRR, Transport);
                        not_impl ->
                            generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?NOTIMP)}, OPTRR, Transport);
                        refused ->
                            generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?REFUSED)}, OPTRR, Transport);
                        {ok, ResponseList} ->
                            %% If there was an OPT RR, it will be concated at the end of the process
                            NewRec = lists:foldl(
                                fun(CurrentRecord, #dns_rec{header = CurrHeader, anlist = CurrAnList, nslist = CurrNSList, arlist = CurrArList} = CurrRec) ->
                                    case CurrentRecord of
                                        {aa, Flag} ->
                                            CurrRec#dns_rec{header = inet_dns:make_header(CurrHeader, aa, Flag)};
                                        {Section, CurrDomain, CurrType, CurrData} ->
                                            RR = inet_dns:make_rr([
                                                {data, CurrData},
                                                {domain, CurrDomain},
                                                {type, CurrType},
                                                {ttl, get_dns_response_ttl()},
                                                {class, Class}]),
                                            case Section of
                                                answer -> CurrRec#dns_rec{anlist = CurrAnList ++ [RR]};
                                                authority -> CurrRec#dns_rec{nslist = CurrNSList ++ [RR]};
                                                additional -> CurrRec#dns_rec{arlist = CurrArList ++ [RR]}
                                            end
                                    end
                                end, DNSRec#dns_rec{}, ResponseList),
                            generate_answer(NewRec#dns_rec{arlist = NewRec#dns_rec.arlist}, OPTRR, Transport)
                    end
            end;
        _ ->
            {error, uprocessable}
    end.


%% validate_query/1
%% ====================================================================
%% @doc Checks if a query is valid.
%% @end
%% ====================================================================
-spec validate_query(#dns_rec{}) -> ok | form_error | bad_version.
%% ====================================================================
validate_query(DNSRec) ->
    case DNSRec of
        #dns_rec{qdlist = [#dns_query{class = ?C_IN}], anlist = [], nslist = [], arlist = []} ->
            % The record includes a question section and no OPT RR section - ok
            ok;
        #dns_rec{qdlist = [#dns_query{class = ?C_IN}], anlist = [], nslist = [], arlist = [#dns_rr_opt{version = Version}]} ->
            % The record includes a question section and an OPT RR section - check EDNS version
            case Version of
                0 -> ok;
                _ -> bad_version
            end;
        #dns_rec{} ->
            % Other - not ok
            form_error
    end.


%% call_handler_module/3
%% ====================================================================
%% @doc Calls the handler module to handle the query or returns not_impl if
%% this kind of query is not accepted by the server.
%% @end
%% ====================================================================
-spec call_handler_module(HandlerModule :: atom(), Domain :: binary(), Type :: atom()) ->
    {ok, [binary()]} | serv_fail | nx_domain | not_impl | refused.
%% ====================================================================
call_handler_module(HandlerModule, Domain, Type) ->
    case type_to_fun(Type) of
        not_impl ->
            not_impl;
        Fun ->
            erlang:apply(HandlerModule, Fun, [Domain])
    end.<commit_msg>VFS-914, fix error with dns class mathing<commit_after>%% ===================================================================
%% @author <NAME>
%% @copyright (C): 2014 ACK CYFRONET AGH
%% This software is released under the MIT license
%% cited in 'LICENSE.txt'.
%% @end
%% ===================================================================
%% @doc: This module handles DNS protocol requests (TCP and UDP) and allows using
%% custom query handlers.
%% @end
%% ===================================================================
-module(dns_server).

-include("dns/dns.hrl").
-include("logging.hrl").

%% Local name of the process waiting for dns udp messages
-define(DNS_UDP_LISTENER, dns_udp).

%% Module name of dns tcp ranch listener
-define(DNS_TCP_LISTENER, dns_tcp).

%% Timeout to wait for DNS listeners to start. After it, they are assumed to have failed to start.
-define(LISTENERS_START_TIMEOUT, 10000).


%% ====================================================================
%% API
%% ====================================================================
-export([start/8, stop/1, handle_query/2, start_listening/5]).

% Functions useful in qury handler modules
-export([answer_record/3, authority_record/3, additional_record/3, authoritative_answer_flag/1]).

% Server configuration (in runtime)
-export([set_handler_module/1, get_handler_module/0, set_dns_response_ttl/1, get_dns_response_ttl/0, set_max_edns_udp_size/1, get_max_edns_udp_size/0]).

-export([test/1, test/2]).

test(Domain) ->
    test(Domain, {8, 8, 8, 8}).

test(Domain, NS) ->
    Query = inet_dns:encode(
        #dns_rec{
            header = #dns_header{
                id = crypto:rand_uniform(1, 16#FFFF),
                opcode = 'query',
                rd = true
            },
            qdlist = [#dns_query{
                domain = Domain,
                type = soa,
                class = in
            }],
            arlist = [{dns_rr_opt, ".", opt, 1280, 0, 0, 0, <<>>}]
        }),
    {ok, Socket} = gen_udp:open(0, [binary, {active, false}]),
    gen_udp:send(Socket, NS, 53, Query),
    {ok, {NS, 53, Reply}} = gen_udp:recv(Socket, 65535),
    inet_dns:decode(Reply).


%% start/8
%% ====================================================================
%% @doc Starts a DNS server. The server will listen on chosen port (UDP and TCP).
%% QueryHandlerModule must conform to dns_query_handler_behaviour.
%% The server is started in a new process.
%% OnFailureFun is evalueated by the process if the server fails to start.
%% @end
%% ====================================================================
-spec start(SupervisorName :: atom(), DNSPort :: integer(), QueryHandlerModule :: atom(), DNSResponseTTL :: integer(), EdnsMaxUdpSize :: integer(),
    TCPNumAcceptors :: integer(), TCPTimeout :: integer(), OnFailureFun :: function()) -> ok | {error, Reason :: term()}.
%% ====================================================================
start(SupervisorName, DNSPort, QueryHandlerModule, DNSResponseTTL, EdnsMaxUdpSize, TCPNumAcceptors, TCPTimeout, OnFailureFun) ->
    set_handler_module(QueryHandlerModule),
    set_dns_response_ttl(DNSResponseTTL),
    set_max_edns_udp_size(EdnsMaxUdpSize),
    % Listeners start is done in another process.
    % This is because this function is often called during the init of the supervisor process
    % that supervises the listeners - which causes a deadlock.
    case proc_lib:start(?MODULE, start_listening, [SupervisorName, DNSPort, TCPNumAcceptors, TCPTimeout, OnFailureFun], ?LISTENERS_START_TIMEOUT) of
        ok -> ok;
        {error, Reason} -> {error, Reason}
    end.


%% stop/1
%% ====================================================================
%% @doc Stops the DNS server and cleans up.
%% @end
%% ====================================================================
-spec stop(SupervisorName :: atom()) -> ok.
%% ====================================================================
stop(SupervisorName) ->
    % Cleaning up must be done in another process. This is because this function is often called during the init of the supervisor process
    % that supervises the listeners - which causes a deadlock.
    spawn(fun() -> clear_children_and_listeners(SupervisorName, true) end),
    ok.


%% handle_query/2
%% ====================================================================
%% @doc Handles a DNS request. This function is called from dns_upd_handler or dns_tcp_handler after
%% they have received a request. After evaluation, the response is sent back to the client.
%% @end
%% ====================================================================
-spec handle_query(Packet :: binary(), Transport :: udp | tcp) -> {ok, Packet :: binary()} | {error, term()}.
%% ====================================================================
handle_query(Packet, Transport) ->
    case inet_dns:decode(Packet) of
        {ok, #dns_rec{header = Header, qdlist = QDList, anlist = _AnList, nslist = _NSList, arlist = ARList} = DNSRecWithAdditionalSection} ->
            % Detach OPT RR from the DNS query record an proceed with processing it - the OPT RR will be added during answer generation
            DNSRec = DNSRecWithAdditionalSection#dns_rec{arlist = []},
            OPTRR = case ARList of
                        [] -> undefined;
                        [#dns_rr_opt{} = OptRR] -> OptRR
                    end,
            case validate_query(DNSRec) of
                form_error ->
                    generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?FORMERR)}, OPTRR, Transport);
                bad_version ->
                    generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?BADVERS)}, OPTRR, Transport);
                ok ->
                    HandlerModule = get_handler_module(),
                    [#dns_query{domain = Domain, type = Type, class = Class}] = QDList,
                    case call_handler_module(HandlerModule, string:to_lower(Domain), Type) of
                        serv_fail ->
                            generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?SERVFAIL)}, OPTRR, Transport);
                        nx_domain ->
                            generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?NXDOMAIN)}, OPTRR, Transport);
                        not_impl ->
                            generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?NOTIMP)}, OPTRR, Transport);
                        refused ->
                            generate_answer(DNSRec#dns_rec{header = inet_dns:make_header(Header, rcode, ?REFUSED)}, OPTRR, Transport);
                        {ok, ResponseList} ->
                            %% If there was an OPT RR, it will be concated at the end of the process
                            NewRec = lists:foldl(
                                fun(CurrentRecord, #dns_rec{header = CurrHeader, anlist = CurrAnList, nslist = CurrNSList, arlist = CurrArList} = CurrRec) ->
                                    case CurrentRecord of
                                        {aa, Flag} ->
                                            CurrRec#dns_rec{header = inet_dns:make_header(CurrHeader, aa, Flag)};
                                        {Section, CurrDomain, CurrType, CurrData} ->
                                            RR = inet_dns:make_rr([
                                                {data, CurrData},
                                                {domain, CurrDomain},
                                                {type, CurrType},
                                                {ttl, get_dns_response_ttl()},
                                                {class, Class}]),
                                            case Section of
                                                answer -> CurrRec#dns_rec{anlist = CurrAnList ++ [RR]};
                                                authority -> CurrRec#dns_rec{nslist = CurrNSList ++ [RR]};
                                                additional -> CurrRec#dns_rec{arlist = CurrArList ++ [RR]}
                                            end
                                    end
                                end, DNSRec#dns_rec{}, ResponseList),
                            generate_answer(NewRec#dns_rec{arlist = NewRec#dns_rec.arlist}, OPTRR, Transport)
                    end
            end;
        _ ->
            {error, uprocessable}
    end.


%% validate_query/1
%% ====================================================================
%% @doc Checks if a query is valid.
%% @end
%% ====================================================================
-spec validate_query(#dns_rec{}) -> ok | form_error | bad_version.
%% ====================================================================
validate_query(DNSRec) ->
    case DNSRec of
        #dns_rec{qdlist = [#dns_query{class = Class}], anlist = [], nslist = [], arlist = []}
            when Class =:= ?C_IN orelse Class =:= in ->
            % The record includes a question section and no OPT RR section - ok
            ok;
        #dns_rec{qdlist = [#dns_query{class = Class}], anlist = [], nslist = [], arlist = [#dns_rr_opt{version = Version}]}
            when Class =:= ?C_IN orelse Class =:= in ->
            % The record includes a question section and an OPT RR section - check EDNS version
            case Version of
                0 -> ok;
                _ -> bad_version
            end;
        #dns_rec{} ->
            % Other - not ok
            form_error
    end.


%% call_handler_module/3
%% ====================================================================
%% @doc Calls the handler module to handle the query or returns not_impl if
%% this kind of query is not accepted by the server.
%% @end
%% ====================================================================
-spec call_handler_module(HandlerModule :: atom(), Domain :: binary(), Type :: atom()) ->
    {ok, [binary()]} | serv_fail | nx_domain | not_impl | refused.
%% ====================================================================
call_handler_module(HandlerModule, Domain, Type) ->
    case type_to_fun(Type) of
        not_impl ->
            not_impl;
        Fun ->
            erlang:apply(HandlerModule, Fun, [Domain])
    end.
{'fraction_non_alphanumeric': 0.22710942809981552, 'fraction_numerical': 0.005485969511603068, 'mean_word_length': 2.456208053691275, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 89, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 60, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2812878', 'n_tokens_mistral': 5572, 'n_tokens_neox': 5182, 'n_words': 1401}