couchbase 3.0.0.alpha.3 → 3.0.0.alpha.4

data/ext/couchbase/io/dns_client.hxx

@@ -0,0 +1,215 @@
+/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
+/*
+ *     Copyright 2020 Couchbase, Inc.
+ *
+ *   Licensed 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.
+ */
+
+#pragma once
+
+#include <memory>
+#include <sstream>
+
+#include <io/dns_codec.hxx>
+#include <io/dns_config.hxx>
+
+#include <asio/read.hpp>
+
+namespace couchbase::io::dns
+{
+class dns_client
+{
+  public:
+    struct dns_srv_response {
+        struct address {
+            std::string hostname;
+            std::uint16_t port;
+        };
+        std::error_code ec;
+        std::vector<address> targets{};
+    };
+
+    class dns_srv_command : public std::enable_shared_from_this<dns_srv_command>
+    {
+      public:
+        dns_srv_command(asio::io_context& ctx,
+                        const std::string& name,
+                        const std::string& service,
+                        const asio::ip::address& address,
+                        std::uint16_t port)
+          : deadline_(ctx)
+          , udp_(ctx)
+          , tcp_(ctx)
+          , address_(address)
+          , port_(port)
+        {
+            static std::string protocol{ "_tcp" };
+            dns_message request{};
+            question_record qr;
+            qr.klass = resource_class::in;
+            qr.type = resource_type::srv;
+            qr.name.labels.push_back(service);
+            qr.name.labels.push_back(protocol);
+            std::string label;
+            std::istringstream name_stream(name);
+            while (std::getline(name_stream, label, '.')) {
+                qr.name.labels.push_back(label);
+            }
+            request.questions.emplace_back(qr);
+            send_buf_ = dns_codec::encode(request);
+        }
+
+        template<class Handler>
+        void execute(std::chrono::milliseconds timeout, Handler&& handler)
+        {
+            asio::ip::udp::endpoint endpoint(address_, port_);
+            udp_.open(endpoint.protocol());
+            udp_.async_send_to(
+              asio::buffer(send_buf_),
+              endpoint,
+              [self = shared_from_this(), handler = std::forward<Handler>(handler)](std::error_code ec1,
+                                                                                    std::size_t /* bytes_transferred */) mutable {
+                  if (ec1 == asio::error::operation_aborted) {
+                      self->deadline_.cancel();
+                      return handler({ std::make_error_code(error::common_errc::ambiguous_timeout) });
+                  }
+                  if (ec1) {
+                      self->deadline_.cancel();
+                      return handler({ ec1 });
+                  }
+
+                  asio::ip::udp::endpoint sender_endpoint;
+                  self->recv_buf_.resize(512);
+                  self->udp_.async_receive_from(
+                    asio::buffer(self->recv_buf_),
+                    sender_endpoint,
+                    [self, handler = std::forward<Handler>(handler)](std::error_code ec2, std::size_t bytes_transferred) mutable {
+                        self->deadline_.cancel();
+                        if (ec2) {
+                            return handler({ ec2 });
+                        }
+                        self->recv_buf_.resize(bytes_transferred);
+                        dns_message message = dns_codec::decode(self->recv_buf_);
+                        if (message.header.flags.tc == truncation::yes) {
+                            self->udp_.close();
+                            return self->retry_with_tcp(std::forward<Handler>(handler));
+                        }
+                        dns_srv_response resp{ ec2 };
+                        resp.targets.reserve(message.answers.size());
+                        for (const auto& answer : message.answers) {
+                            resp.targets.emplace_back(
+                              dns_srv_response::address{ fmt::format("{}", fmt::join(answer.target.labels, ".")), answer.port });
+                        }
+                        return handler(resp);
+                    });
+              });
+            deadline_.expires_after(timeout);
+            deadline_.async_wait([self = shared_from_this()](std::error_code ec) {
+                if (ec == asio::error::operation_aborted) {
+                    return;
+                }
+                self->udp_.cancel();
+                self->tcp_.cancel();
+            });
+        }
+
+      private:
+        template<class Handler>
+        void retry_with_tcp(Handler&& handler)
+        {
+            asio::ip::tcp::no_delay no_delay(true);
+            std::error_code ignore_ec;
+            tcp_.set_option(no_delay, ignore_ec);
+            asio::ip::tcp::endpoint endpoint(address_, port_);
+            tcp_.async_connect(
+              endpoint, [self = shared_from_this(), handler = std::forward<Handler>(handler)](std::error_code ec1) mutable {
+                  if (ec1) {
+                      self->deadline_.cancel();
+                      return handler({ ec1 });
+                  }
+                  auto send_size = static_cast<uint16_t>(self->send_buf_.size());
+                  self->send_buf_.insert(self->send_buf_.begin(), std::uint8_t(send_size & 0xffU));
+                  self->send_buf_.insert(self->send_buf_.begin(), std::uint8_t(send_size >> 8U));
+                  asio::async_write(
+                    self->tcp_,
+                    asio::buffer(self->send_buf_),
+                    [self, handler = std::forward<Handler>(handler)](std::error_code ec2, std::size_t /* bytes_transferred */) mutable {
+                        if (ec2) {
+                            self->deadline_.cancel();
+                            if (ec2 == asio::error::operation_aborted) {
+                                ec2 = std::make_error_code(error::common_errc::ambiguous_timeout);
+                            }
+                            return handler({ ec2 });
+                        }
+                        asio::async_read(self->tcp_,
+                                         asio::buffer(&self->recv_buf_size_, sizeof(self->recv_buf_size_)),
+                                         [self, handler = std::forward<Handler>(handler)](std::error_code ec3,
+                                                                                          std::size_t /* bytes_transferred */) mutable {
+                                             if (ec3) {
+                                                 self->deadline_.cancel();
+                                                 return handler({ ec3 });
+                                             }
+                                             self->recv_buf_size_ = ntohs(self->recv_buf_size_);
+                                             self->recv_buf_.resize(self->recv_buf_size_);
+                                             asio::async_read(
+                                               self->tcp_,
+                                               asio::buffer(self->recv_buf_),
+                                               [self, handler = std::forward<Handler>(handler)](std::error_code ec4,
+                                                                                                std::size_t bytes_transferred) mutable {
+                                                   self->deadline_.cancel();
+                                                   if (ec4) {
+                                                       return handler({ ec4 });
+                                                   }
+                                                   self->recv_buf_.resize(bytes_transferred);
+                                                   dns_message message = dns_codec::decode(self->recv_buf_);
+                                                   dns_srv_response resp{ ec4 };
+                                                   resp.targets.reserve(message.answers.size());
+                                                   for (const auto& answer : message.answers) {
+                                                       resp.targets.emplace_back(dns_srv_response::address{
+                                                         fmt::format("{}", fmt::join(answer.target.labels, ".")), answer.port });
+                                                   }
+                                                   return handler(resp);
+                                               });
+                                         });
+                    });
+              });
+        }
+
+        asio::steady_timer deadline_;
+        asio::ip::udp::socket udp_;
+        asio::ip::tcp::socket tcp_;
+
+        asio::ip::address address_;
+        std::uint16_t port_;
+
+        std::vector<uint8_t> send_buf_;
+        std::uint16_t recv_buf_size_{ 0 };
+        std::vector<uint8_t> recv_buf_;
+    };
+
+    explicit dns_client(asio::io_context& ctx)
+      : ctx_(ctx)
+    {
+    }
+
+    template<class Handler>
+    void query_srv(const std::string& name, const std::string& service, Handler&& handler)
+    {
+        dns_config& config = dns_config::get();
+        auto cmd = std::make_shared<dns_srv_command>(ctx_, name, service, config.address(), config.port());
+        cmd->execute(config.timeout(), std::forward<Handler>(handler));
+    }
+
+    asio::io_context& ctx_;
+};
+} // namespace couchbase::io::dns

data/ext/couchbase/io/dns_codec.hxx

@@ -0,0 +1,207 @@
+/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
+/*
+ *     Copyright 2020 Couchbase, Inc.
+ *
+ *   Licensed 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.
+ */
+
+#pragma once
+
+#include <cstring>
+
+#include <arpa/inet.h>
+
+#include <io/dns_message.hxx>
+
+namespace couchbase::io::dns
+{
+class dns_codec
+{
+  public:
+    static dns_message decode(const std::vector<uint8_t>& payload)
+    {
+        dns_message message{};
+        std::size_t offset = 0;
+
+        std::memcpy(&message.header.id, payload.data() + offset, sizeof(std::uint16_t));
+        offset += sizeof(std::uint16_t);
+        message.header.id = ntohs(message.header.id);
+
+        uint16_t flags = 0;
+        std::memcpy(&flags, payload.data() + offset, sizeof(std::uint16_t));
+        offset += sizeof(std::uint16_t);
+        message.header.flags.decode(ntohs(flags));
+
+        std::memcpy(&message.header.question_records, payload.data() + offset, sizeof(std::uint16_t));
+        offset += sizeof(std::uint16_t);
+        message.header.question_records = ntohs(message.header.question_records);
+
+        std::memcpy(&message.header.answer_records, payload.data() + offset, sizeof(std::uint16_t));
+        offset += sizeof(std::uint16_t);
+        message.header.answer_records = ntohs(message.header.answer_records);
+
+        std::memcpy(&message.header.authority_records, payload.data() + offset, sizeof(std::uint16_t));
+        offset += sizeof(std::uint16_t);
+        message.header.authority_records = ntohs(message.header.authority_records);
+
+        std::memcpy(&message.header.additional_records, payload.data() + offset, sizeof(std::uint16_t));
+        offset += sizeof(std::uint16_t);
+        message.header.additional_records = ntohs(message.header.additional_records);
+
+        for (std::uint16_t idx = 0; idx < message.header.question_records; ++idx) {
+            question_record qr;
+            qr.name = get_name(payload, offset);
+
+            std::uint16_t val = 0;
+            std::memcpy(&val, payload.data() + offset, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+            val = ntohs(val);
+            qr.type = static_cast<resource_type>(val);
+
+            std::memcpy(&val, payload.data() + offset, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+            val = ntohs(val);
+            qr.klass = static_cast<resource_class>(val);
+
+            message.questions.emplace_back(qr);
+        }
+
+        message.answers.reserve(message.header.answer_records);
+        for (std::uint16_t idx = 0; idx < message.header.answer_records; ++idx) {
+            srv_record ar;
+            ar.name = get_name(payload, offset);
+
+            std::uint16_t val = 0;
+            std::memcpy(&val, payload.data() + offset, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+            val = ntohs(val);
+            ar.type = static_cast<resource_type>(val);
+
+            std::memcpy(&val, payload.data() + offset, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+            val = ntohs(val);
+            ar.klass = static_cast<resource_class>(val);
+
+            std::memcpy(&ar.ttl, payload.data() + offset, sizeof(std::uint32_t));
+            offset += static_cast<std::uint16_t>(4U);
+            ar.ttl = ntohl(ar.ttl);
+
+            std::uint16_t size = 0;
+            std::memcpy(&size, payload.data() + offset, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+            size = ntohs(size);
+
+            if (ar.klass != resource_class::in || ar.type != resource_type::srv) {
+                // ignore everything except SRV answers
+                offset += size;
+                continue;
+            }
+
+            std::memcpy(&val, payload.data() + offset, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+            ar.priority = ntohs(val);
+
+            std::memcpy(&val, payload.data() + offset, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+            ar.weight = ntohs(val);
+
+            std::memcpy(&val, payload.data() + offset, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+            ar.port = ntohs(val);
+
+            ar.target = get_name(payload, offset);
+
+            message.answers.emplace_back(ar);
+        }
+        return message;
+    }
+
+    static std::vector<uint8_t> encode(const dns_message& message)
+    {
+        std::vector<std::uint8_t> payload;
+        payload.resize(message.request_size(), 0);
+        std::size_t offset = 0;
+
+        // write header
+        {
+            uint16_t val;
+
+            val = htons(message.header.id);
+            std::memcpy(payload.data() + offset, &val, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+
+            val = htons(message.header.flags.encode());
+            std::memcpy(payload.data() + offset, &val, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+
+            val = htons(static_cast<std::uint16_t>(message.questions.size()));
+            std::memcpy(payload.data() + offset, &val, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t) + 3 * sizeof(std::uint16_t); // answer, authority, additional are all zeros
+        }
+
+        // write body
+        for (const auto& question : message.questions) {
+            for (const auto& label : question.name.labels) {
+                payload[offset] = static_cast<std::uint8_t>(label.size());
+                ++offset;
+                std::memcpy(payload.data() + offset, label.data(), label.size());
+                offset += label.size();
+            }
+            payload[offset] = '\0';
+            ++offset;
+
+            uint16_t val;
+
+            val = htons(static_cast<std::uint16_t>(question.type));
+            std::memcpy(payload.data() + offset, &val, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+
+            val = htons(static_cast<std::uint16_t>(question.klass));
+            std::memcpy(payload.data() + offset, &val, sizeof(std::uint16_t));
+            offset += sizeof(std::uint16_t);
+        }
+        return payload;
+    }
+
+  private:
+    static resource_name get_name(const std::vector<std::uint8_t>& payload, std::size_t& offset)
+    {
+        resource_name name{};
+        std::optional<std::size_t> save_offset{};
+        while (true) {
+            std::uint8_t len = payload[offset];
+            if (len == 0) {
+                offset += 1;
+                if (save_offset) {
+                    // restore offset after pointer jump
+                    offset = *save_offset;
+                }
+                return name;
+            }
+            if ((len & 0b1100'0000U) != 0) {
+                std::uint16_t ptr = 0;
+                std::memcpy(&ptr, payload.data() + offset, sizeof(std::uint16_t));
+                ptr = ntohs(ptr);
+                ptr &= 0b0011'1111'1111'1111U;
+                // store old offset and jump to pointer
+                save_offset = offset + sizeof(std::uint16_t);
+                offset = ptr;
+            } else {
+                std::string label(payload.data() + offset + 1, payload.data() + offset + 1 + len);
+                name.labels.emplace_back(label);
+                offset += static_cast<std::uint16_t>(1U + len);
+            }
+        }
+    }
+};
+} // namespace couchbase::io::dns

data/ext/couchbase/io/dns_config.hxx

@@ -0,0 +1,116 @@
+/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
+/*
+ *     Copyright 2020 Couchbase, Inc.
+ *
+ *   Licensed 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.
+ */
+
+#pragma once
+
+#include <unistd.h>
+
+#include <string>
+#include <fstream>
+
+#include <asio/ip/address.hpp>
+
+#include <timeout_defaults.hxx>
+
+namespace couchbase::io::dns
+{
+class dns_config
+{
+  public:
+    static inline constexpr auto default_resolv_conf_path = "/etc/resolv.conf";
+    static inline constexpr auto default_host = "8.8.8.8";
+    static inline constexpr std::uint16_t default_port = 53;
+
+    [[nodiscard]] const asio::ip::address& address() const
+    {
+        return address_;
+    }
+
+    [[nodiscard]] std::uint16_t port() const
+    {
+        return port_;
+    }
+
+    [[nodiscard]] std::chrono::milliseconds timeout() const
+    {
+        return timeout_;
+    }
+
+    static dns_config& get()
+    {
+        static dns_config instance{};
+
+        instance.initialize();
+
+        return instance;
+    }
+
+  private:
+    void initialize()
+    {
+        if (!initialized_) {
+            load_resolv_conf(default_resolv_conf_path);
+            std::error_code ec;
+            address_ = asio::ip::address::from_string(host_, ec);
+            if (ec) {
+                host_ = default_host;
+                address_ = asio::ip::address::from_string(host_, ec);
+            }
+            initialized_ = true;
+        }
+    }
+
+    void load_resolv_conf(const char* conf_path)
+    {
+        if (access(conf_path, R_OK) == 0) {
+            std::ifstream conf(conf_path);
+            while (conf.good()) {
+                std::string line;
+                std::getline(conf, line);
+                if (line.empty()) {
+                    continue;
+                }
+                std::size_t offset = 0;
+                while (line[offset] == ' ') {
+                    ++offset;
+                }
+                if (line[offset] == '#') {
+                    continue;
+                }
+                std::size_t space = line.find(' ', offset);
+                if (space == std::string::npos || space == offset || line.size() < space + 2) {
+                    continue;
+                }
+                std::string keyword = line.substr(offset, space);
+                if (keyword != "nameserver") {
+                    continue;
+                }
+                offset = space + 1;
+                space = line.find(' ', offset);
+                host_ = line.substr(offset, space);
+                break;
+            }
+        }
+    }
+
+    std::atomic_bool initialized_{ false };
+    std::string host_{ default_host };
+    asio::ip::address address_{};
+    std::uint16_t port_{ default_port };
+    std::chrono::milliseconds timeout_{ timeout_defaults::dns_srv_timeout };
+};
+} // namespace couchbase::io::dns