- add more mlock and zeroize

- remove mlock and zeroize from bench

benches/bench.py

@@ -7,7 +7,6 @@
 import shutil
 import io
 import numpy as np
-from zeroize import zeroize1, mlock, munlock
 
 
 def hash(bytes_in):
@@ -101,18 +100,14 @@ def encrypt(block_len):
     cipher_meta = CipherMeta.Ring(RingAlgorithm.AES256GCM)
     key_len = cipher_meta.key_len()
     key = bytearray(key_len)
-    mlock(key)
     cipher_meta.generate_key(key)
-    munlock(key)
     cipher = Cipher(cipher_meta, key)
 
     plaintext_len = block_len
     ciphertext_len = cipher.ciphertext_len(plaintext_len)
     buf = np.array([0] * ciphertext_len, dtype=np.uint8)
-    mlock(buf)
 
     plaintext = bytearray(os.urandom(plaintext_len))
-    mlock(plaintext)
     cipher.copy_slice(plaintext, buf[:plaintext_len])
     aad = b"AAD"
 
@@ -128,26 +123,20 @@ def encrypt(block_len):
 
     average = sum(deltas, 0) / len(deltas)
     print(f"| {block_len/1024/1024} | {average:.5f} |")
-    zeroize1(plaintext)
-    munlock(plaintext)
 
 
 def encrypt_from(block_len):
     cipher_meta = CipherMeta.Ring(RingAlgorithm.AES256GCM)
     key_len = cipher_meta.key_len()
     key = bytearray(key_len)
-    mlock(key)
     cipher_meta.generate_key(key)
-    munlock(key)
     cipher = Cipher(cipher_meta, key)
 
     plaintext_len = block_len
     ciphertext_len = cipher.ciphertext_len(plaintext_len)
     buf = np.array([0] * ciphertext_len, dtype=np.uint8)
-    mlock(buf)
 
     plaintext = bytearray(os.urandom(plaintext_len))
-    mlock(plaintext)
     aad = b"AAD"
 
     deltas = []
@@ -163,9 +152,6 @@ def encrypt_from(block_len):
     average = sum(deltas, 0) / len(deltas)
     print(f"| {block_len/1024/1024} | {average:.5f} |")
 
-    zeroize1(plaintext)
-    munlock(plaintext)
-
 
 def encrypt_file(path_in, path_out):
     chunk_len = 256 * 1024
@@ -175,15 +161,12 @@ def encrypt_file(path_in, path_out):
     cipher_meta = CipherMeta.Ring(RingAlgorithm.AES256GCM)
     key_len = cipher_meta.key_len()
     key = bytearray(key_len)
-    mlock(key)
     cipher_meta.generate_key(key)
-    munlock(key)
     cipher = Cipher(cipher_meta, key)
 
     plaintext_len = chunk_len
     ciphertext_len = cipher.ciphertext_len(plaintext_len)
     buf = np.array([0] * ciphertext_len, dtype=np.uint8)
-    mlock(buf)
 
     aad = b"AAD"
 
@@ -211,26 +194,19 @@ def encrypt_file(path_in, path_out):
     filesize = get_file_size(path_in)
     print(f"| {(filesize / 1024 / 1024):.5g} | {average:.5f} |")
 
-    zeroize1(buf)
-    munlock(buf)
-
 
 def decrypt(block_len):
     cipher_meta = CipherMeta.Ring(RingAlgorithm.AES256GCM)
     key_len = cipher_meta.key_len()
     key = bytearray(key_len)
-    mlock(key)
     cipher_meta.generate_key(key)
-    munlock(key)
     cipher = Cipher(cipher_meta, key)
 
     plaintext_len = block_len
     ciphertext_len = cipher.ciphertext_len(plaintext_len)
     buf = np.array([0] * ciphertext_len, dtype=np.uint8)
-    mlock(buf)
 
     plaintext = bytearray(os.urandom(plaintext_len))
-    mlock(plaintext)
     aad = b"AAD"
 
     deltas = []
@@ -251,28 +227,19 @@ def decrypt(block_len):
     average = sum(deltas, 0) / len(deltas)
     print(f"| {block_len/1024/1024} | {average:.5f} |")
 
-    zeroize1(plaintext)
-    zeroize1(buf)
-    munlock(plaintext)
-    munlock(buf)
-
 
 def decrypt_from(block_len):
     cipher_meta = CipherMeta.Ring(RingAlgorithm.AES256GCM)
     key_len = cipher_meta.key_len()
     key = bytearray(key_len)
-    mlock(key)
     cipher_meta.generate_key(key)
-    munlock(key)
     cipher = Cipher(cipher_meta, key)
 
     plaintext_len = block_len
     ciphertext_len = cipher.ciphertext_len(plaintext_len)
     buf = np.array([0] * ciphertext_len, dtype=np.uint8)
-    mlock(buf)
 
     plaintext = bytearray(os.urandom(plaintext_len))
-    mlock(plaintext)
     aad = b"AAD"
 
     deltas = []
@@ -294,11 +261,6 @@ def decrypt_from(block_len):
     average = sum(deltas, 0) / len(deltas)
     print(f"| {block_len/1024/1024} | {average:.5f} |")
 
-    zeroize1(plaintext)
-    zeroize1(buf)
-    munlock(plaintext)
-    munlock(buf)
-
 
 def decrypt_file(plaintext_file, ciphertext_file):
     cipher_meta = CipherMeta.Ring(RingAlgorithm.AES256GCM)
@@ -353,23 +315,17 @@ def decrypt_file(plaintext_file, ciphertext_file):
     filesize = get_file_size(plaintext_file)
     print(f"| {(filesize / 1024 / 1024):.5g} | {average:.5f} |")
 
-    zeroize1(buf)
-    munlock(buf)
-
 
 def encrypt_speed_per_mb(block_len):
     cipher_meta = CipherMeta.Ring(RingAlgorithm.AES256GCM)
     key_len = cipher_meta.key_len()
     key = bytearray(key_len)
-    mlock(key)
     cipher_meta.generate_key(key)
-    munlock(key)
     cipher = Cipher(cipher_meta, key)
 
     plaintext_len = block_len
     ciphertext_len = cipher.ciphertext_len(plaintext_len)
     buf = np.array([0] * ciphertext_len, dtype=np.uint8)
-    mlock(buf)
 
     plaintext = bytearray(os.urandom(plaintext_len))
     cipher.copy_slice(plaintext, buf[:plaintext_len])

tests/test_rencrypt.py

@@ -87,30 +87,29 @@ def test_encrypt_aes(self):
         buf = np.array([0] * ciphertext_len, dtype=np.uint8)
         mlock(buf)
 
-        aad = b"AAD"
-
         # put some plaintext in the buffer, it would be ideal if you can directly collect the data into the buffer without allocating new memory
         # but for the sake of example we will allocate and copy the data
         plaintext = bytearray(os.urandom(plaintext_len))
         mlock(plaintext)
+        aad = b"AAD"
         # cipher.copy_slice is slighlty faster than buf[:plaintext_len] = plaintext, especially for large plaintext, because it copies the data in parallel
         # cipher.copy_slice takes bytes as input, cipher.copy_slice takes bytearray
         cipher.copy_slice(plaintext, buf)
         # encrypt it, this will encrypt in-place the data in the buffer
         ciphertext_len = cipher.encrypt(buf, plaintext_len, 42, aad)
-        cipertext = buf[:ciphertext_len]
-        mlock(cipertext)
+        buf[:ciphertext_len]
 
         # decrypt it
         # if you need to copy ciphertext to buffer, we don't need to do it now as it's already in the buffer
         # cipher.copy_slice(ciphertext, buf[:len(ciphertext)])
         plaintext_len = cipher.decrypt(buf, ciphertext_len, 42, aad)
         plaintext2 = buf[:plaintext_len]
-        mlock(plaintext2)
         self.assertEqual(plaintext, plaintext2)
 
         zeroize1(plaintext)
         zeroize1(buf)
+        munlock(plaintext)
+        munlock(buf)
 
     def test_encrypt_chacha(self):
         cipher_meta = CipherMeta.Ring(RingAlgorithm.ChaCha20Poly1305)
@@ -127,19 +126,16 @@ def test_encrypt_chacha(self):
         buf = np.array([0] * ciphertext_len, dtype=np.uint8)
         mlock(buf)
 
-        aad = b"AAD"
-
         # put some plaintext in the buffer, it would be ideal if you can directly collect the data into the buffer without allocating new memory
         # but for the sake of example we will allocate and copy the data
         plaintext = bytearray(os.urandom(plaintext_len))
         mlock(plaintext)
+        aad = b"AAD"
         # cipher.copy_slice is slighlty faster than buf[:plaintext_len] = plaintext, especially for large plaintext, because it copies the data in parallel
         # cipher.copy_slice takes bytes as input, cipher.copy_slice takes bytearray
         cipher.copy_slice(plaintext, buf)
         # encrypt it, this will encrypt in-place the data in the buffer
         ciphertext_len = cipher.encrypt(buf, plaintext_len, 42, aad)
-        cipertext = buf[:ciphertext_len]
-        # you can do something with the ciphertext
 
         # decrypt it
         # if you need to copy ciphertext to buffer, we don't need to do it now as it's already in the buffer
@@ -148,24 +144,30 @@ def test_encrypt_chacha(self):
         plaintext2 = buf[:plaintext_len]
         self.assertEqual(plaintext, plaintext2)
 
+        zeroize1(plaintext)
+        zeroize1(buf)
+        munlock(plaintext)
         munlock(buf)
 
     def test_encrypt_from_aes(self):
         cipher_meta = CipherMeta.Ring(RingAlgorithm.AES256GCM)
         key_len = cipher_meta.key_len()
         key = bytearray(key_len)
+        mlock(key)
         cipher_meta.generate_key(key)
+        munlock(key)
         cipher = Cipher(cipher_meta, key)
 
         # we create a buffer based on plaintext block len of 4096
         # the actual buffer needs to be a bit larger as the ciphertext also includes the tag and nonce
         plaintext_len = 4096
         ciphertext_len = cipher.ciphertext_len(plaintext_len)
         buf = np.array([0] * ciphertext_len, dtype=np.uint8)
-
-        aad = b"AAD"
+        mlock(buf)
 
         plaintext = bytearray(os.urandom(plaintext_len))
+        mlock(plaintext)
+        aad = b"AAD"
 
         # encrypt it, after this will have the ciphertext in the buffer
         ciphertext_len = cipher.encrypt_from(plaintext, buf, 42, aad)
@@ -175,23 +177,31 @@ def test_encrypt_from_aes(self):
         plaintext_len = cipher.decrypt_from(cipertext, buf, 42, aad)
         plaintext2 = buf[:plaintext_len]
         self.assertEqual(plaintext, plaintext2)
+
+        zeroize1(plaintext)
+        zeroize1(buf)
+        munlock(plaintext)
+        munlock(buf)
 
-    def test_encrypt_from_cha(self):
+    def test_encrypt_from_chacha(self):
         cipher_meta = CipherMeta.Ring(RingAlgorithm.ChaCha20Poly1305)
         key_len = cipher_meta.key_len()
         key = bytearray(key_len)
+        mlock(key)
         cipher_meta.generate_key(key)
+        munlock(key)
         cipher = Cipher(cipher_meta, key)
 
         # we create a buffer based on plaintext block len of 4096
         # the actual buffer needs to be a bit larger as the ciphertext also includes the tag and nonce
         plaintext_len = 4096
         ciphertext_len = cipher.ciphertext_len(plaintext_len)
         buf = np.array([0] * ciphertext_len, dtype=np.uint8)
-
-        aad = b"AAD"
+        mlock(buf)
 
         plaintext = bytearray(os.urandom(plaintext_len))
+        mlock(plaintext)
+        aad = b"AAD"
 
         # encrypt it, after this will have the ciphertext in the buffer
         ciphertext_len = cipher.encrypt_from(plaintext, buf, 42, aad)
@@ -201,6 +211,11 @@ def test_encrypt_from_cha(self):
         plaintext_len = cipher.decrypt_from(cipertext, buf, 42, aad)
         plaintext2 = buf[:plaintext_len]
         self.assertEqual(plaintext, plaintext2)
+
+        zeroize1(plaintext)
+        zeroize1(buf)
+        munlock(plaintext)
+        munlock(buf)
 
 
     def test_encrypt_file_aes(self):
@@ -214,14 +229,17 @@ def test_encrypt_file_aes(self):
         cipher_meta = CipherMeta.Ring(RingAlgorithm.AES256GCM)
         key_len = cipher_meta.key_len()
         key = bytearray(key_len)
+        mlock(key)
         cipher_meta.generate_key(key)
+        munlock(key)
         cipher = Cipher(cipher_meta, key)
 
         # we create a buffer based on plaintext block len of 4096
         # the actual buffer needs to be a bit larger as the ciphertext also includes the tag and nonce
         plaintext_len = chunk_len
         ciphertext_len = cipher.ciphertext_len(plaintext_len)
         buf = np.array([0] * ciphertext_len, dtype=np.uint8)
+        mlock(buf)
 
         aad = b"AAD"
 
@@ -245,6 +263,9 @@ def test_encrypt_file_aes(self):
             file_out.flush()
 
         compare_files_by_hash(fin, tmp)
+
+        zeroize1(buf)
+        munlock(buf)
 
     def test_encrypt_file_chacha(self):
         tmp_dir = create_directory_in_home("Cipher_tmp")
@@ -257,15 +278,17 @@ def test_encrypt_file_chacha(self):
         cipher_meta = CipherMeta.Ring(RingAlgorithm.ChaCha20Poly1305)
         key_len = cipher_meta.key_len()
         key = bytearray(key_len)
+        mlock(key)
         cipher_meta.generate_key(key)
+        munlock(key)
         cipher = Cipher(cipher_meta, key)
 
         # we create a buffer based on plaintext block len of 4096
         # the actual buffer needs to be a bit larger as the ciphertext also includes the tag and nonce
         plaintext_len = chunk_len
         ciphertext_len = cipher.ciphertext_len(plaintext_len)
         buf = np.array([0] * ciphertext_len, dtype=np.uint8)
-
+        mlock(buf)
         aad = b"AAD"
 
         # encrypt
@@ -288,6 +311,9 @@ def test_encrypt_file_chacha(self):
             file_out.flush()
 
         compare_files_by_hash(fin, tmp)
+
+        zeroize1(buf)
+        munlock(buf)
 
 
 if __name__ == "__main__":