From 9186f0223ab3c35e220435415d2cd2312e795579 Mon Sep 17 00:00:00 2001
From: Elichai Turkel <elichai.turkel@gmail.com>
Date: Thu, 4 Jul 2019 18:09:55 -0400
Subject: [PATCH] Added preallocation size functions and added a test for the
 preallocation

---
 src/context.rs | 26 +++++++++++++++++++++-----
 src/lib.rs     | 26 +++++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index 81f0308a9..9c5ac4b4d 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -84,7 +84,7 @@ mod std_only {
 
     impl<C: Context> Secp256k1<C> {
         fn gen_new() -> Secp256k1<C> {
-            let buf = vec![0u8; Self::preallocate_size()].into_boxed_slice();
+            let buf = vec![0u8; Self::preallocate_size_gen()].into_boxed_slice();
             let ptr = Box::into_raw(buf);
             Secp256k1 {
                 ctx: unsafe { ffi::secp256k1_context_preallocated_create(ptr as *mut c_void, C::FLAGS) },
@@ -170,14 +170,14 @@ unsafe impl<'buf> Context for AllPreallocated<'buf> {
 
 impl<'buf, C: Context + 'buf> Secp256k1<C> {
     fn preallocated_gen_new(buf: &'buf mut [u8]) -> Result<Secp256k1<C>, Error> {
-        if buf.len() < Self::preallocate_size() {
+        if buf.len() < Self::preallocate_size_gen() {
             return Err(Error::NotEnoughMemory);
         }
         Ok(Secp256k1 {
             ctx: unsafe {
                 ffi::secp256k1_context_preallocated_create(
                     buf.as_mut_ptr() as *mut c_void,
-                    AllPreallocated::FLAGS)
+                    C::FLAGS)
             },
             phantom: PhantomData,
             buf: buf as *mut [u8],
@@ -190,18 +190,34 @@ impl<'buf> Secp256k1<AllPreallocated<'buf>> {
     pub fn preallocated_new(buf: &'buf mut [u8]) -> Result<Secp256k1<AllPreallocated<'buf>>, Error> {
         Secp256k1::preallocated_gen_new(buf)
     }
+    /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for a context
+    pub fn preallocate_size() -> usize {
+        Self::preallocate_size_gen()
+    }
 }
 
 impl<'buf> Secp256k1<SignOnlyPreallocated<'buf>> {
     /// Creates a new Secp256k1 context that can only be used for signing
-    pub fn preallocated_new(buf: &'buf mut [u8]) -> Result<Secp256k1<SignOnlyPreallocated<'buf>>, Error> {
+    pub fn preallocated_signing_only(buf: &'buf mut [u8]) -> Result<Secp256k1<SignOnlyPreallocated<'buf>>, Error> {
         Secp256k1::preallocated_gen_new(buf)
     }
+
+    /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for the context
+    #[inline]
+    pub fn preallocate_signing_size() -> usize {
+        Self::preallocate_size_gen()
+    }
 }
 
 impl<'buf> Secp256k1<VerifyOnlyPreallocated<'buf>> {
     /// Creates a new Secp256k1 context that can only be used for verification
-    pub fn preallocated_new(buf: &'buf mut [u8]) -> Result<Secp256k1<VerifyOnlyPreallocated<'buf>>, Error> {
+    pub fn preallocated_verification_only(buf: &'buf mut [u8]) -> Result<Secp256k1<VerifyOnlyPreallocated<'buf>>, Error> {
         Secp256k1::preallocated_gen_new(buf)
     }
+
+    /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for the context
+    #[inline]
+    pub fn preallocate_verification_size() -> usize {
+        Self::preallocate_size_gen()
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index d6ace6ae4..05f2fdc83 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -570,7 +570,7 @@ impl<C: Context> Secp256k1<C> {
     }
 
     /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for a context
-    pub fn preallocate_size() -> usize {
+    pub(crate) fn preallocate_size_gen() -> usize {
         unsafe { ffi::secp256k1_context_preallocated_size(C::FLAGS) }
     }
 
@@ -695,6 +695,30 @@ mod tests {
         });
     }
 
+    #[test]
+    fn test_preallocation() {
+        let mut buf_ful = vec![0u8; Secp256k1::preallocate_size()];
+        let mut buf_sign = vec![0u8; Secp256k1::preallocate_signing_size()];
+        let mut buf_vfy = vec![0u8; Secp256k1::preallocate_verification_size()];
+//
+        let full = Secp256k1::preallocated_new(&mut buf_ful).unwrap();
+        let sign = Secp256k1::preallocated_signing_only(&mut buf_sign).unwrap();
+        let vrfy = Secp256k1::preallocated_verification_only(&mut buf_vfy).unwrap();
+
+//        drop(buf_vfy); // The buffer can't get dropped before the context.
+//        println!("{:?}", buf_ful[5]); // Can't even read the data thanks to the borrow checker.
+
+        let (sk, pk) = full.generate_keypair(&mut thread_rng());
+        let msg = Message::from_slice(&[2u8; 32]).unwrap();
+        // Try signing
+        assert_eq!(sign.sign(&msg, &sk), full.sign(&msg, &sk));
+        let sig = full.sign(&msg, &sk);
+
+        // Try verifying
+        assert!(vrfy.verify(&msg, &sig, &pk).is_ok());
+        assert!(full.verify(&msg, &sig, &pk).is_ok());
+    }
+
     #[test]
     fn capabilities() {
         let sign = Secp256k1::signing_only();