/*
 * SRT - Secure, Reliable, Transport
 * Copyright (c) 2018 Haivision Systems Inc.
 * 
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 * 
 */

/*****************************************************************************
written by
   Haivision Systems Inc.
 *****************************************************************************/

#include "platform_sys.h"

#include <cstring>
#include <string>
#include <sstream>
#include <iterator>

#include "udt.h"
#include "utilities.h"
#include <haicrypt.h>
#include "crypto.h"
#include "logging.h"
#include "core.h"
#include "api.h"

using namespace srt_logging;

#define SRT_MAX_KMRETRY     10
 
//#define SRT_CMD_KMREQ       3           /* HaiCryptTP SRT Keying Material */
//#define SRT_CMD_KMRSP       4           /* HaiCryptTP SRT Keying Material ACK */
#define SRT_CMD_KMREQ_SZ    HCRYPT_MSG_KM_MAX_SZ          /* */
#if     SRT_CMD_KMREQ_SZ > SRT_CMD_MAXSZ
#error  SRT_CMD_MAXSZ too small
#endif
/*      Key Material Request (Network Order)
        See HaiCryptTP SRT (hcrypt_xpt_srt.c)
*/

// 10* HAICRYPT_DEF_KM_PRE_ANNOUNCE
const int SRT_CRYPT_KM_PRE_ANNOUNCE SRT_ATR_UNUSED = 0x10000;

namespace srt_logging
{
std::string KmStateStr(SRT_KM_STATE state)
{
    switch (state)
    {
#define TAKE(val) case SRT_KM_S_##val : return #val
        TAKE(UNSECURED);
        TAKE(SECURED);
        TAKE(SECURING);
        TAKE(NOSECRET);
        TAKE(BADSECRET);
#ifdef ENABLE_AEAD_API_PREVIEW
        TAKE(BADCRYPTOMODE);
#endif
#undef TAKE
    default:
        {
            char buf[256];
#if defined(_MSC_VER) && _MSC_VER < 1900
            _snprintf(buf, sizeof(buf) - 1, "??? (%d)", state);
#else
            snprintf(buf, sizeof(buf), "??? (%d)", state);
#endif
            return buf;
        }
    }
}
} // namespace

using srt_logging::KmStateStr;

void srt::CCryptoControl::globalInit()
{
#ifdef SRT_ENABLE_ENCRYPTION
    // We need to force the Cryspr to be initialized during startup to avoid the
    // possibility of multiple threads initializing the same static data later on.
    HaiCryptCryspr_Get_Instance();
#endif
}

bool srt::CCryptoControl::isAESGCMSupported()
{
#ifdef SRT_ENABLE_ENCRYPTION
    return HaiCrypt_IsAESGCM_Supported() != 0;
#else
    return false;
#endif
}

#if ENABLE_LOGGING
std::string srt::CCryptoControl::FormatKmMessage(std::string hdr, int cmd, size_t srtlen)
{
    std::ostringstream os;
    os << hdr << ": cmd=" << cmd << "(" << (cmd == SRT_CMD_KMREQ ? "KMREQ":"KMRSP") <<") len="
        << size_t(srtlen*sizeof(int32_t)) << " KmState: SND="
        << KmStateStr(m_SndKmState)
        << " RCV=" << KmStateStr(m_RcvKmState);
    return os.str();
}
#endif

void srt::CCryptoControl::updateKmState(int cmd, size_t srtlen SRT_ATR_UNUSED)
{
    if (cmd == SRT_CMD_KMREQ)
    {
        if ( SRT_KM_S_UNSECURED == m_SndKmState)
        {
            m_SndKmState = SRT_KM_S_SECURING;
        }
        LOGP(cnlog.Note, FormatKmMessage("sendSrtMsg", cmd, srtlen));
    }
    else
    {
        LOGP(cnlog.Note, FormatKmMessage("sendSrtMsg", cmd, srtlen));
    }
}

void srt::CCryptoControl::createFakeSndContext()
{
    if (!m_iSndKmKeyLen)
        m_iSndKmKeyLen = 16;

    if (!createCryptoCtx((m_hSndCrypto), m_iSndKmKeyLen, HAICRYPT_CRYPTO_DIR_TX, false))
    {
        HLOGC(cnlog.Debug, log << "Error: Can't create fake crypto context for sending - sending will return ERROR!");
        m_hSndCrypto = 0;
    }
}

#ifdef SRT_ENABLE_ENCRYPTION

int srt::CCryptoControl::processSrtMsg_KMREQ(
        const uint32_t* srtdata, size_t bytelen, int hsv, unsigned srtv,
        uint32_t pw_srtdata_out[], size_t& w_srtlen)
{
    unsigned char* kmdata = reinterpret_cast<unsigned char*>(pw_srtdata_out);

    // The side that has received KMREQ is always an HSD_RESPONDER, regardless of
    // what has called this function. The HSv5 handshake only enforces bidirectional
    // connection.

    // They actually mean the same thing in HSv5, but removal of bidirectional can be
    // only done when the compat with <1.3.0 is allowed to be broken.
    const bool bidirectional = hsv > CUDT::HS_VERSION_UDT4;
    const bool kmx_update = m_hRcvCrypto;
    SRT_KM_STATE failure_state = m_KmSecret.len == 0 ? SRT_KM_S_NOSECRET : SRT_KM_S_BADSECRET;
    bool bUseGCM = false;

    // TRY-BLOCK, with THROW done by "goto Error".
    {
        if (bytelen % sizeof(uint32_t) != 0 || bytelen > HCRYPT_MSG_KM_MAX_SZ)
        {
            LOGC(cnlog.Error, log << "processSrtMsg_KMREQ: size of the KM (" << bytelen << ") is too high, must be < " << HCRYPT_MSG_KM_MAX_SZ);
            goto Error;
        }

        // Default successful settings: original length and contents
        w_srtlen = bytelen/sizeof(srtdata[SRT_KMR_KMSTATE]);

        /* All 32-bit msg fields swapped on reception
         * But HaiCrypt expect network order message
         * Re-swap to cancel it.
         */
        HtoNLA((pw_srtdata_out), srtdata, w_srtlen);

        if (!kmx_update) // Only in initial/handshake
        {
            // If this is NOT changed anywhere later, failure_state value will be used
            m_RcvKmState = SRT_KM_S_SECURING;
        }

        if ((m_iCryptoMode == CSrtConfig::CIPHER_MODE_AUTO && kmdata[HCRYPT_MSG_KM_OFS_CIPHER] == HCRYPT_CIPHER_AES_GCM)
                || (m_iCryptoMode == CSrtConfig::CIPHER_MODE_AES_GCM))
            bUseGCM = true;

        m_bUseGcm153 = srtv <= SrtVersion(1, 5, 3);

        // INITIAL ACTIONS (first time KMREQ received):
        // If encryption is on (we know that by having m_KmSecret nonempty), create
        // the crypto context (if bidirectional, create for both sending and receiving).
        // Both crypto contexts should be set with the same length of the key.
        // Report SRT_CMD_NONE if no response is to be sent or SRT_CMD_KMRSP if the
        // response is filled and ready (including erroneous).
        if (bytelen <= HCRYPT_MSG_KM_OFS_SALT)  //Sanity on message
        {
            LOGC(cnlog.Error, log << "processSrtMsg_KMREQ: size of the KM (" << bytelen << ") is too small, must be >" << HCRYPT_MSG_KM_OFS_SALT);
            goto Error;
        }

        size_t sek_len = hcryptMsg_KM_GetSekLen(kmdata);
        if (sek_len == 0)
        {
            LOGC(cnlog.Error, log << "processSrtMsg_KMREQ: Received SEK is empty - REJECTING!");
            goto Error;
        }

        bool new_keylen = m_iSndKmKeyLen != sek_len;

        if (kmx_update && new_keylen)
        {
            LOGC(cnlog.Error, log << "processSrtMsg_KMREQ: KMX refresh came in with a DIFFERENT KEY LEN: " << sek_len);
            goto Error;
        }

        HLOGC(cnlog.Debug, log << "KMREQ: getting SEK and creating receiver crypto sek_len=" << sek_len
                << " local=" << m_iSndKmKeyLen
                << (new_keylen ? " - OVERRIDE with received SEK len" : ""));

        // Write the key length.
        // Overwrite the key length anyway - it doesn't make sense to somehow
        // keep the original setting because it will only make KMX impossible.
        m_iSndKmKeyLen = m_iRcvKmKeyLen = sek_len;

        // This is checked only now so that the SRTO_PBKEYLEN return always the correct value,
        // even if encryption is not possible because Agent didn't set a password, or supplied
        // a wrong password.
        if (!hasPassphrase())  //We have a shared secret <==> encryption is on
        {
            LOGC(cnlog.Warn, log << "processSrtMsg_KMREQ: Agent does not declare encryption - won't decrypt incoming packets!");
            goto Error;
        }

        if (!createCryptoCtx((m_hRcvCrypto), m_iRcvKmKeyLen, HAICRYPT_CRYPTO_DIR_RX, bUseGCM))
        {
            LOGC(cnlog.Error, log << "processSrtMsg_KMREQ: Can't create RCV CRYPTO CTX - must reject...");
            failure_state = SRT_KM_S_NOSECRET;
            goto Error;
        }

        if (!kmx_update)
        {
            // Deduce resulting mode.
            m_iCryptoMode = bUseGCM ? CSrtConfig::CIPHER_MODE_AES_GCM : CSrtConfig::CIPHER_MODE_AES_CTR;
            HLOGC(cnlog.Debug, log << "processSrtMsg_KMREQ: created RX ENC with KeyLen=" << m_iRcvKmKeyLen);
        }

        // We have both sides set with password, so both are pending for security
        int rc = HaiCrypt_Rx_Process(m_hRcvCrypto, kmdata, bytelen, NULL, NULL, 0);
        LOGC(cnlog.Note, log << FormatKmMessage("processSrtMsg_KMREQ", SRT_CMD_KMREQ, bytelen) << " result: " << rc);

        // Since now, when CCryptoControl::decrypt() encounters an error, it will print it, ONCE,
        // until the next KMREQ is received as a key regeneration.
        m_bErrorReported = false;

        if (rc >= HAICRYPT_OK)
        {
            m_RcvKmState = SRT_KM_S_SECURED;
            HLOGC(cnlog.Debug, log << "KMREQ/rcv: (snd) Rx process successful - SECURED.");
        }
        else
        {
            switch(rc)
            {
            case HAICRYPT_ERROR_WRONG_SECRET: //Unmatched shared secret to decrypt wrapped key
                failure_state = SRT_KM_S_BADSECRET;
                //Send status KMRSP message to tel error
                LOGC(cnlog.Warn, log << "KMREQ/rcv: (snd) Rx process failure - BADSECRET");
                break;
            case HAICRYPT_ERROR_CIPHER:
#ifdef ENABLE_AEAD_API_PREVIEW
                failure_state = SRT_KM_S_BADCRYPTOMODE;
#else
                failure_state = SRT_KM_S_BADSECRET; // Use "bad secret" as a fallback.
#endif
                LOGC(cnlog.Warn, log << "KMREQ/rcv: (snd) Rx process failure - BADCRYPTOMODE");
                break;
            case HAICRYPT_ERROR: //Other errors
            default:
                failure_state = SRT_KM_S_NOSECRET;
                LOGC(cnlog.Warn, log << "KMREQ/rcv: (snd) Rx process failure (IPE) - NOSECRET");
                break;
            }

            if (!kmx_update)
            {
                // Only initially, set this also to SND state.
                m_SndKmState = failure_state;
            }
            goto Error;
        }

        // Configure the sender context also, if it succeeded to configure the
        // receiver context and we are using bidirectional mode.
        if (bidirectional)
        {
            // Note: 'bidirectional' means that we want a bidirectional key update,
            // which happens only and exclusively with HSv5 handshake - not when the
            // usual key update through UMSG_EXT+SRT_CMD_KMREQ was done (which is used
            // in HSv4 versions also to initialize the first key, unlike HSv5).
            if (m_RcvKmState == SRT_KM_S_SECURED)
            {
                if (m_SndKmState == SRT_KM_S_SECURING && !m_hSndCrypto)
                {
                    m_iSndKmKeyLen = m_iRcvKmKeyLen;
                    if (HaiCrypt_Clone(m_hRcvCrypto, HAICRYPT_CRYPTO_DIR_TX, &m_hSndCrypto) != HAICRYPT_OK)
                    {
                        LOGC(cnlog.Error, log << "processSrtMsg_KMREQ: Can't create SND CRYPTO CTX - WILL NOT SEND-ENCRYPT correctly!");
                        m_SndKmState = failure_state;
                    }
                    else
                    {
                        m_SndKmState = SRT_KM_S_SECURED;
                    }

                    LOGC(cnlog.Note, log << FormatKmMessage("processSrtMsg_KMREQ", SRT_CMD_KMREQ, bytelen)
                            << " SndKeyLen=" << m_iSndKmKeyLen
                            << " TX CRYPTO CTX CLONED FROM RX"
                        );

                    // Write the KM message into the field from which it will be next sent.
                    memcpy((m_SndKmMsg[0].Msg), kmdata, bytelen);
                    m_SndKmMsg[0].MsgLen = bytelen;
                    m_SndKmMsg[0].iPeerRetry = 0; // Don't start sending them upon connection :)
                }
                else
                {
                    HLOGC(cnlog.Debug, log << "processSrtMsg_KMREQ: NOT cloning RX to TX crypto: already in "
                            << KmStateStr(m_SndKmState) << " state");
                }
            }
            else
            {
                HLOGP(cnlog.Debug, "processSrtMsg_KMREQ: NOT SECURED - not replaying failed security association to TX CRYPTO CTX");
            }
        }
        else
        {
            HLOGC(cnlog.Debug, log << "processSrtMsg_KMREQ: NOT REPLAYING the key update to TX CRYPTO CTX.");
        }

        if (m_hRcvCrypto != NULL)
            HaiCrypt_UpdateGcm153(m_hRcvCrypto, m_bUseGcm153);
        if (m_hSndCrypto != NULL)
            HaiCrypt_UpdateGcm153(m_hSndCrypto, m_bUseGcm153);

        // NOTE: The "loop" gets exit with return HERE.
        return SRT_CMD_KMRSP;
    }
Error: // CATCH POINT

    // NOTE: This is set, but if we return NONE, this will not be
    // taken into account anyhow. The state is returned this way only for
    // unit tests so that they recognize the result of the call.
    w_srtlen = 1;
    pw_srtdata_out[SRT_KMR_KMSTATE] = failure_state;

    if (!kmx_update)
    {
        // Set the appropriate error, if it wasn't already set before 
        m_RcvKmState = failure_state;
        if (bidirectional && hasPassphrase())
        {
            // If the Forward KMX process has failed, the reverse-KMX process was not done at all.
            // This will lead to incorrect object configuration and will fail to properly declare
            // the transmission state.
            // Create the "fake crypto" with the passphrsae you currently have.
            createFakeSndContext();
        }

        return SRT_CMD_KMRSP;
    }
    return SRT_CMD_NONE;
}

int srt::CCryptoControl::processSrtMsg_KMRSP(const uint32_t* srtdata, size_t len, unsigned srtv)
{
    // Validate the wire-supplied length before using it:
    //  - oversize would overflow the fixed-size stack buffer below;
    //  - non-word-aligned or too-small payloads are malformed by protocol and would
    //    feed uninitialised stack into downstream key-matching logic.
    if (len > SRT_CMD_MAXSZ
        || len < sizeof(uint32_t)
        || (len % sizeof(uint32_t)) != 0)
    {
        LOGC(cnlog.Error, log << "processSrtMsg_KMRSP: malformed len " << len
                              << " (must be a non-zero multiple of " << sizeof(uint32_t)
                              << ", up to " << SRT_CMD_MAXSZ << ") - rejecting");
        return SRT_CMD_NONE;
    }

    /* All 32-bit msg fields (if present) swapped on reception
     * But HaiCrypt expect network order message
     * Re-swap to cancel it.
     */
    uint32_t srtd[SRTDATA_MAXSIZE];
    size_t srtlen = len/sizeof(uint32_t);
    HtoNLA(srtd, srtdata, srtlen);

    int retstatus = -1;

    // Since now, when CCryptoControl::decrypt() encounters an error, it will print it, ONCE,
    // until the next KMREQ is received as a key regeneration.
    m_bErrorReported = false;

    if (srtlen == 1) // Error report. Set accordingly.
    {
        SRT_KM_STATE peerstate = SRT_KM_STATE(srtd[SRT_KMR_KMSTATE]); /* Bad or no passphrase */
        m_SndKmMsg[0].iPeerRetry = 0;
        m_SndKmMsg[1].iPeerRetry = 0;

        switch (peerstate)
        {
        case SRT_KM_S_BADSECRET:
            m_SndKmState = m_RcvKmState = SRT_KM_S_BADSECRET;
            retstatus = -1;
            break;

            // Default embraces two cases:
            // NOSECRET: this KMRSP was sent by secured Peer, but Agent supplied no password.
            // UNSECURED: this KMRSP was sent by unsecure Peer because Agent sent KMREQ.

        case SRT_KM_S_NOSECRET:
            // This means that the peer did not set the password, while Agent did.
            m_RcvKmState = SRT_KM_S_UNSECURED;
            m_SndKmState = SRT_KM_S_NOSECRET;
            retstatus = -1;
            break;

        case SRT_KM_S_UNSECURED:
            // This means that KMRSP was sent without KMREQ, to inform the Agent,
            // that the Peer, unlike Agent, does use password. Agent can send then,
            // but can't decrypt what Peer would send.
            m_RcvKmState = SRT_KM_S_NOSECRET;
            m_SndKmState = SRT_KM_S_UNSECURED;
            retstatus = 0;
            break;
#ifdef ENABLE_AEAD_API_PREVIEW
        case SRT_KM_S_BADCRYPTOMODE:
            // The peer expects to use a different cryptographic mode (e.g. AES-GCM, not AES-CTR).
            m_RcvKmState = SRT_KM_S_BADCRYPTOMODE;
            m_SndKmState = SRT_KM_S_BADCRYPTOMODE;
            retstatus = -1;
            break;
#endif

        default:
            LOGC(cnlog.Fatal, log << "processSrtMsg_KMRSP: IPE: unknown peer error state: "
                    << KmStateStr(peerstate) << " (" << int(peerstate) << ")");
            m_RcvKmState = SRT_KM_S_NOSECRET;
            m_SndKmState = SRT_KM_S_NOSECRET;
            retstatus = -1; //This is IPE
            break;
        }

        LOGC(cnlog.Warn, log << "processSrtMsg_KMRSP: received failure report. STATE: " << KmStateStr(m_RcvKmState));
    }
    else
    {
        HLOGC(cnlog.Debug, log << "processSrtMsg_KMRSP: received key response len=" << len);
        // XXX INSECURE << ": [" << FormatBinaryString((uint8_t*)srtd, len) << "]";
        bool key1 = getKmMsg_acceptResponse(0, srtd, len);
        bool key2 = true;
        if ( !key1 )
            key2 = getKmMsg_acceptResponse(1, srtd, len); // <--- NOTE SEQUENCING!

        if (key1 || key2)
        {
            m_SndKmState = m_RcvKmState = SRT_KM_S_SECURED;
            HLOGC(cnlog.Debug, log << "processSrtMsg_KMRSP: KM response matches " << (key1 ? "EVEN" : "ODD") << " key");
            retstatus = 1;
        }
        else
        {
            retstatus = -1;
            LOGC(cnlog.Error, log << "processSrtMsg_KMRSP: IPE??? KM response key matches no key");
            /* XXX INSECURE
            LOGC(cnlog.Error, log << "processSrtMsg_KMRSP: KM response: [" << FormatBinaryString((uint8_t*)srtd, len)
                << "] matches no key 0=[" << FormatBinaryString((uint8_t*)m_SndKmMsg[0].Msg, m_SndKmMsg[0].MsgLen)
                << "] 1=[" << FormatBinaryString((uint8_t*)m_SndKmMsg[1].Msg, m_SndKmMsg[1].MsgLen) << "]");
                */

            m_SndKmState = m_RcvKmState = SRT_KM_S_BADSECRET;
        }
        HLOGC(cnlog.Debug, log << "processSrtMsg_KMRSP: key[0]: len=" << m_SndKmMsg[0].MsgLen << " retry=" << m_SndKmMsg[0].iPeerRetry
            << "; key[1]: len=" << m_SndKmMsg[1].MsgLen << " retry=" << m_SndKmMsg[1].iPeerRetry);

        m_bUseGcm153 = srtv <= SrtVersion(1, 5, 3);
        if (m_hRcvCrypto != NULL)
            HaiCrypt_UpdateGcm153(m_hRcvCrypto, m_bUseGcm153);
        if (m_hSndCrypto != NULL)
            HaiCrypt_UpdateGcm153(m_hSndCrypto, m_bUseGcm153);
    }

    LOGP(cnlog.Note, FormatKmMessage("processSrtMsg_KMRSP", SRT_CMD_KMRSP, len));

    return retstatus;
}

#else

int srt::CCryptoControl::processSrtMsg_KMREQ(
        const uint32_t*, size_t, int, unsigned, // ignore input
        uint32_t pw_srtdata_out[], size_t& w_srtlen)
{
    // It's ok that this is reported as error because this happens in a scenario,
    // when non-encryption-enabled SRT application is contacted by encryption-enabled SRT
    // application which tries to make a security association.
    LOGC(cnlog.Warn, log << "processSrtMsg_KMREQ: Encryption not enabled at compile time - must reject...");
    m_RcvKmState = SRT_KM_S_NOSECRET;
    pw_srtdata_out[SRT_KMR_KMSTATE] = m_RcvKmState;
    w_srtlen = 1;

    return SRT_CMD_KMRSP;
}

int srt::CCryptoControl::processSrtMsg_KMRSP(const uint32_t*, size_t, unsigned)
{
    LOGP(cnlog.Error, "processSrtMsg_KMRSP: Encryption not enabled at compile time; not expected to receive SRT_CMD_KMRSP");
    return SRT_CMD_NONE;
}
#endif

void srt::CCryptoControl::sendKeysToPeer(CUDT* sock SRT_ATR_UNUSED, int iSRTT SRT_ATR_UNUSED)
{
    sync::ScopedLock lck(m_mtxLock);
    if (!m_hSndCrypto || m_SndKmState == SRT_KM_S_UNSECURED)
    {
        HLOGC(cnlog.Debug, log << "sendKeysToPeer: NOT sending/regenerating keys: "
                << (m_hSndCrypto ? "CONNECTION UNSECURED" : "NO TX CRYPTO CTX created"));
        return;
    }
#ifdef SRT_ENABLE_ENCRYPTION
    srt::sync::steady_clock::time_point now = srt::sync::steady_clock::now();
    /*
     * Crypto Key Distribution to peer:
     * If...
     * - we want encryption; and
     * - we have not tried more than CSRTCC_MAXRETRY times (peer may not be SRT); and
     * - and did not get answer back from peer; and
     * - last sent Keying Material req should have been replied (RTT*1.5 elapsed);
     * then (re-)send handshake request.
     */
    if (((m_SndKmMsg[0].iPeerRetry > 0) || (m_SndKmMsg[1].iPeerRetry > 0))
        && ((m_SndKmLastTime + srt::sync::microseconds_from((iSRTT * 3)/2)) <= now))
    {
        for (int ki = 0; ki < 2; ki++)
        {
            if (m_SndKmMsg[ki].iPeerRetry > 0 && m_SndKmMsg[ki].MsgLen > 0)
            {
                m_SndKmMsg[ki].iPeerRetry--;
                HLOGC(cnlog.Debug, log << "sendKeysToPeer: SENDING ki=" << ki << " len=" << m_SndKmMsg[ki].MsgLen
                        << " retry(updated)=" << m_SndKmMsg[ki].iPeerRetry);
                m_SndKmLastTime = now;
                sock->sendSrtMsg(SRT_CMD_KMREQ, (uint32_t *)m_SndKmMsg[ki].Msg, m_SndKmMsg[ki].MsgLen / sizeof(uint32_t));
            }
        }
    }
#endif
}

void srt::CCryptoControl::regenCryptoKm(CUDT* sock SRT_ATR_UNUSED, bool bidirectional SRT_ATR_UNUSED)
{
#ifdef SRT_ENABLE_ENCRYPTION
    sync::ScopedLock lck(m_mtxLock);
    if (!m_hSndCrypto)
        return;

    void *out_p[2];
    size_t out_len_p[2];
    int nbo = HaiCrypt_Tx_ManageKeys(m_hSndCrypto, out_p, out_len_p, 2);
    int sent = 0;

    HLOGC(cnlog.Debug, log << "regenCryptoKm: regenerating crypto keys nbo=" << nbo <<
            " THEN=" << (sock ? "SEND" : "KEEP") << " DIR=" << (bidirectional ? "BOTH" : "SENDER"));

    for (int i = 0; i < nbo && i < 2; i++)
    {
        /*
         * New connection keying material
         * or regenerated after crypto_cfg.km_refresh_rate_pkt packets .
         * Send to peer
         */
        // XXX Need to make it clearer and less hardcoded values
        int kix = hcryptMsg_KM_GetKeyIndex((unsigned char *)(out_p[i]));
        int ki = kix & 0x1;
        if ((out_len_p[i] != m_SndKmMsg[ki].MsgLen)
                ||  (0 != memcmp(out_p[i], m_SndKmMsg[ki].Msg, m_SndKmMsg[ki].MsgLen))) 
        {

            uint8_t* oldkey SRT_ATR_UNUSED = m_SndKmMsg[ki].Msg;
            HLOGC(cnlog.Debug, log << "new key[" << ki << "] index=" << kix
                    << " OLD=[" << m_SndKmMsg[ki].MsgLen << "]"
                    << FormatBinaryString(m_SndKmMsg[ki].Msg, m_SndKmMsg[ki].MsgLen)
                    << " NEW=[" << out_len_p[i] << "]"
                    << FormatBinaryString((const uint8_t*)out_p[i], out_len_p[i]));

            /* New Keying material, send to peer */
            memcpy((m_SndKmMsg[ki].Msg), out_p[i], out_len_p[i]);
            m_SndKmMsg[ki].MsgLen = out_len_p[i];
            m_SndKmMsg[ki].iPeerRetry = SRT_MAX_KMRETRY;  

            if (bidirectional && !sock)
            {
                // "Send" this key also to myself, just to be applied to the receiver crypto,
                // exactly the same way how this key is interpreted on the peer side into its receiver crypto
                int rc = HaiCrypt_Rx_Process(m_hRcvCrypto, m_SndKmMsg[ki].Msg, m_SndKmMsg[ki].MsgLen, NULL, NULL, 0);
                if ( rc < 0 )
                {
                    LOGC(cnlog.Fatal, log << "regenCryptoKm: IPE: applying key generated in snd crypto into rcv crypto: failed code=" << rc);
                    // The party won't be able to decrypt incoming data!
                    // Not sure if anything has to be reported.
                }
            }

            if (sock)
            {
                HLOGC(cnlog.Debug, log << "regenCryptoKm: SENDING ki=" << ki << " len=" << m_SndKmMsg[ki].MsgLen
                        << " retry(updated)=" << m_SndKmMsg[ki].iPeerRetry);
                sock->sendSrtMsg(SRT_CMD_KMREQ, (uint32_t *)m_SndKmMsg[ki].Msg, m_SndKmMsg[ki].MsgLen / sizeof(uint32_t));
                sent++;
            }
        }
        else if (out_len_p[i] == 0)
        {
            HLOGC(cnlog.Debug, log << "no key[" << ki << "] index=" << kix << ": not generated");
        }
        else
        {
            HLOGC(cnlog.Debug, log << "no key[" << ki << "] index=" << kix << ": key unchanged");
        }
    }

    HLOGC(cnlog.Debug, log << "regenCryptoKm: key[0]: len=" << m_SndKmMsg[0].MsgLen << " retry=" << m_SndKmMsg[0].iPeerRetry
            << "; key[1]: len=" << m_SndKmMsg[1].MsgLen << " retry=" << m_SndKmMsg[1].iPeerRetry);

    if (sent)
        m_SndKmLastTime = srt::sync::steady_clock::now();
#endif
}

srt::CCryptoControl::CCryptoControl(SRTSOCKET id)
    : m_SocketID(id)
    , m_iSndKmKeyLen(0)
    , m_iRcvKmKeyLen(0)
    , m_SndKmState(SRT_KM_S_UNSECURED)
    , m_RcvKmState(SRT_KM_S_UNSECURED)
    , m_CurrentKey(EK_NOENC)
    , m_KmRefreshRatePkt(0)
    , m_KmPreAnnouncePkt(0)
    , m_iCryptoMode(CSrtConfig::CIPHER_MODE_AUTO)
    , m_bUseGcm153(false)
    , m_bErrorReported(false)
{
    m_KmSecret.len = 0;
    //send
    m_SndKmMsg[0].MsgLen = 0;
    m_SndKmMsg[0].iPeerRetry = 0;
    m_SndKmMsg[1].MsgLen = 0;
    m_SndKmMsg[1].iPeerRetry = 0;
    m_hSndCrypto = NULL;
    //recv
    m_hRcvCrypto = NULL;
}

bool srt::CCryptoControl::init(HandshakeSide side, const CSrtConfig& cfg, bool bidirectional SRT_ATR_UNUSED, bool bUseGcm153 SRT_ATR_UNUSED)
{
    // NOTE: initiator creates m_hSndCrypto. When bidirectional,
    // it creates also m_hRcvCrypto with the same key length.
    // Acceptor creates nothing - it will create appropriate
    // contexts when receiving KMREQ from the initiator.

    HLOGC(cnlog.Debug, log << "CCryptoControl::init: HS SIDE:"
        << (side == HSD_INITIATOR ? "INITIATOR" : "RESPONDER")
        << " DIRECTION:" << (bidirectional ? "BOTH" : (side == HSD_INITIATOR) ? "SENDER" : "RECEIVER"));

    // Set UNSECURED state as default
    m_RcvKmState = SRT_KM_S_UNSECURED;
    m_iCryptoMode = cfg.iCryptoMode;
    m_bUseGcm153 = bUseGcm153;

#ifdef SRT_ENABLE_ENCRYPTION
    if (!cfg.bTSBPD && m_iCryptoMode == CSrtConfig::CIPHER_MODE_AUTO)
        m_iCryptoMode = CSrtConfig::CIPHER_MODE_AES_CTR;
    const bool bUseGCM = m_iCryptoMode == CSrtConfig::CIPHER_MODE_AES_GCM;

    if (bUseGCM && !isAESGCMSupported())
    {
        LOGC(cnlog.Warn, log << "CCryptoControl: AES GCM is not supported by the crypto service provider.");
        return false;
    }
#endif

    // Set security-pending state, if a password was set.
    m_SndKmState = hasPassphrase() ? SRT_KM_S_SECURING : SRT_KM_S_UNSECURED;

    m_KmPreAnnouncePkt = cfg.uKmPreAnnouncePkt;
    m_KmRefreshRatePkt = cfg.uKmRefreshRatePkt;

    if (side == HSD_INITIATOR)
    {
        if (hasPassphrase())
        {
#ifdef SRT_ENABLE_ENCRYPTION
            if (m_iSndKmKeyLen == 0)
            {
                HLOGC(cnlog.Debug, log << "CCryptoControl::init: PBKEYLEN still 0, setting default 16");
                m_iSndKmKeyLen = 16;
            }

            bool ok = createCryptoCtx((m_hSndCrypto), m_iSndKmKeyLen, HAICRYPT_CRYPTO_DIR_TX, bUseGCM);
            HLOGC(cnlog.Debug, log << "CCryptoControl::init: creating SND crypto context: " << ok);

            if (ok && bidirectional)
            {
                m_iRcvKmKeyLen = m_iSndKmKeyLen;
                const int st = HaiCrypt_Clone(m_hSndCrypto, HAICRYPT_CRYPTO_DIR_RX, &m_hRcvCrypto);
                HLOGC(cnlog.Debug, log << "CCryptoControl::init: creating CLONED RCV crypto context: status=" << st);
                ok = st == 0;
            }

            // Note: this is sanity check, it should never happen.
            if (!ok)
            {
                m_SndKmState = SRT_KM_S_NOSECRET; // wanted to secure, but error occurred.
                if (bidirectional)
                    m_RcvKmState = SRT_KM_S_NOSECRET;

                return false;
            }

            regenCryptoKm(
                NULL, // Do not send the key (the KM msg will be attached to the HSv5 handshake)
                bidirectional // replicate the key to the receiver context, if bidirectional
            );

            m_iCryptoMode = bUseGCM ? CSrtConfig::CIPHER_MODE_AES_GCM : CSrtConfig::CIPHER_MODE_AES_CTR;
#else
            // This error would be a consequence of setting the passphrase, while encryption
            // is turned off at compile time. Setting the password itself should be not allowed
            // so this could only happen as a consequence of an IPE.
            LOGC(cnlog.Error, log << "CCryptoControl::init: IPE: encryption not supported");
            return false;
#endif
        }
        else
        {
            HLOGC(cnlog.Debug, log << "CCryptoControl::init: CAN'T CREATE crypto: key length for SND = " << m_iSndKmKeyLen);
        }
    }
    else
    {
        HLOGC(cnlog.Debug, log << "CCryptoControl::init: NOT creating crypto contexts - will be created upon reception of KMREQ");
    }

    return true;
}

void srt::CCryptoControl::close() 
{
    /* Wipeout secrets */
    sync::ScopedLock lck(m_mtxLock);
    memset(&m_KmSecret, 0, sizeof(m_KmSecret));
}

std::string srt::CCryptoControl::CONID() const
{
    if (m_SocketID == 0)
        return "";

    std::ostringstream os;
    os << "@" << m_SocketID << ":";

    return os.str();
}

#ifdef SRT_ENABLE_ENCRYPTION

#if ENABLE_HEAVY_LOGGING
namespace srt {
static std::string CryptoFlags(int flg)
{
    using namespace std;

    vector<string> f;
    if (flg & HAICRYPT_CFG_F_CRYPTO)
        f.push_back("crypto");
    if (flg & HAICRYPT_CFG_F_TX)
        f.push_back("TX");
    if (flg & HAICRYPT_CFG_F_FEC)
        f.push_back("fec");

    ostringstream os;
    copy(f.begin(), f.end(), ostream_iterator<string>(os, "|"));
    return os.str();
}
} // namespace srt
#endif // ENABLE_HEAVY_LOGGING

bool srt::CCryptoControl::createCryptoCtx(HaiCrypt_Handle& w_hCrypto, size_t keylen, HaiCrypt_CryptoDir cdir, bool bAESGCM)
{
    if (w_hCrypto)
    {
        // XXX You can check here if the existing handle represents
        // a correctly defined crypto. But this doesn't seem to be
        // necessary - the whole CCryptoControl facility seems to be valid only
        // within the frames of one connection.
        return true;
    }

    if ((m_KmSecret.len <= 0) || (keylen <= 0))
    {
        LOGC(cnlog.Error, log << CONID() << "cryptoCtx: IPE missing secret (" << m_KmSecret.len << ") or key length (" << keylen << ")");
        return false;
    }

    HaiCrypt_Cfg crypto_cfg;
    memset(&crypto_cfg, 0, sizeof(crypto_cfg));
#if 0//test key refresh (fast rate)
    m_KmRefreshRatePkt = 2000;
    m_KmPreAnnouncePkt = 500;
#endif
    crypto_cfg.flags = HAICRYPT_CFG_F_CRYPTO | (cdir == HAICRYPT_CRYPTO_DIR_TX ? HAICRYPT_CFG_F_TX : 0) | (bAESGCM ? HAICRYPT_CFG_F_GCM : 0);
    crypto_cfg.xport = HAICRYPT_XPT_SRT;
    crypto_cfg.cryspr = HaiCryptCryspr_Get_Instance();
    crypto_cfg.key_len = (size_t)keylen;
    crypto_cfg.data_max_len = HAICRYPT_DEF_DATA_MAX_LENGTH;    //MTU
    crypto_cfg.km_tx_period_ms = 0;//No HaiCrypt KM inject period, handled in SRT;
    crypto_cfg.km_refresh_rate_pkt = m_KmRefreshRatePkt == 0 ? HAICRYPT_DEF_KM_REFRESH_RATE : m_KmRefreshRatePkt;
    crypto_cfg.km_pre_announce_pkt = m_KmPreAnnouncePkt == 0 ? SRT_CRYPT_KM_PRE_ANNOUNCE : m_KmPreAnnouncePkt;
    crypto_cfg.secret = m_KmSecret;

    HLOGC(cnlog.Debug, log << "CRYPTO CFG: flags=" << CryptoFlags(crypto_cfg.flags) << " xport=" << crypto_cfg.xport << " cryspr=" << crypto_cfg.cryspr
        << " keylen=" << crypto_cfg.key_len << " passphrase_length=" << crypto_cfg.secret.len);

    if (HaiCrypt_Create(&crypto_cfg, (&w_hCrypto)) != HAICRYPT_OK)
    {
        LOGC(cnlog.Error, log << CONID() << "cryptoCtx: could not create " << (cdir == HAICRYPT_CRYPTO_DIR_TX ? "tx" : "rx") << " crypto ctx");
        return false;
    }

    HLOGC(cnlog.Debug, log << CONID() << "cryptoCtx: CREATED crypto for dir=" << (cdir == HAICRYPT_CRYPTO_DIR_TX ? "tx" : "rx") << " keylen=" << keylen);

    return true;
}
#else
bool srt::CCryptoControl::createCryptoCtx(HaiCrypt_Handle&, size_t, HaiCrypt_CryptoDir, bool)
{
    return false;
}
#endif // SRT_ENABLE_ENCRYPTION


srt::EncryptionStatus srt::CCryptoControl::encrypt(CPacket& w_packet SRT_ATR_UNUSED)
{
#ifdef SRT_ENABLE_ENCRYPTION
    // Encryption not enabled - do nothing.
    if ( getSndCryptoFlags() == EK_NOENC )
        return ENCS_CLEAR;

    // Note that in case of GCM the header has to zero Retransmitted Packet Flag (R).
    // If TSBPD is disabled, timestamp also has to be zeroed.
    int rc = HaiCrypt_Tx_Data(m_hSndCrypto, ((uint8_t*)w_packet.getHeader()), ((uint8_t*)w_packet.m_pcData), w_packet.getLength());
    if (rc < 0)
    {
        return ENCS_FAILED;
    }
    else if (rc > 0)
    {
        // XXX what happens if the encryption is said to be "succeeded",
        // but the length is 0? Shouldn't this be treated as unwanted?
        w_packet.setLength(rc);
    }

    return ENCS_CLEAR;
#else
    return ENCS_NOTSUP;
#endif
}

srt::EncryptionStatus srt::CCryptoControl::decrypt(CPacket& w_packet SRT_ATR_UNUSED)
{
#ifdef SRT_ENABLE_ENCRYPTION
    if (w_packet.getMsgCryptoFlags() == EK_NOENC)
    {
        HLOGC(cnlog.Debug, log << "CPacket::decrypt: packet not encrypted");
        return ENCS_CLEAR; // not encrypted, no need do decrypt, no flags to be modified
    }

    if (m_RcvKmState == SRT_KM_S_UNSECURED)
    {
        if (m_KmSecret.len != 0)
        {
            // We were unaware that the peer has set password,
            // but now here we are.
            m_RcvKmState = SRT_KM_S_SECURING;
            LOGC(cnlog.Note, log << "SECURITY UPDATE: Peer has surprised Agent with encryption, but KMX is pending - current packet size="
                    << w_packet.getLength() << " dropped");
            return ENCS_FAILED;
        }
        else
        {
            // Peer has set a password, but Agent did not,
            // which means that it will be unable to decrypt
            // sent payloads anyway.
            m_RcvKmState = SRT_KM_S_NOSECRET;
            LOGP(cnlog.Warn, "SECURITY FAILURE: Agent has no PW, but Peer sender has declared one, can't decrypt");
            // This only informs about the state change; it will be also caught by the condition below
        }
    }

    if (m_RcvKmState != SRT_KM_S_SECURED)
    {
        // If not "secured", it means that it won't be able to decrypt packets,
        // so there's no point in even trying to send them to HaiCrypt_Rx_Data.
        // Actually the current conditions concerning m_hRcvCrypto are such that this object
        // is created in case of SRT_KM_S_BADSECRET, so it will simply fail to decrypt,
        // but with SRT_KM_S_NOSECRET m_hRcvCrypto is not even created (is NULL), which
        // will then cause an error to be reported, misleadingly. Simply don't try to
        // decrypt anything as long as you are not sure that the connection is secured.

        // This problem will occur every time a packet comes in, it's worth reporting,
        // but not with every single packet arriving. Print it once and turn off the flag;
        // it will be restored at the next attempt of KMX.
        if (!m_bErrorReported)
        {
            m_bErrorReported = true;
            LOGC(cnlog.Error, log << "SECURITY STATUS: " << KmStateStr(m_RcvKmState) << " - can't decrypt w_packet.");
        }
        HLOGC(cnlog.Debug, log << "Packet still not decrypted, status=" << KmStateStr(m_RcvKmState)
                << " - dropping size=" << w_packet.getLength());
        return ENCS_FAILED;
    }

    const int rc = HaiCrypt_Rx_Data(m_hRcvCrypto, ((uint8_t *)w_packet.getHeader()), ((uint8_t *)w_packet.m_pcData), w_packet.getLength());
    if (rc <= 0)
    {
        LOGC(cnlog.Note, log << "decrypt ERROR: HaiCrypt_Rx_Data failure=" << rc << " - returning failed decryption");
        // -1: decryption failure
        // 0: key not received yet
        return ENCS_FAILED;
    }
    // Otherwise: rc == decrypted text length.
    w_packet.setLength(rc); /* In case clr txt size is different from cipher txt */

    // Decryption succeeded. Update flags.
    w_packet.setMsgCryptoFlags(EK_NOENC);

    HLOGC(cnlog.Debug, log << "decrypt: successfully decrypted, resulting length=" << rc);
    return ENCS_CLEAR;
#else
    return ENCS_NOTSUP;
#endif
}


srt::CCryptoControl::~CCryptoControl()
{
#ifdef SRT_ENABLE_ENCRYPTION
    close();
    if (m_hSndCrypto)
    {
        HaiCrypt_Close(m_hSndCrypto);
    }

    if (m_hRcvCrypto)
    {
        HaiCrypt_Close(m_hRcvCrypto);
    }
#endif
}
