From 55e0dcd84f4883be6deb06810a2bc012da70c9f9 Mon Sep 17 00:00:00 2001
From: sinu <65924192+sinui0@users.noreply.github.com>
Date: Thu, 1 Feb 2024 12:44:02 -0800
Subject: [PATCH 1/8] refactor: chunk KOS extend message

---
 ot/mpz-ot-core/src/kos/mod.rs      |  5 ++++
 ot/mpz-ot-core/src/kos/msgs.rs     | 45 ++++++++++++++++++++++++++++--
 ot/mpz-ot-core/src/kos/receiver.rs |  2 +-
 ot/mpz-ot-core/src/kos/sender.rs   | 13 ++-------
 ot/mpz-ot/src/kos/mod.rs           |  3 ++
 ot/mpz-ot/src/kos/receiver.rs      | 14 ++++++++--
 ot/mpz-ot/src/kos/sender.rs        | 33 +++++++++++++++++++---
 7 files changed, 95 insertions(+), 20 deletions(-)

diff --git a/ot/mpz-ot-core/src/kos/mod.rs b/ot/mpz-ot-core/src/kos/mod.rs
index 438539b2..34d73c84 100644
--- a/ot/mpz-ot-core/src/kos/mod.rs
+++ b/ot/mpz-ot-core/src/kos/mod.rs
@@ -28,6 +28,11 @@ 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.
+pub fn extension_matrix_size(count: usize) -> usize {
+    count * CSP / 8
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/ot/mpz-ot-core/src/kos/msgs.rs b/ot/mpz-ot-core/src/kos/msgs.rs
index afdf2b63..febc0232 100644
--- a/ot/mpz-ot-core/src/kos/msgs.rs
+++ b/ot/mpz-ot-core/src/kos/msgs.rs
@@ -18,6 +18,7 @@ use crate::msgs::Derandomize;
 #[allow(missing_docs)]
 pub enum Message<BaseMsg> {
     BaseMsg(BaseMsg),
+    StartExtend(StartExtend),
     Extend(Extend),
     Check(Check),
     Derandomize(Derandomize),
@@ -33,15 +34,55 @@ impl<BaseMsg> From<MessageError<BaseMsg>> for std::io::Error {
     }
 }
 
-/// Extension message sent by the receiver.
+/// Extension message sent by the receiver to agree upon the number of OTs to set up.
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-pub struct Extend {
+pub struct StartExtend {
     /// The number of OTs to set up.
     pub count: usize,
+}
+
+/// Extension message sent by the receiver.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct Extend {
     /// The receiver's extension vectors.
     pub us: Vec<u8>,
 }
 
+impl Extend {
+    /// Splits the message into chunks.
+    pub fn into_chunks(self, chunk_size: usize) -> ExtendChunks {
+        ExtendChunks {
+            chunk_size,
+            us: self.us,
+        }
+    }
+}
+
+/// Iterator over the chunks of an extension message.
+pub struct ExtendChunks {
+    chunk_size: usize,
+    us: Vec<u8>,
+}
+
+impl Iterator for ExtendChunks {
+    type Item = Extend;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let remaining = self.us.len();
+        if remaining == 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 })
+    }
+}
+
 /// Values for the correlation check sent by the receiver.
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[allow(missing_docs)]
diff --git a/ot/mpz-ot-core/src/kos/receiver.rs b/ot/mpz-ot-core/src/kos/receiver.rs
index b7b8fb59..6bbb56af 100644
--- a/ot/mpz-ot-core/src/kos/receiver.rs
+++ b/ot/mpz-ot-core/src/kos/receiver.rs
@@ -196,7 +196,7 @@ impl Receiver<state::Extension> {
         );
         self.state.unchecked_choices.extend(choices);
 
-        Ok(Extend { count, us })
+        Ok(Extend { us })
     }
 
     /// Performs the correlation check for all outstanding OTS.
diff --git a/ot/mpz-ot-core/src/kos/sender.rs b/ot/mpz-ot-core/src/kos/sender.rs
index 8e692727..4bd7a65a 100644
--- a/ot/mpz-ot-core/src/kos/sender.rs
+++ b/ot/mpz-ot-core/src/kos/sender.rs
@@ -1,5 +1,6 @@
 use crate::{
     kos::{
+        extension_matrix_size,
         msgs::{Check, Ciphertexts, Extend, SenderPayload},
         Aes128Ctr, Rng, RngSeed, SenderConfig, SenderError, CSP, SSP,
     },
@@ -126,17 +127,9 @@ impl Sender<state::Extension> {
         const NROWS: usize = CSP;
         let row_width = count / 8;
 
-        let Extend {
-            us,
-            count: receiver_count,
-        } = extend;
+        let Extend { us } = extend;
 
-        // Make sure the number of OTs to extend matches the receiver's setup message.
-        if receiver_count != count {
-            return Err(SenderError::CountMismatch(receiver_count, count));
-        }
-
-        if us.len() != NROWS * row_width {
+        if us.len() != extension_matrix_size(count) {
             return Err(SenderError::InvalidExtend);
         }
 
diff --git a/ot/mpz-ot/src/kos/mod.rs b/ot/mpz-ot/src/kos/mod.rs
index 39aca16b..3e267259 100644
--- a/ot/mpz-ot/src/kos/mod.rs
+++ b/ot/mpz-ot/src/kos/mod.rs
@@ -18,6 +18,9 @@ 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;
+
 /// 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>(
     sink: &'a mut Si,
diff --git a/ot/mpz-ot/src/kos/receiver.rs b/ot/mpz-ot/src/kos/receiver.rs
index 1562e76a..22e10625 100644
--- a/ot/mpz-ot/src/kos/receiver.rs
+++ b/ot/mpz-ot/src/kos/receiver.rs
@@ -3,7 +3,8 @@ use futures::SinkExt;
 use itybity::{FromBitIterator, IntoBitIterator};
 use mpz_core::{cointoss, prg::Prg, Block, ProtocolMessage};
 use mpz_ot_core::kos::{
-    msgs::Message, receiver_state as state, Receiver as ReceiverCore, ReceiverConfig, CSP, SSP,
+    msgs::{Message, StartExtend},
+    receiver_state as state, Receiver as ReceiverCore, ReceiverConfig, CSP, SSP,
 };
 
 use enum_try_as_inner::EnumTryAsInner;
@@ -15,7 +16,9 @@ use utils_aio::{
     stream::{ExpectStreamExt, IoStream},
 };
 
-use super::{into_base_sink, into_base_stream, ReceiverError, ReceiverVerifyError};
+use super::{
+    into_base_sink, into_base_stream, ReceiverError, ReceiverVerifyError, EXTEND_CHUNK_SIZE,
+};
 use crate::{
     OTError, OTReceiver, OTSender, OTSetup, RandomOTReceiver, VerifiableOTReceiver,
     VerifiableOTSender,
@@ -105,7 +108,12 @@ where
         let (cointoss_sender, cointoss_commitment) = cointoss::Sender::new(vec![seed]).send();
 
         // Send the extend message and cointoss commitment
-        sink.feed(Message::Extend(extend)).await?;
+        sink.feed(Message::StartExtend(StartExtend { count }))
+            .await?;
+        let mut chunks = extend.into_chunks(EXTEND_CHUNK_SIZE);
+        while let Some(extend) = chunks.next() {
+            sink.feed(Message::Extend(extend)).await?;
+        }
         sink.feed(Message::CointossCommit(cointoss_commitment))
             .await?;
         sink.flush().await?;
diff --git a/ot/mpz-ot/src/kos/sender.rs b/ot/mpz-ot/src/kos/sender.rs
index 407ddadd..be7b759c 100644
--- a/ot/mpz-ot/src/kos/sender.rs
+++ b/ot/mpz-ot/src/kos/sender.rs
@@ -4,7 +4,9 @@ use futures_util::SinkExt;
 use itybity::IntoBits;
 use mpz_core::{cointoss, prg::Prg, Block, ProtocolMessage};
 use mpz_ot_core::kos::{
-    msgs::Message, sender_state as state, Sender as SenderCore, SenderConfig, CSP, SSP,
+    extension_matrix_size,
+    msgs::{Extend, Message, StartExtend},
+    sender_state as state, Sender as SenderCore, SenderConfig, CSP, SSP,
 };
 use rand::{thread_rng, Rng};
 use rand_core::{RngCore, SeedableRng};
@@ -138,13 +140,36 @@ where
         let mut ext_sender =
             std::mem::replace(&mut self.state, State::Error).try_into_extension()?;
 
-        // Receive extend message from the receiver.
-        let extend = stream
+        let StartExtend {
+            count: receiver_count,
+        } = stream
             .expect_next()
             .await?
-            .try_into_extend()
+            .try_into_start_extend()
             .map_err(SenderError::from)?;
 
+        if count != receiver_count {
+            return Err(SenderError::ConfigError(
+                "sender and receiver count mismatch".to_string(),
+            ));
+        }
+
+        let expected_us = extension_matrix_size(count);
+        let mut extend = Extend {
+            us: Vec::with_capacity(expected_us),
+        };
+
+        // Receive extension matrix from the receiver.
+        while extend.us.len() < expected_us {
+            let Extend { us: chunk } = stream
+                .expect_next()
+                .await?
+                .try_into_extend()
+                .map_err(SenderError::from)?;
+
+            extend.us.extend(chunk);
+        }
+
         // Receive coin toss commitments from the receiver.
         let commitment = stream.expect_next().await?.try_into_cointoss_commit()?;
 

From 186e54e0476acee70e1a02fef91dbad9f5ab0931 Mon Sep 17 00:00:00 2001
From: sinu <65924192+sinui0@users.noreply.github.com>
Date: Thu, 1 Feb 2024 12:48:57 -0800
Subject: [PATCH 2/8] clippy

---
 ot/mpz-ot/src/kos/receiver.rs | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/ot/mpz-ot/src/kos/receiver.rs b/ot/mpz-ot/src/kos/receiver.rs
index 22e10625..59c02760 100644
--- a/ot/mpz-ot/src/kos/receiver.rs
+++ b/ot/mpz-ot/src/kos/receiver.rs
@@ -110,8 +110,7 @@ where
         // Send the extend message and cointoss commitment
         sink.feed(Message::StartExtend(StartExtend { count }))
             .await?;
-        let mut chunks = extend.into_chunks(EXTEND_CHUNK_SIZE);
-        while let Some(extend) = chunks.next() {
+        for extend in extend.into_chunks(EXTEND_CHUNK_SIZE) {
             sink.feed(Message::Extend(extend)).await?;
         }
         sink.feed(Message::CointossCommit(cointoss_commitment))

From f19eb5227c4510d7d0f7d90d7c6bb57abbd473a7 Mon Sep 17 00:00:00 2001
From: sinu <65924192+sinui0@users.noreply.github.com>
Date: Fri, 9 Feb 2024 15:47:42 -0800
Subject: [PATCH 3/8] fix chunking

---
 ot/mpz-ot-core/src/kos/mod.rs  | 11 ++++++++++-
 ot/mpz-ot-core/src/kos/msgs.rs | 19 +++++++------------
 ot/mpz-ot/Cargo.toml           |  1 +
 ot/mpz-ot/src/kos/mod.rs       | 11 +++++++++--
 ot/mpz-ot/src/kos/receiver.rs  |  6 ++++--
 ot/mpz-ot/src/kos/sender.rs    | 12 +++++-------
 6 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/ot/mpz-ot-core/src/kos/mod.rs b/ot/mpz-ot-core/src/kos/mod.rs
index 34d73c84..41f9f0a8 100644
--- a/ot/mpz-ot-core/src/kos/mod.rs
+++ b/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
 }
diff --git a/ot/mpz-ot-core/src/kos/msgs.rs b/ot/mpz-ot-core/src/kos/msgs.rs
index febc0232..372723ab 100644
--- a/ot/mpz-ot-core/src/kos/msgs.rs
+++ b/ot/mpz-ot-core/src/kos/msgs.rs
@@ -53,7 +53,7 @@ impl Extend {
     pub fn into_chunks(self, chunk_size: usize) -> ExtendChunks {
         ExtendChunks {
             chunk_size,
-            us: self.us,
+            us: self.us.into_iter(),
         }
     }
 }
@@ -61,25 +61,20 @@ impl Extend {
 /// 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<_>>(),
+            })
+        }
     }
 }
 
diff --git a/ot/mpz-ot/Cargo.toml b/ot/mpz-ot/Cargo.toml
index ec171c99..96d15c8e 100644
--- a/ot/mpz-ot/Cargo.toml
+++ b/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 }
diff --git a/ot/mpz-ot/src/kos/mod.rs b/ot/mpz-ot/src/kos/mod.rs
index 3e267259..d43542fb 100644
--- a/ot/mpz-ot/src/kos/mod.rs
+++ b/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>(
diff --git a/ot/mpz-ot/src/kos/receiver.rs b/ot/mpz-ot/src/kos/receiver.rs
index 59c02760..34aa7c6e 100644
--- a/ot/mpz-ot/src/kos/receiver.rs
+++ b/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)
         })
diff --git a/ot/mpz-ot/src/kos/sender.rs b/ot/mpz-ot/src/kos/sender.rs
index be7b759c..e84b0f36 100644
--- a/ot/mpz-ot/src/kos/sender.rs
+++ b/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();

From 0fae323d5412ff7b65edfda105300b0f3d59eb5b Mon Sep 17 00:00:00 2001
From: "sinu.eth" <65924192+sinui0@users.noreply.github.com>
Date: Fri, 9 Feb 2024 15:48:44 -0800
Subject: [PATCH 4/8] Update ot/mpz-ot-core/src/kos/msgs.rs

Co-authored-by: dan <themighty1@users.noreply.github.com>
---
 ot/mpz-ot-core/src/kos/msgs.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ot/mpz-ot-core/src/kos/msgs.rs b/ot/mpz-ot-core/src/kos/msgs.rs
index 372723ab..5fd96223 100644
--- a/ot/mpz-ot-core/src/kos/msgs.rs
+++ b/ot/mpz-ot-core/src/kos/msgs.rs
@@ -49,7 +49,7 @@ pub struct Extend {
 }
 
 impl Extend {
-    /// Splits the message into chunks.
+    /// Returns an iterator over the chunks of the message.
     pub fn into_chunks(self, chunk_size: usize) -> ExtendChunks {
         ExtendChunks {
             chunk_size,

From f6d9e3b7ae2bab453d34ffaa1c191b3791644fb9 Mon Sep 17 00:00:00 2001
From: sinu <65924192+sinui0@users.noreply.github.com>
Date: Fri, 9 Feb 2024 15:49:31 -0800
Subject: [PATCH 5/8] clippy

---
 ot/mpz-ot-core/src/kos/msgs.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ot/mpz-ot-core/src/kos/msgs.rs b/ot/mpz-ot-core/src/kos/msgs.rs
index 5fd96223..7ee1bc13 100644
--- a/ot/mpz-ot-core/src/kos/msgs.rs
+++ b/ot/mpz-ot-core/src/kos/msgs.rs
@@ -69,7 +69,7 @@ impl Iterator for ExtendChunks {
 
     fn next(&mut self) -> Option<Self::Item> {
         if self.us.len() == 0 {
-            return None;
+            None
         } else {
             Some(Extend {
                 us: self.us.by_ref().take(self.chunk_size).collect::<Vec<_>>(),

From efb5f305acf920939f948a35aebfbf8592594442 Mon Sep 17 00:00:00 2001
From: sinu <65924192+sinui0@users.noreply.github.com>
Date: Fri, 9 Feb 2024 15:55:27 -0800
Subject: [PATCH 6/8] pad before rounding

---
 ot/mpz-ot-core/src/kos/mod.rs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/ot/mpz-ot-core/src/kos/mod.rs b/ot/mpz-ot-core/src/kos/mod.rs
index 41f9f0a8..a2c268e1 100644
--- a/ot/mpz-ot-core/src/kos/mod.rs
+++ b/ot/mpz-ot-core/src/kos/mod.rs
@@ -30,11 +30,11 @@ pub(crate) type Aes128Ctr = ctr::Ctr64LE<aes::Aes128>;
 
 /// 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;
+pub fn pad_ot_count(mut count: usize) -> usize {
     // Add OTs for the KOS extension check.
-    count + CSP + SSP
+    count += CSP + SSP;
+    // Round up the OTs to extend to the nearest multiple of 64 (matrix transpose optimization).
+    (count + 63) & !63
 }
 
 /// Returns the size in bytes of the extension matrix for a given number of OTs.

From 4ea2509b653cc647d4e1c998aeaca62d2fb997f4 Mon Sep 17 00:00:00 2001
From: sinu <65924192+sinui0@users.noreply.github.com>
Date: Mon, 12 Feb 2024 11:37:39 -0800
Subject: [PATCH 7/8] update comment

---
 ot/mpz-ot/src/kos/receiver.rs | 2 +-
 ot/mpz-ot/src/kos/sender.rs   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/ot/mpz-ot/src/kos/receiver.rs b/ot/mpz-ot/src/kos/receiver.rs
index 34aa7c6e..1d48a349 100644
--- a/ot/mpz-ot/src/kos/receiver.rs
+++ b/ot/mpz-ot/src/kos/receiver.rs
@@ -95,7 +95,7 @@ where
 
         let count = pad_ot_count(count);
 
-        // Extend the OTs, adding padding for the consistency check.
+        // Extend the OTs.
         let (mut ext_receiver, extend) = Backend::spawn(move || {
             let extend = ext_receiver.extend(count);
 
diff --git a/ot/mpz-ot/src/kos/sender.rs b/ot/mpz-ot/src/kos/sender.rs
index e84b0f36..91358b90 100644
--- a/ot/mpz-ot/src/kos/sender.rs
+++ b/ot/mpz-ot/src/kos/sender.rs
@@ -175,7 +175,7 @@ where
         // Receive coin toss commitments from the receiver.
         let commitment = stream.expect_next().await?.try_into_cointoss_commit()?;
 
-        // Extend the OTs, adding padding for the consistency check.
+        // Extend the OTs.
         let mut ext_sender =
             Backend::spawn(move || ext_sender.extend(count, extend).map(|_| ext_sender)).await?;
 

From 9f021ac9cc69640d929acee304889e267737a0a0 Mon Sep 17 00:00:00 2001
From: sinu <65924192+sinui0@users.noreply.github.com>
Date: Mon, 12 Feb 2024 11:43:17 -0800
Subject: [PATCH 8/8] require extend count to be a multiple of 64

---
 ot/mpz-ot-core/src/kos/error.rs    |  4 ++++
 ot/mpz-ot-core/src/kos/receiver.rs |  9 ++++++---
 ot/mpz-ot-core/src/kos/sender.rs   | 11 +++++++----
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/ot/mpz-ot-core/src/kos/error.rs b/ot/mpz-ot-core/src/kos/error.rs
index b91e4eb7..3e230367 100644
--- a/ot/mpz-ot-core/src/kos/error.rs
+++ b/ot/mpz-ot-core/src/kos/error.rs
@@ -4,6 +4,8 @@
 pub enum SenderError {
     #[error("invalid state: expected {0}")]
     InvalidState(String),
+    #[error("invalid count, must be a multiple of 64: {0}")]
+    InvalidCount(usize),
     #[error("count mismatch: expected {0}, got {1}")]
     CountMismatch(usize, usize),
     #[error("id mismatch: expected {0}, got {1}")]
@@ -22,6 +24,8 @@ pub enum SenderError {
 pub enum ReceiverError {
     #[error("invalid state: expected {0}")]
     InvalidState(String),
+    #[error("invalid count, must be a multiple of 64: {0}")]
+    InvalidCount(usize),
     #[error("count mismatch: expected {0}, got {1}")]
     CountMismatch(usize, usize),
     #[error("id mismatch: expected {0}, got {1}")]
diff --git a/ot/mpz-ot-core/src/kos/receiver.rs b/ot/mpz-ot-core/src/kos/receiver.rs
index 6bbb56af..89c28c9d 100644
--- a/ot/mpz-ot-core/src/kos/receiver.rs
+++ b/ot/mpz-ot-core/src/kos/receiver.rs
@@ -117,6 +117,8 @@ impl Receiver<state::Extension> {
 
     /// Perform the IKNP OT extension.
     ///
+    /// The provided count _must_ be a multiple of 64, otherwise an error will be returned.
+    ///
     /// # Sacrificial OTs
     ///
     /// Performing the consistency check sacrifices 256 OTs, so be sure to
@@ -132,7 +134,7 @@ impl Receiver<state::Extension> {
     ///
     /// # Arguments
     ///
-    /// * `count` - The number of OTs to extend.
+    /// * `count` - The number of OTs to extend (must be a multiple of 64).
     pub fn extend(&mut self, count: usize) -> Result<Extend, ReceiverError> {
         if self.state.extended {
             return Err(ReceiverError::InvalidState(
@@ -140,8 +142,9 @@ impl Receiver<state::Extension> {
             ));
         }
 
-        // Round up the OTs to extend to the nearest multiple of 64 (matrix transpose optimization).
-        let count = (count + 63) & !63;
+        if count % 64 != 0 {
+            return Err(ReceiverError::InvalidCount(count));
+        }
 
         const NROWS: usize = CSP;
         let row_width = count / 8;
diff --git a/ot/mpz-ot-core/src/kos/sender.rs b/ot/mpz-ot-core/src/kos/sender.rs
index 4bd7a65a..1e0ceee1 100644
--- a/ot/mpz-ot-core/src/kos/sender.rs
+++ b/ot/mpz-ot-core/src/kos/sender.rs
@@ -97,6 +97,8 @@ impl Sender<state::Extension> {
 
     /// Perform the IKNP OT extension.
     ///
+    /// The provided count _must_ be a multiple of 64, otherwise an error will be returned.
+    ///
     /// # Sacrificial OTs
     ///
     /// Performing the consistency check sacrifices 256 OTs, so be sure to extend enough to
@@ -112,8 +114,8 @@ impl Sender<state::Extension> {
     ///
     /// # Arguments
     ///
-    /// * `count` - The number of additional OTs to extend
-    /// * `extend` - The receiver's setup message
+    /// * `count` - The number of additional OTs to extend (must be a multiple of 64).
+    /// * `extend` - The receiver's setup message.
     pub fn extend(&mut self, count: usize, extend: Extend) -> Result<(), SenderError> {
         if self.state.extended {
             return Err(SenderError::InvalidState(
@@ -121,8 +123,9 @@ impl Sender<state::Extension> {
             ));
         }
 
-        // Round up the OTs to extend to the nearest multiple of 64 (matrix transpose optimization).
-        let count = (count + 63) & !63;
+        if count % 64 != 0 {
+            return Err(SenderError::InvalidCount(count));
+        }
 
         const NROWS: usize = CSP;
         let row_width = count / 8;