<commit_before>%% -------------------------------------------------------------------
%%
%% Copyright (c) 2011 <NAME>. All Rights Reserved.
%%
%% This file is provided to you 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(dnsxd_ds_server).
-include("dnsxd.hrl").
-behaviour(gen_server).

%% API
-export([start_link/0, ets_memory/0]).

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

%% zone management
-export([load_zone/1, reload_zone/1, delete_zone/1, zone_loaded/1]).

%% querying
-export([zone_for_name/1, get_zone/1, find_zone/1, get_key/1]).

%% llq
-export([new_llq/3, msg_llq/2]).

-define(SERVER, ?MODULE).

-define(TAB_INDEX, dnsxd_ds_index).
-define(TAB_CAT, dnsxd_ds_catalog).
-define(TAB_SW, dnsxd_ds_serials).
-define(TAB_LLQ, dnsxd_ds_llq).

-record(state, {llq_count = 0}).
-record(index, {hash, soa = false, count = 1, domain}).
-record(sw, {zonename, ref, serials}).
-record(llq_ref, {id, pid, zonename}).

%%%===================================================================
%%% API
%%%===================================================================

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

load_zone(#dnsxd_zone{} = Zone) ->
    gen_server:call(?SERVER, {load_zone, Zone}).

reload_zone(#dnsxd_zone{} = Zone) ->
    gen_server:call(?SERVER, {reload_zone, Zone}).

delete_zone(ZoneName) ->
    gen_server:call(?SERVER, {delete_zone, ZoneName}).

zone_for_name(Name) ->
    Labels = lists:reverse(dns:dname_to_labels(dns:dname_to_lower(Name))),
    zone_for_name(undefined, <<>>, Labels).

zone_for_name(undefined, _PrevHash, []) -> undefined;
zone_for_name(LastDom, _PrevHash, []) -> LastDom;
zone_for_name(LastDom, PrevHash, [Label|Labels]) ->
    Hash = crypto:sha([Label, PrevHash]),
    case ets:lookup(?TAB_INDEX, Hash) of
	[#index{hash = Hash, soa = false, count = Count}] when Count > 0 ->
	    zone_for_name(LastDom, Hash, Labels);
	[#index{hash = Hash, soa = true, count = Count, domain = Dom}]
	  when Count > 0 ->
	    zone_for_name(Dom, Hash, Labels);
	_ ->
	    if LastDom =:= <<>> -> undefined;
	       true -> LastDom end
    end.

get_zone(Name) ->
    case ets:lookup(?TAB_CAT, Name) of
	[Result] -> Result;
	[] -> undefined
    end.

find_zone(Name) ->
    case zone_for_name(Name) of
	undefined -> undefined;
	SOAName -> get_zone(SOAName)
    end.

get_key(KeyName) ->
    [KeyLabel|ZoneLabels] = dns:dname_to_labels(dns:dname_to_lower(KeyName)),
    ZoneName = join_labels(ZoneLabels),
    case ets:lookup(?TAB_CAT, ZoneName) of
	[#dnsxd_zone{tsig_keys = Keys}] ->
	    case lists:keyfind(KeyLabel, #dnsxd_tsig_key.name, Keys) of
		#dnsxd_tsig_key{} = Key -> {ZoneName, Key};
		false -> undefined
	    end;
	[] -> undefined
    end.

new_llq(Pid, MsgCtx, #dns_message{} = Msg) ->
    gen_server:call(?SERVER, {new_llq, Pid, MsgCtx, Msg}).

msg_llq(MsgCtx, #dns_message{
	  additional = [#dns_optrr{data = [#dns_opt_llq{id = Id}]}]
	 } = Msg) ->
    case ets:lookup(?TAB_LLQ, Id) of
	[#llq_ref{id = Id, pid = Pid}] when is_pid(Pid) ->
	    dnsxd_llq_server:handle_msg(Pid, MsgCtx, Msg);
	[] -> {error, nosuchllq}
    end.

ets_memory() ->
    Tabs = [?TAB_INDEX, ?TAB_CAT, ?TAB_SW, ?TAB_LLQ],
    WordSize = erlang:system_info(wordsize),
    Fun = fun(Tab, Acc) ->
		  TabSize = ets:info(Tab, memory) * WordSize,
		  {{Tab, TabSize}, TabSize + Acc}
	  end,
    {TabSizes, Total} = lists:mapfoldl(Fun, 0, Tabs),
    [{total, Total}|TabSizes].

%%%===================================================================
%%% gen_server callbacks
%%%===================================================================

init([]) ->
    process_flag(trap_exit, true),
    ?TAB_INDEX = ets:new(?TAB_INDEX, [named_table, {keypos, #index.hash}]),
    ?TAB_CAT = ets:new(?TAB_CAT, [named_table, {keypos, #dnsxd_zone.name}]),
    ?TAB_LLQ = ets:new(?TAB_LLQ, [named_table, {keypos, #llq_ref.id}]),
    ?TAB_SW = ets:new(?TAB_SW, [named_table, {keypos, #sw.zonename}]),
    {ok, #state{}}.

handle_call({load_zone, #dnsxd_zone{name = ZoneName} = Zone}, _From, State) ->
    Reply = case zone_loaded(ZoneName) of
		true -> {error, loaded};
		false -> insert_zone(Zone)
	    end,
    {reply, Reply, State};
handle_call({reload_zone, #dnsxd_zone{} = Zone}, _From, State) ->
    Reply =  insert_zone(Zone),
    {reply, Reply, State};
handle_call({delete_zone, ZoneName}, _From, State) ->
    Reply = case zone_loaded(ZoneName) of
		true ->
		    ok = remove_from_index(ZoneName),
		    true = ets:delete(?TAB_CAT, ZoneName),
		    ok = cancel_serial_change(ZoneName);
		false ->
		    {error, not_loaded}
	    end,
    {reply, Reply, State};
handle_call({new_llq, Pid, MsgCtx,
	     #dns_message{questions = [#dns_query{name = NameM} = Q],
			  additional = [#dns_optrr{dnssec = ClientDNSSEC}|_]
			 } = Msg}, _From, #state{llq_count = Count} = State) ->
    Name = dns:dname_to_lower(NameM),
    {ok, Opts} = dnsxd:llq_opts(),
    MaxLLQ = proplists:get_value(max_llq, Opts, 500),
    Reply = case Count >= MaxLLQ of
		true ->
		    {ok, servfull};
		false ->
		    case find_zone(Name) of
			#dnsxd_zone{name = ZoneName} = Zone ->
			    Id = new_llqid(),
			    DNSSEC = ClientDNSSEC andalso
				Zone#dnsxd_zone.dnssec_enabled,
			    {ok, LLQPid} = dnsxd_llq_server:start_link(
					  Pid, Id, ZoneName, MsgCtx, Q, DNSSEC),
			    LLQRef = #llq_ref{id = Id, pid = LLQPid,
					      zonename = ZoneName},
			    ets:insert(?TAB_LLQ, LLQRef),
			    ok = dnsxd_llq_server:handle_msg(LLQPid, MsgCtx,
							     Msg);
			undefined ->
			    {ok, bad_zone}
		    end
	    end,
    {reply, Reply, State};
handle_call(Request, _From, State) ->
    ?DNSXD_ERR("Stray call:~n~p", [Request]),<commit_msg>Fix creation of new LLQs<commit_after>%% -------------------------------------------------------------------
%%
%% Copyright (c) 2011 <NAME>. All Rights Reserved.
%%
%% This file is provided to you 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(dnsxd_ds_server).
-include("dnsxd.hrl").
-behaviour(gen_server).

%% API
-export([start_link/0, ets_memory/0]).

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

%% zone management
-export([load_zone/1, reload_zone/1, delete_zone/1, zone_loaded/1]).

%% querying
-export([zone_for_name/1, get_zone/1, find_zone/1, get_key/1]).

%% llq
-export([new_llq/3, msg_llq/2]).

-define(SERVER, ?MODULE).

-define(TAB_INDEX, dnsxd_ds_index).
-define(TAB_CAT, dnsxd_ds_catalog).
-define(TAB_SW, dnsxd_ds_serials).
-define(TAB_LLQ, dnsxd_ds_llq).

-record(state, {llq_count = 0}).
-record(index, {hash, soa = false, count = 1, domain}).
-record(sw, {zonename, ref, serials}).
-record(llq_ref, {id, pid, zonename}).

%%%===================================================================
%%% API
%%%===================================================================

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

load_zone(#dnsxd_zone{} = Zone) ->
    gen_server:call(?SERVER, {load_zone, Zone}).

reload_zone(#dnsxd_zone{} = Zone) ->
    gen_server:call(?SERVER, {reload_zone, Zone}).

delete_zone(ZoneName) ->
    gen_server:call(?SERVER, {delete_zone, ZoneName}).

zone_for_name(Name) ->
    Labels = lists:reverse(dns:dname_to_labels(dns:dname_to_lower(Name))),
    zone_for_name(undefined, <<>>, Labels).

zone_for_name(undefined, _PrevHash, []) -> undefined;
zone_for_name(LastDom, _PrevHash, []) -> LastDom;
zone_for_name(LastDom, PrevHash, [Label|Labels]) ->
    Hash = crypto:sha([Label, PrevHash]),
    case ets:lookup(?TAB_INDEX, Hash) of
	[#index{hash = Hash, soa = false, count = Count}] when Count > 0 ->
	    zone_for_name(LastDom, Hash, Labels);
	[#index{hash = Hash, soa = true, count = Count, domain = Dom}]
	  when Count > 0 ->
	    zone_for_name(Dom, Hash, Labels);
	_ ->
	    if LastDom =:= <<>> -> undefined;
	       true -> LastDom end
    end.

get_zone(Name) ->
    case ets:lookup(?TAB_CAT, Name) of
	[Result] -> Result;
	[] -> undefined
    end.

find_zone(Name) ->
    case zone_for_name(Name) of
	undefined -> undefined;
	SOAName -> get_zone(SOAName)
    end.

get_key(KeyName) ->
    [KeyLabel|ZoneLabels] = dns:dname_to_labels(dns:dname_to_lower(KeyName)),
    ZoneName = join_labels(ZoneLabels),
    case ets:lookup(?TAB_CAT, ZoneName) of
	[#dnsxd_zone{tsig_keys = Keys}] ->
	    case lists:keyfind(KeyLabel, #dnsxd_tsig_key.name, Keys) of
		#dnsxd_tsig_key{} = Key -> {ZoneName, Key};
		false -> undefined
	    end;
	[] -> undefined
    end.

new_llq(Pid, MsgCtx, #dns_message{} = Msg) ->
    gen_server:call(?SERVER, {new_llq, Pid, MsgCtx, Msg}).

msg_llq(MsgCtx, #dns_message{
	  additional = [#dns_optrr{data = [#dns_opt_llq{id = Id}]}]
	 } = Msg) ->
    case ets:lookup(?TAB_LLQ, Id) of
	[#llq_ref{id = Id, pid = Pid}] when is_pid(Pid) ->
	    dnsxd_llq_server:handle_msg(Pid, MsgCtx, Msg);
	[] -> {error, nosuchllq}
    end.

ets_memory() ->
    Tabs = [?TAB_INDEX, ?TAB_CAT, ?TAB_SW, ?TAB_LLQ],
    WordSize = erlang:system_info(wordsize),
    Fun = fun(Tab, Acc) ->
		  TabSize = ets:info(Tab, memory) * WordSize,
		  {{Tab, TabSize}, TabSize + Acc}
	  end,
    {TabSizes, Total} = lists:mapfoldl(Fun, 0, Tabs),
    [{total, Total}|TabSizes].

%%%===================================================================
%%% gen_server callbacks
%%%===================================================================

init([]) ->
    process_flag(trap_exit, true),
    ?TAB_INDEX = ets:new(?TAB_INDEX, [named_table, {keypos, #index.hash}]),
    ?TAB_CAT = ets:new(?TAB_CAT, [named_table, {keypos, #dnsxd_zone.name}]),
    ?TAB_LLQ = ets:new(?TAB_LLQ, [named_table, {keypos, #llq_ref.id}]),
    ?TAB_SW = ets:new(?TAB_SW, [named_table, {keypos, #sw.zonename}]),
    {ok, #state{}}.

handle_call({load_zone, #dnsxd_zone{name = ZoneName} = Zone}, _From, State) ->
    Reply = case zone_loaded(ZoneName) of
		true -> {error, loaded};
		false -> insert_zone(Zone)
	    end,
    {reply, Reply, State};
handle_call({reload_zone, #dnsxd_zone{} = Zone}, _From, State) ->
    Reply =  insert_zone(Zone),
    {reply, Reply, State};
handle_call({delete_zone, ZoneName}, _From, State) ->
    Reply = case zone_loaded(ZoneName) of
		true ->
		    ok = remove_from_index(ZoneName),
		    true = ets:delete(?TAB_CAT, ZoneName),
		    ok = cancel_serial_change(ZoneName);
		false ->
		    {error, not_loaded}
	    end,
    {reply, Reply, State};
handle_call({new_llq, Pid, MsgCtx,
	     #dns_message{questions = [#dns_query{name = NameM} = Q],
			  additional = [#dns_optrr{dnssec = ClientDNSSEC}|_]
			 } = Msg}, _From, #state{llq_count = Count} = State) ->
    Name = dns:dname_to_lower(NameM),
    Opts = dnsxd:llq_opts(),
    MaxLLQ = proplists:get_value(max_llq, Opts, 500),
    Reply = case Count >= MaxLLQ of
		true ->
		    {ok, servfull};
		false ->
		    case find_zone(Name) of
			#dnsxd_zone{name = ZoneName} = Zone ->
			    Id = new_llqid(),
			    DNSSEC = ClientDNSSEC andalso
				Zone#dnsxd_zone.dnssec_enabled,
			    {ok, LLQPid} = dnsxd_llq_server:start_link(
					  Pid, Id, ZoneName, MsgCtx, Q, DNSSEC),
			    LLQRef = #llq_ref{id = Id, pid = LLQPid,
					      zonename = ZoneName},
			    ets:insert(?TAB_LLQ, LLQRef),
			    ok = dnsxd_llq_server:handle_msg(LLQPid, MsgCtx,
							     Msg);
			undefined ->
			    {ok, bad_zone}
		    end
	    end,
    {reply, Reply, State};
handle_call(Request, _From, State) ->
    ?DNSXD_ERR("Stray call:~n~p", [Request]),
{'fraction_non_alphanumeric': 0.21699232655644332, 'fraction_numerical': 0.005379321256229729, 'mean_word_length': 3.671840354767184, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 107, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6117502', 'n_tokens_mistral': 4945, 'n_tokens_neox': 4578, 'n_words': 966}