AuroraRuntime/Source/IO/TLS/TLSContext.cpp

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

    File: TLSContext.cpp
    Date: 2022-8-24
    Author: Reece
***/
#include "TLS.hpp"
#include "TLSContext.hpp"
#include <Source/IO/Protocol/ProtocolStack.hpp>
#include <Aurora/IO/Net/NetExperimental.hpp>
#include <Source/IO/Net/AuNetSocket.hpp>

namespace Aurora::IO::TLS
{
    static mbedtls_entropy_context gEntropy;
    static mbedtls_ctr_drbg_context gCtrDrbg;
    static bool gTlsReady {};

    void TLSInit()
    {
        ::mbedtls_ctr_drbg_init(&gCtrDrbg);
        ::mbedtls_entropy_init(&gEntropy);

        int ret;

        if ((ret = ::mbedtls_ctr_drbg_seed(&gCtrDrbg,
                                           ::mbedtls_entropy_func,
                                           &gEntropy,
                                           (const unsigned char *)"ReeceWasHere",
                                           12)) != 0)
        {
            SysPushErrorNet("{}", ret);
            return;
        }

        gTlsReady = true;
    }

    static int TLSContextRecv(void *ctx,
                              unsigned char *buf,
                              size_t len)
    {
        return ((TLSContext *)ctx)->Read(buf, len);
    }

    static int TLSContextSend(void *ctx,
                              const unsigned char *buf,
                              size_t len)
    {
        return ((TLSContext *)ctx)->Write(buf, len);
    }

    TLSContext::TLSContext(const TLSMeta &meta) :
        channelRecv_(this),
        channelSend_(this),
        meta_(meta)
    {
        this->recvStack_ = AuMakeShared<Protocol::ProtocolStack>();
        this->sendStack_ = AuMakeShared<Protocol::ProtocolStack>();
    }

    TLSContext::TLSContext(const AuSPtr<Protocol::IProtocolStack> &pSendStack,
                           const AuSPtr<Protocol::IProtocolStack> &pRecvStack,
                           const TLSMeta &meta) :
        channelRecv_(this),
        channelSend_(this),
        recvStack_(AuStaticCast<Protocol::ProtocolStack>(pRecvStack)),
        sendStack_(AuStaticCast<Protocol::ProtocolStack>(pSendStack)),
        meta_(meta)
    {
    }

    TLSContext::~TLSContext()
    {
        this->Destroy();
    }

    //
    // mbedtls nonblocking interface
    // 
    //

    int TLSContext::Write(const void *pIn, AuUInt length)
    {
        return this->sendStack_->pDrainBuffer->Write(pIn, length);
    }

    int TLSContext::Read(void *pOut, AuUInt length)
    {
        auto tempReadBuffer = this->channelRecv_.pReadInByteBuffer.lock();
        if (!tempReadBuffer)
        {
            //SysPushErrorNet();
            return MBEDTLS_ERR_SSL_WANT_READ;
        }

        auto toRead = length;// AuMin<AuUInt>(length, this->channelRecv_.uBytesReadAvail - this->channelRecv_.uBytesRead);
        auto uBytesRead = tempReadBuffer->Read(pOut, toRead);
        if (!uBytesRead)
        {
            return MBEDTLS_ERR_SSL_WANT_READ;
        }

        this->channelRecv_.bHasRead = true;
        auto old = this->channelRecv_.uBytesRead;
        this->channelRecv_.uBytesRead += uBytesRead;
        return uBytesRead;
    }

    bool TLSContext::CheckCertificate(const AuMemoryViewRead &read)
    {
        if (!this->meta_.pCertPin)
        {
            return true;
        }

        return this->meta_.pCertPin->CheckCertificate(read);
    }

    //
    // tls context
    // 
    //

    void TLSContext::Init()
    {
        int ret;

        if (!this->sendStack_)
        {
            return;
        }

        if (!this->recvStack_)
        {
            return;
        }

        SysAssert(this->sendStack_->AddInterceptorEx(this->GetSendInterceptor(), this->meta_.uOutPageSize));
        SysAssert(this->recvStack_->AddInterceptorEx(this->GetRecvInterceptor(), this->meta_.uOutPageSize));
    
        ::mbedtls_ssl_init(&ssl);
        ::mbedtls_ssl_config_init(&conf);
        ::mbedtls_x509_crt_init(&cacert);

        if ((ret = ::mbedtls_ssl_config_defaults(&conf,
            this->meta_.bIsClient ? MBEDTLS_SSL_IS_CLIENT : MBEDTLS_SSL_IS_SERVER,
            this->meta_.transportProtocol == AuNet::ETransportProtocol::eProtocolUDP ? MBEDTLS_SSL_TRANSPORT_DATAGRAM : MBEDTLS_SSL_TRANSPORT_STREAM,
            MBEDTLS_SSL_PRESET_DEFAULT)) != 0)
        {
            SysPushErrorNet("{}", ret);
            return;
        }

        ::mbedtls_ssl_conf_authmode(&conf, MBEDTLS_SSL_VERIFY_OPTIONAL);
        ::mbedtls_ssl_conf_ca_cb(&conf, [](void *p_ctx,
                                 mbedtls_x509_crt const *child,
                                 mbedtls_x509_crt **candidate_cas) -> int
        {
            
            return ((TLSContext *)p_ctx)->CheckCertificate({ child->raw.p, child->raw.len }) ? 0 : -1;
        }, this);

        ::mbedtls_ssl_conf_rng(&conf, mbedtls_ctr_drbg_random, &gCtrDrbg);
        ::mbedtls_ssl_conf_dbg(&conf, [](void *, int, const char *as, int, const char *ad)
        {
            //AuLogDbg("{} <--> {}", as, ad);
        }, nullptr);

        if ((ret = ::mbedtls_ssl_setup(&ssl, &conf)) != 0)
        {
            SysPushErrorNet("{}", ret);
            return;
        }

        if (this->meta_.sSNIServerName.size())
        {
            if ((ret = ::mbedtls_ssl_set_hostname(&ssl, this->meta_.sSNIServerName.c_str())) != 0)
            {
                SysPushErrorNet("{}", ret);
                return;
            }
        }

        ::mbedtls_ssl_set_bio(&ssl, this, TLSContextSend, TLSContextRecv, nullptr);
        
        if (this->meta_.cipherSuites.size())
        {
            this->cipherSuites_.reserve(this->meta_.cipherSuites.size());
            for (const auto &cipher : this->meta_.cipherSuites)
            {
                this->cipherSuites_.push_back(cipher);
            }
        }
        else
        {
            auto &defaultCiphers = GetDefaultCipherSuites();
            this->cipherSuites_.reserve(defaultCiphers.size());
            for (const auto &cipher : defaultCiphers)
            {
                this->cipherSuites_.push_back(cipher);
            }
        }

        this->cipherSuites_.push_back(0);
        ((mbedtls_ssl_config *)ssl.private_conf/*fuck yourself*/)->private_ciphersuite_list = this->cipherSuites_.data();
    }

    void TLSContext::Destroy()
    {
        ::mbedtls_ssl_free(&ssl);
        ::mbedtls_ssl_config_free(&conf);
        ::mbedtls_x509_crt_free(&cacert);

        this->Attach({});
    }

    void TLSContext::OnClose()
    {
        this->bIsDead = true;

        if (auto pSocket = this->wpSocket_.lock())
        {
            pSocket->Shutdown();
        }

        AuResetMember(this->meta_);
    }

    void TLSContext::OnFatal()
    {
        this->bIsDead = true;
        this->bIsFatal = true;

        if (auto pSocket = this->wpSocket_.lock())
        {
            AuDynamicCast<AuNet::Socket, AuNet::ISocket>(pSocket)->SendErrorBeginShutdown(AuNet::ENetworkError::eTLSError);
        }

        AuResetMember(this->meta_);
    }

    //
    // public api
    // 
    //

    AuSPtr<Protocol::IProtocolStack> TLSContext::ToReadStack()
    {
        return this->recvStack_;
    }
    
    AuSPtr<Protocol::IProtocolStack> TLSContext::ToWriteStack()
    {
        return this->sendStack_;
    }

    AuSPtr<Protocol::IProtocolInterceptorEx> TLSContext::GetRecvInterceptor()
    {
        return AuSPtr<Protocol::IProtocolInterceptorEx>(AuSharedFromThis(), &this->channelRecv_);
    }

    AuSPtr<Protocol::IProtocolInterceptorEx> TLSContext::GetSendInterceptor()
    {
        return AuSPtr<Protocol::IProtocolInterceptorEx>(AuSharedFromThis(), &this->channelSend_);
    }
    
    void TLSContext::Attach(const AuSPtr<Net::ISocket> &pSocket)
    {
        if (!pSocket)
        {
            if (auto pOldSocket = this->wpSocket_.lock())
            {
                this->wpSocket_.reset();
                pOldSocket->ToChannel()->SpecifyRecvProtocol({});
                pOldSocket->ToChannel()->SpecifySendProtocol({});
                // TODO (Reece): Shutdown hook
            }
            return;
        }

        this->wpSocket_ = pSocket;
        pSocket->ToChannel()->SpecifyRecvProtocol(ToReadStack());
        pSocket->ToChannel()->SpecifySendProtocol(ToWriteStack());
        // TODO (Reece): Shutdown hook
    }

    void TLSContext::StartHandshake()
    {
        this->channelRecv_.TryHandshake();
    }

    void TLSContext::StartClose()
    {
    }

    bool TLSContext::HasCompletedHandshake()
    {
        return this->channelRecv_.HasCompletedHandshake();
    }

    bool TLSContext::HasEnded()
    {
        return this->bIsDead;
    }

    bool TLSContext::HasFailed()
    {
        return this->bIsFatal;
    }

    int TLSContext::GetFatalErrorCode()
    {
        return this->iFatalError;
    }

    AUKN_SYM AuSPtr<ITLSContext> NewTLSContext(const TLSMeta &meta)
    {
        auto pTlsContext = AuMakeShared<TLSContext>(meta);
        if (!pTlsContext)
        {
            return {};
        }

        pTlsContext->Init();

        return pTlsContext;
    }

    AUKN_SYM AuSPtr<ITLSContext> NewTLSContextEx(const AuSPtr<Protocol::IProtocolStack> &pSendStack,
                                                 const AuSPtr<Protocol::IProtocolStack> &pRecvStack,
                                                 const TLSMeta &meta)
    {
        auto pTlsContext = AuMakeShared<TLSContext>(pSendStack, pRecvStack, meta);
        if (!pTlsContext)
        {
            return {};
        }

        pTlsContext->Init();

        return pTlsContext;
    }
}