From b4b52a98589255c7bc281edb426697d5a42e0a66 Mon Sep 17 00:00:00 2001
From: Elichai Turkel <elichai.turkel@gmail.com>
Date: Thu, 4 Jul 2019 17:38:03 -0400
Subject: [PATCH] Moved the context specific traits/enums into a separate file
 with `std` gate

---
 src/context.rs | 122 +++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs     | 101 ++++++++++------------------------------
 2 files changed, 145 insertions(+), 78 deletions(-)
 create mode 100644 src/context.rs

diff --git a/src/context.rs b/src/context.rs
new file mode 100644
index 0000000..92d09cf
--- /dev/null
+++ b/src/context.rs
@@ -0,0 +1,122 @@
+use core::marker::PhantomData;
+use {ffi, types::{c_uint, c_void}, Error, Secp256k1, };
+
+#[cfg(feature = "std")]
+pub use self::std_only::*;
+
+/// A trait for all kinds of Context's that let's you define the exact flags and a function to deallocate memory.
+/// * DO NOT * implement it for your own types.
+pub unsafe trait Context {
+    /// Flags for the ffi.
+    const FLAGS: c_uint;
+    /// A constant description of the context.
+    const DESCRIPTION: &'static str;
+    /// A function to deallocate the memory when the context is dropped.
+    fn deallocate(ptr: *mut [u8]);
+}
+
+/// Marker trait for indicating that an instance of `Secp256k1` can be used for signing.
+pub trait Signing: Context {}
+
+/// Marker trait for indicating that an instance of `Secp256k1` can be used for verification.
+pub trait Verification: Context {}
+
+
+#[cfg(feature = "std")]
+mod std_only {
+    use super::*;
+
+    /// Represents the set of capabilities needed for signing.
+    pub enum SignOnly {}
+
+    /// Represents the set of capabilities needed for verification.
+    pub enum VerifyOnly {}
+
+    /// Represents the set of all capabilities.
+    pub enum All {}
+
+    impl Signing for SignOnly {}
+    impl Signing for All {}
+
+    impl Verification for VerifyOnly {}
+    impl Verification for All {}
+
+    unsafe impl Context for SignOnly {
+        const FLAGS: c_uint = ffi::SECP256K1_START_SIGN;
+        const DESCRIPTION: &'static str = "signing only";
+
+        fn deallocate(ptr: *mut [u8]) {
+            let _ = unsafe { Box::from_raw(ptr) };
+        }
+    }
+
+    unsafe impl Context for VerifyOnly {
+        const FLAGS: c_uint = ffi::SECP256K1_START_VERIFY;
+        const DESCRIPTION: &'static str = "verification only";
+
+        fn deallocate(ptr: *mut [u8]) {
+            let _ = unsafe { Box::from_raw(ptr) };
+        }
+    }
+
+    unsafe impl Context for All {
+        const FLAGS: c_uint = VerifyOnly::FLAGS | SignOnly::FLAGS;
+        const DESCRIPTION: &'static str = "all capabilities";
+
+        fn deallocate(ptr: *mut [u8]) {
+            let _ = unsafe { Box::from_raw(ptr) };
+        }
+    }
+
+    impl<C: Context> Secp256k1<C> {
+        fn gen_new() -> Secp256k1<C> {
+            let buf = vec![0u8; Self::preallocate_size()].into_boxed_slice();
+            let ptr = Box::into_raw(buf);
+            Secp256k1 {
+                ctx: unsafe { ffi::secp256k1_context_preallocated_create(ptr as *mut c_void, C::FLAGS) },
+                phantom: PhantomData,
+                buf: ptr,
+            }
+        }
+    }
+
+    impl Secp256k1<All> {
+        /// Creates a new Secp256k1 context with all capabilities
+        pub fn new() -> Secp256k1<All> {
+            Secp256k1::gen_new()
+        }
+    }
+
+    impl Secp256k1<SignOnly> {
+        /// Creates a new Secp256k1 context that can only be used for signing
+        pub fn signing_only() -> Secp256k1<SignOnly> {
+            Secp256k1::gen_new()
+        }
+    }
+
+    impl Secp256k1<VerifyOnly> {
+        /// Creates a new Secp256k1 context that can only be used for verification
+        pub fn verification_only() -> Secp256k1<VerifyOnly> {
+            Secp256k1::gen_new()
+        }
+    }
+
+    impl Default for Secp256k1<All> {
+        fn default() -> Self {
+            Self::new()
+        }
+    }
+
+    impl<C: Context> Clone for Secp256k1<C> {
+        fn clone(&self) -> Secp256k1<C> {
+            let buf = vec![0u8; unsafe { (&*self.buf).len() }].into_boxed_slice();
+            let ptr = Box::into_raw(buf);
+            Secp256k1 {
+                ctx: unsafe { ffi::secp256k1_context_preallocated_create(ptr as *mut c_void, C::FLAGS) },
+                phantom: PhantomData,
+                buf: ptr,
+            }
+        }
+    }
+
+}
diff --git a/src/lib.rs b/src/lib.rs
index 0947ce7..d6ace6a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -148,6 +148,7 @@ use core::{fmt, ptr, str};
 #[macro_use]
 mod macros;
 mod types;
+mod context;
 pub mod constants;
 pub mod ecdh;
 pub mod ffi;
@@ -157,6 +158,7 @@ pub mod recovery;
 
 pub use key::SecretKey;
 pub use key::PublicKey;
+pub use context::*;
 use core::marker::PhantomData;
 use core::ops::Deref;
 
@@ -463,6 +465,9 @@ pub enum Error {
     InvalidRecoveryId,
     /// Invalid tweak for add_*_assign or mul_*_assign
     InvalidTweak,
+    /// Didn't pass enough memory to context creation with preallocated memory
+    NotEnoughMemory,
+
 }
 
 impl Error {
@@ -475,6 +480,7 @@ impl Error {
             Error::InvalidSecretKey => "secp: malformed or out-of-range secret key",
             Error::InvalidRecoveryId => "secp: bad recovery id",
             Error::InvalidTweak => "secp: bad tweak",
+            Error::NotEnoughMemory => "secp: not enough memory allocated",
         }
     }
 }
@@ -491,48 +497,21 @@ impl std::error::Error for Error {
     fn description(&self) -> &str { self.as_str() }
 }
 
-/// Marker trait for indicating that an instance of `Secp256k1` can be used for signing.
-pub trait Signing {}
-
-/// Marker trait for indicating that an instance of `Secp256k1` can be used for verification.
-pub trait Verification {}
-
-/// Represents the set of capabilities needed for signing.
-pub struct SignOnly {}
-
-/// Represents the set of capabilities needed for verification.
-pub struct VerifyOnly {}
-
-/// Represents the set of all capabilities.
-pub struct All {}
-
-impl Signing for SignOnly {}
-impl Signing for All {}
-
-impl Verification for VerifyOnly {}
-impl Verification for All {}
 
 /// The secp256k1 engine, used to execute all signature operations
-pub struct Secp256k1<C> {
+pub struct Secp256k1<C: Context> {
     ctx: *mut ffi::Context,
-    phantom: PhantomData<C>
+    phantom: PhantomData<C>,
+    buf: *mut [u8],
 }
 
 // The underlying secp context does not contain any references to memory it does not own
-unsafe impl<C> Send for Secp256k1<C> {}
+unsafe impl<C: Context> Send for Secp256k1<C> {}
 // The API does not permit any mutation of `Secp256k1` objects except through `&mut` references
-unsafe impl<C> Sync for Secp256k1<C> {}
+unsafe impl<C: Context> Sync for Secp256k1<C> {}
 
-impl<C> Clone for Secp256k1<C> {
-    fn clone(&self) -> Secp256k1<C> {
-        Secp256k1 {
-            ctx: unsafe { ffi::secp256k1_context_clone(self.ctx) },
-            phantom: self.phantom
-        }
-    }
-}
 
-impl<C> PartialEq for Secp256k1<C> {
+impl<C: Context> PartialEq for Secp256k1<C> {
     fn eq(&self, _other: &Secp256k1<C>) -> bool { true }
 }
 
@@ -566,60 +545,21 @@ impl Deref for SerializedSignature {
 
 impl Eq for SerializedSignature {}
 
-impl<C> Eq for Secp256k1<C> { }
+impl<C: Context> Eq for Secp256k1<C> { }
 
-impl<C> Drop for Secp256k1<C> {
+impl<C: Context> Drop for Secp256k1<C> {
     fn drop(&mut self) {
-        unsafe { ffi::secp256k1_context_destroy(self.ctx); }
+        C::deallocate(self.buf)
     }
 }
 
-impl fmt::Debug for Secp256k1<SignOnly> {
+impl<C: Context> fmt::Debug for Secp256k1<C> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "<secp256k1 context {:?}, signing only>", self.ctx)
+        write!(f, "<secp256k1 context {:?}, {}>", self.ctx, C::DESCRIPTION)
     }
 }
 
-impl fmt::Debug for Secp256k1<VerifyOnly> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "<secp256k1 context {:?}, verification only>", self.ctx)
-    }
-}
-
-impl fmt::Debug for Secp256k1<All> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "<secp256k1 context {:?}, all capabilities>", self.ctx)
-    }
-}
-
-impl Secp256k1<All> {
-    /// Creates a new Secp256k1 context with all capabilities
-    pub fn new() -> Secp256k1<All> {
-        Secp256k1 { ctx: unsafe { ffi::secp256k1_context_create(ffi::SECP256K1_START_SIGN | ffi::SECP256K1_START_VERIFY) }, phantom: PhantomData }
-    }
-}
-
-impl Default for Secp256k1<All> {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl Secp256k1<SignOnly> {
-    /// Creates a new Secp256k1 context that can only be used for signing
-    pub fn signing_only() -> Secp256k1<SignOnly> {
-        Secp256k1 { ctx: unsafe { ffi::secp256k1_context_create(ffi::SECP256K1_START_SIGN) }, phantom: PhantomData }
-    }
-}
-
-impl Secp256k1<VerifyOnly> {
-    /// Creates a new Secp256k1 context that can only be used for verification
-    pub fn verification_only() -> Secp256k1<VerifyOnly> {
-        Secp256k1 { ctx: unsafe { ffi::secp256k1_context_create(ffi::SECP256K1_START_VERIFY) }, phantom: PhantomData }
-    }
-}
-
-impl<C> Secp256k1<C> {
+impl<C: Context> Secp256k1<C> {
 
     /// Getter for the raw pointer to the underlying secp256k1 context. This
     /// shouldn't be needed with normal usage of the library. It enables
@@ -629,6 +569,11 @@ impl<C> Secp256k1<C> {
         &self.ctx
     }
 
+    /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for a context
+    pub fn preallocate_size() -> usize {
+        unsafe { ffi::secp256k1_context_preallocated_size(C::FLAGS) }
+    }
+
     /// (Re)randomizes the Secp256k1 context for cheap sidechannel resistance;
     /// see comment in libsecp256k1 commit d2275795f by Gregory Maxwell. Requires
     /// compilation with "rand" feature.