プログラマーとDNS (3)

#include <WinSock2.h>
#pragma comment(lib, "ws2_32.lib")
#include <ws2tcpip.h>
#include <Windows.h>
#include <tchar.h>

int _tmain(int argc, TCHAR** argv)
{
    // Winsockの初期化処理
    WSADATA wsa;
    if (WSAStartup(MAKEWORD(2, 2), &wsa) != 0)
        return -1;

    // 問い合わせる
    auto cacheserver = "8.8.8.8";
    auto port = "53";
    auto fqdn = "www.microsoft.com"s;
    u_short id = 0x1234;
    u_short type = TYPE_A;
    Response response;

    bool ret = query(cacheserver, port, fqdn, type, id, response);
    if (!ret) {
        // エラー: response.headerをチェック
    }
    else {
        // typeに応じてresponse.answer,response.authority,response.additionをチェック
    }

    // Winsockの終了処理
    if (WSACleanup() != 0)
        return -1;

    return 0;
}

#pragma pack(push, 1)
struct Header final
{
    u_short id;       // メッセージを識別する任意の16ビットの数値
    u_short flags;    // 各種フラグ
    u_short qdcount;  // 問い合わせの数
    u_short ancount;  // 回答レコードの数
    u_short nscount;  // 権威レコードの数
    u_short arcount;  // 付加情報レコードの数
};
#pragma pack(pop)

// ヘッダーをストリームに出力する
void writeHeader(std::ostream& ostr, Header& header)
{
    header.toBigEndian();
    ostr.write(reinterpret_cast<const char*>(&header), sizeof(Header));
    header.toLittleEndian();
}

// FQDNをエンコードしてストリームに出力する
void writeFqdn(std::ostream& ostr, const std::string& str)
{
    if (str.length() + 1 > Limits::MAX_NAME)
        throw InvalidFormat();

    size_t end = str.length();
    size_t beg = 0, cur = 0;
    do {
        // ドットを検索
        cur = str.find('.', beg);
        if (cur == std::string::npos) {
            cur = end; // ドットがもうないので省略されたものとみなす
        }

        // ラベル長とラベルの出力
        auto val = cur - beg;
        if (val > Limits::MAX_LABEL) {
            throw InvalidFormat();
        }
        ostr << static_cast<u_char>(val);
        for (auto i = beg; i < cur; ++i) {
            ostr << str[i];
        }

        // 終端文字の出力
        if (cur == end) {
            ostr << static_cast<u_char>(0);
            return;
        }

        beg = cur + 1;
	  } while (beg < end);

    throw InvalidFormat();
}

// u_short値をストリームに出力する
void writeUShort(std::ostream& ostr, u_short val)
{
    val = ::htons(val);
    ostr.write(reinterpret_cast<const char*>(&val), sizeof(u_short));
}

// 問い合わせメッセージを作成する
std::string buildMessage(const std::string& fqdn, u_short type, u_short id)
{
    // ヘッダーの初期化
    Header header{ 0 };
    header.id = id;
    header.flags = RD; // 再帰的問い合わせ(決め打ち)
    header.qdcount = 1;

    // メッセージの作成
    std::ostringstream request;
    writeHeader(request, header);
    writeFqdn(request, fqdn);
    writeUShort(request, type);
    writeUShort(request, CLASS_IN);

    return request.str();
}

std::string sendMessage(const char* cacheserver, const char* port, const std::string& request)
{
    // ソケットの作成
    addrinfo* addr, hints;
    memset(&hints, 0, sizeof(addrinfo));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_DGRAM;
    hints.ai_protocol = IPPROTO_UDP;
    hints.ai_flags = AI_NUMERICHOST | AI_NUMERICSERV;
    if (getaddrinfo(cacheserver, port, &hints, &addr) != 0) {
        throw SocketError();
    }

    auto sock = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol);
    if (sock == INVALID_SOCKET) {
        auto err = WSAGetLastError();
        freeaddrinfo(addr);
        throw SocketError(err);
    }
    
    // リクエストの送信
    auto sendlen = sendto(sock, request.c_str(), (int)request.length(), 0, addr->ai_addr, addr->ai_addrlen);
    if (sendlen == SOCKET_ERROR) {
        auto err = WSAGetLastError();
        freeaddrinfo(addr);
        closesocket(sock);
        throw SocketError(err);
    }

    // レスポンスの受信
    std::string response(Limits::MAX_UDPDATA, 0);
    auto addrlen = addr->ai_addrlen;
    auto recvlen = recvfrom(sock, response.data(), response.size(), 0, addr->ai_addr, reinterpret_cast<int*>(&addrlen));
    if (recvlen == SOCKET_ERROR) {
        auto err = WSAGetLastError();
        freeaddrinfo(addr);
        closesocket(sock);
        throw SocketError(err);
    }

    response.resize(recvlen);

    freeaddrinfo(addr);
    closesocket(sock);
    
    return response;
}

bool query(const char* cacheserver, const char* port, const std::string& fqdn, u_short type, u_short id, Response& response)
{
    auto rawrequest = buildMessage(fqdn, type, id);
    auto rawresponse = sendMessage(cacheserver, port, rawrequest);
    if (rawresponse.size() == 0)
        return false;

    return parseMessage(id, rawresponse, response);
}

struct Response
{
    Header header;
    std::string qname;
    u_short qtype;
    u_short qclass;
    std::vector<ResourceRecord> answer;
    std::vector<ResourceRecord> authority;
    std::vector<ResourceRecord> addition;
};

bool parseMessage(u_short id, const std::string& rawresponse, Response& response)
{
    auto pTop = rawresponse.c_str(); // 応答メッセージの先頭を表す
    auto pEnd = pTop + rawresponse.size(); // 応答メッセージの終端を表す
    auto pCur = pTop; // 次に解析すべき位置を示す

    // ヘッダーの解析
    response.header = *reinterpret_cast<const Header*>(pCur); // ヘッダー部をコピー
    pCur += sizeof(Header); // ヘッダー分のポインタを進める

    response.header.toLittleEndian(); // ヘッダーをリトルエンディアンにする
    if (response.header.id != id)
        return false; // メッセージ識別子が一致していない
    if (!response.header.isSet(QR) || !response.header.isSet(RD) || !response.header.isSet(RA))
        throw InvalidFormat(); // 応答でない || 再帰的問い合わせでない || 再帰的問い合わせに対応していない
    if (response.header.getRCode() != 0)
        return false; // 問い合わせが正常に処理されなかった

    // 問い合わせ部の解析
    // QNAME
    pCur = parseName(pTop, pEnd, pCur, response.qname);
    // QTYPE
    pCur = parseUShort(pTop, pEnd, pCur, response.qtype);
    // QCLASS
    pCur = parseUShort(pTop, pEnd, pCur, response.qclass);

// ドメイン名の読み込み(メッセージ圧縮を解除するけど、デコードは行わない)
const char* parseName(const char* pTop, const char* pEnd, const char* pCur, std::ostream& name)
{
    while (pCur < pEnd) {
        u_char ch = static_cast<u_char>(*pCur++);
        if (ch == 0) { // ラベル長が0 => ドメイン名の終端
            name << ch;
            return pCur;
        }
        else if ((ch & 0xC0) == 0) { // ラベル長(0x3F以内)
            name << ch;
            for (u_char i = 0; i < ch; ++i) // ラベル長分読み込む
                if (pCur < pEnd)
                    name << *pCur++;
                else
                    throw InvalidFormat(); // ここで終端に達するのはエラー
        }
        else if ((ch & 0xC0) == 0xC0) { // 上位2ビットが11の場合はメッセージ圧縮されている
            auto ch2 = static_cast<u_char>(*pCur++); // 次の1オクテットも読み込む
            u_short offset = ((ch & 0x3F) << 8) + ch2; // chの下位6ビットとch2がpTopからのオフセットを差している
            auto unused = parseName(pTop, pEnd, pTop + offset, name); // pTop+offsetで示される位置のドメイン名を読み込む

            return pCur; // 参照先ではなく元のpCurを返す
        }
        else {
            throw NotImplemented();
        }
    }

    throw InvalidFormat(); // ここに来たらエラー
}


// ドメイン名の読み込み(std::stringで受け取るラッパー関数)
const char* parseName(const char* pTop, const char* pEnd, const char* pCur, std::string& name)
{
    std::ostringstream ostr;
    auto ret = parseName(pTop, pEnd, pCur, ostr);
    name = ostr.str();
    return ret;
}

// 2オクテット値の読み込み＆リトルエンディアンに変換
const char* parseUShort(const char* pTop, const char* pEnd, const char* pCur, u_short& value)
{
    if (pCur + 1 >= pEnd)
        throw InvalidFormat(); // ここで終端に達するのはエラー
    value = static_cast<u_char>(*pCur++);
    value = (value << 8) + static_cast<u_char>(*pCur++);
    return pCur;
}

// 4オクテット値の読み込み＆リトルエンディアンに変換
const char* parseUInt(const char* pTop, const char* pEnd, const char* pCur, u_int& value)
{
    if (pCur + 3 >= pEnd)
        throw InvalidFormat(); // ここで終端に達するのはエラー
    value = static_cast<u_char>(*pCur++);
    value = (value << 8) + static_cast<u_char>(*pCur++);
    value = (value << 8) + static_cast<u_char>(*pCur++);
    value = (value << 8) + static_cast<u_char>(*pCur++);
    return pCur;
}

struct ResourceRecord
{
    std::string name;
    u_short type;
    u_short cls;
    u_int ttl;
    std::string resource;
};

// リソースレコードの解析
const char* parseResourceRecord(const char* pTop, const char* pEnd, const char* pCur, ResourceRecord& rr)
{
    // NAME
    pCur = parseName(pTop, pEnd, pCur, rr.name);

    // TYPE
    pCur = parseUShort(pTop, pEnd, pCur, rr.type);

    // CLASS
    pCur = parseUShort(pTop, pEnd, pCur, rr.cls);

    // TTL
    pCur = parseUInt(pTop, pEnd, pCur, rr.ttl);

    // RDLENGTH
    u_short rdlen;
    pCur = parseUShort(pTop, pEnd, pCur, rdlen);

    auto const pRData = pCur;

    // RDATA
    if ((rr.type == TYPE_A && rdlen == 4) || (rr.type == TYPE_AAAA && rdlen == 16)) {
        // AレコードはIPv4アドレス、AAAAレコードはIPv6アドレスをリソースとして持つ
        // IPアドレスはバイト列として読み込むだけで、文字列表記に直さない
        std::ostringstream ostr;
        for (u_short i = 0; i < rdlen; ++i) {
            ostr << *pCur++;
        }
        rr.resource = ostr.str();
    }
    else if (rr.type == TYPE_CNAME || rr.type == TYPE_NS) {
        // CNAME、NSレコードはドメイン名をリソースとして持つ
        pCur = parseName(pTop, pEnd, pCur, rr.resource);
    }
    else if (rr.type == TYPE_MX) {
        // MXレコードは2オクテットの優先順位とドメイン名をリソースとして持つ
        std::ostringstream mxstr;
        u_short temp;
        pCur = parseUShort(pTop, pEnd, pCur, temp);
        mxstr.write(reinterpret_cast<const char*>(&temp), sizeof(u_short));
        pCur = parseName(pTop, pEnd, pCur, mxstr); // 圧縮メッセージの展開はしておく
        rr.resource = mxstr.str();
    }
    else if (rr.type == TYPE_SOA) {
        // SOAレコードは MNAME,RNAME,SERIAL,REFRESH,RETRY,EXPIRE,MINIMUM
        std::ostringstream soastr;
        u_int temp;
        pCur = parseName(pTop, pEnd, pCur, soastr); // 圧縮メッセージの展開はしておく
        pCur = parseName(pTop, pEnd, pCur, soastr); // 圧縮メッセージの展開はしておく
        pCur = parseUInt(pTop, pEnd, pCur, temp);
        soastr.write(reinterpret_cast<const char*>(&temp), sizeof(u_int));
        pCur = parseUInt(pTop, pEnd, pCur, temp);
        soastr.write(reinterpret_cast<const char*>(&temp), sizeof(u_int));
        pCur = parseUInt(pTop, pEnd, pCur, temp);
        soastr.write(reinterpret_cast<const char*>(&temp), sizeof(u_int));
        pCur = parseUInt(pTop, pEnd, pCur, temp);
        soastr.write(reinterpret_cast<const char*>(&temp), sizeof(u_int));
        pCur = parseUInt(pTop, pEnd, pCur, temp);
        soastr.write(reinterpret_cast<const char*>(&temp), sizeof(u_int));
        rr.resource = soastr.str();
    }
    else {
        // 未対応のものはRDLENGTH分だけRDATAを受け取り解析はしない
        std::ostringstream ostr;
        for (u_short i = 0; i < rdlen; ++i) {
            ostr << *pCur++;
        }
        rr.resource = ostr.str();
    }

    // リソース部を解析した結果、RDLENGTHと不整合があればエラー
    if (pCur - pRData != rdlen)
        throw InvalidFormat();

    return pCur;
}

    // 回答部
    for (u_short answer = 0; answer < response.header.ancount; ++answer) {
        ResourceRecord rr;
        pCur = parseResourceRecord(pTop, pEnd, pCur, rr);
        response.answer.push_back(rr);
    }
    // 権威部
    for (u_short authority = 0; authority < response.header.nscount; ++authority) {
        ResourceRecord rr;
        pCur = parseResourceRecord(pTop, pEnd, pCur, rr);
        response.authority.push_back(rr);
    }
    // 付加情報部
    for (u_short addition = 0; addition < response.header.arcount; ++addition) {
        ResourceRecord rr;
        pCur = parseResourceRecord(pTop, pEnd, pCur, rr);
        response.addition.push_back(rr);
    }
    // 正しく終端まで来てなければエラー
    if (pCur != pEnd)
        throw InvalidFormat();

    return true;
}