From cd90fbea50c1ab127ca0454d21670adc83db8cac Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Wed, 29 Nov 2023 13:19:17 -0700
Subject: [PATCH] BoxedUint: add `cond_map` and `cond_and_then`

Unfortunately `BoxedUint` can't impl `subtle::ConditionallySelectable`
due to a supertrait bound on `Copy`. See dalek-cryptography/subtle#94

This bound is required by `CtOption::map` and `CtOption::and_then` which
are important for writing constant-time code.

As a workaround which makes it still possible to leverate `CtOption`,
this adds special `BoxedUint`-specialized combinators that are able to
work around this issue by generating a placeholder (zero) value to pass
to the provided callbacks in the event `CtOption` is none.

This requires branching on the output of `CtOption` (which is
unavoidable without an upstream fix in `subtle` itself), but still
ensures that the provided callback function is called with a `BoxedUint`
of a matching number of limbs regardless of whether the `CtOption` is
some or none, which is the best we can do for now (and really quite
close to what `subtle` is doing under the hood anyway).
---
 src/uint/boxed.rs    |   1 +
 src/uint/boxed/ct.rs | 120 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 121 insertions(+)
 create mode 100644 src/uint/boxed/ct.rs

diff --git a/src/uint/boxed.rs b/src/uint/boxed.rs
index 23b268b1..4b73779b 100644
--- a/src/uint/boxed.rs
+++ b/src/uint/boxed.rs
@@ -6,6 +6,7 @@ mod bit_and;
 mod bit_or;
 mod bits;
 mod cmp;
+mod ct;
 mod div;
 pub(crate) mod encoding;
 mod inv_mod;
diff --git a/src/uint/boxed/ct.rs b/src/uint/boxed/ct.rs
new file mode 100644
index 00000000..7e2baea5
--- /dev/null
+++ b/src/uint/boxed/ct.rs
@@ -0,0 +1,120 @@
+//! Constant-time helper functions.
+//!
+//! These largely exist as a workaround for the `Copy` bound on [`subtle::ConditionallySelectable`].
+
+use super::BoxedUint;
+use subtle::CtOption;
+
+impl BoxedUint {
+    /// Conditional `map`: workaround which provides a [`CtOption::map`]-like API.
+    ///
+    /// Ensures both functions are called regardless of whether the first returns some/none with an
+    /// argument whose precision matches `self`.
+    ///
+    /// Workaround due to `Copy` in [`subtle::ConditionallySelectable`] supertrait bounds.
+    pub fn cond_map<C, F, T>(&self, condition: C, f: F) -> CtOption<T>
+    where
+        C: Fn(&Self) -> CtOption<Self>,
+        F: Fn(Self) -> T,
+    {
+        let placeholder = Self::zero_with_precision(self.bits_precision());
+        let cond_val = condition(self);
+        let is_some = cond_val.is_some();
+
+        let value = Option::<Self>::from(cond_val).unwrap_or(placeholder);
+        debug_assert_eq!(self.bits_precision(), value.bits_precision());
+        CtOption::new(f(value), is_some)
+    }
+
+    /// Conditional `and_then`: workaround which provides a [`CtOption::and_then`]-like API.
+    ///
+    /// Ensures both functions are called regardless of whether they return some/none with an
+    /// argument whose precision matches `self`.
+    ///
+    /// Workaround due to `Copy` in [`subtle::ConditionallySelectable`] supertrait bounds.
+    pub fn cond_and_then<C, F>(&self, condition: C, f: F) -> CtOption<Self>
+    where
+        C: Fn(&Self) -> CtOption<Self>,
+        F: Fn(Self) -> CtOption<Self>,
+    {
+        let cond_val = condition(self);
+        let mut is_some = cond_val.is_some();
+
+        let placeholder = Self::zero_with_precision(self.bits_precision());
+        let value = Option::<Self>::from(cond_val).unwrap_or(placeholder);
+        debug_assert_eq!(self.bits_precision(), value.bits_precision());
+
+        let cond_val = f(value);
+        is_some &= cond_val.is_some();
+
+        let placeholder = Self::zero_with_precision(self.bits_precision());
+        let value = Option::from(cond_val).unwrap_or(placeholder);
+        CtOption::new(value, is_some)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::BoxedUint;
+    use subtle::CtOption;
+
+    #[test]
+    fn cond_map_some() {
+        let n = BoxedUint::one();
+
+        let ret = n
+            .cond_map(
+                |n| CtOption::new(n.clone(), 1u8.into()),
+                |n| n.wrapping_add(&BoxedUint::one()),
+            )
+            .unwrap();
+
+        assert_eq!(ret, BoxedUint::from(2u8));
+    }
+
+    #[test]
+    fn cond_map_none() {
+        let n = BoxedUint::one();
+
+        let ret = n.cond_map(
+            |n| CtOption::new(n.clone(), 0u8.into()),
+            |n| n.wrapping_add(&BoxedUint::one()),
+        );
+
+        assert!(bool::from(ret.is_none()));
+    }
+
+    #[test]
+    fn cond_and_then_all_some() {
+        let n = BoxedUint::one();
+
+        let ret = n
+            .cond_and_then(
+                |n| CtOption::new(n.clone(), 1u8.into()),
+                |n| CtOption::new(n.wrapping_add(&BoxedUint::one()), 1u8.into()),
+            )
+            .unwrap();
+
+        assert_eq!(ret, BoxedUint::from(2u8));
+    }
+
+    macro_rules! cond_and_then_none_test {
+        ($name:ident, $a:expr, $b:expr) => {
+            #[test]
+            fn $name() {
+                let n = BoxedUint::one();
+
+                let ret = n.cond_and_then(
+                    |n| CtOption::new(n.clone(), $a.into()),
+                    |n| CtOption::new(n.wrapping_add(&BoxedUint::one()), $b.into()),
+                );
+
+                assert!(bool::from(ret.is_none()));
+            }
+        };
+    }
+
+    cond_and_then_none_test!(cond_and_then_none_some, 0, 1);
+    cond_and_then_none_test!(cond_and_then_some_none, 1, 0);
+    cond_and_then_none_test!(cond_and_then_none_none, 0, 0);
+}