fix chunking

ot/mpz-ot-core/src/kos/mod.rs

@@ -28,7 +28,16 @@ pub(crate) type RngSeed = <Rng as SeedableRng>::Seed;
 /// AES-128 CTR used for encryption.
 pub(crate) type Aes128Ctr = ctr::Ctr64LE<aes::Aes128>;
 
-/// Returns the size in bytes of the extension message for a given number of OTs.
+/// Pads the number of OTs to accomodate for the KOS extension check and
+/// the extension matrix transpose optimization.
+pub fn pad_ot_count(count: usize) -> usize {
+    // Round up the OTs to extend to the nearest multiple of 64 (matrix transpose optimization).
+    let count = (count + 63) & !63;
+    // Add OTs for the KOS extension check.
+    count + CSP + SSP
+}
+
+/// Returns the size in bytes of the extension matrix for a given number of OTs.
 pub fn extension_matrix_size(count: usize) -> usize {
     count * CSP / 8
 }

ot/mpz-ot-core/src/kos/msgs.rs

@@ -53,33 +53,28 @@ impl Extend {
     pub fn into_chunks(self, chunk_size: usize) -> ExtendChunks {
         ExtendChunks {
             chunk_size,
-            us: self.us,
+            us: self.us.into_iter(),
         }
     }
 }
 
 /// Iterator over the chunks of an extension message.
 pub struct ExtendChunks {
     chunk_size: usize,
-    us: Vec<u8>,
+    us: <Vec<u8> as IntoIterator>::IntoIter,
 }
 
 impl Iterator for ExtendChunks {
     type Item = Extend;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let remaining = self.us.len();
-        if remaining == 0 {
+        if self.us.len() == 0 {
             return None;
-        }
-
-        let us = if remaining <= self.chunk_size {
-            std::mem::take(&mut self.us)
         } else {
-            self.us.split_off(remaining - self.chunk_size)
-        };
-
-        Some(Extend { us })
+            Some(Extend {
+                us: self.us.by_ref().take(self.chunk_size).collect::<Vec<_>>(),
+            })
+        }
     }
 }
 

ot/mpz-ot/Cargo.toml

@@ -32,6 +32,7 @@ itybity.workspace = true
 enum-try-as-inner.workspace = true
 opaque-debug.workspace = true
 serde = { workspace = true, optional = true }
+cfg-if.workspace = true
 
 [dev-dependencies]
 rstest = { workspace = true }

ot/mpz-ot/src/kos/mod.rs

@@ -18,8 +18,15 @@ pub use mpz_ot_core::kos::{
 };
 use utils_aio::{sink::IoSink, stream::IoStream};
 
-/// The size of the chunks used to send the extension matrix, 1MB.
-const EXTEND_CHUNK_SIZE: usize = 1024 * 1024;
+// If we're testing we use a smaller chunk size to make sure the chunking code paths are tested.
+cfg_if::cfg_if! {
+    if #[cfg(test)] {
+        pub(crate) const EXTEND_CHUNK_SIZE: usize = 1024;
+    } else {
+        /// The size of the chunks used to send the extension matrix, 4MB.
+        pub(crate) const EXTEND_CHUNK_SIZE: usize = 4 * 1024 * 1024;
+    }
+}
 
 /// Converts a sink of KOS messages into a sink of base OT messages.
 pub(crate) fn into_base_sink<'a, Si: IoSink<msgs::Message<T>> + Send + Unpin, T: Send + 'a>(

ot/mpz-ot/src/kos/receiver.rs

@@ -4,7 +4,7 @@ use itybity::{FromBitIterator, IntoBitIterator};
 use mpz_core::{cointoss, prg::Prg, Block, ProtocolMessage};
 use mpz_ot_core::kos::{
     msgs::{Message, StartExtend},
-    receiver_state as state, Receiver as ReceiverCore, ReceiverConfig, CSP, SSP,
+    pad_ot_count, receiver_state as state, Receiver as ReceiverCore, ReceiverConfig, CSP,
 };
 
 use enum_try_as_inner::EnumTryAsInner;
@@ -93,9 +93,11 @@ where
         let mut ext_receiver =
             std::mem::replace(&mut self.state, State::Error).try_into_extension()?;
 
+        let count = pad_ot_count(count);
+
         // Extend the OTs, adding padding for the consistency check.
         let (mut ext_receiver, extend) = Backend::spawn(move || {
-            let extend = ext_receiver.extend(count + CSP + SSP);
+            let extend = ext_receiver.extend(count);
 
             (ext_receiver, extend)
         })

ot/mpz-ot/src/kos/sender.rs

@@ -6,7 +6,7 @@ use mpz_core::{cointoss, prg::Prg, Block, ProtocolMessage};
 use mpz_ot_core::kos::{
     extension_matrix_size,
     msgs::{Extend, Message, StartExtend},
-    sender_state as state, Sender as SenderCore, SenderConfig, CSP, SSP,
+    pad_ot_count, sender_state as state, Sender as SenderCore, SenderConfig, CSP,
 };
 use rand::{thread_rng, Rng};
 use rand_core::{RngCore, SeedableRng};
@@ -140,6 +140,8 @@ where
         let mut ext_sender =
             std::mem::replace(&mut self.state, State::Error).try_into_extension()?;
 
+        let count = pad_ot_count(count);
+
         let StartExtend {
             count: receiver_count,
         } = stream
@@ -174,12 +176,8 @@ where
         let commitment = stream.expect_next().await?.try_into_cointoss_commit()?;
 
         // Extend the OTs, adding padding for the consistency check.
-        let mut ext_sender = Backend::spawn(move || {
-            ext_sender
-                .extend(count + CSP + SSP, extend)
-                .map(|_| ext_sender)
-        })
-        .await?;
+        let mut ext_sender =
+            Backend::spawn(move || ext_sender.extend(count, extend).map(|_| ext_sender)).await?;
 
         // Execute coin toss protocol for consistency check.
         let seed: Block = thread_rng().gen();