AuroraRuntime/Source/IO/Net/AuNetResolver.NT.cpp

/***
    Copyright (C) 2022 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: AuNetResolver.NT.cpp
    Date: 2022-8-17
    Author: Reece
***/
#include "Networking.hpp"
#include "AuNetResolver.NT.hpp"
#include "AuNetEndpoint.hpp"

namespace Aurora::IO::Net
{
    NetResolver::NetResolver(const AuSPtr<INetWorker> &pWorker) :
        pWorker_(pWorker)
    {
    
    }

    NetResolver::NetResolver(const AuSPtr<INetWorker> &pWorker,
                             const AuSPtr<AuAsync::PromiseCallback<AuList<IPAddress>, NetError>> &pCompletion) :
        pCompletion_(pCompletion),
        pWorker_(pWorker)
    {

    }

    NetResolver::~NetResolver()
    {
    }

    bool NetResolver::Start()
    {
        ADDRINFOEXW infoEx {};
        int iInfoEx { 0 };

        if (!this->pEvent)
        {
            return false;
        }

        if (this->bA && this->bAAAA)
        {
            infoEx.ai_family = AF_UNSPEC;
            //infoEx.ai_flags = AI_ALL;
        }
        else if (this->bA && !this->bAAAA)
        {
            infoEx.ai_family = AF_INET;
        }
        else if (this->bAAAA)
        {
            infoEx.ai_family = AF_INET6;
            //infoEx.ai_flags = AI_V4MAPPED; // beats returning nothing...
        }
        else
        {
            return false;
        }

        int iRet {};
        HANDLE hHandle = NULL;

        if (pGetAddrInfoExW)
        {
            if (pGetAddrInfoExCancel)
            {
                iRet = pGetAddrInfoExW(AuLocale::ConvertFromUTF8(this->hostname).data(),
                                       nullptr,
                                       NS_DNS,
                                       nullptr,
                                       &infoEx,
                                       &this->resultHandle_,
                                       nullptr,
                                       &this->overlapped,
                                       NULL,
                                       &hHandle);
            }
            else
            {
                auto pShared = this->SharedFromThis();

                auto pThread = AuThreads::ThreadUnique(AuThreads::ThreadInfo(
                    AuMakeShared<AuThreads::IThreadVectorsFunctional>(AuThreads::IThreadVectorsFunctional::OnEntry_t(std::bind([=]()
                    {
                        int iRet {};
                        if ((iRet = pGetAddrInfoExW(AuLocale::ConvertFromUTF8(pShared->hostname).data(),
                                                    nullptr,
                                                    NS_DNS,
                                                    nullptr,
                                                    &infoEx,
                                                    &pShared->resultHandle_,
                                                    nullptr,
                                                    nullptr,
                                                    nullptr,
                                                    nullptr)) != ERROR_SUCCESS)
                        {
                            pShared->bForceError_ = true;
                            pShared->uOsError = iRet;
                            pShared->error_ = pShared->ToError();
                        }

                        pShared->pEvent->Set();
                    })),
                    AuThreads::IThreadVectorsFunctional::OnExit_t{}),
                    "Legacy Win32 DNS Request"
                ));

                if (!pThread)
                {
                    return false;
                }

                pThread->Run();
                pThread->Detach();

                iRet = ERROR_IO_PENDING;
            }
        }
        else if (pgetaddrinfo)
        {
            auto pShared = this->SharedFromThis();

            auto pThread = AuThreads::ThreadUnique(AuThreads::ThreadInfo(
                AuMakeShared<AuThreads::IThreadVectorsFunctional>(AuThreads::IThreadVectorsFunctional::OnEntry_t(std::bind([=]()
                {
                    int iRet {};
                    ADDRINFOA infoA { 0 };
                    PADDRINFOA pInfoRet {};
                    infoA.ai_family = infoEx.ai_family;

                    iRet = pgetaddrinfo(pShared->hostname.c_str(),
                                        nullptr,
                                        &infoA,
                                        &pInfoRet);

                    if (iRet == 0)
                    {
                        auto pCurrent = pInfoRet;

                        while (pCurrent)
                        {
                            NetEndpoint endpoint;
                            AuMemcpy(endpoint.hint, pCurrent->ai_addr, pCurrent->ai_addrlen);
                            DeoptimizeEndpoint(endpoint);

                            if (!AuTryInsert(pShared->processedIps_, endpoint.ip))
                            {
                                pShared->bForceError_ = true;
                                break;
                            }

                            pCurrent = pCurrent->ai_next;
                        }
                    }
                    else
                    {
                        pShared->bForceError_ = true;
                        pShared->uOsError = iRet;
                        pShared->error_ = pShared->ToError();
                    }

                    if (pInfoRet && pfreeaddrinfo)
                    {
                        pfreeaddrinfo(pInfoRet);
                    }

                    pShared->pEvent->Set();
                })),
                AuThreads::IThreadVectorsFunctional::OnExit_t{}),
                "Legacy Win32 DNS Request"
            ));

            if (!pThread)
            {
                return false;
            }

            pThread->Run();
            pThread->Detach();

            iRet = ERROR_IO_PENDING;
        }
        else
        {
            return false;
        }

        this->hName_ = hHandle;
        return this->FinishOperationEx(AuSharedFromThis(),
                                       this->pWorker_,
                                       iRet == 0,
                                       iRet);
    }

    void NetResolver::SetCompletion(const AuSPtr<AuAsync::PromiseCallback<AuList<IPAddress>, NetError>> &callback)
    {
        this->pCompletion_ = callback;
    }

    void NetResolver::OnOverlappedComplete()
    {
        if (AuExchange(this->bHasCompleted_, true))
        {
            return;
        }
        
        if (this->Unpack())
        {
            if (this->pCompletion_)
            {
                try
                {
                    this->pCompletion_->OnSuccess((void *)&this->processedIps_);
                }
                catch (...)
                {
                    SysPushErrorCatch();
                }
            }
        }
        else
        {
            if (this->pCompletion_)
            {
                try
                {
                    if (this->error_.netError == ENetworkError::eEnumInvalid)
                    {
                        this->error_ = ENetworkError::eAsyncError;
                    }
                    this->pCompletion_->OnFailure((void *)&this->error_);
                }
                catch (...)
                {
                    SysPushErrorCatch();
                }
            }
        }
    }

    void NetResolver::OnOverlappedFailure(const NetError &error)
    {
        this->error_ = error;
        if (this->pCompletion_)
        {
            try
            {
                this->pCompletion_->OnFailure((void *)&this->error_);
            }
            catch (...)
            {
                SysPushErrorCatch();
            }
        }
        this->bHasCompleted_ = true;
    }

    bool NetResolver::HasFinished()
    {
        return this->bHasCompleted_;
    }

    void NetResolver::Cancel()
    {
        auto hHandle = AuExchange(this->hName_, (HANDLE)NULL);
        if (hHandle != NULL)
        {
            if (pGetAddrInfoExCancel)
            {
                pGetAddrInfoExCancel(&hHandle);
            }
        }

        while (!this->HasComplete())
        {
            IOYield();
        }
    }

    const NetError &NetResolver::GetError()
    {
        return this->error_;
    }

    bool NetResolver::Unpack()
    {
        auto pCurrent = this->resultHandle_;

        if (this->bForceError_)
        {
            return false;
        }

        while (pCurrent)
        {
            NetEndpoint endpoint;
            AuMemcpy(endpoint.hint, pCurrent->ai_addr, pCurrent->ai_addrlen);
            DeoptimizeEndpoint(endpoint);

            if (!AuTryInsert(this->processedIps_, endpoint.ip))
            {
                return false;
            }

            pCurrent = pCurrent->ai_next;
        }

        if (this->resultHandle_ && pFreeAddrInfoExW)
        {
            pFreeAddrInfoExW(this->resultHandle_);
        }

        return true;
    }
}