Work in progress

dnstest.erl

#!/usr/bin/env escript
%% -*- erlang -*-
%%! -smp enable
-module(dnstest).

-mode(compile).

-record(dnsmsg, {id, qr, opcode, aa, tc, rd, ra, z, rcode,
                 questions, answers, authorities, extras, packet}).
-record(dnsquestion, {name, type, class}).
-record(dnsrr, {name, type, class, ttl, rdata}).

main([]) ->
    %R = dnsparse(<<0, 15, 128, 5, 0, 2, 0, 0, 0, 0, 0, 0, 5, "kitsu", 2, "ru", 0, 0, 0, 0, 0, 3, "www", (192+0), 12, 0, 1, 0, 1>>),
    %R = dnsparse(<<0,3,129,128,0,1,0,1,0,0,0,0,6,103,111,111,103,108,101,3,99,111,109,0,0,28,0,1>>),
    %io:format("Test: ~s~n", [pretty_print(R)]),
    {ok, ServerSocket} = gen_udp:open(53, [binary, {active, true}, {recbuf, 16384}, {sndbuf, 16384}, {ip, {127, 0, 0, 1}}]),
    {ok, LocalSocket} = gen_udp:open(0, [binary, {active, true}, {recbuf, 16384}, {sndbuf, 16384}]),
    io:format("DNS Started up!~n"),
    handle(ServerSocket, LocalSocket, none, none).

handle(ServerSocket, LocalSocket, ClientIP, ClientPort) ->
    receive
        {udp, ServerSocket, IP, Port, Packet} ->
            io:format("Request from ~w:~w = ~p~n", [IP, Port, Packet]),
            try dnsparse(Packet) of
                Msg ->
                    io:format("~s~n", [pretty_print(Msg)])
            catch
                Error ->
                    io:format("Error: ~p~n", [Error])
            end,
            gen_udp:send(LocalSocket, {8, 8, 8, 8}, 53, Packet),
            handle(ServerSocket, LocalSocket, IP, Port);
        {udp, LocalSocket, IP, Port, Packet} ->
            io:format("Response from ~w:~w = ~p~n", [IP, Port, Packet]),
            try dnsparse(Packet) of
                Msg ->
                    io:format("~s~n", [pretty_print(Msg)])
            catch
                Error ->
                    io:format("Error: ~p~n", [Error])
            end,
            gen_udp:send(ServerSocket, ClientIP, ClientPort, Packet),
            handle(ServerSocket, LocalSocket, ClientIP, ClientPort);
        Unknown ->
            io:format("Unexpected message ~p~n", [Unknown]),
            handle(ServerSocket, LocalSocket, ClientIP, ClientPort)
    end.

dnsparse(<<ID:16, QR:1, Opcode:4, AA:1, TC:1, RD:1, RA:1, Z:3, RCODE:4,
           QDCOUNT:16, ANCOUNT:16, NSCOUNT:16, ARCOUNT:16, RestH/binary>> = Packet) ->
    {Questions, RestQ} = dnsquestions(RestH, Packet, QDCOUNT),
    {Answers, RestA} = case TC of 0 -> dnsrrs(RestQ, Packet, ANCOUNT); _ -> {truncated, RestQ} end,
    {Authorities, RestNS} = case TC of 0 -> dnsrrs(RestA, Packet, NSCOUNT); _ -> {truncated, RestA} end,
    {Extras, _RestAR} = case TC of 0 -> dnsrrs(RestNS, Packet, ARCOUNT); _ -> {truncated, RestNS} end,
    #dnsmsg{id=ID, qr=QR, opcode=Opcode, aa=AA, tc=TC, rd=RD, ra=RA, z=Z, rcode=RCODE,
            questions=Questions, answers=Answers, authorities=Authorities, extras=Extras,
            packet=Packet}.

%% ===================================================================
%% Parses a series of dns labels
%% ===================================================================

dnsname(Bin, Packet) ->
    {Acc, Rest} = dnsname(Bin, Packet, [], 0, 0, unknown),
    {lists:reverse(Acc), Rest}.

dnsname(Bin, Packet, Acc, Total, Jumps, Ending) ->
    case Bin of
        _ when Total >= 255 ->
            throw(dns_name_too_long);
        _ when Jumps >= 16 ->
            throw(dns_name_compression_loop);
        <<0, Rest/binary>> ->
            {Acc, case Ending of unknown -> Rest; _ -> Ending end};
        <<2#00:2, Length:6, Label:Length/binary, Rest/binary>> ->
            dnsname(Rest, Packet, [Label | Acc], Total + 1 + Length, Jumps, Ending);
        <<2#11:2, Offset:14, Rest/binary>> when Offset < byte_size(Packet) ->
            <<_:Offset/binary, BinSuffix/binary>> = Packet,
            dnsname(BinSuffix, Packet, Acc, Total, Jumps + 1, case Ending of unknown -> Rest; _ -> Ending end);
        <<2#01:2, _:6, _/binary>> ->
            throw(dns_name_unsupported);
        <<2#10:2, _:6, _/binary>> ->
            throw(dns_name_unsupported);
        _ ->
            throw(dns_name_truncated)
    end.

%% ===================================================================
%% Parses a series of dns questions
%% ===================================================================

dnsquestion(Bin, Packet) ->
    {QNAME, Rest} = dnsname(Bin, Packet),
    case Rest of
        <<QTYPE:16, QCLASS:16, Rest2/binary>> ->
            {#dnsquestion{name=QNAME, type=QTYPE, class=QCLASS}, Rest2};
        _ ->
            throw(dns_question_truncated)
    end.

dnsquestions(Bin, Packet, N) ->
    {Acc, Rest} = dnsquestions(Bin, Packet, N, []),
    {lists:reverse(Acc), Rest}.

dnsquestions(Rest, _, 0, Acc) ->
    {Acc, Rest};
dnsquestions(Bin, Packet, N, Acc) ->
    {Q, Rest} = dnsquestion(Bin, Packet),
    dnsquestions(Rest, Packet, N - 1, [Q | Acc]).

%% ===================================================================
%% Parses a series of dns RRs
%% ===================================================================

dnsrr(Bin, Packet) ->
    {NAME, Rest} = dnsname(Bin, Packet),
    case Rest of
        <<TYPE:16, CLASS:16, TTL:32, RDLENGTH:16, RDATA:RDLENGTH/binary, Rest2/binary>> ->
            {#dnsrr{name=NAME, type=TYPE, class=CLASS, ttl=TTL, rdata=RDATA}, Rest2};
        _ ->
            throw(dns_rr_truncated)
    end.

dnsrrs(Bin, Packet, N) ->
    {Acc, Rest} = dnsrrs(Bin, Packet, N, []),
    {lists:reverse(Acc), Rest}.

dnsrrs(Rest, _, 0, Acc) ->
    {Acc, Rest};
dnsrrs(Bin, Packet, N, Acc) ->
    {RR, Rest} = dnsrr(Bin, Packet),
    dnsrrs(Rest, Packet, N - 1, [RR | Acc]).

%% ===================================================================
%% Pretty printing for records
%% ===================================================================

pretty_print(Record) -> 
    io_lib_pretty:print(Record, fun pretty_print/2).

pretty_print(dnsmsg, N) ->
    case record_info(size, dnsmsg) - 1 of
        N -> record_info(fields, dnsmsg);
        _ -> no
    end;
pretty_print(dnsquestion, N) ->
    case record_info(size, dnsquestion) - 1 of
        N -> record_info(fields, dnsquestion);
        _ -> no
    end;
pretty_print(_, _) ->
    no.