[ARMORY-30]

- Improve KDF compute duration prediction - Cover prediction code - Add WalletCreationParams, feed custom unlock times at wallet creation
goatpig · goatpig · May 31, 2023 · May 31, 2023 · May 31, 2023 · 4f7b60bab7a4caf9a1a4ce74f1f9d8f939529fb2
commit 4f7b60bab7a4caf9a1a4ce74f1f9d8f939529fb2
diff --git a/cppForSwig/BridgeAPI/WalletManager.cpp b/cppForSwig/BridgeAPI/WalletManager.cpp
@@ -156,8 +156,9 @@ shared_ptr<WalletContainer> WalletManager::createNewWallet(
    unique_ptr<ClearTextSeed> seed(new ClearTextSeed_Armory135(root,
       ClearTextSeed_Armory135::LegacyType::Armory200));
    auto wallet = AssetWallet_Single::createFromSeed(move(seed),
-      pass, controlPass,
-      path_, lookup);
+      WalletCreationParams{
+         pass, controlPass, path_, lookup
+      });
    return addWallet(wallet, wallet->getMainAccountID());
 }
 
@@ -940,8 +941,9 @@ shared_ptr<AssetWallet_Single> Armory135Header::migrate(
       unique_ptr<ClearTextSeed> seed(new ClearTextSeed_Armory135(
          decryptedRoot, chaincodeCopy));
       wallet = AssetWallet_Single::createFromSeed(move(seed),
-         privKeyPass, controlPass,
-         folder, highestIndex);
+         WalletCreationParams{
+            privKeyPass, controlPass, folder, (uint32_t)highestIndex
+         });
    }
 
    //main account id, check it matches armory wallet id

diff --git a/cppForSwig/EncryptionUtils.h b/cppForSwig/EncryptionUtils.h
@@ -162,13 +162,14 @@ class KdfRomix
 
    /////////////////////////////////////////////////////////////////////////////
    // Default max-memory reqt will 
-   void computeKdfParams(double   targetComputeSec=0.25, 
-                         uint32_t maxMemReqtsBytes=DEFAULT_KDF_MAX_MEMORY);
+   void computeKdfParams(double targetComputeSec=0.25,
+      uint32_t maxMemReqtsBytes=DEFAULT_KDF_MAX_MEMORY,
+      bool verbose = false);
 
    /////////////////////////////////////////////////////////////////////////////
-   void usePrecomputedKdfParams(uint32_t memReqts, 
-                                uint32_t numIter, 
-                                SecureBinaryData salt);
+   void usePrecomputedKdfParams(uint32_t memReqts,
+      uint32_t numIter,
+      SecureBinaryData salt);
 
    /////////////////////////////////////////////////////////////////////////////
    void printKdfParams(void);

diff --git a/cppForSwig/KDF.cpp b/cppForSwig/KDF.cpp
@@ -31,15 +31,18 @@ KdfRomix::KdfRomix(uint32_t memReqts, uint32_t numIter, SecureBinaryData salt) :
 }
 
 /////////////////////////////////////////////////////////////////////////////
-void KdfRomix::computeKdfParams(double targetComputeSec, uint32_t maxMemReqts)
+void KdfRomix::computeKdfParams(
+   double targetComputeMSec,
+   uint32_t maxMemReqts,
+   bool verbose)
 {
    // Create a random salt, even though this is probably unnecessary:
    // the variation in numIter and memReqts is probably effective enough
    salt_ = CryptoPRNG::generateRandom(32);
 
-   // If target compute is 0s, then this method really only generates 
+   // If target compute is 0s, then this method really only generates
    // a random salt, and sets the other params to default minimum.
-   if (targetComputeSec == 0)
+   if (targetComputeMSec <= 4)
    {
       numIterations_ = 1;
       memoryReqtBytes_ = 1024;
@@ -48,64 +51,68 @@ void KdfRomix::computeKdfParams(double targetComputeSec, uint32_t maxMemReqts)
 
 
    // Here, we pick the largest memory reqt that allows the executing system
-   // to compute the KDF is less than the target time.  A maximum can be 
+   // to compute the KDF is less than the target time.  A maximum can be
    // specified, in case the target system is likely to be memory-limited
    // more than compute-speed limited
-   auto&& testKey = SecureBinaryData::fromString("This is an example key to test KDF iteration speed");
+
+   //12 bytes test key
+   auto testKey = SecureBinaryData::fromString("- Test Key -");
 
    // Start the search for a memory value at 1kB
    memoryReqtBytes_ = 1024;
-   double approxSec = 0;
-   while (approxSec <= targetComputeSec / 4 && memoryReqtBytes_ < maxMemReqts)
+   uint64_t approxMSec = 0;
+   while (approxMSec <= targetComputeMSec / 4 &&
+      memoryReqtBytes_ < maxMemReqts)
    {
       memoryReqtBytes_ *= 2;
 
       sequenceCount_ = memoryReqtBytes_ / hashOutputBytes_;
       lookupTable_.resize(memoryReqtBytes_);
 
-      TIMER_RESTART("KDF_Mem_Search");
+      auto start = chrono::system_clock::now();
       testKey = DeriveKey_OneIter(testKey);
-      TIMER_STOP("KDF_Mem_Search");
-      approxSec = TIMER_READ_SEC("KDF_Mem_Search");
+      auto end = chrono::system_clock::now();
+      approxMSec = chrono::duration_cast<chrono::milliseconds>(
+         end - start).count();
    }
 
-   // Recompute here, in case we didn't enter the search above 
+   // Recompute here, in case we didn't enter the search above
    sequenceCount_ = memoryReqtBytes_ / hashOutputBytes_;
    lookupTable_.resize(memoryReqtBytes_);
 
 
-   // Depending on the search above (or if a low max memory was chosen, 
+   // Depending on the search above (or if a low max memory was chosen,
    // we may need to do multiple iterations to achieve the desired compute
    // time on this system.
-   double allItersSec = 0;
+   uint64_t allItersMSec = 0;
    uint32_t numTest = 1;
-   while (allItersSec < 0.02)
+   while (allItersMSec < 100)
    {
       numTest *= 2;
-      TIMER_RESTART("KDF_Time_Search");
+      auto start = chrono::system_clock::now();
       for (uint32_t i = 0; i < numTest; i++)
       {
-         auto&& _testKey = SecureBinaryData::fromString(
-            "This is an example key to test KDF iteration speed");
-
+         auto&& _testKey = testKey;
          _testKey = DeriveKey_OneIter(_testKey);
       }
-      TIMER_STOP("KDF_Time_Search");
-      allItersSec = TIMER_READ_SEC("KDF_Time_Search");
+      auto end = chrono::system_clock::now();
+      allItersMSec = chrono::duration_cast<chrono::milliseconds>(
+         end - start).count();
    }
 
-   double perIterSec = allItersSec / numTest;
-   numIterations_ = (uint32_t)(targetComputeSec / (perIterSec + 0.0005));
-   numIterations_ = (numIterations_ < 1 ? 1 : numIterations_);
-   //cout << "System speed test results    :  " << endl;
-   //cout << "   Total test of the KDF took:  " << allItersSec*1000 << " ms" << endl;
-   //cout << "                   to execute:  " << numTest << " iterations" << endl;
-   //cout << "   Target computation time is:  " << targetComputeSec*1000 << " ms" << endl;
-   //cout << "   Setting numIterations to:    " << numIterations_ << endl;
+   uint64_t perIterMSec = allItersMSec / numTest;
+   numIterations_ = (uint32_t)(targetComputeMSec / (perIterMSec + 1));
+   numIterations_ = (numIterations_ < 1 ? 1 : numIterations_ + 1);
+   if (verbose)
+   {
+      cout << "System speed test results    :  " << endl;
+      cout << "   Total test of the KDF took:  " << allItersMSec << " ms" << endl;
+      cout << "                   to execute:  " << numTest << " iterations" << endl;
+      cout << "   Target computation time is:  " << targetComputeMSec << " ms" << endl;
+      cout << "   Setting numIterations to:    " << numIterations_ << endl;
+   }
 }
 
-
-
 /////////////////////////////////////////////////////////////////////////////
 void KdfRomix::usePrecomputedKdfParams(uint32_t memReqts,
    uint32_t numIter,
@@ -206,5 +213,5 @@ SecureBinaryData KdfRomix::DeriveKey(SecureBinaryData const & password)
    for (uint32_t i = 0; i < numIterations_; i++)
       masterKey = DeriveKey_OneIter(masterKey);
 
-   return SecureBinaryData(masterKey);
+   return masterKey;
 }
diff --git a/cppForSwig/Wallets/AssetEncryption.cpp b/cppForSwig/Wallets/AssetEncryption.cpp
@@ -25,6 +25,23 @@ using namespace Armory::Wallets::Encryption;
 KeyDerivationFunction::~KeyDerivationFunction()
 {}
 
+////////////////////////////////////////////////////////////////////////////////
+KeyDerivationFunction_Romix::KeyDerivationFunction_Romix(
+   uint32_t unlockTime_ms) :
+   KeyDerivationFunction(), salt_((initialize(unlockTime_ms)))
+{}
+
+////
+KeyDerivationFunction_Romix::KeyDerivationFunction_Romix(
+   unsigned iterations, unsigned memTarget, SecureBinaryData salt) :
+   KeyDerivationFunction(),
+   iterations_(iterations), memTarget_(memTarget), salt_(move(salt))
+{}
+
+////
+KeyDerivationFunction_Romix::~KeyDerivationFunction_Romix()
+{}
+
 ////////////////////////////////////////////////////////////////////////////////
 BinaryData KeyDerivationFunction_Romix::computeID() const
 {
@@ -39,10 +56,10 @@ BinaryData KeyDerivationFunction_Romix::computeID() const
 }
 
 ////////////////////////////////////////////////////////////////////////////////
-BinaryData KeyDerivationFunction_Romix::initialize()
+BinaryData KeyDerivationFunction_Romix::initialize(uint32_t unlockTime_ms)
 {
    KdfRomix kdf;
-   kdf.computeKdfParams(0);
+   kdf.computeKdfParams(unlockTime_ms);
    iterations_ = kdf.getNumIterations();
    memTarget_ = kdf.getMemoryReqtBytes();
    return kdf.getSalt();
@@ -53,7 +70,7 @@ SecureBinaryData KeyDerivationFunction_Romix::deriveKey(
    const SecureBinaryData& rawKey) const
 {
    KdfRomix kdfObj(memTarget_, iterations_, salt_);
-   return move(kdfObj.DeriveKey(rawKey));
+   return kdfObj.DeriveKey(rawKey);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -95,7 +112,7 @@ shared_ptr<KeyDerivationFunction> KeyDerivationFunction::deserialize(
          SecureBinaryData salt(move(brr.get_BinaryData(len)));
 
          kdfPtr = make_shared<KeyDerivationFunction_Romix>(
-            iterations, memTarget, salt);
+            iterations, memTarget, move(salt));
          break;
       }
 
@@ -157,6 +174,22 @@ unsigned KeyDerivationFunction_Romix::memTarget() const
    return memTarget_;
 }
 
+////////
+unsigned KeyDerivationFunction_Romix::iterations() const
+{
+   return iterations_;
+}
+
+////////
+void KeyDerivationFunction_Romix::prettyPrint() const
+{
+   cout << "KDF Parameters:" << endl;
+   cout << "   HashFunction : " << "sha512" << endl;
+   cout << "   Memory/thread: " << memTarget_ << " bytes" << endl;
+   cout << "   NumIterations: " << iterations_ << endl;
+   cout << "   Salt         : " << salt_.toHexStr() << endl;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////
 //// Cipher

diff --git a/cppForSwig/Wallets/AssetEncryption.h b/cppForSwig/Wallets/AssetEncryption.h
@@ -45,6 +45,7 @@ namespace Armory
          struct ClearTextEncryptionKey;
          class ClearTextAssetData;
 
+         ///////////////////////////////////////////////////////////////////////
          class KeyDerivationFunction
          {
          public:
@@ -67,7 +68,7 @@ namespace Armory
                deserialize(const BinaryDataRef&);
          };
 
-         ///////////////////////////////////////////////////////////////////////
+         ////////
          class KeyDerivationFunction_Romix : public KeyDerivationFunction
          {
          private:
@@ -78,25 +79,23 @@ namespace Armory
 
          private:
             BinaryData computeID(void) const;
-            BinaryData initialize(void);
+            BinaryData initialize(uint32_t);
 
          public:
-            KeyDerivationFunction_Romix() :
-               KeyDerivationFunction(),
-               salt_(std::move(initialize()))
-            {}
+            KeyDerivationFunction_Romix(uint32_t);
+            KeyDerivationFunction_Romix(unsigned, unsigned, SecureBinaryData);
+            ~KeyDerivationFunction_Romix(void) override;
 
-            KeyDerivationFunction_Romix(unsigned iterations, unsigned memTarget,
-               SecureBinaryData& salt) :
-               KeyDerivationFunction(),
-               iterations_(iterations), memTarget_(memTarget), salt_(salt)
-            {}
+            //overrides
+            SecureBinaryData deriveKey(const SecureBinaryData&) const override;
+            bool isSame(KeyDerivationFunction* const) const override;
+            BinaryData serialize(void) const override;
+            const BinaryData& getId(void) const override;
 
-            SecureBinaryData deriveKey(const SecureBinaryData& rawKey) const;
-            bool isSame(KeyDerivationFunction* const) const;
-            BinaryData serialize(void) const;
-            const BinaryData& getId(void) const;
+            //locals
             unsigned memTarget(void) const;
+            unsigned iterations(void) const;
+            void prettyPrint(void) const;
          };
 
          ///////////////////////////////////////////////////////////////////////

diff --git a/cppForSwig/Wallets/AuthorizedPeers.cpp b/cppForSwig/Wallets/AuthorizedPeers.cpp
@@ -26,7 +26,7 @@ using namespace Armory::Seeds;
 
 ////////////////////////////////////////////////////////////////////////////////
 AuthorizedPeers::AuthorizedPeers(
-   const string& datadir, const string& filename, 
+   const string& datadir, const string& filename,
    const PassphraseLambda& passLbd)
 {
    auto path = datadir;
@@ -157,10 +157,10 @@ void AuthorizedPeers::loadWallet(
 
 ////////////////////////////////////////////////////////////////////////////////
 void AuthorizedPeers::createWallet(
-   const string& baseDir, const string& filename, 
+   const string& baseDir, const string& filename,
    const PassphraseLambda& passLbd)
 {
-   //Default peers wallet password. Asset wallets always encrypt private keys, 
+   //Default peers wallet password. Asset wallets always encrypt private keys,
    //have to provide a password at creation.
    auto&& password = SecureBinaryData::fromString(PEERS_WALLET_PASSWORD);
    auto&& controlPassphrase = passLbd({});
@@ -175,7 +175,9 @@ void AuthorizedPeers::createWallet(
       wallet_ = AssetWallet_Single::createFromSeed(
          make_unique<ClearTextSeed_BIP32>(
             CryptoPRNG::generateRandom(32), SeedType::BIP32_Virgin),
-         password, controlPassphrase, baseDir);
+         WalletCreationParams{
+            password, controlPassphrase, baseDir, 0, 100, 250
+         });
       auto wltSingle = dynamic_pointer_cast<AssetWallet_Single>(wallet_);
 
       auto rootBip32 = dynamic_pointer_cast<AssetEntry_BIP32Root>(

diff --git a/cppForSwig/Wallets/Seeds/Backups.cpp b/cppForSwig/Wallets/Seeds/Backups.cpp
@@ -11,6 +11,7 @@
 #include "BtcUtils.h"
 #include "../WalletIdTypes.h"
 #include "Seeds.h"
+#include "Wallets.h"
 #include "protobuf/BridgeProto.pb.h"
 extern "C" {
 #include <trezor-crypto/bip39.h>
@@ -1024,8 +1025,8 @@ unique_ptr<Backup_Base58> Helpers::getBase58BackupString(
 
 ////////////////////////////// -- restore methods -- ///////////////////////////
 shared_ptr<AssetWallet> Helpers::restoreFromBackup(
-   unique_ptr<WalletBackup> backup, const std::string& homedir,
-   const UserPrompt& callback)
+   unique_ptr<WalletBackup> backup, const UserPrompt& callback,
+   const WalletCreationParams& params)
 {
    unique_ptr<ClearTextSeed> seed = nullptr;
    auto bType = backup->type();
@@ -1074,7 +1075,9 @@ shared_ptr<AssetWallet> Helpers::restoreFromBackup(
    }
 
    //prompt for passwords
-   SecureBinaryData privkey, control;
+   BinaryDataRef pass = params.passphrase.getRef();
+   BinaryDataRef control = params.controlPassphrase.getRef();
+   if (pass.empty())
    {
       BridgeProto::RestorePrompt prompt;
       prompt.set_get_passphrases(true);
@@ -1083,13 +1086,16 @@ shared_ptr<AssetWallet> Helpers::restoreFromBackup(
       if (!reply.success())
          throw RestoreUserException("user did not provide a passphrase");
 
-      privkey = SecureBinaryData::fromString(reply.passphrases().privkey());
-      control = SecureBinaryData::fromString(reply.passphrases().control());
+      pass.setRef(reply.passphrases().privkey());
+      control.setRef(reply.passphrases().control());
    }
 
+   WalletCreationParams paramsCopy{ pass, control,
+      params.folder, params.lookup,
+      params.publicUnlockDuration_ms, params.privateUnlockDuration_ms };
+
    //return wallet
-   return AssetWallet_Single::createFromSeed(
-      std::move(seed), privkey, control, homedir);
+   return AssetWallet_Single::createFromSeed(std::move(seed), paramsCopy);
 }
 
 ////////