[DOH] Since v8.5.0, a change in API behavior causes curl_multi_poll() or curl_multi_perform() to return error code #14414

#include <cstdio>
#include <exception>
#include <iostream>
#include <iterator>
#include <memory>
#include <ostream>
#include <sstream>
#include <mutex>
#include <future>
#include <functional>
#include <thread>
#include <atomic>
#include <chrono>
#include <string>
#include <utility>
#include <list>
#include <vector>
#include <string.h>
#include <assert.h>
#include <curl/curl.h>

static bool globalInitCalled = false;

void doInit()
{
    static bool initialized = ([] {
        curl_global_init(CURL_GLOBAL_ALL);
        globalInitCalled = true;
        return true;
    })();

    // silence unused warning
    (void)initialized;
}

enum Method
{
    HEAD,
    GET,
    POST,
    PUT,
};

class Session
{
public:
    enum State
    {
        NEW,
        PERFORMING,
        ABORTED,
        HTTP_OK,
        HTTP_ERR,
        OTHER_ERR
    };

    Session(Method method, const std::string& url, const std::string& dnsServer, std::shared_ptr<std::atomic_bool> flag)
        : state_(NEW)
        , statusStr_("new")
        , origUrl_(url)
    {
        doInit();
        curl_ = curl_easy_init();

        // Set URL
        setopt(CURLOPT_URL, url.c_str());

        // Set method
        switch (method)
        {
            case HEAD:
                setopt(CURLOPT_NOBODY, 1L);
                break;
            case GET:
                setopt(CURLOPT_HTTPGET, 1L);
                break;
            case POST:
                setopt(CURLOPT_POST, 1L);
                break;
            case PUT:
                setopt(CURLOPT_CUSTOMREQUEST, "PUT");
                setopt(CURLOPT_UPLOAD, 1L);
                break;
        }

        // encoding
        setopt(CURLOPT_ACCEPT_ENCODING, "");
        setopt(CURLOPT_SSL_VERIFYPEER, 0L);

        if (flag)
        {
            setAbortFlag(flag);
        }

        if (!dnsServer.empty())
        {
            enableDoH(dnsServer);
        }
    }

    ~Session()
    {
        if (curl_ != nullptr)
        {
            curl_easy_cleanup(curl_);
            curl_ = nullptr;
        }
    }

    void setAbortFlag(std::shared_ptr<std::atomic_bool> f)
    {
        assert(state_ == NEW);
        isAborted_ = [f]()
        {
            bool ret = f->load();
            if (ret)
            {
                std::cout << "> abort flag set." << std::endl;
            }

            return f->load();
        };
    }

    void enableDoH(const std::string& v)
    {
        assert(state_ == NEW);

        std::cout << "enableDoH: " << v << std::endl;
        setopt(CURLOPT_DOH_URL, v.c_str());
        setopt(CURLOPT_DOH_SSL_VERIFYPEER, 0);
        setopt(CURLOPT_DOH_SSL_VERIFYHOST, 0);
    }

    bool checkAborted()
    {
        return isAborted_();
    }

    void start()
    {
        assert(state_ == NEW);
        setopt(CURLOPT_PRIVATE, this);
        state_ = PERFORMING;
        statusStr_ = "performing";
    }

    void abort()
    {
        assert(state_ == PERFORMING);
        state_ = ABORTED;
        statusStr_ = "aborted";
    }

    State getState() const
    {
        return state_;
    }

    bool inProgress() const
    {
        return state_ == PERFORMING;
    }

    bool failed() const
    {
        return state_ != NEW &&
               state_ != PERFORMING &&
               state_ != HTTP_OK;
    }

    const std::string& getStatusStr() const
    {
        return statusStr_;
    }

    const std::string& url() const
    {
        return origUrl_;
    }

    CURL* handle() const
    {
        return curl_;
    }

    void finish(CURLcode res)
    {
        if (res == CURLE_OK)
        {
            state_ = HTTP_OK;
            statusStr_ = "OK";
        }
        else
        {
            state_ = OTHER_ERR;
            statusStr_ = "error";
        }
    }

private:
    template <typename T>
    void setopt(CURLoption opt, T&& v) const
    {
        curl_easy_setopt(curl_, opt, v);
    }

private:
    State state_;
    std::string statusStr_;
    std::string origUrl_;
    CURL* curl_ = nullptr;

    std::function<bool()> isAborted_ = []() {
        return false;
    };
};

class RequestError: public std::exception
{
public:
    RequestError(std::unique_ptr<Session> session)
        : session_(std::make_shared<std::unique_ptr<Session>>(std::move(session)))
    {
    }
    const char* what() const noexcept override
    {
        return response().getStatusStr().c_str();
    }

    const Session& response() const
    {
        return **session_;
    }

private:
    std::shared_ptr<std::unique_ptr<Session>> session_;
};

inline std::ostream& operator<<(std::ostream& out, const RequestError& err)
{
    return out << err.what();
}

using Continuation = std::function<void(std::future<std::unique_ptr<Session>>)>;

class Pool {
public:
    Pool()
    {
        doInit();
        curlm_ = curl_multi_init();
    }

    ~Pool()
    {
        std::cout << "Pool destruction." << std::endl;
        if (curlm_ != nullptr)
        {
            curl_multi_cleanup(curlm_);
            curlm_ = nullptr;
        }
    }

    std::unique_ptr<Session> create(Method method, const std::string& url, const std::string& dnsServer, std::shared_ptr<std::atomic_bool> flag)
    {
        return std::make_unique<Session>(method, url, dnsServer, flag);
    }

    void add(Session& s)
    {
        s.start();
        check(curl_multi_add_handle(curlm_, s.handle()), "curl_multi_add_handle");
    }

    void remove(Session& s)
    {
        check(curl_multi_remove_handle(curlm_, s.handle()), "curl_multi_remove_handle");
    }

    void abortIfInProgress(Session& s)
    {
        if (s.getState() == Session::PERFORMING)
        {
            std::cout << ">> abort session" << std::endl;
            remove(s);
            s.abort();
        }
    }

    bool wait(int timeoutMs = 20)
    {
        auto ec = curl_multi_poll(curlm_, NULL, 0, timeoutMs, NULL);
        check(ec, "curl_multi_poll");
        return ec == CURLM_OK;
    }

    void wakeup()
    {
        check(curl_multi_wakeup(curlm_), "curl_multi_wakeup");
    }

    void check(CURLMcode ec, const char* msg)
    {
        if (ec != CURLM_OK)
        {
            std::stringstream err;
            err << "[curl error " << ec << "], " << msg << ": " << curl_multi_strerror(ec);
            std::cout << err.str();
        }
    }

    int step()
    {
        int active = 0;
        CURLMcode mc = curl_multi_perform(curlm_, &active);
        if (!mc)
        {
            CURLMsg* msg = 0;

            do {
                int nmsg = 0;
                msg = curl_multi_info_read(curlm_, &nmsg);
                if (msg && (msg->msg == CURLMSG_DONE))
                {
                    Session* session = nullptr;
                    curl_easy_getinfo(msg->easy_handle, CURLINFO_PRIVATE, &session);
                    assert(session);
                    session->finish(msg->data.result);
                    remove(*session);
                }

                std::cout << "-- active: " << active << ", number of message: " << nmsg << std::endl;
            }
            while(msg);
        }
        else
        {
            check(mc, "curl_multi_perform");
        }

        return active;
    }
private:
    CURLM* curlm_ = nullptr;
};

// Client Inteface
class IClient
{
public:
    virtual ~IClient() = default;

    virtual std::unique_ptr<Session> get(const std::string& url, const std::string& dnsServer, std::shared_ptr<std::atomic_bool> flag) = 0;
};

class AsyncClient: public IClient
{
public:
    AsyncClient()
    {
        // create thread
        thread_ = std::thread{[&]()
        {
            try
            {
                std::unique_lock<std::mutex> lock(mutex_);

                while (true)
                {
                    // Add incoming requests
                    for (auto& p : incoming_)
                    {
                        pool_.add(*std::get<0>(p));
                        processing_.push_back(std::move(p));
                    }
                    incoming_.clear();

                    // Check for aborted sessions
                    int numOfInprogressRequests = 0;
                    for (auto& p : processing_)
                    {
                        std::unique_ptr<Session>& session = std::get<0>(p);
                        if (aborted_ || session->checkAborted())
                        {
                            pool_.abortIfInProgress(*session);  // this will remove the handle and set abort state
                        }
                        else
                        {
                            ++numOfInprogressRequests;
                        }
                    }

                    // Do I/O
                    int active = 0;
                    if (numOfInprogressRequests > 0)
                    {
                        active = pool_.step();
                    }

                    // Handle completed requests
                    for (auto& p : processing_)
                    {
                        std::unique_ptr<Session>& session = std::get<0>(p);
                        Continuation& then = std::get<1>(p);

                        if (active < 1)
                        {
                            // abort in progress sessions that are not completed if curl is not active anymore.
                            pool_.abortIfInProgress(*session);
                        }

                        if (!session->inProgress())
                        {
                            std::promise<std::unique_ptr<Session>> prom;
                            if (!session->failed())
                            {
                                prom.set_value(std::move(session));
                            }
                            else
                            {
                                std::cout << " aborted: " << session->url() << ". status: " << session->getStatusStr() << std::endl;
                                prom.set_exception(std::make_exception_ptr(RequestError{std::move(session)}));
                            }

                            // trigger callback
                            then(prom.get_future());
                        }
                    }

                    // Remove handled requests(the session has been moved!)
                    processing_.remove_if([](auto& p)
                    {
                        return !std::get<0>(p);
                    });

                    // It's possible to have extra active requests due to DNS over HTTPS activity
                    assert(processing_.size() <= active);

                    // Exit thread if we're shutting down and all requests are handled
                    if (shutdown_ && active < 1)
                    {
                        assert(processing_.empty());
                        lock.unlock();
                        return;
                    }

                    lock.unlock();
                    if (pool_.wait())
                    {
                        lock.lock();
                    }
                    else
                    {
                        std::cout << "work thread failed to poll." << std::endl;
                        // break;
                        lock.lock();
                    }
                }

                std::cout << "work thread exiting ..." << std::endl;
            }
            catch (const std::exception& err)
            {
                std::cerr << "Uncaught exception in async client's work thread: " <<
                          err.what() << std::endl;
                std::terminate();
            }
        }};
    }

    ~AsyncClient() override
    {
        std::cout << "AsyncClient destruction." << std::endl;
        shutdown();
    }

    std::unique_ptr<Session> get(const std::string& url, const std::string& dnsServer, std::shared_ptr<std::atomic_bool> flag) override
    {
        auto prom = std::make_shared<std::promise<std::unique_ptr<Session>>>();
        auto then = [prom](std::future<std::unique_ptr<Session>> res) {
            try
            {
                prom->set_value(res.get());
            }
            catch (...)
            {
                prom->set_exception(std::current_exception());
            }
        };

        getAsync(url, dnsServer, flag, then);
        return prom->get_future().get();
    }

    void getAsync(const std::string& url, const std::string& dnsServer, std::shared_ptr<std::atomic_bool> flag, Continuation then)
    {

        std::unique_lock<std::mutex> lock(mutex_);
        if (shutdown_)
        {
            std::promise<std::unique_ptr<Session>> prom;
            prom.set_exception(std::make_exception_ptr(
                               std::runtime_error("can't perform request on a shut down client")));
            then(prom.get_future());
            return;
        }

        try
        {
            incoming_.emplace_back(pool_.create(Method::GET, url, dnsServer, flag), then);
        }
        catch (const std::exception& e)
        {
            std::promise<std::unique_ptr<Session>> prom;
            prom.set_exception(std::make_exception_ptr(
                std::runtime_error(std::string("Can't perform request due to error: ") + e.what())));
            then(prom.get_future());
        }
    }

    void shutdown()
    {
        std::unique_lock<std::mutex> lock(mutex_);
        if (!shutdown_)
        {
            // set flags
            shutdown_ = true;
            aborted_ = true;

            pool_.wakeup();
            lock.unlock();

            thread_.join();
        }
    }

private:

    using SessionContinuationPair = std::tuple<std::unique_ptr<Session>, Continuation>;

    std::mutex mutex_;

    std::list<SessionContinuationPair> incoming_;
    std::list<SessionContinuationPair> processing_;
    bool shutdown_ = false;
    bool aborted_ = false;

    std::thread thread_;
    Pool pool_;
};

int main(int argc, char* argv[])
{
    std::cout << "=== curl test ===" << std::endl;

    auto url = "https://httpbin.org/delay/5";
    auto dnsServer = "https://1.1.1.1/dns-query";

    auto test_abort_func = [=](IClient& cli, const std::string& dnsServer, int abortAfterMs) {
        auto abortFlag = std::make_shared<std::atomic_bool>(false);
        auto th = std::thread{[&]()
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(abortAfterMs));
            abortFlag->store(true);
        }};

        std::cout << "\n>> start a request that will be aborted." << std::endl;
        try
        {
            auto session = cli.get(url, dnsServer, abortFlag);
            std::cout << ">> abort test " << abortAfterMs << " ms: " << session->getStatusStr() << std::endl;
        }
        catch(RequestError& error)
        {
            std::cout << ">> abort test " << error.response().url() << ". status: " << error.response().getStatusStr() << std::endl;
        }

        std::cout << "\n>> start another request." << std::endl;
        auto flag = std::make_shared<std::atomic_bool>(false);
        cli.get(url, "", flag);
        th.join();
    };

    // test async client with DOH
    std::unique_ptr<IClient> client(new AsyncClient());
    test_abort_func(*client, dnsServer, 500);

    return 0;
}