/*
 * .============.
 * // M A K E / \
 * // C++ DEV / \
 * // E A S Y / \/ \
 * ++ ----------. \/\ .
 * \\ \ \ /\ /
 * \\ \ \ /
 * \\ \ \ /
 * -============'
 *
 * Copyright (c) 2018 Hevake and contributors, all rights reserved.
 *
 * This file is part of cpp-tbox (https://github.com/cpp-main/cpp-tbox)
 * Use of this source code is governed by MIT license that can be found
 * in the LICENSE file in the root of the source tree. All contributing
 * project authors may be found in the CONTRIBUTORS.md file in the root
 * of the source tree.
 */
#include "dns_request.h"

#include <unistd.h>
#include <sstream>
#include <map>
#include <memory>

#include <tbox/base/assert.h>
#include <tbox/base/wrapped_recorder.h>
#include <tbox/util/serializer.h>
#include <tbox/util/string.h>
#include <tbox/util/fs.h>

#undef MODULE_ID
#define MODULE_ID "tbox.dns"

namespace tbox {
namespace network {

namespace {

enum DNS_TYPE {
 DNS_TYPE_A = 1, //!< 期望获得查询名的IP地址
 DNS_TYPE_NS = 2, //!< 一个授权的域名服务器
 DNS_TYPE_CNAME = 5, //!< 规范名称
 DNS_TYPE_PTR = 12, //!< 指针记录
 DNS_TYPE_HINFO = 13, //!< 主机信息
 DNS_TYPE_MX = 15, //!< 邮件交换记录
 DNS_TYPE_AXFR = 252, //!< 对区域转换的请求
 DNS_TYPE_ANY = 255, //!< 对所有记录的请求
};

enum DNS_CLASS {
 DNS_CLASS_IN = 1, //!< 互联网地址
};

enum DNS_RC {
 DNS_RC_FORMAT_ERROR = 1, //!< 数据包格式错误
 DNS_RC_SERVER_FAIL = 2, //!< 服务器的错误
 DNS_RC_NAME_ERROR = 3, //!< 名称不存在
 DNS_RC_NOT_IMPL = 4, //!< 服务器不支持查询
 DNS_RC_REFUSED = 5, //!< 服务器拒绝
};

/// 将domain写入到缓冲
/**
 * "www.baidu.com\x0" --> "\x03www0円x5baidu0円x3com0円x0"
 */
void AppendDomain(util::Serializer &dump, const std::string domain)
{
 std::vector<std::string> str_vec;
 util::string::Split(domain, ".", str_vec);

 for (auto &seg : str_vec) {
 dump << uint8_t(seg.length());
 dump.append(seg.data(), seg.length());
 }
 dump << uint8_t(0);
}

/// 从缓冲中提取domain,与AppendDomain()相反
std::string FetchDomain(util::Deserializer &parser)
{
 std::ostringstream oss;
 bool first = true;
 for (;;) {
 uint8_t len = 0;
 parser >> len;
 if (len == 0)
 break;

 if (!first)
 oss << '.';
 first = false;

 //! 处理压缩的字串
 if ((len & 0xc0) == 0xc0) {
 uint8_t offset_low = 0;
 parser >> offset_low;
 uint16_t offset = (len & 0x3f) << 8 | offset_low;
 util::Deserializer sub_parser(parser);
 sub_parser.set_pos(offset);
 oss << FetchDomain(sub_parser);
 break;
 } else {
 char str[len + 1];
 parser.fetch(str, len);
 str[len] = '0円';
 oss << str;
 }
 }
 return oss.str();
}

}

DnsRequest::DnsRequest(event::Loop *wp_loop) :
 udp_(wp_loop),
 timeout_monitor_(wp_loop)
{
 init();
}

DnsRequest::DnsRequest(event::Loop *wp_loop, const IPAddressVec &dns_ip_vec) :
 udp_(wp_loop),
 timeout_monitor_(wp_loop),
 dns_ip_vec_(dns_ip_vec)
{
 init();
}

DnsRequest::~DnsRequest() {
 udp_.disable();
}

void DnsRequest::init()
{
 using namespace std::placeholders;
 udp_.setRecvCallback(std::bind(&DnsRequest::onUdpRecv, this, _1, _2, _3));

 timeout_monitor_.initialize(std::chrono::seconds(1), 5);
 timeout_monitor_.setCallback(std::bind(&DnsRequest::onRequestTimeout, this, _1));
}

void DnsRequest::setDnsIPAddresses(const IPAddressVec &dns_ip_vec)
{
 dns_ip_vec_ = dns_ip_vec;
}

DnsRequest::ReqId DnsRequest::request(const DomainName &domain, const Callback &cb)
{
 if (dns_ip_vec_.empty()) {
 LogWarn("dns srv ip not specify");
 return false;
 }

 ReqId req_id = ++req_id_alloc_;
 std::vector<uint8_t> send_buff;

 util::Serializer dump(send_buff);

 uint16_t flags = 0x0100; //!< QR:请求, OPCODE:标准查询, RA:期望递归
 uint16_t qd_count = 1;
 uint16_t an_count = 0;
 uint16_t ns_count = 0;
 uint16_t ar_count = 0;

 dump << req_id
 << flags
 << qd_count
 << an_count
 << ns_count
 << ar_count;

 AppendDomain(dump, domain.toString());

 uint16_t dns_type = DNS_TYPE_A;
 uint16_t dns_class = DNS_CLASS_IN;

 dump << dns_type
 << dns_class;

#if 0
 std::string hex_str = util::string::RawDataToHexStr(send_buff.data(), send_buff.size());
 LogTrace("send data: %s", hex_str.c_str());
#endif

 for (auto &ip : dns_ip_vec_) {
 SockAddr dns_srv(ip, 53);
 udp_.send(send_buff.data(), send_buff.size(), dns_srv);
 }

 addRequest(req_id, cb);
 return req_id;
}

bool DnsRequest::cancel(ReqId req_id)
{
 return deleteRequest(req_id);
}

bool DnsRequest::isRunning(ReqId req_id) const
{
 auto iter = requests_.find(req_id);
 return iter != requests_.end();
}

void DnsRequest::onUdpRecv(const void *data_ptr, size_t data_size, const SockAddr &from)
{
#if 0
 std::string hex_str = util::string::RawDataToHexStr(data_ptr, data_size);
 LogTrace("recv from %s : %s", from.toString().c_str(), hex_str.c_str());
#endif

 if (requests_.empty())
 return;

 RECORD_SCOPE();
 util::Deserializer parser(data_ptr, data_size);

 uint16_t req_id, flags;
 parser >> req_id >> flags;

 Request *req = findRequest(req_id);
 if (req == nullptr)
 return;

 if ((flags & 0x8000) == 0) //! 必须是回复包
 return;

 //! 检查flags
 uint8_t rcode = flags & 0x000f;

 Result result;

 if (rcode == 0) { //! 正常
 uint16_t qd_count, an_count, ns_count, ar_count;
 parser >> qd_count >> an_count >> ns_count >> ar_count;

#if 0
 LogTrace("id:%d, flags:%04x, qd_count:%d, an_count:%d, ns_count:%d, ar_count:%d",
 id, flags, qd_count, an_count, an_count, ns_count, ar_count);
#endif

 //! 解析Question字段
 for (uint16_t i = 0; i < qd_count; ++i) {
 FetchDomain(parser);
 uint16_t dns_type, dns_class;
 parser >> dns_type >> dns_class;
 }

 for (uint16_t i = 0; i < an_count; ++i) {
 FetchDomain(parser);
 uint16_t an_type, an_class, an_len;
 uint32_t an_ttl;
 parser >> an_type >> an_class >> an_ttl >> an_len;

#if 0
 LogTrace("type:%d, class:%d, ttl:%d, len:%d", an_type, an_class, an_ttl, an_len);
#endif
 if (an_type == DNS_TYPE_A) {
 uint32_t ip_value;
 auto old_endian = parser.setEndian(util::Endian::kLittle);
 parser >> ip_value;
 parser.setEndian(old_endian);
 A a = { an_ttl, IPAddress(ip_value) };
 result.a_vec.push_back(a);

 } else if (an_type == DNS_TYPE_CNAME) {
 std::string domain = FetchDomain(parser);
 CNAME cname = { an_ttl, DomainName(domain) };
 result.cname_vec.push_back(cname);

 } else {
 LogNotice("unknow type:%d", an_type);
 parser.skip(an_len);
 }
 }
 } else {
 //! 出现异常
 if (rcode == DNS_RC_NAME_ERROR) {
 //! 如果是域名自身的问题,则直接返回失败
 result.status = Result::Status::kDomainError;
 } else if (rcode == DNS_RC_FORMAT_ERROR) {
 LogNotice("dns packet error"); //! 不应该发生
 result.status = Result::Status::kFail;
 } else {
 //! 如果是服务器的问题,则略过当前数据,等待其它服务器的数据
 ++req->response_count;
 if (req->response_count < dns_ip_vec_.size())
 return;
 result.status = Result::Status::kAllDnsFail;
 }
 }

 if (req->cb)
 req->cb(result);

 deleteRequest(req_id);
 (void)from;
}

void DnsRequest::onRequestTimeout(ReqId req_id)
{
 auto req = findRequest(req_id);
 if (req == nullptr)
 return;

 Result result;
 result.status = Result::Status::kTimeout;

 if (req->cb)
 req->cb(result);

 deleteRequest(req_id);
}

void DnsRequest::addRequest(ReqId req_id, const Callback &cb)
{
 if (requests_.empty())
 udp_.enable();

 Request req;
 req.cb = cb;

 requests_[req_id] = req;
 timeout_monitor_.add(req_id);
}

DnsRequest::Request* DnsRequest::findRequest(ReqId req_id)
{
 auto iter = requests_.find(req_id);
 if (iter != requests_.end())
 return &(iter->second);
 return nullptr;
}

bool DnsRequest::deleteRequest(ReqId req_id)
{
 auto iter = requests_.find(req_id);
 if (iter != requests_.end()) {
 requests_.erase(iter);
 if (requests_.empty())
 udp_.disable();
 return true;
 }
 return false;
}

}
}