Allow passing workspace to Montgomery_Int

Improves DH and RSA by 5-15% depending on param sizes.

At larger sizes (3072+) doesn't make much difference since the
cost of allocation is relatively small compared to the work.

3 files changed, 140 insertions, 34 deletions

diff --git a/src/lib/math/numbertheory/monty.cpp b/src/lib/math/numbertheory/monty.cpp
index 76575a88c..6ab847ead 100644
--- a/src/lib/math/numbertheory/monty.cpp
+++ b/src/lib/math/numbertheory/monty.cpp
@@ -32,10 +32,13 @@ BigInt Montgomery_Params::inv_mod_p(const BigInt& x) const
    return ct_inverse_mod_odd_modulus(x, p());
    }
 
-BigInt Montgomery_Params::redc(const BigInt& x) const
+BigInt Montgomery_Params::redc(const BigInt& x, secure_vector<word>& ws) const
    {
    const size_t output_size = 2*m_p_words + 2;
-   std::vector<word> ws(output_size);
+
+   if(ws.size() < output_size)
+      ws.resize(output_size);
+
    BigInt z = x;
    z.grow_to(output_size);
 
@@ -43,14 +46,17 @@ BigInt Montgomery_Params::redc(const BigInt& x) const
                      m_p.data(), m_p_words, m_p_dash,
                      ws.data(), ws.size());
 
-   secure_scrub_memory(ws.data(), ws.size() * sizeof(word));
    return z;
    }
 
-BigInt Montgomery_Params::mul(const BigInt& x, const BigInt& y) const
+BigInt Montgomery_Params::mul(const BigInt& x, const BigInt& y,
+                              secure_vector<word>& ws) const
    {
    const size_t output_size = 2*m_p_words + 2;
-   std::vector<word> ws(output_size);
+
+   if(ws.size() < output_size)
+      ws.resize(output_size);
+
    BigInt z(BigInt::Positive, output_size);
    bigint_mul(z.mutable_data(), z.size(),
               x.data(), x.size(), x.sig_words(),
@@ -61,14 +67,16 @@ BigInt Montgomery_Params::mul(const BigInt& x, const BigInt& y) const
                      m_p.data(), m_p_words, m_p_dash,
                      ws.data(), ws.size());
 
-   secure_scrub_memory(ws.data(), ws.size() * sizeof(word));
    return z;
    }
 
-BigInt Montgomery_Params::mul(const BigInt& x, const secure_vector<word>& y) const
+BigInt Montgomery_Params::mul(const BigInt& x,
+                              const secure_vector<word>& y,
+                              secure_vector<word>& ws) const
    {
    const size_t output_size = 2*m_p_words + 2;
-   std::vector<word> ws(output_size);
+   if(ws.size() < output_size)
+      ws.resize(output_size);
    BigInt z(BigInt::Positive, output_size);
 
    bigint_mul(z.mutable_data(), z.size(),
@@ -80,14 +88,42 @@ BigInt Montgomery_Params::mul(const BigInt& x, const secure_vector<word>& y) con
                      m_p.data(), m_p_words, m_p_dash,
                      ws.data(), ws.size());
 
-   secure_scrub_memory(ws.data(), ws.size() * sizeof(word));
    return z;
    }
 
-BigInt Montgomery_Params::sqr(const BigInt& x) const
+void Montgomery_Params::mul_by(BigInt& x,
+                               const secure_vector<word>& y,
+                               secure_vector<word>& ws) const
    {
    const size_t output_size = 2*m_p_words + 2;
-   std::vector<word> ws(output_size);
+
+   if(ws.size() < 2*output_size)
+      ws.resize(2*output_size);
+
+   word* z_data = &ws[0];
+   word* ws_data = &ws[output_size];
+
+   bigint_mul(z_data, output_size,
+              x.data(), x.size(), x.sig_words(),
+              y.data(), y.size(), y.size(),
+              ws_data, output_size);
+
+   bigint_monty_redc(z_data,
+                     m_p.data(), m_p_words, m_p_dash,
+                     ws_data, output_size);
+
+   if(x.size() < output_size)
+      x.grow_to(output_size);
+   copy_mem(x.mutable_data(), z_data, output_size);
+   }
+
+BigInt Montgomery_Params::sqr(const BigInt& x, secure_vector<word>& ws) const
+   {
+   const size_t output_size = 2*m_p_words + 2;
+
+   if(ws.size() < output_size)
+      ws.resize(output_size);
+
    BigInt z(BigInt::Positive, output_size);
 
    bigint_sqr(z.mutable_data(), z.size(),
@@ -98,16 +134,48 @@ BigInt Montgomery_Params::sqr(const BigInt& x) const
                      m_p.data(), m_p_words, m_p_dash,
                      ws.data(), ws.size());
 
-   secure_scrub_memory(ws.data(), ws.size() * sizeof(word));
    return z;
    }
 
+void Montgomery_Params::square_this(BigInt& x,
+                                    secure_vector<word>& ws) const
+   {
+   const size_t output_size = 2*m_p_words + 2;
+
+   if(ws.size() < 2*output_size)
+      ws.resize(2*output_size);
+
+   word* z_data = &ws[0];
+   word* ws_data = &ws[output_size];
+
+   bigint_sqr(z_data, output_size,
+              x.data(), x.size(), x.sig_words(),
+              ws_data, output_size);
+
+   bigint_monty_redc(z_data,
+                     m_p.data(), m_p_words, m_p_dash,
+                     ws_data, output_size);
+
+   if(x.size() < output_size)
+      x.grow_to(output_size);
+   copy_mem(x.mutable_data(), z_data, output_size);
+   }
+
 Montgomery_Int::Montgomery_Int(const std::shared_ptr<const Montgomery_Params> params,
                                const BigInt& v,
                                bool redc_needed) :
-   m_params(params),
-   m_v(redc_needed ? m_params->mul(v % m_params->p(), m_params->R2()) : v)
-   {}
+   m_params(params)
+   {
+   if(redc_needed == false)
+      {
+      m_v = v;
+      }
+   else
+      {
+      secure_vector<word> ws;
+      m_v = m_params->mul(v % m_params->p(), m_params->R2(), ws);
+      }
+   }
 
 void Montgomery_Int::fix_size()
    {
@@ -154,7 +222,8 @@ bool Montgomery_Int::is_zero() const
 
 BigInt Montgomery_Int::value() const
    {
-   return m_params->redc(m_v);
+   secure_vector<word> ws;
+   return m_params->redc(m_v, ws);
    }
 
 Montgomery_Int Montgomery_Int::operator+(const Montgomery_Int& other) const
@@ -191,36 +260,53 @@ Montgomery_Int& Montgomery_Int::operator-=(const Montgomery_Int& other)
 
 Montgomery_Int Montgomery_Int::operator*(const Montgomery_Int& other) const
    {
-   return Montgomery_Int(m_params, m_params->mul(m_v, other.m_v), false);
+   secure_vector<word> ws;
+   return Montgomery_Int(m_params, m_params->mul(m_v, other.m_v, ws), false);
    }
 
-Montgomery_Int& Montgomery_Int::operator*=(const Montgomery_Int& other)
+Montgomery_Int& Montgomery_Int::mul_by(const Montgomery_Int& other,
+                                       secure_vector<word>& ws)
    {
-   m_v = m_params->mul(m_v, other.m_v);
+   m_v = m_params->mul(m_v, other.m_v, ws);
    return (*this);
    }
 
-Montgomery_Int& Montgomery_Int::operator*=(const secure_vector<word>& other)
+Montgomery_Int& Montgomery_Int::mul_by(const secure_vector<word>& other,
+                                       secure_vector<word>& ws)
    {
-   m_v = m_params->mul(m_v, other);
+   m_params->mul_by(m_v, other, ws);
+   //m_v = m_params->mul(m_v, other, ws);
    return (*this);
    }
 
-Montgomery_Int& Montgomery_Int::square_this()
+Montgomery_Int& Montgomery_Int::operator*=(const Montgomery_Int& other)
+   {
+   secure_vector<word> ws;
+   return mul_by(other, ws);
+   }
+
+Montgomery_Int& Montgomery_Int::operator*=(const secure_vector<word>& other)
+   {
+   secure_vector<word> ws;
+   return mul_by(other, ws);
+   }
+
+Montgomery_Int& Montgomery_Int::square_this(secure_vector<word>& ws)
    {
-   m_v = m_params->sqr(m_v);
+   m_params->square_this(m_v, ws);
    return (*this);
    }
 
-Montgomery_Int Montgomery_Int::square() const
+Montgomery_Int Montgomery_Int::square(secure_vector<word>& ws) const
    {
-   const BigInt v = m_params->sqr(m_v);
+   const BigInt v = m_params->sqr(m_v, ws);
    return Montgomery_Int(m_params, v, false);
    }
 
 Montgomery_Int Montgomery_Int::multiplicative_inverse() const
    {
-   const BigInt iv = m_params->mul(m_params->inv_mod_p(m_v), m_params->R3());
+   secure_vector<word> ws;
+   const BigInt iv = m_params->mul(m_params->inv_mod_p(m_v), m_params->R3(), ws);
    return Montgomery_Int(m_params, iv, false);
    }
 
diff --git a/src/lib/math/numbertheory/monty.h b/src/lib/math/numbertheory/monty.h
index ea08789e1..137e0b967 100644
--- a/src/lib/math/numbertheory/monty.h
+++ b/src/lib/math/numbertheory/monty.h
@@ -68,9 +68,15 @@ class Montgomery_Int final
 
       Montgomery_Int& operator*=(const secure_vector<word>& other);
 
-      Montgomery_Int square() const;
+      Montgomery_Int& mul_by(const Montgomery_Int& other,
+                             secure_vector<word>& ws);
 
-      Montgomery_Int& square_this();
+      Montgomery_Int& mul_by(const secure_vector<word>& other,
+                             secure_vector<word>& ws);
+
+      Montgomery_Int square(secure_vector<word>& ws) const;
+
+      Montgomery_Int& square_this(secure_vector<word>& ws);
 
       Montgomery_Int multiplicative_inverse() const;
 
@@ -110,13 +116,26 @@ class Montgomery_Params final
 
       size_t p_words() const { return m_p_words; }
 
-      BigInt redc(const BigInt& x) const;
+      BigInt redc(const BigInt& x,
+                  secure_vector<word>& ws) const;
+
+      BigInt mul(const BigInt& x,
+                 const BigInt& y,
+                 secure_vector<word>& ws) const;
+
+      BigInt mul(const BigInt& x,
+                 const secure_vector<word>& y,
+                 secure_vector<word>& ws) const;
 
-      BigInt mul(const BigInt& x, const BigInt& y) const;
+      void mul_by(BigInt& x,
+                  const secure_vector<word>& y,
+                  secure_vector<word>& ws) const;
 
-      BigInt mul(const BigInt& x, const secure_vector<word>& y) const;
+      BigInt sqr(const BigInt& x,
+                 secure_vector<word>& ws) const;
 
-      BigInt sqr(const BigInt& x) const;
+      void square_this(BigInt& x,
+                       secure_vector<word>& ws) const;
 
       BigInt inv_mod_p(const BigInt& x) const;
 
diff --git a/src/lib/math/numbertheory/monty_exp.cpp b/src/lib/math/numbertheory/monty_exp.cpp
index 6b7af3b09..657686044 100644
--- a/src/lib/math/numbertheory/monty_exp.cpp
+++ b/src/lib/math/numbertheory/monty_exp.cpp
@@ -92,19 +92,20 @@ BigInt Montgomery_Exponentation_State::exponentiation(const BigInt& scalar) cons
    Montgomery_Int x(m_params, m_params->R1(), false);
 
    secure_vector<word> e_bits(m_params->p_words());
+   secure_vector<word> ws;
 
    for(size_t i = exp_nibbles; i > 0; --i)
       {
       for(size_t j = 0; j != m_window_bits; ++j)
          {
-         x.square_this();
+         x.square_this(ws);
          }
 
       const uint32_t nibble = scalar.get_substring(m_window_bits*(i-1), m_window_bits);
 
       const_time_lookup(e_bits, m_g, nibble);
 
-      x *= e_bits;
+      x.mul_by(e_bits, ws);
       }
 
    return x.value();