erlang dns proxy, with long ttl and persistant etc table cache.

dns_proxy_srv.erl

%%%-------------------------------------------------------------------
%%% @author  <[email protected]>
%%% @copyright (C) 2013, 
%%% @doc
%%%
%%% @end
%%% Created : 18 Apr 2013 by  <[email protected]>
%%%-------------------------------------------------------------------
-module(dns_proxy_srv).

-behaviour(gen_server).

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

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

-compile([export_all]).

-define(SERVER, ?MODULE). 

-record(state, {sock,table}).

-include_lib("stdlib/include/ms_transform.hrl"). 
-include_lib("kernel/src/inet_dns.hrl").

-define(DEBUG(Term), io:format("debug: ~p~n", [Term])).
-define(DEBUG(What, Term), io:format("debug ~p: ~p~n", [What, Term])).

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

%%--------------------------------------------------------------------
%% @doc
%% Starts the server
%%
%% @spec start_link() -> {ok, Pid} | ignore | {error, Error}
%% @end
%%--------------------------------------------------------------------
start_link() ->
    gen_server:start_link({local, ?SERVER}, ?MODULE, [], []).

test() ->
    io:format("ddddddddddddddddddebug: ~p~n",
	      [ets:fun2ms(fun(#dns_rr{type=cname,domain="www.baidu.com"}=R) when data =:= "www.baidu.com" -> R end)]).

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

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Initializes the server
%%
%% @spec init(Args) -> {ok, State} |
%%                     {ok, State, Timeout} |
%%                     ignore |
%%                     {stop, Reason}
%% @end
%%--------------------------------------------------------------------
init([]) ->
    Tid = ets:new(resolve_table, [bag, public, named_table, {keypos, #dns_rr.domain}]),
    case gen_udp:open(53, [binary, {active, true}]) of
	{ok, Sock} ->
	    {ok, #state{sock=Sock, table=Tid}};
	{error, Reason} ->
	    {stop, Reason}
    end.		

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Handling call messages
%%
%% @spec handle_call(Request, From, State) ->
%%                                   {reply, Reply, State} |
%%                                   {reply, Reply, State, Timeout} |
%%                                   {noreply, State} |
%%                                   {noreply, State, Timeout} |
%%                                   {stop, Reason, Reply, State} |
%%                                   {stop, Reason, State}
%% @end
%%--------------------------------------------------------------------
handle_call(_Request, _From, State) ->
    Reply = ok,
    {reply, Reply, State}.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Handling cast messages
%%
%% @spec handle_cast(Msg, State) -> {noreply, State} |
%%                                  {noreply, State, Timeout} |
%%                                  {stop, Reason, State}
%% @end
%%--------------------------------------------------------------------
handle_cast(_Msg, State) ->
    {noreply, State}.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Handling all non call/cast messages
%%
%% @spec handle_info(Info, State) -> {noreply, State} |
%%                                   {noreply, State, Timeout} |
%%                                   {stop, Reason, State}
%% @end
%%--------------------------------------------------------------------
handle_info({udp, Sock, FromIP, FromPort, Data}, #state{sock=Sock,table=Tid} = State) ->
    io:format("got packet from ~p:~p~n", [FromIP, FromPort]),
    handle_dns_packet(Data, fun(D, normal) ->
				    gen_udp:send(Sock, FromIP, FromPort, D),
				    save_dns_result_to_table(Tid, D);
			       (D, cached) ->
				    gen_udp:send(Sock, FromIP, FromPort, D)
			    end, Tid),
    {noreply, State};
handle_info({test_ok, From}, #state{sock=Sock,table=Tid} = State) ->
    From ! ets:select(Tid, ets:fun2ms(fun(#dns_rr{type=cname,domain="www.baidu.com",data=Cname}=R) ->
					 Cname
				      end)),
					      
    {noreply, State};
handle_info(_Info, State) ->
    io:format("unhandled message: ~p~n", [_Info]),
    {noreply, State}.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% This function is called by a gen_server when it is about to
%% terminate. It should be the opposite of Module:init/1 and do any
%% necessary cleaning up. When it returns, the gen_server terminates
%% with Reason. The return value is ignored.
%%
%% @spec terminate(Reason, State) -> void()
%% @end
%%--------------------------------------------------------------------
terminate(_Reason, #state{sock=Sock}) ->
    gen_udp:close(Sock),
    ok.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Convert process state when code is changed
%%
%% @spec code_change(OldVsn, State, Extra) -> {ok, NewState}
%% @end
%%--------------------------------------------------------------------
code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%%%===================================================================
%%% Internal functions
%%%===================================================================
handle_dns_packet(Data, SendFunc, Tid) ->
    {ok, Packet} = inet_dns:decode(Data),
    #dns_rec{header=Header, qdlist=Questions, anlist=Answers,
	     nslist=Authorities, arlist=Resources} = Packet,
    handle_dns_header(Header),
    case Header#dns_header.qr of
	true ->
	    %% this is response
	    handle_dns_header(Header);
	false ->
	    %% this is a request
	    case query_table_by_dns_query(Tid, Questions) of
		[] ->
		    handle_dns_request(Packet, SendFunc);
		Cached ->
		    io:format("!!! found in cache! items=~p~n", [length(Cached)]),
		    SendFunc(inet_dns:encode(dns_rec_fill_answer(Packet, Cached)),
			     cached)
	    end
    end,

    handle_dns_questions(Questions),
    %% io:format("got dns packet ~p~n", [Packet]),
    ok.

handle_dns_request(#dns_rec{header=#dns_header{qr=false,opcode='query'},
			    qdlist=Queries} = Packet,
		   SendFunc) ->
    QuerySendFunc =
	fun() ->
		Qdlist = lists:filter(fun(#dns_query{type=aaaa}) -> false;
					  (_)                     -> true
				       end,
				       Packet#dns_rec.qdlist),
		case Qdlist of
		    [] ->
			SendFunc(inet_dns:encode(dns_rec_requst_to_response(Packet)),
				 normal);
		    _  ->
			query_google_and_send_response(Packet, SendFunc)
		end
	end,
    spawn(QuerySendFunc).

dns_rec_requst_to_response(Packet) ->
    _OldHeader = Packet#dns_rec.header,
    Packet#dns_rec{header=_OldHeader#dns_header{qr=true}}.

dns_rec_fill_answer(Packet, Answers) when is_list(Answers) ->
    Packet#dns_rec{anlist=Answers}.

query_google_and_send_response(Packet, SendFunc) ->
    {ok, S} = gen_udp:open(0, [binary]),
    gen_udp:send(S, "8.8.8.8", 53, inet_dns:encode(Packet)),
    receive
	{udp, S, {8,8,8,8}, 53, Data} ->
	    io:format("query google ok, sent response!~n"),
	    SendFunc(Data, normal)
    after 1000 ->
	    io:format("query time out.~n")
    end,
    gen_udp:close(S).

		   
handle_dns_header(#dns_header{id=Id,qr=RespFlag,opcode=OpCode,rcode=RCode}) ->
    ok.
	    

handle_dns_questions([#dns_query{domain=Domain,type=Type,class=Class}|Rest]) ->
    io:format("query type:~p ~p~n", [Type, Domain]),
    handle_dns_questions(Rest);
handle_dns_questions([]) ->
    ok.

handle_dns_answers([#dns_rr{domain=Domain,type=Type,class=Class,
			    cnt=Count,ttl=TTL,data=Data,tm=Time,
			    bm=_,func=_}|Rest]) ->
    io:format("dns record ~p ~p ~p~n", [Type, Class, Domain]),
    handle_dns_answers(Rest);
handle_dns_answers([]) ->
    ok.

    
    
save_dns_result_to_table(T, Data)    when is_binary(Data) ->
    {ok, Packet} = inet_dns:decode(Data),
    save_dns_result_to_table(T, Packet);
save_dns_result_to_table(T, Packet) when is_record(Packet, dns_rec) ->
    case Packet#dns_rec.anlist of
	[] ->
	    ok;
	Records ->
	    Records1 = lists:map(fun(R) -> dns_rr_set_ttl(R, 1000) end,
				 Records),
	    ets:insert(T, Records1)
    end.

query_table_by_dns_query(T, Query) when is_record(Query, dns_query) ->
    #dns_query{domain=Domain,type=Type,class=Class} = Query,
    Result = table_query(T, Domain),
    fill_cname_query(T, Result);
query_table_by_dns_query(T, [Query|Rest]=Queries) when is_list(Queries),
						       is_record(Query, dns_query) ->
    #dns_query{domain=Domain,type=Type,class=Class} = Query,
    lists:flatten([query_table_by_dns_query(T, Query),
		   query_table_by_dns_query(T, Rest)]);
query_table_by_dns_query(_, []) ->
    [].


fill_cname_query(Tid, [R|Rest]) ->
    #dns_rr{type=Type, domain=Domain, data=Data} = R,
    case Type of
	cname ->
	    [R|fill_cname_query(Tid, table_query(Tid, Data))] ++
		fill_cname_query(Tid, Rest);
	Other ->
	    [R|fill_cname_query(Tid, Rest)]
    end;
fill_cname_query(_, []) ->
    [].

%% make ttl longer and store to a `bag`
dns_rr_set_ttl(Query, TTL) when is_record(Query, dns_rr), is_integer(TTL) ->
    Query#dns_rr{ttl=TTL}.

table_query(Tid, Domain) ->
    %% fun2ms is bad here
    A = ets:select(Tid, [{{dns_rr,Domain,'_','_','_','_','_','_','_','_'},[],['$_']}]),
    A.