Store node encryption key with user data, fix encryption bug when using additional data

7 files changed, 57 insertions, 39 deletions

diff --git a/include/odhtdb/Database.hpp b/include/odhtdb/Database.hpp
index 87389eb..04c3b62 100644
--- a/include/odhtdb/Database.hpp
+++ b/include/odhtdb/Database.hpp
@@ -188,11 +188,11 @@ namespace odhtdb
         // Username has to be either unique or if it's the same as existing one, then password has to match.
         // Node has to be unique for the user.
         // Throws DatabaseStorageWrongPassword or SqlExecException on failure (if username is not unique in node).
-        void storeUserPasswordEncrypted(const Hash &nodeHash, const std::string &username, const std::string &password, const Signature::KeyPair &keyPair);
+        void storeNodeInfoForUserEncrypted(const DatabaseNode &nodeInfo, const std::string &username, const std::string &password, const Signature::KeyPair &keyPair);
         
-        // Returns nodes, public key and private key of encrypted user.
+        // Returns nodes, node encryption key, public key and private key of encrypted user.
         // Throws DatabaseStorageWrongPassword if password for the stored user is wrong.
-        MapHash<Signature::KeyPair> getStoredUserNodeDataDecrypted(const std::string &username, const std::string &password);
+        MapHash<StoredNodeInfo> getStoredNodeUserInfoDecrypted(const std::string &username, const std::string &password) const;
         
         std::vector<OwnedMemory> getUserGroups(const Hash &nodeHash, const Signature::PublicKey &userPublicKey) const;
         
diff --git a/include/odhtdb/DatabaseStorage.hpp b/include/odhtdb/DatabaseStorage.hpp
index 264ab57..a618420 100644
--- a/include/odhtdb/DatabaseStorage.hpp
+++ b/include/odhtdb/DatabaseStorage.hpp
@@ -10,6 +10,7 @@
 #include "OwnedMemory.hpp"
 #include "DatabaseOperation.hpp"
 #include "DatabaseOrder.hpp"
+#include "DatabaseNode.hpp"
 #include "sql/SqlQuery.hpp"
 #include "sql/SqlExec.hpp"
 #include <vector>
@@ -72,6 +73,12 @@ namespace odhtdb
     using FetchNodeUserActionGapsCallbackFunc = std::function<void(const DataView userPublicKey, u64 start, u64 range)>;
     using FetchNodeUserLatestActionCounterCallbackFunc = std::function<void(const DataView userPublicKey, u64 latestActionCounter)>;
     
+    struct StoredNodeInfo
+    {
+        std::shared_ptr<OwnedMemory> nodeEncryptionKey;
+        std::shared_ptr<Signature::KeyPair> userKeyPair;
+    };
+    
     class DatabaseStorage
     {
     public:
@@ -119,12 +126,13 @@ namespace odhtdb
         // Username has to be either unique or if it's the same as existing one, then password has to match.
         // Node has to be unique for the user.
         // Throws DatabaseStorageWrongPassword or SqlExecException on failure (if username is not unique in node).
-        void storeUserPasswordEncrypted(const Hash &nodeHash, const std::string &username, const std::string &password, const Signature::KeyPair &keyPair);
+        void storeNodeInfoForUserEncrypted(const DatabaseNode &nodeInfo, const std::string &username, const std::string &password, const Signature::KeyPair &keyPair);
         
-        // Returns nodes, public key and private key of encrypted user.
+        // Returns nodes, node encryption key, public key and private key of encrypted user.
         // Throws DatabaseStorageWrongPassword if password for the stored user is wrong.
         // Throws DatabaseStorageNoSuchStoredUser if user doesn't exist.
-        MapHash<Signature::KeyPair> getStoredUserNodeDataDecrypted(const std::string &username, const std::string &password);
+        // Otherwise throw DatabaseStorageException on other errors.
+        MapHash<StoredNodeInfo> getStoredNodeUserInfoDecrypted(const std::string &username, const std::string &password) const;
         
         // Returns true and node decryption key if node exists and we have the decryption key, 
         // otherwise return false and OwnedMemory with data set to nullptr
@@ -158,7 +166,7 @@ namespace odhtdb
         void setNodeAddDataDecryptedData(i64 rowId, const DataView &decryptedData);
         
         // Throws DatabaseStorageNoSuchStoredUser or DatabaseStorageWrongPassword
-        i64 getStoredUserId(const std::string &username, const DataView &hashedPassword);
+        i64 getStoredUserId(const std::string &username, const DataView &hashedPassword) const;
     private:
         Database *database;
         sqlite3 *sqliteDb;
diff --git a/include/odhtdb/Encryption.hpp b/include/odhtdb/Encryption.hpp
index 2457630..b2ae67e 100644
--- a/include/odhtdb/Encryption.hpp
+++ b/include/odhtdb/Encryption.hpp
@@ -32,7 +32,7 @@ namespace odhtdb
         DISABLE_COPY(Encryption)
     public:
         // Throws EncryptionException on failure (or std::bad_alloc on failed memory allocation)
-        Encryption(const DataView &data, const DataView &additionalData = DataView(), const DataView &key = DataView());
+        Encryption(const DataView &data, const DataView &key = DataView());
         ~Encryption();
         
         DataView getKey() const;
diff --git a/src/Database.cpp b/src/Database.cpp
index fff74f4..df11b3c 100644
--- a/src/Database.cpp
+++ b/src/Database.cpp
@@ -90,10 +90,10 @@ namespace odhtdb
     }
     
     Database::Database(const char *bootstrapNodeAddr, u16 port, const boost::filesystem::path &storageDir, DatabaseCallbackFuncs callbackFuncs) : 
+        databaseStorage(this, storageDir),
         onCreateNodeCallbackFunc(callbackFuncs.createNodeCallbackFunc),
         onAddNodeCallbackFunc(callbackFuncs.addNodeCallbackFunc),
         onAddUserCallbackFunc(callbackFuncs.addUserCallbackFunc),
-        databaseStorage(this, storageDir),
         shuttingDown(false)
     {
         node.run(port , {
@@ -443,7 +443,7 @@ namespace odhtdb
         serializer.add(newActionCounter);
         
         DataView encryptionKey(nodeInfo.getNodeEncryptionKey()->data, ENCRYPTION_KEY_BYTE_SIZE);
-        Encryption encryptedBody(dataToAdd, DataView(), encryptionKey);
+        Encryption encryptedBody(dataToAdd, encryptionKey);
         OwnedMemory requestData = combine(serializer, encryptedBody);
         string signedRequestData = userToPerformActionWith.getPrivateKey().sign(requestData.getView());
         OwnedMemory stagedAddObject = combine(userToPerformActionWith.getPublicKey(), signedRequestData);
@@ -643,14 +643,14 @@ namespace odhtdb
         return databaseStorage.storeUserWithoutNodes(username, password);
     }
     
-    void Database::storeUserPasswordEncrypted(const Hash &nodeHash, const string &username, const string &password, const Signature::KeyPair &keyPair)
+    void Database::storeNodeInfoForUserEncrypted(const DatabaseNode &nodeInfo, const string &username, const string &password, const Signature::KeyPair &keyPair)
     {
-        return databaseStorage.storeUserPasswordEncrypted(nodeHash, username, password, keyPair);
+        return databaseStorage.storeNodeInfoForUserEncrypted(nodeInfo, username, password, keyPair);
     }
         
-    MapHash<Signature::KeyPair> Database::getStoredUserNodeDataDecrypted(const string &username, const string &password)
+    MapHash<StoredNodeInfo> Database::getStoredNodeUserInfoDecrypted(const string &username, const string &password) const
     {
-        return databaseStorage.getStoredUserNodeDataDecrypted(username, password);
+        return databaseStorage.getStoredNodeUserInfoDecrypted(username, password);
     }
     
     vector<OwnedMemory> Database::getUserGroups(const Hash &nodeHash, const Signature::PublicKey &userPublicKey) const
diff --git a/src/DatabaseStorage.cpp b/src/DatabaseStorage.cpp
index 7b316f7..016c498 100644
--- a/src/DatabaseStorage.cpp
+++ b/src/DatabaseStorage.cpp
@@ -117,7 +117,7 @@ namespace odhtdb
                             "CREATE TABLE IF NOT EXISTS NodeUserActionGap(id INTEGER PRIMARY KEY, nodeUserId INTEGER NOT NULL, start INTEGER NOT NULL, range INTEGER NOT NULL, FOREIGN KEY(nodeUserId) REFERENCES NodeUser(id));"
                             
                             "CREATE TABLE IF NOT EXISTS NodeEncryptedUser(id INTEGER PRIMARY KEY, username BLOB UNIQUE NOT NULL, nonce BLOB NOT NULL, encryptedPassword BLOB NOT NULL);"
-                            "CREATE TABLE IF NOT EXISTS NodeEncryptedUserData(node INTEGER NOT NULL, usernameId INTEGER NOT NULL, userPublicKey BLOB NOT NULL, nonce BLOB NOT NULL, userPrivateKeyEncrypted BLOB NOT NULL, FOREIGN KEY(usernameId) REFERENCES NodeEncryptedUser(id), FOREIGN KEY(node, userPublicKey) REFERENCES NodeUser(node, publicKey));"
+                            "CREATE TABLE IF NOT EXISTS NodeEncryptedUserData(node INTEGER NOT NULL, usernameId INTEGER NOT NULL, userPublicKey BLOB NOT NULL, nonce BLOB NOT NULL, nodeAndUserPrivateKeyEncrypted BLOB NOT NULL, FOREIGN KEY(usernameId) REFERENCES NodeEncryptedUser(id), FOREIGN KEY(node, userPublicKey) REFERENCES NodeUser(node, publicKey));"
                             
                             "CREATE UNIQUE INDEX IF NOT EXISTS UniqueUserInNode ON NodeUser(node, publicKey);"
                             "CREATE UNIQUE INDEX IF NOT EXISTS UniqueUserGroupAssoc ON NodeUserGroupAssoc(node, userPublicKey, groupId);"
@@ -851,12 +851,12 @@ namespace odhtdb
         DataView usernameView((void*)username.data(), username.size());
         
         DataView saltView((void*)passwordSalt, PASSWORD_SALT_LEN);
-        Encryption encryptedPassword(saltView, {}, hashedPasswordView);
+        Encryption encryptedPassword(saltView, hashedPasswordView);
         SqlExec addUserSql(sqliteDb, "INSERT INTO NodeEncryptedUser(username, nonce, encryptedPassword) VALUES(?, ?, ?)");
         addUserSql.execWithArgs({ usernameView, encryptedPassword.getNonce(), encryptedPassword.getCipherText() });
     }
     
-    i64 DatabaseStorage::getStoredUserId(const string &username, const DataView &hashedPassword)
+    i64 DatabaseStorage::getStoredUserId(const string &username, const DataView &hashedPassword) const
     {
         DataView usernameView((void*)username.data(), username.size());
         SqlQuery userQuery(sqliteDb, "SELECT id, nonce, encryptedPassword FROM NodeEncryptedUser WHERE username = ?", { usernameView });
@@ -882,12 +882,16 @@ namespace odhtdb
         }
     }
     
-    void DatabaseStorage::storeUserPasswordEncrypted(const Hash &nodeHash, const string &username, const string &password, const Signature::KeyPair &keyPair)
+    void DatabaseStorage::storeNodeInfoForUserEncrypted(const DatabaseNode &nodeInfo, const string &username, const string &password, const Signature::KeyPair &keyPair)
     {
         OwnedMemory hashedPassword = hashPassword(DataView((void*)password.data(), password.size()), DataView((void*)passwordSalt, PASSWORD_SALT_LEN));
         DataView hashedPasswordView(hashedPassword.data, hashedPassword.size);
         DataView privateKeyView((void*)keyPair.getPrivateKey().getData(), PRIVATE_KEY_NUM_BYTES);
-        Encryption encryptedPrivateKey(privateKeyView, {}, hashedPasswordView);
+        
+        sibs::SafeSerializer serializer;
+        serializer.add((const u8*)privateKeyView.data, privateKeyView.size);
+        serializer.add((const u8*)nodeInfo.getNodeEncryptionKey()->data, nodeInfo.getNodeEncryptionKey()->size);
+        Encryption encryptedNodeUserPrivateKey(DataView(serializer.getBuffer().data(), serializer.getBuffer().size()), hashedPasswordView);
         
         SqlTransaction transaction(sqliteDb);
         DataView usernameView((void*)username.data(), username.size());
@@ -900,7 +904,7 @@ namespace odhtdb
         catch(DatabaseStorageNoSuchStoredUser &e)
         {
             DataView saltView((void*)passwordSalt, PASSWORD_SALT_LEN);
-            Encryption encryptedPassword(saltView, {}, hashedPasswordView);
+            Encryption encryptedPassword(saltView, hashedPasswordView);
             
             SqlExec addUserSql(sqliteDb, "INSERT INTO NodeEncryptedUser(username, nonce, encryptedPassword) VALUES(?, ?, ?)");
             addUserSql.execWithArgs({ usernameView, encryptedPassword.getNonce(), encryptedPassword.getCipherText() });
@@ -912,25 +916,25 @@ namespace odhtdb
             encryptedUserRowId = userQuery.getInt64(0);
         }
         
-        SqlExec sqlExec(sqliteDb, "INSERT INTO NodeEncryptedUserData(node, usernameId, userPublicKey, nonce, userPrivateKeyEncrypted) VALUES(?, ?, ?, ?, ?)");
+        SqlExec sqlExec(sqliteDb, "INSERT INTO NodeEncryptedUserData(node, usernameId, userPublicKey, nonce, nodeAndUserPrivateKeyEncrypted) VALUES(?, ?, ?, ?, ?)");
         sqlExec.execWithArgs({
-            DataView(nodeHash.getData(), nodeHash.getSize()),
+            DataView(nodeInfo.getRequestHash()->getData(), nodeInfo.getRequestHash()->getSize()),
             encryptedUserRowId,
             DataView((void*)keyPair.getPublicKey().getData(), PUBLIC_KEY_NUM_BYTES),
-            encryptedPrivateKey.getNonce(),
-            encryptedPrivateKey.getCipherText()
+            encryptedNodeUserPrivateKey.getNonce(),
+            encryptedNodeUserPrivateKey.getCipherText()
         });
         transaction.commit();
     }
     
-    MapHash<Signature::KeyPair> DatabaseStorage::getStoredUserNodeDataDecrypted(const string &username, const string &password)
+    MapHash<StoredNodeInfo> DatabaseStorage::getStoredNodeUserInfoDecrypted(const string &username, const string &password) const
     {
         OwnedMemory hashedPassword = hashPassword(DataView((void*)password.data(), password.size()), DataView((void*)passwordSalt, PASSWORD_SALT_LEN));
         DataView hashedPasswordView(hashedPassword.data, hashedPassword.size);
         i64 encryptedUserRowId = getStoredUserId(username, hashedPasswordView);
         
-        MapHash<Signature::KeyPair> result;
-        SqlQuery query(sqliteDb, "SELECT node, userPublicKey, nonce, userPrivateKeyEncrypted FROM NodeEncryptedUserData WHERE usernameId = ?", { encryptedUserRowId });
+        MapHash<StoredNodeInfo> result;
+        SqlQuery query(sqliteDb, "SELECT node, userPublicKey, nonce, nodeAndUserPrivateKeyEncrypted FROM NodeEncryptedUserData WHERE usernameId = ?", { encryptedUserRowId });
         while(query.next())
         {
             Hash nodeHash;
@@ -940,14 +944,18 @@ namespace odhtdb
             const DataView storedUserPublicKey = query.getBlob(1);
             Signature::PublicKey userPublicKey((const char*)storedUserPublicKey.data, storedUserPublicKey.size);
             const DataView storedNonce = query.getBlob(2);
-            const DataView storedPrivateKeyEncrypted = query.getBlob(3);
+            const DataView storedNodeUserPrivateKeyEncrypted = query.getBlob(3);
             try
             {
                 // TODO: We can bypass decrypting several times by storing several nodes private key user the same field (as a blob)
-                Decryption decryptedStoredPrivateKey(storedPrivateKeyEncrypted, storedNonce, hashedPasswordView);
-                Signature::PrivateKey userPrivateKey((const char*)decryptedStoredPrivateKey.getDecryptedText().data, decryptedStoredPrivateKey.getDecryptedText().size);
-                Signature::KeyPair keyPair(userPublicKey, userPrivateKey);
-                result[nodeHash] = keyPair;
+                Decryption decryptedStoredNodeUserPrivateKey(storedNodeUserPrivateKeyEncrypted, storedNonce, hashedPasswordView);
+                if(decryptedStoredNodeUserPrivateKey.getDecryptedText().size != PRIVATE_KEY_NUM_BYTES + ENCRYPTION_KEY_BYTE_SIZE)
+                    throw DatabaseStorageException("Encrypted data size is of unexpected size");
+                Signature::PrivateKey userPrivateKey((const char*)decryptedStoredNodeUserPrivateKey.getDecryptedText().data, PRIVATE_KEY_NUM_BYTES);
+                shared_ptr<Signature::KeyPair> keyPair = make_shared<Signature::KeyPair>(userPublicKey, userPrivateKey);
+                shared_ptr<OwnedMemory> nodeEncryptionKey = make_shared<OwnedMemory>(new u8[ENCRYPTION_KEY_BYTE_SIZE], ENCRYPTION_KEY_BYTE_SIZE);
+                memcpy(nodeEncryptionKey->data, (char*)decryptedStoredNodeUserPrivateKey.getDecryptedText().data + PRIVATE_KEY_NUM_BYTES, ENCRYPTION_KEY_BYTE_SIZE);
+                result[nodeHash] = { nodeEncryptionKey, keyPair };
             }
             catch(DecryptionException &e)
             {
diff --git a/src/Encryption.cpp b/src/Encryption.cpp
index ff37270..d4763b8 100644
--- a/src/Encryption.cpp
+++ b/src/Encryption.cpp
@@ -7,7 +7,7 @@ namespace odhtdb
 {
     static_assert(ENCRYPTION_CHECKSUM_BYTE_SIZE == crypto_aead_xchacha20poly1305_ietf_ABYTES, "Encryption checksum key size has changed for some reason, oops...");
     
-    Encryption::Encryption(const DataView &data, const DataView &additionalData, const DataView &_key)
+    Encryption::Encryption(const DataView &data, const DataView &_key)
     {
         cipherTextLength = crypto_aead_xchacha20poly1305_ietf_ABYTES + data.size;
         cipherText = new unsigned char[cipherTextLength];   
@@ -22,7 +22,7 @@ namespace odhtdb
             generateKey(key);
         
         randombytes_buf(nonce, ENCRYPTION_NONCE_BYTE_SIZE);
-        if(crypto_aead_xchacha20poly1305_ietf_encrypt(cipherText, &cipherTextLength, (const unsigned char*)data.data, data.size, (const unsigned char*)additionalData.data, additionalData.size, nullptr, nonce, key) < 0)
+        if(crypto_aead_xchacha20poly1305_ietf_encrypt(cipherText, &cipherTextLength, (const unsigned char*)data.data, data.size, nullptr, 0, nullptr, nonce, key) < 0)
             throw EncryptionException("Failed to encrypt data");
     }
     
diff --git a/tests/main.cpp b/tests/main.cpp
index 3f9bbf9..e4438e7 100644
--- a/tests/main.cpp
+++ b/tests/main.cpp
@@ -205,10 +205,10 @@ int main()
         if(database.doesStoredUserExist(username))
             fail("Expected stored to not exist until it has been added");
         
-        database.storeUserPasswordEncrypted(*databaseNode.getRequestHash(), username, password, *adminUserKey);
+        database.storeNodeInfoForUserEncrypted(databaseNode, username, password, *adminUserKey);
         try
         {
-            database.storeUserPasswordEncrypted(*databaseNode.getRequestHash(), username, password, localUserKeyPair);
+            database.storeNodeInfoForUserEncrypted(databaseNode, username, password, localUserKeyPair);
             fail("Expected store user password to fail since we have already stored an user in the node");
         }
         catch(SqlExecException &e)
@@ -219,15 +219,17 @@ int main()
         if(!database.doesStoredUserExist(username))
             fail("Expected stored to exist after it has been added");
         
-        auto nodeUserData = database.getStoredUserNodeDataDecrypted(username, password);
+        auto nodeUserData = database.getStoredNodeUserInfoDecrypted(username, password);
         assertEquals((size_t)1, nodeUserData.size());
         auto userDataIt = nodeUserData.find(*databaseNode.getRequestHash());
         if(userDataIt == nodeUserData.end())
             fail("Expected stored node hash to match node hash");
-        if(userDataIt->second.getPublicKey() != adminUserKey->getPublicKey())
+        if(userDataIt->second.userKeyPair->getPublicKey() != adminUserKey->getPublicKey())
             fail("Expected stored public key to match admin user public key");
-        if(userDataIt->second.getPrivateKey() != adminUserKey->getPrivateKey())
+        if(userDataIt->second.userKeyPair->getPrivateKey() != adminUserKey->getPrivateKey())
             fail("Expected stored private key to match admin user private key");
+        if(userDataIt->second.nodeEncryptionKey->getView() != databaseCreateResponse->getNodeEncryptionKey()->getView())
+            fail("Expected stored node encryption key to match node encryption key");
         
         try
         {