<commit_before>// Copyright (c) 2009-2012 The Bitcoin Developers
// Distributed under the MIT/X11 software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "crypter.h"

#include "script.h"
#include "scrypt.h"

#include <string>
#include <vector>
#include <boost/foreach.hpp>
#include <openssl/opensslv.h>     // For using openssl 1.0 and 1.1 branches.
#include <openssl/aes.h>
#include <openssl/evp.h>
#ifdef WIN32
#include <windows.h>
#endif

bool CCrypter::SetKeyFromPassphrase(const SecureString& strKeyData, const std::vector<unsigned char>& chSalt, const unsigned int nRounds, const unsigned int nDerivationMethod)
{
    if (nRounds < 1 || chSalt.size() != WALLET_CRYPTO_SALT_SIZE)
        return false;

    int i = 0;
    if (nDerivationMethod == 0)
    {
        i = EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha512(), &chSalt[0],
                          (unsigned char *)&strKeyData[0], strKeyData.size(), nRounds, chKey, chIV);
    }

    if (nDerivationMethod == 1)
    {
        // Passphrase conversion
        uint256 scryptHash = scrypt_salted_multiround_hash((const void*)strKeyData.c_str(), strKeyData.size(), &chSalt[0], 8, nRounds);

        i = EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha512(), &chSalt[0],
                          (unsigned char *)&scryptHash, sizeof scryptHash, nRounds, chKey, chIV);
        OPENSSL_cleanse(&scryptHash, sizeof scryptHash);
    }

    if (i != (int)WALLET_CRYPTO_KEY_SIZE)
    {
        OPENSSL_cleanse(chKey, sizeof(chKey));
        OPENSSL_cleanse(chIV, sizeof(chIV));
        return false;
    }

    fKeySet = true;
    return true;
}

bool CCrypter::SetKey(const CKeyingMaterial& chNewKey, const std::vector<unsigned char>& chNewIV)
{
    if (chNewKey.size() != WALLET_CRYPTO_KEY_SIZE || chNewIV.size() != WALLET_CRYPTO_KEY_SIZE)
        return false;

    memcpy(&chKey[0], &chNewKey[0], sizeof chKey);
    memcpy(&chIV[0], &chNewIV[0], sizeof chIV);

    fKeySet = true;
    return true;
}

bool CCrypter::Encrypt(const CKeyingMaterial& vchPlaintext, std::vector<unsigned char> &vchCiphertext)
{
    if (!fKeySet)
        return false;

    // max ciphertext len for a n bytes of plaintext is
    // n + AES_BLOCK_SIZE - 1 bytes
    int nLen = vchPlaintext.size();
    int nCLen = nLen + AES_BLOCK_SIZE, nFLen = 0;
    vchCiphertext = std::vector<unsigned char> (nCLen);

#if OPENSSL_VERSION_NUMBER < 0x10100000L
    EVP_CIPHER_CTX ctx;
#else
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
#endif

    bool fOk = true;

#if OPENSSL_VERSION_NUMBER < 0x10100000L
    EVP_CIPHER_CTX_init(&ctx);
    if (fOk) fOk = EVP_EncryptInit_ex(&ctx, EVP_aes_256_cbc(), NULL, chKey, chIV);
    if (fOk) fOk = EVP_EncryptUpdate(&ctx, &vchCiphertext[0], &nCLen, &vchPlaintext[0], nLen);
    if (fOk) fOk = EVP_EncryptFinal_ex(&ctx, (&vchCiphertext[0])+nCLen, &nFLen);
    EVP_CIPHER_CTX_cleanup(&ctx);
#else
    EVP_CIPHER_CTX_init(ctx);
    if (fOk) fOk = EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, chKey, chIV);
    if (fOk) fOk = EVP_EncryptUpdate(ctx, &vchCiphertext[0], &nCLen, &vchPlaintext[0], nLen);
    if (fOk) fOk = EVP_EncryptFinal_ex(ctx, (&vchCiphertext[0])+nCLen, &nFLen);
    EVP_CIPHER_CTX_free(ctx);
#endif

    if (!fOk) return false;

    vchCiphertext.resize(nCLen + nFLen);
    return true;
}

bool CCrypter::Decrypt(const std::vector<unsigned char>& vchCiphertext, CKeyingMaterial& vchPlaintext)
{
    if (!fKeySet)
        return false;

    // plaintext will always be equal to or lesser than length of ciphertext
    int nLen = vchCiphertext.size();
    int nPLen = nLen, nFLen = 0;

    vchPlaintext = CKeyingMaterial(nPLen);

#if OPENSSL_VERSION_NUMBER < 0x10100000L
    EVP_CIPHER_CTX ctx;
#else
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
#endif

    bool fOk = true;

#if OPENSSL_VERSION_NUMBER < 0x10100000L
    EVP_CIPHER_CTX_init(&ctx);
    if (fOk) fOk = EVP_DecryptInit_ex(&ctx, EVP_aes_256_cbc(), NULL, chKey, chIV);
    if (fOk) fOk = EVP_DecryptUpdate(&ctx, &vchPlaintext[0], &nPLen, &vchCiphertext[0], nLen);
    if (fOk) fOk = EVP_DecryptFinal_ex(&ctx, (&vchPlaintext[0])+nPLen, &nFLen);
    EVP_CIPHER_CTX_cleanup(&ctx);
#else
    EVP_CIPHER_CTX_init(ctx);
    if (fOk) fOk = EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, chKey, chIV);
    if (fOk) fOk = EVP_DecryptUpdate(ctx, &vchPlaintext[0], &nPLen, &vchCiphertext[0], nLen);
    if (fOk) fOk = EVP_DecryptFinal_ex(ctx, (&vchPlaintext[0])+nPLen, &nFLen);
    EVP_CIPHER_CTX_free(ctx);
#endif

    if (!fOk) return false;

    vchPlaintext.resize(nPLen + nFLen);
    return true;
}

bool EncryptSecret(const CKeyingMaterial& vMasterKey, const CKeyingMaterial &vchPlaintext, const uint256& nIV, std::vector<unsigned char> &vchCiphertext)
{
    CCrypter cKeyCrypter;
    std::vector<unsigned char> chIV(WALLET_CRYPTO_KEY_SIZE);
    memcpy(&chIV[0], &nIV, WALLET_CRYPTO_KEY_SIZE);
    if(!cKeyCrypter.SetKey(vMasterKey, chIV))
        return false;
    return cKeyCrypter.Encrypt(*((const CKeyingMaterial*)&vchPlaintext), vchCiphertext);
}

bool DecryptSecret(const CKeyingMaterial& vMasterKey, const std::vector<unsigned char>& vchCiphertext, const uint256& nIV, CKeyingMaterial& vchPlaintext)
{
    CCrypter cKeyCrypter;
    std::vector<unsigned char> chIV(WALLET_CRYPTO_KEY_SIZE);
    memcpy(&chIV[0], &nIV, WALLET_CRYPTO_KEY_SIZE);
    if(!cKeyCrypter.SetKey(vMasterKey, chIV))
        return false;
    return cKeyCrypter.Decrypt(vchCiphertext, *((CKeyingMaterial*)&vchPlaintext));
}

bool CCryptoKeyStore::SetCrypted()
{
    LOCK(cs_KeyStore);
    if (fUseCrypto)
        return true;
    if (!mapKeys.empty())
        return false;
    fUseCrypto = true;
    return true;
}

bool CCryptoKeyStore::LockKeyStore()
{
    if (!SetCrypted())
        return false;

    {
        LOCK(cs_KeyStore);
        vMasterKey.clear();
    }

    NotifyStatusChanged(this);
    return true;
}

bool CCryptoKeyStore::Unlock(const CKeyingMaterial& vMasterKeyIn)
{
    if (fDebug)
        LogPrintf("CCryptoKeyStore::Unlock()\n");

    {
        LOCK(cs_KeyStore);
        if (!SetCrypted())
            return false;

        int nUnlocked = 0;

        CryptedKeyMap::const_iterator mi = mapCryptedKeys.begin();
        for (; mi != mapCryptedKeys.end(); ++mi)
        {
            const CPubKey &vchPubKey = (*mi).second.first;
            const std::vector<unsigned char> &vchCryptedSecret = (*mi).second.second;
            CSecret vchSecret;

            if (vchCryptedSecret.size() < 1) // key was recieved from stealth/anon txn with wallet locked, will be expanded after this
            {
                if (fDebug)
                    LogPrintf("Skipping unexpanded key %s.\n", vchPubKey.GetHash().ToString().c_str());
                continue;
            };

            if (!DecryptSecret(vMasterKeyIn, vchCryptedSecret, vchPubKey.GetHash(), vchSecret))
            {
                LogPrintf("DecryptSecret() failed.\n");
                return false;
            };

            if (vchSecret.size() != 32)
                return false;

            CKey key;
            key.Set(vchSecret.begin(), vchSecret.end(), vchPubKey.IsCompressed());

            if (key.GetPubKey() != vchPubKey)
            {
                LogPrintf("Unlock failed: PubKey mismatch %s.\n", vchPubKey.GetHash().ToString().c_str());
                return false;
            };

            nUnlocked++;
            break;
        };

        if (nUnlocked < 1) // at least 1 key must pass the test
        {
            LogPrintf("Unlock failed: No keys unlocked.\n");
            return false;
        };

        vMasterKey = vMasterKeyIn;
    }

    NotifyStatusChanged(this);

    return true;
}

bool CCryptoKeyStore::AddKey(const CKey& key)
{
    return CCryptoKeyStore::AddKeyPubKey(key, key.GetPubKey());
}

bool CCryptoKeyStore::AddKeyPubKey(const CKey& key, const CPubKey &pubkey)
{
    {
        LOCK(cs_KeyStore);
        if (!IsCrypted())
            return CBasicKeyStore::AddKeyPubKey(key, pubkey);
        if (IsLocked())
            return false;
        std::vector<unsigned char> vchCryptedSecret;
        CKeyingMaterial vchSecret(key.begin(), key.end());
        if (!EncryptSecret(vMasterKey, vchSecret, pubkey.GetHash(), vchCryptedSecret))
            return false;

        // -- NOTE: this is CWallet::AddCryptedKey
        if (!AddCryptedKey(pubkey, vchCryptedSecret))
            return false;
    }
    return true;
}

bool CCryptoKeyStore::AddCryptedKey(const CPubKey &vchPubKey, const std::vector<unsigned char> &vchCryptedSecret)
{
    {
        LOCK(cs_KeyStore);
        if (!SetCrypted())
            return false;

        mapCryptedKeys[vchPubKey.GetID()] = make_pair(vchPubKey, vchCryptedSecret);
    }
    return true;
}

bool CCryptoKeyStore::GetKey(const CKeyID &address, CKey& keyOut) const
{
    {
        LOCK(cs_KeyStore);
        if (!IsCrypted())
            return CBasicKeyStore::GetKey(address, keyOut);

        CryptedKeyMap::const_iterator mi = mapCryptedKeys.find(address);
        if (mi != mapCryptedKeys.end())
        {
            const CPubKey &vchPubKey = (*mi).second.first;
            const std::vector<unsigned char> &vchCryptedSecret = (*mi).second.second;
            CKeyingMaterial vchSecret;
            if (!DecryptSecret(vMasterKey, vchCryptedSecret, vchPubKey.GetHash(), vchSecret))
                return false;
            if (vchSecret.size() != 32)
                return false;
            keyOut.Set(vchSecret.begin(), vchSecret.end(), vchPubKey.IsCompressed());
            return true;
        }
    }
    return false;
}

bool CCryptoKeyStore::GetPubKey(const CKeyID &address, CPubKey& vchPubKeyOut) const
{
    {
        LOCK(cs_KeyStore);
        if (!IsCrypted())
            return CKeyStore::GetPubKey(address, vchPubKeyOut);

        CryptedKeyMap::const_iterator mi = mapCryptedKeys.find(address);
        if (mi != mapCryptedKeys.end())
        {
            vchPubKeyOut = (*mi).second.first;
            return true;
        }
    }
    return false;
}

bool CCryptoKeyStore::EncryptKeys(CKeyingMaterial& vMasterKeyIn)
{
    {
        LOCK(cs_KeyStore);
        if (!mapCryptedKeys.empty() || IsCrypted())
            return false;

        fUseCrypto = true;
        BOOST_FOREACH(KeyMap::value_type& mKey, mapKeys)
        {
            const CKey &key = mKey.second;
            CPubKey vchPubKey = key.GetPubKey();
            CKeyingMaterial vchSecret(key.begin(), key.end());
            std::vector<unsigned char> vchCryptedSecret;
            if (!EncryptSecret(vMasterKeyIn, vchSecret, vchPubKey.GetHash(), vchCryptedSecret))
                return false;
            if (!AddCryptedKey(vchPubKey, vchCryptedSecret))
                return false;
        }
        mapKeys.clear();
    }
    return true;
}
<commit_msg>map c keys<commit_after>// Copyright (c) 2009-2012 The Bitcoin Developers
// Distributed under the MIT/X11 software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "crypter.h"

#include "script.h"
#include "scrypt.h"

#include <string>
#include <vector>
#include <boost/foreach.hpp>
#include <openssl/opensslv.h>     // For using openssl 1.0 and 1.1 branches.
#include <openssl/aes.h>
#include <openssl/evp.h>
#ifdef WIN32
#include <windows.h>
#endif

bool CCrypter::SetKeyFromPassphrase(const SecureString& strKeyData, const std::vector<unsigned char>& chSalt, const unsigned int nRounds, const unsigned int nDerivationMethod)
{
    if (nRounds < 1 || chSalt.size() != WALLET_CRYPTO_SALT_SIZE)
        return false;

    int i = 0;
    if (nDerivationMethod == 0)
    {
        i = EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha512(), &chSalt[0],
                          (unsigned char *)&strKeyData[0], strKeyData.size(), nRounds, chKey, chIV);
    }

    if (nDerivationMethod == 1)
    {
        // Passphrase conversion
        uint256 scryptHash = scrypt_salted_multiround_hash((const void*)strKeyData.c_str(), strKeyData.size(), &chSalt[0], 8, nRounds);

        i = EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha512(), &chSalt[0],
                          (unsigned char *)&scryptHash, sizeof scryptHash, nRounds, chKey, chIV);
        OPENSSL_cleanse(&scryptHash, sizeof scryptHash);
    }

    if (i != (int)WALLET_CRYPTO_KEY_SIZE)
    {
        OPENSSL_cleanse(chKey, sizeof(chKey));
        OPENSSL_cleanse(chIV, sizeof(chIV));
        return false;
    }

    fKeySet = true;
    return true;
}

bool CCrypter::SetKey(const CKeyingMaterial& chNewKey, const std::vector<unsigned char>& chNewIV)
{
    if (chNewKey.size() != WALLET_CRYPTO_KEY_SIZE || chNewIV.size() != WALLET_CRYPTO_KEY_SIZE)
        return false;

    memcpy(&chKey[0], &chNewKey[0], sizeof chKey);
    memcpy(&chIV[0], &chNewIV[0], sizeof chIV);

    fKeySet = true;
    return true;
}

bool CCrypter::Encrypt(const CKeyingMaterial& vchPlaintext, std::vector<unsigned char> &vchCiphertext)
{
    if (!fKeySet)
        return false;

    // max ciphertext len for a n bytes of plaintext is
    // n + AES_BLOCK_SIZE - 1 bytes
    int nLen = vchPlaintext.size();
    int nCLen = nLen + AES_BLOCK_SIZE, nFLen = 0;
    vchCiphertext = std::vector<unsigned char> (nCLen);

#if OPENSSL_VERSION_NUMBER < 0x10100000L
    EVP_CIPHER_CTX ctx;
#else
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
#endif

    bool fOk = true;

#if OPENSSL_VERSION_NUMBER < 0x10100000L
    EVP_CIPHER_CTX_init(&ctx);
    if (fOk) fOk = EVP_EncryptInit_ex(&ctx, EVP_aes_256_cbc(), NULL, chKey, chIV);
    if (fOk) fOk = EVP_EncryptUpdate(&ctx, &vchCiphertext[0], &nCLen, &vchPlaintext[0], nLen);
    if (fOk) fOk = EVP_EncryptFinal_ex(&ctx, (&vchCiphertext[0])+nCLen, &nFLen);
    EVP_CIPHER_CTX_cleanup(&ctx);
#else
    EVP_CIPHER_CTX_init(ctx);
    if (fOk) fOk = EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, chKey, chIV);
    if (fOk) fOk = EVP_EncryptUpdate(ctx, &vchCiphertext[0], &nCLen, &vchPlaintext[0], nLen);
    if (fOk) fOk = EVP_EncryptFinal_ex(ctx, (&vchCiphertext[0])+nCLen, &nFLen);
    EVP_CIPHER_CTX_free(ctx);
#endif

    if (!fOk) return false;

    vchCiphertext.resize(nCLen + nFLen);
    return true;
}

bool CCrypter::Decrypt(const std::vector<unsigned char>& vchCiphertext, CKeyingMaterial& vchPlaintext)
{
    if (!fKeySet)
        return false;

    // plaintext will always be equal to or lesser than length of ciphertext
    int nLen = vchCiphertext.size();
    int nPLen = nLen, nFLen = 0;

    vchPlaintext = CKeyingMaterial(nPLen);

#if OPENSSL_VERSION_NUMBER < 0x10100000L
    EVP_CIPHER_CTX ctx;
#else
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
#endif

    bool fOk = true;

#if OPENSSL_VERSION_NUMBER < 0x10100000L
    EVP_CIPHER_CTX_init(&ctx);
    if (fOk) fOk = EVP_DecryptInit_ex(&ctx, EVP_aes_256_cbc(), NULL, chKey, chIV);
    if (fOk) fOk = EVP_DecryptUpdate(&ctx, &vchPlaintext[0], &nPLen, &vchCiphertext[0], nLen);
    if (fOk) fOk = EVP_DecryptFinal_ex(&ctx, (&vchPlaintext[0])+nPLen, &nFLen);
    EVP_CIPHER_CTX_cleanup(&ctx);
#else
    EVP_CIPHER_CTX_init(ctx);
    if (fOk) fOk = EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, chKey, chIV);
    if (fOk) fOk = EVP_DecryptUpdate(ctx, &vchPlaintext[0], &nPLen, &vchCiphertext[0], nLen);
    if (fOk) fOk = EVP_DecryptFinal_ex(ctx, (&vchPlaintext[0])+nPLen, &nFLen);
    EVP_CIPHER_CTX_free(ctx);
#endif

    if (!fOk) return false;

    vchPlaintext.resize(nPLen + nFLen);
    return true;
}

bool EncryptSecret(const CKeyingMaterial& vMasterKey, const CKeyingMaterial &vchPlaintext, const uint256& nIV, std::vector<unsigned char> &vchCiphertext)
{
    CCrypter cKeyCrypter;
    std::vector<unsigned char> chIV(WALLET_CRYPTO_KEY_SIZE);
    memcpy(&chIV[0], &nIV, WALLET_CRYPTO_KEY_SIZE);
    if(!cKeyCrypter.SetKey(vMasterKey, chIV))
        return false;
    return cKeyCrypter.Encrypt(*((const CKeyingMaterial*)&vchPlaintext), vchCiphertext);
}

bool DecryptSecret(const CKeyingMaterial& vMasterKey, const std::vector<unsigned char>& vchCiphertext, const uint256& nIV, CKeyingMaterial& vchPlaintext)
{
    CCrypter cKeyCrypter;
    std::vector<unsigned char> chIV(WALLET_CRYPTO_KEY_SIZE);
    memcpy(&chIV[0], &nIV, WALLET_CRYPTO_KEY_SIZE);
    if(!cKeyCrypter.SetKey(vMasterKey, chIV))
        return false;
    return cKeyCrypter.Decrypt(vchCiphertext, *((CKeyingMaterial*)&vchPlaintext));
}

bool CCryptoKeyStore::SetCrypted()
{
    LOCK(cs_KeyStore);
    if (fUseCrypto)
        return true;
    if (!mapKeys.empty())
        return false;
    fUseCrypto = true;
    return true;
}

bool CCryptoKeyStore::LockKeyStore()
{
    if (!SetCrypted())
        return false;

    {
        LOCK(cs_KeyStore);
        vMasterKey.clear();
    }

    NotifyStatusChanged(this);
    return true;
}

bool CCryptoKeyStore::Unlock(const CKeyingMaterial& vMasterKeyIn)
{
    if (fDebug)
        LogPrintf("CCryptoKeyStore::Unlock()\n");

    {
        LOCK(cs_KeyStore);
        if (!SetCrypted())
            return false;

        int nUnlocked = 0;

        CryptedKeyMap::const_iterator mi = mapCryptedKeys.begin();
        for (; mi != mapCryptedKeys.end(); ++mi)
        {
            const CPubKey &vchPubKey = (*mi).second.first;
            const std::vector<unsigned char> &vchCryptedSecret = (*mi).second.second;
            CSecret vchSecret;

            if (vchCryptedSecret.size() < 1) // key was recieved from stealth/anon txn with wallet locked, will be expanded after this
            {
                if (fDebug)
                    LogPrintf("Skipping unexpanded key %s.\n", vchPubKey.GetHash().ToString().c_str());
                continue;
            };

            if (!DecryptSecret(vMasterKeyIn, vchCryptedSecret, vchPubKey.GetHash(), vchSecret))
            {
                LogPrintf("DecryptSecret() failed.\n");
                return false;
            };

            if (vchSecret.size() != 32)
                return false;

            CKey key;
            key.Set(vchSecret.begin(), vchSecret.end(), vchPubKey.IsCompressed());

            if (key.GetPubKey() != vchPubKey)
            {
                LogPrintf("Unlock failed: PubKey mismatch %s.\n", vchPubKey.GetHash().ToString().c_str());
                return false;
            };

            nUnlocked++;
            break;
        };

        if (nUnlocked < 1) // at least 1 key must pass the test
        {
            if (mapCryptedKeys.size() > 0)
            {
                LogPrintf("Unlock failed: No keys unlocked.\n");
                return false;
            };
        };

        vMasterKey = vMasterKeyIn;
    }

    NotifyStatusChanged(this);

    return true;
}

bool CCryptoKeyStore::AddKey(const CKey& key)
{
    return CCryptoKeyStore::AddKeyPubKey(key, key.GetPubKey());
}

bool CCryptoKeyStore::AddKeyPubKey(const CKey& key, const CPubKey &pubkey)
{
    {
        LOCK(cs_KeyStore);
        if (!IsCrypted())
            return CBasicKeyStore::AddKeyPubKey(key, pubkey);
        if (IsLocked())
            return false;
        std::vector<unsigned char> vchCryptedSecret;
        CKeyingMaterial vchSecret(key.begin(), key.end());
        if (!EncryptSecret(vMasterKey, vchSecret, pubkey.GetHash(), vchCryptedSecret))
            return false;

        // -- NOTE: this is CWallet::AddCryptedKey
        if (!AddCryptedKey(pubkey, vchCryptedSecret))
            return false;
    }
    return true;
}

bool CCryptoKeyStore::AddCryptedKey(const CPubKey &vchPubKey, const std::vector<unsigned char> &vchCryptedSecret)
{
    {
        LOCK(cs_KeyStore);
        if (!SetCrypted())
            return false;

        mapCryptedKeys[vchPubKey.GetID()] = make_pair(vchPubKey, vchCryptedSecret);
    }
    return true;
}

bool CCryptoKeyStore::GetKey(const CKeyID &address, CKey& keyOut) const
{
    {
        LOCK(cs_KeyStore);
        if (!IsCrypted())
            return CBasicKeyStore::GetKey(address, keyOut);

        CryptedKeyMap::const_iterator mi = mapCryptedKeys.find(address);
        if (mi != mapCryptedKeys.end())
        {
            const CPubKey &vchPubKey = (*mi).second.first;
            const std::vector<unsigned char> &vchCryptedSecret = (*mi).second.second;
            CKeyingMaterial vchSecret;
            if (!DecryptSecret(vMasterKey, vchCryptedSecret, vchPubKey.GetHash(), vchSecret))
                return false;
            if (vchSecret.size() != 32)
                return false;
            keyOut.Set(vchSecret.begin(), vchSecret.end(), vchPubKey.IsCompressed());
            return true;
        }
    }
    return false;
}

bool CCryptoKeyStore::GetPubKey(const CKeyID &address, CPubKey& vchPubKeyOut) const
{
    {
        LOCK(cs_KeyStore);
        if (!IsCrypted())
            return CKeyStore::GetPubKey(address, vchPubKeyOut);

        CryptedKeyMap::const_iterator mi = mapCryptedKeys.find(address);
        if (mi != mapCryptedKeys.end())
        {
            vchPubKeyOut = (*mi).second.first;
            return true;
        }
    }
    return false;
}

bool CCryptoKeyStore::EncryptKeys(CKeyingMaterial& vMasterKeyIn)
{
    {
        LOCK(cs_KeyStore);
        if (!mapCryptedKeys.empty() || IsCrypted())
            return false;

        fUseCrypto = true;
        BOOST_FOREACH(KeyMap::value_type& mKey, mapKeys)
        {
            const CKey &key = mKey.second;
            CPubKey vchPubKey = key.GetPubKey();
            CKeyingMaterial vchSecret(key.begin(), key.end());
            std::vector<unsigned char> vchCryptedSecret;
            if (!EncryptSecret(vMasterKeyIn, vchSecret, vchPubKey.GetHash(), vchCryptedSecret))
                return false;
            if (!AddCryptedKey(vchPubKey, vchCryptedSecret))
                return false;
        }
        mapKeys.clear();
    }
    return true;
}
<|endoftext|>


<commit_before>// Copyright (c) 2019, Sergey Abbakumov
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "rst/TaskRunner/PollingTaskRunner.h"

#include <utility>

#include "rst/Check/Check.h"
#include "rst/NoDestructor/NoDestructor.h"

namespace chrono = std::chrono;

namespace rst {

PollingTaskRunner::PollingTaskRunner(
    std::function<chrono::milliseconds()>&& time_function)
    : time_function_(std::move(time_function)) {}

PollingTaskRunner::~PollingTaskRunner() = default;

void PollingTaskRunner::PostDelayedTask(std::function<void()>&& task,
                                        const chrono::milliseconds delay) {
  RST_DCHECK(delay.count() >= 0);

  const auto now = time_function_();
  const auto future_time_point = now + delay;
  std::lock_guard<std::mutex> lock(mutex_);
  queue_.emplace(future_time_point, task_id_, std::move(task));
  task_id_++;
}

void PollingTaskRunner::RunPendingTasks() {
  pending_tasks_.clear();

  {
    std::lock_guard<std::mutex> lock(mutex_);

    const auto now = time_function_();
    while (!queue_.empty()) {
      const auto& item = queue_.top();
      if (now < item.time_point)
        break;

      auto task = item.task;
      queue_.pop();
      pending_tasks_.emplace_back(std::move(task));
    }
  }

  for (const auto& task : pending_tasks_)
    task();
}

}  // namespace rst
<commit_msg>Remove unused include<commit_after>// Copyright (c) 2019, Sergey Abbakumov
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "rst/TaskRunner/PollingTaskRunner.h"

#include <utility>

#include "rst/Check/Check.h"

namespace chrono = std::chrono;

namespace rst {

PollingTaskRunner::PollingTaskRunner(
    std::function<chrono::milliseconds()>&& time_function)
    : time_function_(std::move(time_function)) {}

PollingTaskRunner::~PollingTaskRunner() = default;

void PollingTaskRunner::PostDelayedTask(std::function<void()>&& task,
                                        const chrono::milliseconds delay) {
  RST_DCHECK(delay.count() >= 0);

  const auto now = time_function_();
  const auto future_time_point = now + delay;
  std::lock_guard<std::mutex> lock(mutex_);
  queue_.emplace(future_time_point, task_id_, std::move(task));
  task_id_++;
}

void PollingTaskRunner::RunPendingTasks() {
  pending_tasks_.clear();

  {
    std::lock_guard<std::mutex> lock(mutex_);

    const auto now = time_function_();
    while (!queue_.empty()) {
      const auto& item = queue_.top();
      if (now < item.time_point)
        break;

      auto task = item.task;
      queue_.pop();
      pending_tasks_.emplace_back(std::move(task));
    }
  }

  for (const auto& task : pending_tasks_)
    task();
}

}  // namespace rst
<|endoftext|>


<commit_before>#ifndef RFR_MOSTLY_CONTINUOUS_DATA_CONTAINER_WITH_INSTANCES_HPP
#define RFR_MOSTLY_CONTINUOUS_DATA_CONTAINER_WITH_INSTANCES_HPP

#include <vector>
#include <map>
#include <cmath>

#include "rfr/data_containers/data_container_base.hpp"
#include "rfr/data_containers/data_container_utils.hpp"

namespace rfr{ namespace data_containers{

/** \brief A data container for mostly continuous data with instances.
 *
 * Similar to the mostly_continuous_data container, but with the capability
 * to handle instance features.
 */
template<typename num_type = float, typename response_type = float, typename index_type = unsigned int>
class mostly_continuous_data_with_instances : public rfr::data_containers::data_container_base<num_type, response_type, index_type>{
  protected:

	std::vector< std::vector<num_type> > configurations;//!< 2d vector to store the feature values of all configurations
	std::vector< std::vector<num_type> > instances; 	//!< 2d vector to store the feature values of all instances

	std::vector<std::pair<index_type, index_type> > config_instance_pairs;
	std::vector<num_type> response_values;
	std::map<index_type, index_type> categorical_ranges;//!< a map storing the few categorical indices and their range
	index_type response_t;
  public:

	// empty constructor. Use this only if you read the data from a file!
	// the private vectors are not properly initialized! Adding data
	// points via 'add_data_point' may or may not fail!
	mostly_continuous_data_with_instances() { throw std::runtime_error("The empty constructor is not supported by this container.");}

	// if you plan on filling the container with single data points one at a time
	// use this constructor to specify the number of features for configurations and instances
	mostly_continuous_data_with_instances (index_type num_config_f, index_type num_instance_f):
		configurations(num_config_f, std::vector<num_type>(0)),
		instances(num_instance_f, std::vector<num_type>(0)),
		response_t(0){}

	virtual num_type feature  (index_type feature_index, index_type sample_index) const {
		// find out if this is a config feature
		if (feature_index < configurations.size()){
			index_type i = config_instance_pairs[sample_index].first;
			return(configurations[feature_index][i]);
		}

		// otherwise it should be a instance feature
		index_type i = config_instance_pairs[sample_index].second;
		feature_index -= configurations.size();
		return(instances[feature_index][i]);
	}

	virtual std::vector<num_type> features (index_type feature_index, std::vector<index_type> &sample_indices) const {
		std::vector<num_type> rv;
		rv.reserve(sample_indices.size());

		if (feature_index < configurations.size()){
			for (auto i : sample_indices)
				rv.push_back( configurations[feature_index][config_instance_pairs[i].first]);
		}
		else{
			feature_index -= configurations.size();
			for (auto i : sample_indices)
				rv.push_back(instances[feature_index][config_instance_pairs[i].second]);
		}
		return(rv);
	}

	virtual response_type response (index_type sample_index) const{
		return(response_values[sample_index]);
	}

	virtual void add_data_point (num_type*, index_type, response_type){
		throw std::runtime_error("This container does not support adding a data point with this function");
	}

	void add_data_point( index_type config_index, index_type instance_index, response_type r){
		if (config_index >= num_configurations() )
			throw std::runtime_error("Configuration index too large.");
		if (instance_index >= num_instances() )

			throw std::runtime_error("Instance index too large.");
		config_instance_pairs.emplace_back(std::pair<index_type, index_type> (config_index, instance_index));
		response_values.emplace_back(r);
	}

	index_type num_configurations(){
		if (configurations.size() > 0)
			return(configurations[0].size());
		return(0);
	}

	index_type num_instances(){
		if (instances.size() > 0)
			return(instances[0].size());
		return(0);
	}

	index_type add_configuration(num_type* config_features, index_type num_elements){
		if (num_elements != configurations.size())
			throw std::runtime_error("Number of configuration features is not what it should be!");

		for (auto i = 0u; i< num_elements; i++)
			configurations[i].push_back(config_features[i]);
		return(num_configurations()-1);
	}

	index_type add_instance(num_type* instance_features, index_type num_elements){
		if (num_elements != instances.size())
			throw std::runtime_error("Number of instance features is not what it should be!");
		for (auto i = 0u; i< num_elements; i++)
			instances[i].push_back(instance_features[i]);
		return(num_instances()-1);
	}

	virtual std::vector<num_type> retrieve_data_point (index_type index) const {
		std::vector<num_type> vec;
		vec.reserve(num_features());

		for (auto i = 0u; i< num_features(); i++)
			vec.emplace_back(feature(i, index));

		return(vec);
	}

	/** \brief method to query the type of a feature
	 *
	 * As most features are assumed to be numerical, it is actually
	 * beneficial to store only the categorical exceptions in a hash-map.
	 * Type = 0 means continuous, and Type = n >= 1 means categorical with
	 * options \in {1, n}. For consistency, we exclude zero from the categorical
	 * values if anyone wants to add sparse data later on.
	 *
	 * \param feature_index the index of the feature
	 * \return int type of the feature: 0 - numerical value (float or int); n>0 - categorical value with n different values {1,2,...,n}
	 *
	 */
	virtual index_type get_type_of_feature (index_type feature_index) const{
		auto it = categorical_ranges.find(feature_index);
		if ( it == categorical_ranges.end())
			return(0);
		return(it->second);
	}

	void set_type_of_configuration_feature(index_type index, index_type type){
		if (type > 0){
			// check consistency for categorical features
			for (auto &fv: configurations[index]){
				if (!(fv<type))
					throw std::runtime_error("Feature values not consistent with provided type. Data contains a value larger than allowed.");
				if (fv < 0)
					throw std::runtime_error("Feature values contain a negative value, can't make that a categorical feature.");
				// round it properly
				fv = std::round(fv);
			}
			categorical_ranges[index] = type;
		}
		else{
			categorical_ranges.erase(index);
		}
	}

	void set_type_of_instance_feature(index_type index, index_type type){
		if (type > 0){
			// check consistency for categorical features
			for (auto &fv: instances[index]){
				if (!(fv<type))
					throw std::runtime_error("Feature values not consistent with provided type. Data contains a value larger than allowed.");
				if (fv < 0)
					throw std::runtime_error("Feature values contain a negative value, can't make that a categorical feature.");
				// round it properly
				fv = std::round(fv);
			}
			categorical_ranges[index + configurations.size()] = type;
		}
		else{
			categorical_ranges.erase(index + configurations.size());
		}
	}

	virtual void set_type_of_feature(index_type index, index_type type){
		if (index >= num_features())
			throw std::runtime_error("Unknown index specified.");
		if (! (type >= 0))
			throw std::runtime_error("Type value should be >= 0");

		if (index < configurations.size())
			set_type_of_configuration_feature(index, type);

		else
			set_type_of_instance_feature(index-configurations.size(), type);

	}

	virtual index_type num_features() const {return(configurations.size() + instances.size());}

	virtual index_type num_data_points() const {return(config_instance_pairs.size());}

	void check_consistency(){

		/* number of stored values has to be the same for:
		*  1. the response values
		*  2. every (row) vector in feature_values
		*/

		if (config_instance_pairs.size() != num_data_points)
			throw std::runtime_error("config_instance_pairs has the wrong size!");

		if (response_values.size() != num_data_points)
			throw std::runtime_error("response_values has the wrong size!");

		for (auto p: config_instance_pairs){
			if ( (p.first <0 ) || (p.first >= num_configurations()))
				throw std::runtime_error("Invalid configuration index");
			if ( (p.second <0 ) || (p.second >= num_instances()))
				throw std::runtime_error("Invalid instance index");
		}

		// check all values for all features
		for (auto f = 0u; f<num_features(); f++){
			if (get_type_of_feature(f) == 0){
				for (auto n = 0u; n < num_data_points(); n++){
					if (isnan(feature(f,n)))
						throw std::runtime_error("Features contain a NaN!");
				}
			}
			else{
				index_type t = get_type_of_feature(f);
				for (auto n = 0u; n < num_data_points(); n++){
					if (isnan(feature(f,n)))
						throw std::runtime_error("Features contain a NaN!");

					if ((feature(f,n) <0) || (feature(f,n) >= t))
						throw std::runtime_error("A categorical feature has an invalid value!");
				}
			}
		}

		index_type t = get_type_of_response();
		for (auto r: response_values){
			if (isnan(r))
				throw std::runtime_error("Responses contain a NaN!");
		}

	}

	virtual index_type get_type_of_response () const{
		return(response_t);
	}

	virtual void set_type_of_response (index_type resp_t){
		if (resp_t > 0){
			for (auto &rv: response_values){
				if (!(rv < resp_t))
					throw std::runtime_error("Response value not consistent with provided type. Data contains a value larger than allowed.");
				if (rv < 0)
					throw std::runtime_error("Response values contain a negative value, can't make that a categorical value.");
			}
			response_t = resp_t;
		}
	}

	/** \brief method to get instance as set_feature for predict_mean_var_of_mean_response_on_set method in regression forest
	 */
	virtual std::vector<num_type> get_instance_set(){
		std::vector<num_type> set_feature;
		set_feature.reserve( num_instances() * num_features());
		for (auto instance_idx = 0u; instance_idx < num_instances(); ++instance_idx){
				for (auto i = 0u; i <  configurations.size(); ++i){
						set_feature.emplace_back(NAN);
				}
				for (auto i = 0u; i < instances.size(); ++i){
						set_feature.emplace_back(instances[i][instance_idx]);
				}
		}
		return set_feature;
	}
	virtual std::vector<num_type> get_configuration_set(num_type configuration_index){
		std::vector<num_type> features;
		features.reserve(num_features());
		for (auto i = 0u; i < configurations.size(); ++i){
				features.emplace_back(configurations[i][configuration_index]);
		}
		for (auto i = configurations.size(); i < configurations.size() +  instances.size(); ++i){
				features.emplace_back(NAN);
		}
		return features;
	}

};

}}//namespace rfr
#endif
<commit_msg>added a method to retrieve the feauture by configuration an instance idx in data container<commit_after>#ifndef RFR_MOSTLY_CONTINUOUS_DATA_CONTAINER_WITH_INSTANCES_HPP
#define RFR_MOSTLY_CONTINUOUS_DATA_CONTAINER_WITH_INSTANCES_HPP

#include <vector>
#include <map>
#include <cmath>

#include "rfr/data_containers/data_container_base.hpp"
#include "rfr/data_containers/data_container_utils.hpp"

namespace rfr{ namespace data_containers{

/** \brief A data container for mostly continuous data with instances.
 *
 * Similar to the mostly_continuous_data container, but with the capability
 * to handle instance features.
 */
template<typename num_type = float, typename response_type = float, typename index_type = unsigned int>
class mostly_continuous_data_with_instances : public rfr::data_containers::data_container_base<num_type, response_type, index_type>{
  protected:

	std::vector< std::vector<num_type> > configurations;//!< 2d vector to store the feature values of all configurations
	std::vector< std::vector<num_type> > instances; 	//!< 2d vector to store the feature values of all instances

	std::vector<std::pair<index_type, index_type> > config_instance_pairs;
	std::vector<num_type> response_values;
	std::map<index_type, index_type> categorical_ranges;//!< a map storing the few categorical indices and their range
	index_type response_t;
  public:

	// empty constructor. Use this only if you read the data from a file!
	// the private vectors are not properly initialized! Adding data
	// points via 'add_data_point' may or may not fail!
	mostly_continuous_data_with_instances() { throw std::runtime_error("The empty constructor is not supported by this container.");}

	// if you plan on filling the container with single data points one at a time
	// use this constructor to specify the number of features for configurations and instances
	mostly_continuous_data_with_instances (index_type num_config_f, index_type num_instance_f):
		configurations(num_config_f, std::vector<num_type>(0)),
		instances(num_instance_f, std::vector<num_type>(0)),
		response_t(0){}

	virtual num_type feature  (index_type feature_index, index_type sample_index) const {
		// find out if this is a config feature
		if (feature_index < configurations.size()){
			index_type i = config_instance_pairs[sample_index].first;
			return(configurations[feature_index][i]);
		}

		// otherwise it should be a instance feature
		index_type i = config_instance_pairs[sample_index].second;
		feature_index -= configurations.size();
		return(instances[feature_index][i]);
	}

	virtual std::vector<num_type> features (index_type feature_index, std::vector<index_type> &sample_indices) const {
		std::vector<num_type> rv;
		rv.reserve(sample_indices.size());

		if (feature_index < configurations.size()){
			for (auto i : sample_indices)
				rv.push_back( configurations[feature_index][config_instance_pairs[i].first]);
		}
		else{
			feature_index -= configurations.size();
			for (auto i : sample_indices)
				rv.push_back(instances[feature_index][config_instance_pairs[i].second]);
		}
		return(rv);
	}

	virtual response_type response (index_type sample_index) const{
		return(response_values[sample_index]);
	}

	virtual void add_data_point (num_type*, index_type, response_type){
		throw std::runtime_error("This container does not support adding a data point with this function");
	}

	void add_data_point( index_type config_index, index_type instance_index, response_type r){
		if (config_index >= num_configurations() )
			throw std::runtime_error("Configuration index too large.");
		if (instance_index >= num_instances() )

			throw std::runtime_error("Instance index too large.");
		config_instance_pairs.emplace_back(std::pair<index_type, index_type> (config_index, instance_index));
		response_values.emplace_back(r);
	}

	index_type num_configurations(){
		if (configurations.size() > 0)
			return(configurations[0].size());
		return(0);
	}

	index_type num_instances(){
		if (instances.size() > 0)
			return(instances[0].size());
		return(0);
	}

	index_type add_configuration(num_type* config_features, index_type num_elements){
		if (num_elements != configurations.size())
			throw std::runtime_error("Number of configuration features is not what it should be!");

		for (auto i = 0u; i< num_elements; i++)
			configurations[i].push_back(config_features[i]);
		return(num_configurations()-1);
	}

	index_type add_instance(num_type* instance_features, index_type num_elements){
		if (num_elements != instances.size())
			throw std::runtime_error("Number of instance features is not what it should be!");
		for (auto i = 0u; i< num_elements; i++)
			instances[i].push_back(instance_features[i]);
		return(num_instances()-1);
	}

	virtual std::vector<num_type> retrieve_data_point (index_type index) const {
		std::vector<num_type> vec;
		vec.reserve(num_features());

		for (auto i = 0u; i< num_features(); i++)
			vec.emplace_back(feature(i, index));

		return(vec);
	}

	/** \brief method to query the type of a feature
	 *
	 * As most features are assumed to be numerical, it is actually
	 * beneficial to store only the categorical exceptions in a hash-map.
	 * Type = 0 means continuous, and Type = n >= 1 means categorical with
	 * options \in {1, n}. For consistency, we exclude zero from the categorical
	 * values if anyone wants to add sparse data later on.
	 *
	 * \param feature_index the index of the feature
	 * \return int type of the feature: 0 - numerical value (float or int); n>0 - categorical value with n different values {1,2,...,n}
	 *
	 */
	virtual index_type get_type_of_feature (index_type feature_index) const{
		auto it = categorical_ranges.find(feature_index);
		if ( it == categorical_ranges.end())
			return(0);
		return(it->second);
	}

	void set_type_of_configuration_feature(index_type index, index_type type){
		if (type > 0){
			// check consistency for categorical features
			for (auto &fv: configurations[index]){
				if (!(fv<type))
					throw std::runtime_error("Feature values not consistent with provided type. Data contains a value larger than allowed.");
				if (fv < 0)
					throw std::runtime_error("Feature values contain a negative value, can't make that a categorical feature.");
				// round it properly
				fv = std::round(fv);
			}
			categorical_ranges[index] = type;
		}
		else{
			categorical_ranges.erase(index);
		}
	}

	void set_type_of_instance_feature(index_type index, index_type type){
		if (type > 0){
			// check consistency for categorical features
			for (auto &fv: instances[index]){
				if (!(fv<type))
					throw std::runtime_error("Feature values not consistent with provided type. Data contains a value larger than allowed.");
				if (fv < 0)
					throw std::runtime_error("Feature values contain a negative value, can't make that a categorical feature.");
				// round it properly
				fv = std::round(fv);
			}
			categorical_ranges[index + configurations.size()] = type;
		}
		else{
			categorical_ranges.erase(index + configurations.size());
		}
	}

	virtual void set_type_of_feature(index_type index, index_type type){
		if (index >= num_features())
			throw std::runtime_error("Unknown index specified.");
		if (! (type >= 0))
			throw std::runtime_error("Type value should be >= 0");

		if (index < configurations.size())
			set_type_of_configuration_feature(index, type);

		else
			set_type_of_instance_feature(index-configurations.size(), type);

	}

	virtual index_type num_features() const {return(configurations.size() + instances.size());}

	virtual index_type num_data_points() const {return(config_instance_pairs.size());}

	void check_consistency(){

		/* number of stored values has to be the same for:
		*  1. the response values
		*  2. every (row) vector in feature_values
		*/

		if (config_instance_pairs.size() != num_data_points)
			throw std::runtime_error("config_instance_pairs has the wrong size!");

		if (response_values.size() != num_data_points)
			throw std::runtime_error("response_values has the wrong size!");

		for (auto p: config_instance_pairs){
			if ( (p.first <0 ) || (p.first >= num_configurations()))
				throw std::runtime_error("Invalid configuration index");
			if ( (p.second <0 ) || (p.second >= num_instances()))
				throw std::runtime_error("Invalid instance index");
		}

		// check all values for all features
		for (auto f = 0u; f<num_features(); f++){
			if (get_type_of_feature(f) == 0){
				for (auto n = 0u; n < num_data_points(); n++){
					if (isnan(feature(f,n)))
						throw std::runtime_error("Features contain a NaN!");
				}
			}
			else{
				index_type t = get_type_of_feature(f);
				for (auto n = 0u; n < num_data_points(); n++){
					if (isnan(feature(f,n)))
						throw std::runtime_error("Features contain a NaN!");

					if ((feature(f,n) <0) || (feature(f,n) >= t))
						throw std::runtime_error("A categorical feature has an invalid value!");
				}
			}
		}

		index_type t = get_type_of_response();
		for (auto r: response_values){
			if (isnan(r))
				throw std::runtime_error("Responses contain a NaN!");
		}

	}

	virtual index_type get_type_of_response () const{
		return(response_t);
	}

	virtual void set_type_of_response (index_type resp_t){
		if (resp_t > 0){
			for (auto &rv: response_values){
				if (!(rv < resp_t))
					throw std::runtime_error("Response value not consistent with provided type. Data contains a value larger than allowed.");
				if (rv < 0)
					throw std::runtime_error("Response values contain a negative value, can't make that a categorical value.");
			}
			response_t = resp_t;
		}
	}

	/** \brief method to get instance as set_feature for predict_mean_var_of_mean_response_on_set method in regression forest
	 */
	virtual std::vector<num_type> get_instance_set(){
		std::vector<num_type> set_feature;
		set_feature.reserve( num_instances() * num_features());
		for (auto instance_idx = 0u; instance_idx < num_instances(); ++instance_idx){
				for (auto i = 0u; i <  configurations.size(); ++i){
						set_feature.emplace_back(NAN);
				}
				for (auto i = 0u; i < instances.size(); ++i){
						set_feature.emplace_back(instances[i][instance_idx]);
				}
		}
		return set_feature;
	}
	virtual std::vector<num_type> get_configuration_set(num_type configuration_index){
		std::vector<num_type> features;
		features.reserve(num_features());
		for (auto i = 0u; i < configurations.size(); ++i){
				features.emplace_back(configurations[i][configuration_index]);
		}
		for (auto i = 0; i < instances.size(); ++i){
				features.emplace_back(NAN);
		}
		return features;
	}
	virtual std::vector<num_type> get_features_by_configuration_and_instance(num_type configuration_index, num_type instance_index){
		std::vector<num_type> features;
		features.reserve(num_features());
		for (auto i = 0u; i < configurations.size(); ++i){
				features.emplace_back(configurations[i][configuration_index]);
		}
		for (auto i = 0; i < instances.size(); ++i){
				features.emplace_back(instances[i][instance_index]);
		}
		return features;
	}
};

}}//namespace rfr
#endif
<|endoftext|>


<commit_before>// Copyright Joshua Boyce 2010-2012.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
// This file is part of HadesMem.
// <http://www.raptorfactor.com/> <raptorfactor@raptorfactor.com>

#include "hadesmem/process.hpp"

#include <array>
#include <utility>

#include "hadesmem/detail/warning_disable_prefix.hpp"
#include <boost/assert.hpp>
#include "hadesmem/detail/warning_disable_suffix.hpp"

#include "hadesmem/error.hpp"

namespace hadesmem
{

Process::Process(DWORD id)
  : id_(id),
  handle_(Open(id))
{
  CheckWoW64();
}

Process::Process(Process&& other) BOOST_NOEXCEPT
  : id_(other.id_),
  handle_(other.handle_)
{
  other.id_ = 0;
  other.handle_ = nullptr;
}

Process& Process::operator=(Process&& other) BOOST_NOEXCEPT
{
  Cleanup();

  std::swap(this->id_, other.id_);
  std::swap(this->handle_, other.handle_);

  return *this;
}

Process::Process(Process const& other)
  : id_(0),
  handle_(nullptr)
{
  HANDLE new_handle = nullptr;
  if (!DuplicateHandle(GetCurrentProcess(), other.handle_,
    GetCurrentProcess(), &new_handle, 0, FALSE, DUPLICATE_SAME_ACCESS))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("DuplicateHandle failed.") <<
      ErrorCodeWinLast(last_error));
  }

  handle_ = new_handle;
  id_ = other.id_;
}

Process& Process::operator=(Process const& other)
{
  HANDLE new_handle = nullptr;
  if (!DuplicateHandle(GetCurrentProcess(), other.handle_,
    GetCurrentProcess(), &new_handle, 0, FALSE, DUPLICATE_SAME_ACCESS))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("DuplicateHandle failed.") <<
      ErrorCodeWinLast(last_error));
  }

  Cleanup();

  handle_ = new_handle;

  id_ = other.id_;

  return *this;
}

Process::~Process() BOOST_NOEXCEPT
{
  if (handle_ && handle_ != GetCurrentProcess())
  {
    BOOST_VERIFY(CloseHandle(handle_));
  }
}

DWORD Process::GetId() const BOOST_NOEXCEPT
{
  return id_;
}

HANDLE Process::GetHandle() const BOOST_NOEXCEPT
{
  return handle_;
}

void Process::Cleanup()
{
  if (!handle_)
  {
    return;
  }

  BOOST_ASSERT(id_);

  if (!CloseHandle(handle_))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("CloseHandle failed.") <<
      ErrorCodeWinLast(last_error));
  }

  id_ = 0;
  handle_ = nullptr;
}

void Process::CheckWoW64() const
{
  BOOL is_wow64_me = FALSE;
  if (!::IsWow64Process(GetCurrentProcess(), &is_wow64_me))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("Could not detect WoW64 status of current process.") <<
      ErrorCodeWinLast(last_error));
  }

  BOOL is_wow64 = FALSE;
  if (!::IsWow64Process(handle_, &is_wow64))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("Could not detect WoW64 status of target process.") <<
      ErrorCodeWinLast(last_error));
  }

  if (is_wow64_me != is_wow64)
  {
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("Cross-architecture process manipulation is currently "
        "unsupported."));
  }
}

HANDLE Process::Open(DWORD id)
{
  HANDLE handle = ::OpenProcess(PROCESS_CREATE_THREAD |
    PROCESS_QUERY_INFORMATION |
    PROCESS_VM_OPERATION |
    PROCESS_VM_READ |
    PROCESS_VM_WRITE,
    FALSE,
    id);
  if (!handle)
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("OpenProcess failed.") <<
      ErrorCodeWinLast(last_error));
  }

  return handle;
}

void Process::CleanupUnchecked() BOOST_NOEXCEPT
{
  try
  {
    Cleanup();
  }
  catch (std::exception const& e)
  {
    // WARNING: Handle is leaked if 'Cleanup' fails.
    BOOST_ASSERT_MSG(false, boost::diagnostic_information(e).c_str());

    id_ = 0;
    handle_ = nullptr;
  }
}

bool Process::operator==(Process const& other) const BOOST_NOEXCEPT
{
  return this->id_ == other.id_;
}

bool Process::operator!=(Process const& other) const BOOST_NOEXCEPT
{
  return !(*this == other);
}

std::wstring GetPath(Process const& process)
{
  std::array<wchar_t, MAX_PATH> path = { { } };
  DWORD path_len = static_cast<DWORD>(path.size());
  if (!::QueryFullProcessImageName(process.GetHandle(), 0, path.data(),
    &path_len))
  {
      DWORD const last_error = GetLastError();
      BOOST_THROW_EXCEPTION(HadesMemError() <<
        ErrorString("QueryFullProcessImageName failed.") <<
        ErrorCodeWinLast(last_error));
  }

  return path.data();
}

bool IsWoW64(Process const& process)
{
  BOOL is_wow64 = FALSE;
  if (!::IsWow64Process(process.GetHandle(), &is_wow64))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("IsWow64Process failed.") <<
      ErrorCodeWinLast(last_error));
  }

  return is_wow64 != FALSE;
}

}
<commit_msg>* [process] Fix the destructor properly because I'm an idiot. * [process] Minor cleanup.<commit_after>// Copyright Joshua Boyce 2010-2012.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
// This file is part of HadesMem.
// <http://www.raptorfactor.com/> <raptorfactor@raptorfactor.com>

#include "hadesmem/process.hpp"

#include <array>
#include <utility>

#include "hadesmem/detail/warning_disable_prefix.hpp"
#include <boost/assert.hpp>
#include "hadesmem/detail/warning_disable_suffix.hpp"

#include "hadesmem/error.hpp"

namespace hadesmem
{

Process::Process(DWORD id)
  : id_(id),
  handle_(Open(id))
{
  CheckWoW64();
}

Process::Process(Process&& other) BOOST_NOEXCEPT
  : id_(other.id_),
  handle_(other.handle_)
{
  other.id_ = 0;
  other.handle_ = nullptr;
}

Process& Process::operator=(Process&& other) BOOST_NOEXCEPT
{
  Cleanup();

  std::swap(this->id_, other.id_);
  std::swap(this->handle_, other.handle_);

  return *this;
}

Process::Process(Process const& other)
  : id_(0),
  handle_(nullptr)
{
  HANDLE new_handle = Duplicate(other.handle_);

  handle_ = new_handle;
  id_ = other.id_;
}

Process& Process::operator=(Process const& other)
{
  HANDLE new_handle = Duplicate(other.handle_);

  Cleanup();

  handle_ = new_handle;

  id_ = other.id_;

  return *this;
}

Process::~Process() BOOST_NOEXCEPT
{
  CleanupUnchecked();
}

DWORD Process::GetId() const BOOST_NOEXCEPT
{
  return id_;
}

HANDLE Process::GetHandle() const BOOST_NOEXCEPT
{
  return handle_;
}

void Process::Cleanup()
{
  if (!handle_)
  {
    return;
  }

  BOOST_ASSERT(id_);

  if (!::CloseHandle(handle_))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("CloseHandle failed.") <<
      ErrorCodeWinLast(last_error));
  }

  id_ = 0;
  handle_ = nullptr;
}

void Process::CheckWoW64() const
{
  BOOL is_wow64_me = FALSE;
  if (!::IsWow64Process(GetCurrentProcess(), &is_wow64_me))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("Could not detect WoW64 status of current process.") <<
      ErrorCodeWinLast(last_error));
  }

  BOOL is_wow64 = FALSE;
  if (!::IsWow64Process(handle_, &is_wow64))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("Could not detect WoW64 status of target process.") <<
      ErrorCodeWinLast(last_error));
  }

  if (is_wow64_me != is_wow64)
  {
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("Cross-architecture process manipulation is currently "
        "unsupported."));
  }
}

HANDLE Process::Open(DWORD id)
{
  HANDLE handle = ::OpenProcess(PROCESS_CREATE_THREAD |
    PROCESS_QUERY_INFORMATION |
    PROCESS_VM_OPERATION |
    PROCESS_VM_READ |
    PROCESS_VM_WRITE,
    FALSE,
    id);
  if (!handle)
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("OpenProcess failed.") <<
      ErrorCodeWinLast(last_error));
  }

  return handle;
}

void Process::CleanupUnchecked() BOOST_NOEXCEPT
{
  try
  {
    Cleanup();
  }
  catch (std::exception const& e)
  {
    // WARNING: Handle is leaked if 'Cleanup' fails.
    BOOST_ASSERT_MSG(false, boost::diagnostic_information(e).c_str());

    id_ = 0;
    handle_ = nullptr;
  }
}

HANDLE Process::Duplicate(HANDLE handle)
{
  HANDLE new_handle = nullptr;
  if (!::DuplicateHandle(GetCurrentProcess(), handle, GetCurrentProcess(),
    &new_handle, 0, FALSE, DUPLICATE_SAME_ACCESS))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("DuplicateHandle failed.") <<
      ErrorCodeWinLast(last_error));
  }

  return new_handle;
}

bool Process::operator==(Process const& other) const BOOST_NOEXCEPT
{
  return this->id_ == other.id_;
}

bool Process::operator!=(Process const& other) const BOOST_NOEXCEPT
{
  return !(*this == other);
}

std::wstring GetPath(Process const& process)
{
  std::array<wchar_t, MAX_PATH> path = { { } };
  DWORD path_len = static_cast<DWORD>(path.size());
  if (!::QueryFullProcessImageName(process.GetHandle(), 0, path.data(),
    &path_len))
  {
      DWORD const last_error = GetLastError();
      BOOST_THROW_EXCEPTION(HadesMemError() <<
        ErrorString("QueryFullProcessImageName failed.") <<
        ErrorCodeWinLast(last_error));
  }

  return path.data();
}

bool IsWoW64(Process const& process)
{
  BOOL is_wow64 = FALSE;
  if (!::IsWow64Process(process.GetHandle(), &is_wow64))
  {
    DWORD const last_error = GetLastError();
    BOOST_THROW_EXCEPTION(HadesMemError() <<
      ErrorString("IsWow64Process failed.") <<
      ErrorCodeWinLast(last_error));
  }

  return is_wow64 != FALSE;
}

}
<|endoftext|>


<commit_before>// Filename: connectionManager.cxx
// Created by:  jns (07Feb00)
//
////////////////////////////////////////////////////////////////////
//
// PANDA 3D SOFTWARE
// Copyright (c) Carnegie Mellon University.  All rights reserved.
//
// All use of this software is subject to the terms of the revised BSD
// license.  You should have received a copy of this license along
// with this source code in a file named "LICENSE."
//
////////////////////////////////////////////////////////////////////

#include "connectionManager.h"
#include "connection.h"
#include "connectionReader.h"
#include "connectionWriter.h"
#include "netAddress.h"
#include "config_net.h"
#include "lightMutexHolder.h"
#include "trueClock.h"

#if defined(WIN32_VC) || defined(WIN64_VC)
#include <winsock2.h>  // For gethostname()
#endif

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::Constructor
//       Access: Public
//  Description:
////////////////////////////////////////////////////////////////////
ConnectionManager::
ConnectionManager() : _set_mutex("ConnectionManager::_set_mutex")
{
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::Destructor
//       Access: Public, Virtual
//  Description:
////////////////////////////////////////////////////////////////////
ConnectionManager::
~ConnectionManager() {
  // Notify all of our associated readers and writers that we're gone.
  Readers::iterator ri;
  for (ri = _readers.begin(); ri != _readers.end(); ++ri) {
    (*ri)->clear_manager();
  }
  Writers::iterator wi;
  for (wi = _writers.begin(); wi != _writers.end(); ++wi) {
    (*wi)->clear_manager();
  }
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_UDP_connection
//       Access: Public
//  Description: Opens a socket for sending and/or receiving UDP
//               packets.  If the port number is greater than zero,
//               the UDP connection will be opened for listening on
//               the indicated port; otherwise, it will be useful only
//               for sending.
//
//               Use a ConnectionReader and ConnectionWriter to handle
//               the actual communication.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_UDP_connection(int port) {
  Socket_UDP *socket = new Socket_UDP;

  if (port > 0) {
    NetAddress address;
    address.set_any(port);

    if (!socket->OpenForInput(address.get_addr())) {
      net_cat.error()
        << "Unable to bind to port " << port << " for UDP.\n";
      delete socket;
      return PT(Connection)();
    }

    net_cat.info()
      << "Creating UDP connection for port " << port << "\n";

  } else {
    if (!socket->InitNoAddress()) {
      net_cat.error()
        << "Unable to initialize outgoing UDP.\n";
      delete socket;
      return PT(Connection)();
    }

    net_cat.info()
      << "Creating outgoing UDP connection\n";
  }

  PT(Connection) connection = new Connection(this, socket);
  new_connection(connection);
  return connection;
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_TCP_server_rendezvous
//       Access: Public
//  Description: Creates a socket to be used as a rendezvous socket
//               for a server to listen for TCP connections.  The
//               socket returned by this call should only be added to
//               a ConnectionListener (not to a generic
//               ConnectionReader).
//
//               This variant of this method accepts a single port,
//               and will listen to that port on all available
//               interfaces.
//
//               backlog is the maximum length of the queue of pending
//               connections.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_TCP_server_rendezvous(int port, int backlog) {
  NetAddress address;
  address.set_any(port);
  return open_TCP_server_rendezvous(address, backlog);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_TCP_server_rendezvous
//       Access: Public
//  Description: Creates a socket to be used as a rendezvous socket
//               for a server to listen for TCP connections.  The
//               socket returned by this call should only be added to
//               a ConnectionListener (not to a generic
//               ConnectionReader).
//
//               This variant of this method accepts a "hostname",
//               which is usually just an IP address in dotted
//               notation, and a port number.  It will listen on the
//               interface indicated by the IP address.  If the IP
//               address is empty string, it will listen on all
//               interfaces.
//
//               backlog is the maximum length of the queue of pending
//               connections.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_TCP_server_rendezvous(const string &hostname, int port, int backlog) {
  NetAddress address;
  if (hostname.empty()) {
    address.set_any(port);
  } else {
    address.set_host(hostname, port);
  }
  return open_TCP_server_rendezvous(address, backlog);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_TCP_server_rendezvous
//       Access: Public
//  Description: Creates a socket to be used as a rendezvous socket
//               for a server to listen for TCP connections.  The
//               socket returned by this call should only be added to
//               a ConnectionListener (not to a generic
//               ConnectionReader).
//
//               This variant of this method accepts a NetAddress,
//               which allows you to specify a specific interface to
//               listen to.
//
//               backlog is the maximum length of the queue of pending
//               connections.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_TCP_server_rendezvous(const NetAddress &address, int backlog) {
  ostringstream strm;
  if (address.get_ip() == 0) {
    strm << "port " << address.get_port();
  } else {
    strm << address.get_ip_string() << ":" << address.get_port();
  }

  Socket_TCP_Listen *socket = new Socket_TCP_Listen;
  bool okflag = socket->OpenForListen(address.get_addr(), backlog);
  if (!okflag) {
    net_cat.info()
      << "Unable to listen to " << strm.str() << " for TCP.\n";
    delete socket;
    return PT(Connection)();
  }

  net_cat.info()
    << "Listening for TCP connections on " << strm.str() << "\n";

  PT(Connection) connection = new Connection(this, socket);
  new_connection(connection);
  return connection;
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_TCP_client_connection
//       Access: Public
//  Description: Attempts to establish a TCP client connection to a
//               server at the indicated address.  If the connection
//               is not established within timeout_ms milliseconds, a
//               null connection is returned.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_TCP_client_connection(const NetAddress &address, int timeout_ms) {
  Socket_TCP *socket = new Socket_TCP;

  // We always open the connection with non-blocking mode first, so we
  // can implement the timeout.
  bool okflag = socket->ActiveOpenNonBlocking(address.get_addr());
  if (okflag && socket->GetLastError() == LOCAL_CONNECT_BLOCKING) {
    // Now wait for the socket to connect.
    TrueClock *clock = TrueClock::get_global_ptr();
    double start = clock->get_short_time();
    Thread::force_yield();
    Socket_fdset fset;
    fset.setForSocket(*socket);
    int ready = fset.WaitForWrite(true, 0);
    while (ready == 0) {
      double elapsed = clock->get_short_time() - start;
      if (elapsed * 1000.0 > timeout_ms) {
        // Timeout.
        okflag = false;
        break;
      }
      Thread::force_yield();
      fset.setForSocket(*socket);
      ready = fset.WaitForWrite(true, 0);
    }
  }

  if (!okflag) {
    net_cat.error()
      << "Unable to open TCP connection to server "
      << address.get_ip_string() << " on port " << address.get_port() << "\n";
    delete socket;
    return PT(Connection)();
  }

#if !defined(HAVE_THREADS) || !defined(SIMPLE_THREADS)
  // Now we have opened the socket in nonblocking mode.  Unless we're
  // using SIMPLE_THREADS, though, we really want the socket in
  // blocking mode (since that's what we support here).  Change it.
  socket->SetBlocking();

#endif  // SIMPLE_THREADS

  net_cat.info()
    << "Opened TCP connection to server " << address.get_ip_string() << " "
    << " on port " << address.get_port() << "\n";

  PT(Connection) connection = new Connection(this, socket);
  new_connection(connection);
  return connection;
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_TCP_client_connection
//       Access: Public
//  Description: This is a shorthand version of the function to
//               directly establish communications to a named host and
//               port.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_TCP_client_connection(const string &hostname, int port,
                           int timeout_ms) {
  NetAddress address;
  if (!address.set_host(hostname, port)) {
    return PT(Connection)();
  }

  return open_TCP_client_connection(address, timeout_ms);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::close_connection
//       Access: Public
//  Description: Terminates a UDP or TCP socket previously opened.
//               This also removes it from any associated
//               ConnectionReader or ConnectionListeners.
//
//               The socket itself may not be immediately closed--it
//               will not be closed until all outstanding pointers to
//               it are cleared, including any pointers remaining in
//               NetDatagrams recently received from the socket.
//
//               The return value is true if the connection was marked
//               to be closed, or false if close_connection() had
//               already been called (or the connection did not belong
//               to this ConnectionManager).  In neither case can you
//               infer anything about whether the connection has
//               *actually* been closed yet based on the return value.
////////////////////////////////////////////////////////////////////
bool ConnectionManager::
close_connection(const PT(Connection) &connection) {
  if (connection != (Connection *)NULL) {
    connection->flush();
  }

  {
    LightMutexHolder holder(_set_mutex);
    Connections::iterator ci = _connections.find(connection);
    if (ci == _connections.end()) {
      // Already closed, or not part of this ConnectionManager.
      return false;
    }
    _connections.erase(ci);

    Readers::iterator ri;
    for (ri = _readers.begin(); ri != _readers.end(); ++ri) {
      (*ri)->remove_connection(connection);
    }
  }

  Socket_IP *socket = connection->get_socket();

  // We can't *actually* close the connection right now, because
  // there might be outstanding pointers to it.  But we can at least
  // shut it down.  It will be eventually closed when all the
  // pointers let go.

  net_cat.info()
    << "Shutting down connection " << (void *)connection
    << " locally.\n";
  socket->Close();

  return true;
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::get_host_name
//       Access: Public, Static
//  Description: Returns the name of this particular machine on the
//               network, if available, or the empty string if the
//               hostname cannot be determined.
////////////////////////////////////////////////////////////////////
string ConnectionManager::
get_host_name() {
  char temp_buff[1024];
  if (gethostname(temp_buff, 1024) == 0) {
    return string(temp_buff);
  }

  return string();
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::new_connection
//       Access: Protected
//  Description: This internal function is called whenever a new
//               connection is established.  It allows the
//               ConnectionManager to save all of the pointers to open
//               connections so they can't be inadvertently deleted
//               until close_connection() is called.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
new_connection(const PT(Connection) &connection) {
  LightMutexHolder holder(_set_mutex);
  _connections.insert(connection);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::flush_read_connection
//       Access: Protected, Virtual
//  Description: An internal function called by ConnectionWriter only
//               when a write failure has occurred.  This method
//               ensures that all of the read data has been flushed
//               from the pipe before the connection is fully removed.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
flush_read_connection(Connection *connection) {
  Readers readers;
  {
    LightMutexHolder holder(_set_mutex);
    Connections::iterator ci = _connections.find(connection);
    if (ci == _connections.end()) {
      // Already closed, or not part of this ConnectionManager.
      return;
    }
    _connections.erase(ci);

    // Get a copy first, so we can release the lock before traversing.
    readers = _readers;
  }
  Readers::iterator ri;
  for (ri = readers.begin(); ri != readers.end(); ++ri) {
    (*ri)->flush_read_connection(connection);
  }

  Socket_IP *socket = connection->get_socket();
  socket->Close();
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::connection_reset
//       Access: Protected, Virtual
//  Description: An internal function called by the ConnectionReader,
//               ConnectionWriter, or ConnectionListener when a
//               connection has been externally reset.  This adds the
//               connection to the queue of those which have recently
//               been reset.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
connection_reset(const PT(Connection) &connection, bool okflag) {
  if (net_cat.is_info()) {
    if (okflag) {
      net_cat.info()
        << "Connection " << (void *)connection
        << " was closed normally by the other end";

    } else {
      net_cat.info()
        << "Lost connection " << (void *)connection
        << " unexpectedly\n";
    }
  }

  // Turns out we do need to explicitly mark the connection as closed
  // immediately, rather than waiting for the user to do it, since
  // otherwise we'll keep trying to listen for noise on the socket and
  // we'll always hear a "yes" answer.
  close_connection(connection);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::add_reader
//       Access: Protected
//  Description: This internal function is called by ConnectionReader
//               when it is constructed.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
add_reader(ConnectionReader *reader) {
  LightMutexHolder holder(_set_mutex);
  _readers.insert(reader);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::remove_reader
//       Access: Protected
//  Description: This internal function is called by ConnectionReader
//               when it is destructed.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
remove_reader(ConnectionReader *reader) {
  LightMutexHolder holder(_set_mutex);
  _readers.erase(reader);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::add_writer
//       Access: Protected
//  Description: This internal function is called by ConnectionWriter
//               when it is constructed.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
add_writer(ConnectionWriter *writer) {
  LightMutexHolder holder(_set_mutex);
  _writers.insert(writer);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::remove_writer
//       Access: Protected
//  Description: This internal function is called by ConnectionWriter
//               when it is destructed.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
remove_writer(ConnectionWriter *writer) {
  LightMutexHolder holder(_set_mutex);
  _writers.erase(writer);
}
<commit_msg>fix identification of failed TCP connection<commit_after>// Filename: connectionManager.cxx
// Created by:  jns (07Feb00)
//
////////////////////////////////////////////////////////////////////
//
// PANDA 3D SOFTWARE
// Copyright (c) Carnegie Mellon University.  All rights reserved.
//
// All use of this software is subject to the terms of the revised BSD
// license.  You should have received a copy of this license along
// with this source code in a file named "LICENSE."
//
////////////////////////////////////////////////////////////////////

#include "connectionManager.h"
#include "connection.h"
#include "connectionReader.h"
#include "connectionWriter.h"
#include "netAddress.h"
#include "config_net.h"
#include "lightMutexHolder.h"
#include "trueClock.h"

#if defined(WIN32_VC) || defined(WIN64_VC)
#include <winsock2.h>  // For gethostname()
#endif

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::Constructor
//       Access: Public
//  Description:
////////////////////////////////////////////////////////////////////
ConnectionManager::
ConnectionManager() : _set_mutex("ConnectionManager::_set_mutex")
{
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::Destructor
//       Access: Public, Virtual
//  Description:
////////////////////////////////////////////////////////////////////
ConnectionManager::
~ConnectionManager() {
  // Notify all of our associated readers and writers that we're gone.
  Readers::iterator ri;
  for (ri = _readers.begin(); ri != _readers.end(); ++ri) {
    (*ri)->clear_manager();
  }
  Writers::iterator wi;
  for (wi = _writers.begin(); wi != _writers.end(); ++wi) {
    (*wi)->clear_manager();
  }
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_UDP_connection
//       Access: Public
//  Description: Opens a socket for sending and/or receiving UDP
//               packets.  If the port number is greater than zero,
//               the UDP connection will be opened for listening on
//               the indicated port; otherwise, it will be useful only
//               for sending.
//
//               Use a ConnectionReader and ConnectionWriter to handle
//               the actual communication.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_UDP_connection(int port) {
  Socket_UDP *socket = new Socket_UDP;

  if (port > 0) {
    NetAddress address;
    address.set_any(port);

    if (!socket->OpenForInput(address.get_addr())) {
      net_cat.error()
        << "Unable to bind to port " << port << " for UDP.\n";
      delete socket;
      return PT(Connection)();
    }

    net_cat.info()
      << "Creating UDP connection for port " << port << "\n";

  } else {
    if (!socket->InitNoAddress()) {
      net_cat.error()
        << "Unable to initialize outgoing UDP.\n";
      delete socket;
      return PT(Connection)();
    }

    net_cat.info()
      << "Creating outgoing UDP connection\n";
  }

  PT(Connection) connection = new Connection(this, socket);
  new_connection(connection);
  return connection;
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_TCP_server_rendezvous
//       Access: Public
//  Description: Creates a socket to be used as a rendezvous socket
//               for a server to listen for TCP connections.  The
//               socket returned by this call should only be added to
//               a ConnectionListener (not to a generic
//               ConnectionReader).
//
//               This variant of this method accepts a single port,
//               and will listen to that port on all available
//               interfaces.
//
//               backlog is the maximum length of the queue of pending
//               connections.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_TCP_server_rendezvous(int port, int backlog) {
  NetAddress address;
  address.set_any(port);
  return open_TCP_server_rendezvous(address, backlog);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_TCP_server_rendezvous
//       Access: Public
//  Description: Creates a socket to be used as a rendezvous socket
//               for a server to listen for TCP connections.  The
//               socket returned by this call should only be added to
//               a ConnectionListener (not to a generic
//               ConnectionReader).
//
//               This variant of this method accepts a "hostname",
//               which is usually just an IP address in dotted
//               notation, and a port number.  It will listen on the
//               interface indicated by the IP address.  If the IP
//               address is empty string, it will listen on all
//               interfaces.
//
//               backlog is the maximum length of the queue of pending
//               connections.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_TCP_server_rendezvous(const string &hostname, int port, int backlog) {
  NetAddress address;
  if (hostname.empty()) {
    address.set_any(port);
  } else {
    address.set_host(hostname, port);
  }
  return open_TCP_server_rendezvous(address, backlog);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_TCP_server_rendezvous
//       Access: Public
//  Description: Creates a socket to be used as a rendezvous socket
//               for a server to listen for TCP connections.  The
//               socket returned by this call should only be added to
//               a ConnectionListener (not to a generic
//               ConnectionReader).
//
//               This variant of this method accepts a NetAddress,
//               which allows you to specify a specific interface to
//               listen to.
//
//               backlog is the maximum length of the queue of pending
//               connections.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_TCP_server_rendezvous(const NetAddress &address, int backlog) {
  ostringstream strm;
  if (address.get_ip() == 0) {
    strm << "port " << address.get_port();
  } else {
    strm << address.get_ip_string() << ":" << address.get_port();
  }

  Socket_TCP_Listen *socket = new Socket_TCP_Listen;
  bool okflag = socket->OpenForListen(address.get_addr(), backlog);
  if (!okflag) {
    net_cat.info()
      << "Unable to listen to " << strm.str() << " for TCP.\n";
    delete socket;
    return PT(Connection)();
  }

  net_cat.info()
    << "Listening for TCP connections on " << strm.str() << "\n";

  PT(Connection) connection = new Connection(this, socket);
  new_connection(connection);
  return connection;
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_TCP_client_connection
//       Access: Public
//  Description: Attempts to establish a TCP client connection to a
//               server at the indicated address.  If the connection
//               is not established within timeout_ms milliseconds, a
//               null connection is returned.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_TCP_client_connection(const NetAddress &address, int timeout_ms) {
  Socket_TCP *socket = new Socket_TCP;

  // We always open the connection with non-blocking mode first, so we
  // can implement the timeout.
  bool okflag = socket->ActiveOpenNonBlocking(address.get_addr());
  if (okflag && socket->GetLastError() == LOCAL_CONNECT_BLOCKING) {
    // Now wait for the socket to connect.
    TrueClock *clock = TrueClock::get_global_ptr();
    double start = clock->get_short_time();
    Thread::force_yield();
    Socket_fdset fset;
    fset.setForSocket(*socket);
    int ready = fset.WaitForWrite(true, 0);
    while (ready == 0) {
      double elapsed = clock->get_short_time() - start;
      if (elapsed * 1000.0 > timeout_ms) {
        // Timeout.
        okflag = false;
        break;
      }
      Thread::force_yield();
      fset.setForSocket(*socket);
      ready = fset.WaitForWrite(true, 0);
    }
  }

  if (okflag) {
    // So, the connect() operation finished, but did it succeed or fail?
    if (socket->GetPeerName().GetIPAddressRaw() == 0) {
      // No peer means it failed.
      okflag = false;
    }
  }

  if (!okflag) {
    net_cat.error()
      << "Unable to open TCP connection to server "
      << address.get_ip_string() << " on port " << address.get_port() << "\n";
    delete socket;
    return PT(Connection)();
  }

#if !defined(HAVE_THREADS) || !defined(SIMPLE_THREADS)
  // Now we have opened the socket in nonblocking mode.  Unless we're
  // using SIMPLE_THREADS, though, we really want the socket in
  // blocking mode (since that's what we support here).  Change it.
  socket->SetBlocking();

#endif  // SIMPLE_THREADS

  net_cat.info()
    << "Opened TCP connection to server " << address.get_ip_string() << " "
    << " on port " << address.get_port() << "\n";

  PT(Connection) connection = new Connection(this, socket);
  new_connection(connection);
  return connection;
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::open_TCP_client_connection
//       Access: Public
//  Description: This is a shorthand version of the function to
//               directly establish communications to a named host and
//               port.
////////////////////////////////////////////////////////////////////
PT(Connection) ConnectionManager::
open_TCP_client_connection(const string &hostname, int port,
                           int timeout_ms) {
  NetAddress address;
  if (!address.set_host(hostname, port)) {
    return PT(Connection)();
  }

  return open_TCP_client_connection(address, timeout_ms);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::close_connection
//       Access: Public
//  Description: Terminates a UDP or TCP socket previously opened.
//               This also removes it from any associated
//               ConnectionReader or ConnectionListeners.
//
//               The socket itself may not be immediately closed--it
//               will not be closed until all outstanding pointers to
//               it are cleared, including any pointers remaining in
//               NetDatagrams recently received from the socket.
//
//               The return value is true if the connection was marked
//               to be closed, or false if close_connection() had
//               already been called (or the connection did not belong
//               to this ConnectionManager).  In neither case can you
//               infer anything about whether the connection has
//               *actually* been closed yet based on the return value.
////////////////////////////////////////////////////////////////////
bool ConnectionManager::
close_connection(const PT(Connection) &connection) {
  if (connection != (Connection *)NULL) {
    connection->flush();
  }

  {
    LightMutexHolder holder(_set_mutex);
    Connections::iterator ci = _connections.find(connection);
    if (ci == _connections.end()) {
      // Already closed, or not part of this ConnectionManager.
      return false;
    }
    _connections.erase(ci);

    Readers::iterator ri;
    for (ri = _readers.begin(); ri != _readers.end(); ++ri) {
      (*ri)->remove_connection(connection);
    }
  }

  Socket_IP *socket = connection->get_socket();

  // We can't *actually* close the connection right now, because
  // there might be outstanding pointers to it.  But we can at least
  // shut it down.  It will be eventually closed when all the
  // pointers let go.

  net_cat.info()
    << "Shutting down connection " << (void *)connection
    << " locally.\n";
  socket->Close();

  return true;
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::get_host_name
//       Access: Public, Static
//  Description: Returns the name of this particular machine on the
//               network, if available, or the empty string if the
//               hostname cannot be determined.
////////////////////////////////////////////////////////////////////
string ConnectionManager::
get_host_name() {
  char temp_buff[1024];
  if (gethostname(temp_buff, 1024) == 0) {
    return string(temp_buff);
  }

  return string();
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::new_connection
//       Access: Protected
//  Description: This internal function is called whenever a new
//               connection is established.  It allows the
//               ConnectionManager to save all of the pointers to open
//               connections so they can't be inadvertently deleted
//               until close_connection() is called.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
new_connection(const PT(Connection) &connection) {
  LightMutexHolder holder(_set_mutex);
  _connections.insert(connection);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::flush_read_connection
//       Access: Protected, Virtual
//  Description: An internal function called by ConnectionWriter only
//               when a write failure has occurred.  This method
//               ensures that all of the read data has been flushed
//               from the pipe before the connection is fully removed.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
flush_read_connection(Connection *connection) {
  Readers readers;
  {
    LightMutexHolder holder(_set_mutex);
    Connections::iterator ci = _connections.find(connection);
    if (ci == _connections.end()) {
      // Already closed, or not part of this ConnectionManager.
      return;
    }
    _connections.erase(ci);

    // Get a copy first, so we can release the lock before traversing.
    readers = _readers;
  }
  Readers::iterator ri;
  for (ri = readers.begin(); ri != readers.end(); ++ri) {
    (*ri)->flush_read_connection(connection);
  }

  Socket_IP *socket = connection->get_socket();
  socket->Close();
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::connection_reset
//       Access: Protected, Virtual
//  Description: An internal function called by the ConnectionReader,
//               ConnectionWriter, or ConnectionListener when a
//               connection has been externally reset.  This adds the
//               connection to the queue of those which have recently
//               been reset.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
connection_reset(const PT(Connection) &connection, bool okflag) {
  if (net_cat.is_info()) {
    if (okflag) {
      net_cat.info()
        << "Connection " << (void *)connection
        << " was closed normally by the other end";

    } else {
      net_cat.info()
        << "Lost connection " << (void *)connection
        << " unexpectedly\n";
    }
  }

  // Turns out we do need to explicitly mark the connection as closed
  // immediately, rather than waiting for the user to do it, since
  // otherwise we'll keep trying to listen for noise on the socket and
  // we'll always hear a "yes" answer.
  close_connection(connection);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::add_reader
//       Access: Protected
//  Description: This internal function is called by ConnectionReader
//               when it is constructed.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
add_reader(ConnectionReader *reader) {
  LightMutexHolder holder(_set_mutex);
  _readers.insert(reader);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::remove_reader
//       Access: Protected
//  Description: This internal function is called by ConnectionReader
//               when it is destructed.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
remove_reader(ConnectionReader *reader) {
  LightMutexHolder holder(_set_mutex);
  _readers.erase(reader);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::add_writer
//       Access: Protected
//  Description: This internal function is called by ConnectionWriter
//               when it is constructed.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
add_writer(ConnectionWriter *writer) {
  LightMutexHolder holder(_set_mutex);
  _writers.insert(writer);
}

////////////////////////////////////////////////////////////////////
//     Function: ConnectionManager::remove_writer
//       Access: Protected
//  Description: This internal function is called by ConnectionWriter
//               when it is destructed.
////////////////////////////////////////////////////////////////////
void ConnectionManager::
remove_writer(ConnectionWriter *writer) {
  LightMutexHolder holder(_set_mutex);
  _writers.erase(writer);
}
<|endoftext|>