require extend count to be a multiple of 64

privacy-scaling-explorations · Feb 12, 2024 · 9f021ac · 9f021ac
1 parent 4ea2509
commit 9f021ac
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 7 deletions.
diff --git 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
@@ -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,16 +134,17 @@ 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(
                 "extending more than once is currently disabled".to_string(),
             ));
         }
 
-        // 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
@@ -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,17 +114,18 @@ 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(
                 "extending more than once is currently disabled".to_string(),
             ));
         }
 
-        // 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;