From 0d475b1d00a7d67b68e6a1502fdb3363f5932b95 Mon Sep 17 00:00:00 2001
From: Peter Nose <peter.nose@gmail.com>
Date: Mon, 4 Nov 2024 09:38:34 +0100
Subject: [PATCH] secret-sharing/src/poly: Restrict add/sub/mul assign std ops

The AddAssign, SubAssign, and MulAssign functions can now be used
only when the prime field supports zeroization. This ensures that
any leftover data from heap reallocation, when the right-hand-side
polynomial has more coefficients than the left-hand-side, is zeroized.

An alternative solution is to remove these functions, but this
could lead to performance drawbacks.
---
 .changelog/5928.trivial.md            |  0
 secret-sharing/src/poly/univariate.rs | 26 ++++++++++++++++++++------
 2 files changed, 20 insertions(+), 6 deletions(-)
 create mode 100644 .changelog/5928.trivial.md

diff --git a/.changelog/5928.trivial.md b/.changelog/5928.trivial.md
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/secret-sharing/src/poly/univariate.rs b/secret-sharing/src/poly/univariate.rs
index 307caa7499e..7bd5ab3bab1 100644
--- a/secret-sharing/src/poly/univariate.rs
+++ b/secret-sharing/src/poly/univariate.rs
@@ -238,7 +238,7 @@ where
 
 impl<F> AddAssign for Polynomial<F>
 where
-    F: PrimeField,
+    F: PrimeField + Zeroize,
 {
     #[inline]
     fn add_assign(&mut self, rhs: Polynomial<F>) {
@@ -248,9 +248,16 @@ where
 
 impl<F> AddAssign<&Polynomial<F>> for Polynomial<F>
 where
-    F: PrimeField,
+    F: PrimeField + Zeroize,
 {
     fn add_assign(&mut self, rhs: &Polynomial<F>) {
+        if self.a.capacity() < rhs.a.len() {
+            let mut a = Vec::with_capacity(rhs.a.len());
+            a.extend_from_slice(&self.a);
+            self.a.zeroize();
+            self.a = a;
+        }
+
         let min_len = min(self.a.len(), rhs.a.len());
 
         for i in 0..min_len {
@@ -321,7 +328,7 @@ where
 
 impl<F> SubAssign for Polynomial<F>
 where
-    F: PrimeField,
+    F: PrimeField + Zeroize,
 {
     #[inline]
     fn sub_assign(&mut self, rhs: Polynomial<F>) {
@@ -331,9 +338,16 @@ where
 
 impl<F> SubAssign<&Polynomial<F>> for Polynomial<F>
 where
-    F: PrimeField,
+    F: PrimeField + Zeroize,
 {
     fn sub_assign(&mut self, rhs: &Polynomial<F>) {
+        if self.a.capacity() < rhs.a.len() {
+            let mut a = Vec::with_capacity(rhs.a.len());
+            a.extend_from_slice(&self.a);
+            self.a.zeroize();
+            self.a = a;
+        }
+
         let min_len = min(self.a.len(), rhs.a.len());
 
         for i in 0..min_len {
@@ -510,7 +524,7 @@ where
 
 impl<F> Sum for Polynomial<F>
 where
-    F: PrimeField,
+    F: PrimeField + Zeroize,
 {
     fn sum<I: Iterator<Item = Polynomial<F>>>(iter: I) -> Polynomial<F> {
         let mut sum = Polynomial::zero(0);
@@ -521,7 +535,7 @@ where
 
 impl<'a, F> Sum<&'a Polynomial<F>> for Polynomial<F>
 where
-    F: PrimeField,
+    F: PrimeField + Zeroize,
 {
     fn sum<I: Iterator<Item = &'a Polynomial<F>>>(iter: I) -> Polynomial<F> {
         let mut sum = Polynomial::zero(0);