From 15cb4a17f220d517b271fa5804fb7fd7bb8dc718 Mon Sep 17 00:00:00 2001
From: Wes Brown <wesbrown18@gmail.com>
Date: Fri, 3 Nov 2023 16:36:01 -0400
Subject: [PATCH] fix(crypt): Offset and size corrections.

---
 tensorizer/serialization.py | 20 ++++++++------------
 tensorizer/stream_io.py     |  9 +++------
 2 files changed, 11 insertions(+), 18 deletions(-)

diff --git a/tensorizer/serialization.py b/tensorizer/serialization.py
index c80f43a..237cd6c 100644
--- a/tensorizer/serialization.py
+++ b/tensorizer/serialization.py
@@ -2570,13 +2570,15 @@ def encrypt_tensor() -> Tuple[Optional[bytes], Optional[bytes]]:
             tensor_bytes = memoryview(tensor_memory.tobytes())
             setup_end = time.monotonic()
 
+            chunk_repr = ""
+
             for i in range(num_chunks):
                 # We XOR the nonce with the chunk index to avoid nonce reuse.
                 step_nonce = nonce_int ^ i
                 step_nonce_bytes = step_nonce.to_bytes(
                     nacl.secret.SecretBox.NONCE_SIZE, "big", signed=False
                 )
-                plaintext_begin = i * self._crypt_chunk_size
+                plaintext_begin = i * self._cleartext_chunk_size
                 plaintext_end = plaintext_begin + self._cleartext_chunk_size
                 if plaintext_end > tensor_memory.nbytes:
                     plaintext_end = tensor_memory.nbytes
@@ -2584,23 +2586,19 @@ def encrypt_tensor() -> Tuple[Optional[bytes], Optional[bytes]]:
                 chunk = self._lockbox.encrypt(
                     to_encrypt,
                     step_nonce_bytes,
-                )
+                ).ciphertext
                 cryptotext_begin = i * self._crypt_chunk_size
                 cryptotext_end = cryptotext_begin + len(chunk)
                 cryptotext[cryptotext_begin:cryptotext_end] = chunk
                 encryption_header_size = len(chunk) - len(to_encrypt)
-                if i == 0:
-                    print(chunk)
             end = time.monotonic()
             duration_setup_ms = (setup_end - start) * 1000
             duration_ms = (end - start) * 1000
             print(
-                f"Salt: {self._salt} - Key: {self._crypto_key} - Nonce: {nonce}"
-            )
-            print(
-                f"Encryption time: {duration_ms:.2f}ms, "
-                f"setup: {duration_setup_ms:.2f}ms, "
-                f"{cryptotext_end} bytes, {encryption_header_size} header size"
+                f"Pos: {tensor_pos} - Size: {tensor_size} - Encryption time:"
+                f" {duration_ms:.2f}ms, setup: {duration_setup_ms:.2f}ms,"
+                f" {cryptotext_end} bytes, {encryption_header_size} header"
+                f" size, {chunk_repr}"
             )
 
             return nonce, cryptotext[:cryptotext_end]
@@ -2633,8 +2631,6 @@ def commit_header(
             (encrypt_task, crc32_task, sha256_task, commit_header_task)
         )
 
-        encrypt_task = self._computation_pool.submit(encrypt_tensor)
-
         # This task is I/O-bound and has no prerequisites,
         # so it goes into the regular writer pool.
         def write_tensor_data():
diff --git a/tensorizer/stream_io.py b/tensorizer/stream_io.py
index aa8f0b8..80240df 100644
--- a/tensorizer/stream_io.py
+++ b/tensorizer/stream_io.py
@@ -205,9 +205,7 @@ def __init__(
         self._lockbox = nacl.secret.SecretBox(key)
         self._nonce_int = int.from_bytes(nonce, "big", signed=False)
         self._chunk_size = chunk_size
-        self._ciphertext_chunk_sz = (
-            chunk_size + self._lockbox.MACBYTES + self._lockbox.NONCE_SIZE
-        )
+        self._ciphertext_chunk_sz = chunk_size + self._lockbox.MACBYTES
         self._ciphertext_buffer = bytearray(self._ciphertext_chunk_sz)
 
     def __enter__(self):
@@ -233,9 +231,7 @@ def readinto(self, ba: bytearray) -> int:
         num_chunks = goal // self._chunk_size
         if goal % self._chunk_size:
             num_chunks += 1
-        ciphertext_goal = goal + (
-            num_chunks * (self._lockbox.MACBYTES + self._lockbox.NONCE_SIZE)
-        )
+        ciphertext_goal = goal + (num_chunks * self._lockbox.MACBYTES)
 
         step = 0
         while ciphertext_offset < ciphertext_goal:
@@ -260,6 +256,7 @@ def readinto(self, ba: bytearray) -> int:
             ba[plaintext_offset : plaintext_offset + plaintext_sz] = (
                 self._lockbox.decrypt(ciphertext, step_nonce_bytes)
             )
+            step += 1
             ciphertext_offset += ciphertext_read_sz
             plaintext_offset += plaintext_sz