From 6e2e6bdfd8258de84659f652d5bd0631f6808c1b Mon Sep 17 00:00:00 2001 From: Radu Marias Date: Wed, 5 Jun 2024 04:02:11 +0300 Subject: [PATCH] - add more mlock and zeroize - remove mlock and zeroize from bench --- benches/bench.py | 44 --------------------------------- tests/test_rencrypt.py | 56 +++++++++++++++++++++++++++++++----------- 2 files changed, 41 insertions(+), 59 deletions(-) diff --git a/benches/bench.py b/benches/bench.py index b96fb8f..0d2584f 100644 --- a/benches/bench.py +++ b/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]) diff --git a/tests/test_rencrypt.py b/tests/test_rencrypt.py index 5cf65e2..4288b53 100644 --- a/tests/test_rencrypt.py +++ b/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,13 +144,18 @@ 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 @@ -162,10 +163,11 @@ def test_encrypt_from_aes(self): 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,12 +177,19 @@ 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 @@ -188,10 +197,11 @@ def test_encrypt_from_cha(self): 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,7 +229,9 @@ 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 @@ -222,6 +239,7 @@ def test_encrypt_file_aes(self): 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,7 +278,9 @@ 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 @@ -265,7 +288,7 @@ def test_encrypt_file_chacha(self): 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__":