Pull request #480: AG-17827 - fix segfault in ag::dns::AioSocket::on_…

…close Merge in ADGUARD-CORE-LIBS/dns-libs from fix/AG-17827 to dev-2.1 Squashed commit of the following: commit 4d9f0d257cfca6150f9703e306b552cc39c03db4 Author: Max Grupper <[email protected]> Date: Mon Dec 19 10:47:55 2022 +0200 add guard to exchange commit 0df7408856293c83657e55810ea65452f95137fc Author: Max Grupper <[email protected]> Date: Mon Dec 19 09:36:46 2022 +0200 revert commit e4edaf88d70cc12d35c21872b86e4a65f2ca7894 Author: Max Grupper <[email protected]> Date: Fri Dec 16 18:26:42 2022 +0200 try to fix segfault commit 812ac9e9850c22ca27147a9556318dd218b27162 Author: Max Grupper <[email protected]> Date: Fri Dec 16 18:08:55 2022 +0200 try to fix segfault
AdguardTeam · Dec 21, 2022 · 776ece5 · 776ece5
1 parent aed382a
commit 776ece5
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/upstream/upstream_dnscrypt.cpp b/upstream/upstream_dnscrypt.cpp
@@ -38,7 +38,8 @@ DnscryptUpstream::DnscryptUpstream(
         ServerStamp &&stamp, const UpstreamOptions &opts, const UpstreamFactoryConfig &config)
         : Upstream(make_dnscrypt_options(stamp, opts), config)
         , m_log("DNScrypt upstream")
-        , m_stamp(std::move(stamp)) {
+        , m_stamp(std::move(stamp))
+        , m_shutdown_guard(std::make_shared<bool>(true)) {
     static const SodiumInitializer ensure_initialized;
 }
 
@@ -50,15 +51,22 @@ DnscryptUpstream::~DnscryptUpstream() = default;
 
 coro::Task<Upstream::ExchangeResult> DnscryptUpstream::exchange(const ldns_pkt *request_pkt, const DnsMessageInfo *) {
     tracelog_id(m_log, request_pkt, "Started");
+    std::weak_ptr<bool> guard = m_shutdown_guard;
 
     SetupResult result = co_await setup_impl();
+    if (guard.expired()) {
+        co_return make_error(DnsError::AE_SHUTTING_DOWN);
+    }
     if (result.error) {
         co_return result.error;
     }
     if (m_options.timeout < result.rtt) {
         co_return make_error(DnsError::AE_TIMED_OUT, AG_FMT("Certificate fetch took too much time: {}ms", result.rtt.count()));
     }
     auto reply = co_await apply_exchange(*request_pkt, m_options.timeout - result.rtt);
+    if (guard.expired()) {
+        co_return make_error(DnsError::AE_SHUTTING_DOWN);
+    }
     if (reply.has_error()) {
         co_return reply.error();
     }
@@ -90,13 +98,16 @@ coro::Task<DnscryptUpstream::SetupResult> DnscryptUpstream::setup_impl() {
 coro::Task<Upstream::ExchangeResult> DnscryptUpstream::apply_exchange(const ldns_pkt &request_pkt, Millis timeout) {
     Impl local_upstream;
     local_upstream = *m_impl;
+    std::weak_ptr<bool> guard = m_shutdown_guard;
 
     utils::Timer timer;
 
     auto udp_reply_res = co_await local_upstream.udp_client.exchange(
             request_pkt, local_upstream.server_info, this->config().loop,
             timeout, m_config.socket_factory, this->make_socket_parameters());
-
+    if (guard.expired()) {
+        co_return make_error(DnsError::AE_SHUTTING_DOWN);
+    }
     if (udp_reply_res && ldns_pkt_tc(udp_reply_res->packet.get())) {
         tracelog_id(m_log, &request_pkt, "Truncated message was received, retrying over TCP");
         dnscrypt::Client tcp_client(utils::TP_TCP);
@@ -109,6 +120,9 @@ coro::Task<Upstream::ExchangeResult> DnscryptUpstream::apply_exchange(const ldns
 
         auto tcp_reply_res = co_await tcp_client.exchange(request_pkt, local_upstream.server_info,
                 this->config().loop, timeout, m_config.socket_factory, this->make_socket_parameters());
+        if (guard.expired()) {
+            co_return make_error(DnsError::AE_SHUTTING_DOWN);
+        }
         if (tcp_reply_res) {
             co_return std::move(tcp_reply_res->packet);
         } else {

diff --git a/upstream/upstream_dnscrypt.h b/upstream/upstream_dnscrypt.h
@@ -39,6 +39,7 @@ class DnscryptUpstream : public Upstream {
     Logger m_log;
     ServerStamp m_stamp;
     ImplPtr m_impl;
+    std::shared_ptr<bool> m_shutdown_guard;
 };
 
 } // namespace ag::dns