1 files changed, 36 insertions, 6 deletions
diff --git a/src/Encryption.cpp b/src/Encryption.cpp
index c4e6a2c..238861f 100644
--- a/src/Encryption.cpp
+++ b/src/Encryption.cpp
@@ -1,16 +1,26 @@
 #include "../include/Encryption.hpp"
 #include <sodium/crypto_aead_xchacha20poly1305.h>
 #include <sodium/randombytes.h>
-#include <string>
+#include <cstring>
 
 namespace odhtdb
 {
-    Encryption::Encryption(const DataView &data, const DataView &additionalData)
+    Encryption::Encryption(const DataView &data, const DataView &additionalData, const DataView &_key)
     {
-        cipherText = new unsigned char[crypto_aead_xchacha20poly1305_ietf_ABYTES + data.size];        
-        crypto_aead_xchacha20poly1305_ietf_keygen(key);
+        cipherText = new unsigned char[crypto_aead_xchacha20poly1305_ietf_ABYTES + data.size];   
+        cipherTextLength = crypto_aead_xchacha20poly1305_ietf_ABYTES + data.size;
+        
+        if(_key.data)
+        {
+            if(_key.size != KEY_BYTE_SIZE)
+                throw EncryptionException("Encryption key is wrong size");
+            memcpy(key, _key.data, _key.size);
+        }
+        else
+            crypto_aead_xchacha20poly1305_ietf_keygen(key);
+        
         randombytes_buf(nonce, 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, (const unsigned char*)additionalData.data, additionalData.size, nullptr, nonce, key) < 0)
             throw EncryptionException("Failed to encrypt data");
     }
     
@@ -37,6 +47,7 @@ namespace odhtdb
     Decryption::Decryption(const DataView &data, const DataView &nonce, const DataView &key)
     {
         decryptedText = new unsigned char[data.size];
+        decryptedTextLength = data.size;
         
         if(nonce.size < NONCE_BYTE_SIZE)
             throw DecryptionException("Nonce is not big enough");
@@ -44,10 +55,29 @@ namespace odhtdb
         if(key.size < KEY_BYTE_SIZE)
             throw DecryptionException("Key is not big enough");
         
-        if(crypto_aead_xchacha20poly1305_ietf_decrypt(decryptedText, &decryptedTextLength, nullptr, (const unsigned char*)data.data, data.size, nullptr, 0, (const unsigned char*)nonce.data, (const unsigned char*)key.data) != 0)
+        if(crypto_aead_xchacha20poly1305_ietf_decrypt(decryptedText, &decryptedTextLength, nullptr, (const unsigned char*)data.data, data.size, nullptr, 0, (const unsigned char*)nonce.data, (const unsigned char*)key.data) < 0)
             throw DecryptionException("Failed to decrypt data");
     }
     
+    Decryption::Decryption(Decryption &&other)
+    {
+        decryptedText = other.decryptedText;
+        decryptedTextLength = other.decryptedTextLength;
+        
+        other.decryptedText = nullptr;
+        other.decryptedTextLength = 0;
+    }
+    
+    Decryption& Decryption::operator=(Decryption &&other)
+    {
+        decryptedText = other.decryptedText;
+        decryptedTextLength = other.decryptedTextLength;
+        
+        other.decryptedText = nullptr;
+        other.decryptedTextLength = 0;
+        return *this;
+    }
+    
     Decryption::~Decryption()
     {
         delete[](decryptedText);