Implement 'sign[_final]_into()'

These '_into()' variants can be used when the caller knows the expected
buffer size, and has allocated a buffer for the signature themselves.
It requires fewer calls into the underlying library and can save on
heap allocations.

Signed-off-by: arya dradjica <arya@nlnetlabs.nl>

Author: bal-e

cryptoki/src/session/signing_macing.rs

@@ -3,29 +3,38 @@
 //! Signing and authentication functions
 
 use crate::context::Function;
+#[cfg(doc)]
+use crate::error::RvError;
 use crate::error::{Result, Rv};
 use crate::mechanism::Mechanism;
 use crate::object::ObjectHandle;
 use crate::session::Session;
 use cryptoki_sys::*;
 use std::convert::TryInto;
 
+/// # Generating and Verifying Signatures
+///
+/// Several functions are provided for signing data and verifying signatures.
+/// This includes message authentication codes (MACs). The signed data can be
+/// provided in one-shot and streaming modes.
 impl Session {
-    /// Sign data in single-part
+    /// Sign data in one-shot mode.
+    ///
+    /// `data` should be a byte sequence representing the input message. It will
+    /// be signed using the specified key, and the resulting signature will be
+    /// returned in a `Vec`.
+    ///
+    /// Use [`Self::sign_into()`] if (an upper bound for) the size of the
+    /// signature is known, to avoid the heap allocation of `Vec`. Use
+    /// [`Self::sign_init()`] etc. if the input data is being streamed (i.e. it
+    /// is not all immediately available).
     pub fn sign(&self, mechanism: &Mechanism, key: ObjectHandle, data: &[u8]) -> Result<Vec<u8>> {
-        let mut mechanism: CK_MECHANISM = mechanism.into();
         let mut signature_len = 0;
 
-        unsafe {
-            Rv::from(get_pkcs11!(self.client(), C_SignInit)(
-                self.handle(),
-                &mut mechanism as CK_MECHANISM_PTR,
-                key.handle(),
-            ))
-            .into_result(Function::SignInit)?;
-        }
+        // Initialize the signing operation.
+        self.sign_init(mechanism, key)?;
 
-        // Get the output buffer length
+        // Get the output buffer length.
         unsafe {
             Rv::from(get_pkcs11!(self.client(), C_Sign)(
                 self.handle(),
@@ -37,9 +46,10 @@ impl Session {
             .into_result(Function::Sign)?;
         }
 
+        // Allocate the output buffer.
         let mut signature = vec![0; signature_len.try_into()?];
 
-        //TODO: we should add a new error instead of those unwrap!
+        // Perform the actual signing.
         unsafe {
             Rv::from(get_pkcs11!(self.client(), C_Sign)(
                 self.handle(),
@@ -51,11 +61,62 @@ impl Session {
             .into_result(Function::Sign)?;
         }
 
+        // Limit the output buffer to the size of the generated signature.
         signature.truncate(signature_len.try_into()?);
 
         Ok(signature)
     }
 
+    /// Sign data into the given buffer.
+    ///
+    /// `data` should be a byte sequence representing the input message. It will
+    /// be signed using the specified key, and the resulting signature will be
+    /// written to the given output buffer.
+    ///
+    /// The output buffer should be large enough to store the signature.
+    /// If it is large enough, the number of filled bytes is returned
+    /// (i.e. `sig[0..size]` contains the prepared signature). Otherwise,
+    /// [`RvError::BufferTooSmall`] will be returned.
+    ///
+    /// Use [`Self::sign()`] if (an upper bound for) the size of the signature
+    /// is not known. Use [`Self::sign_init()`] etc. if the input data is being
+    /// streamed (i.e. it is not all immediately available).
+    pub fn sign_into(
+        &mut self,
+        mechanism: &Mechanism,
+        key: ObjectHandle,
+        data: &[u8],
+        sig: &mut [u8],
+    ) -> Result<usize> {
+        // The size of the signature buffer, into which 'C_Sign' will write the
+        // size of the generated signature.
+        let sig_buf_len = sig.len().try_into()?;
+        let mut sig_len = sig_buf_len;
+
+        // Initialize the signing operation.
+        self.sign_init(mechanism, key)?;
+
+        // Perform the actual signing.
+        unsafe {
+            Rv::from(get_pkcs11!(self.client(), C_Sign)(
+                self.handle(),
+                data.as_ptr() as *mut u8,
+                data.len().try_into()?,
+                sig.as_mut_ptr(),
+                &mut sig_len,
+            ))
+            .into_result(Function::Sign)?;
+        }
+
+        assert!(
+            sig_len <= sig_buf_len,
+            "'C_Sign' succeeded but increased 'sig_len', possibly indicating out-of-bounds accesses"
+        );
+
+        // NOTE: As checked above, 'sig_len <= sig_buf_len <= usize::MAX'.
+        Ok(sig_len as usize)
+    }
+
     /// Starts new multi-part signing operation
     pub fn sign_init(&self, mechanism: &Mechanism, key: ObjectHandle) -> Result<()> {
         let mut mechanism: CK_MECHANISM = mechanism.into();
@@ -87,12 +148,20 @@ impl Session {
         Ok(())
     }
 
-    /// Finalizes ongoing multi-part signing operation,
-    /// returning the signature
+    /// Complete an ongoing streaming signing operation.
+    ///
+    /// This must be preceded by [`Self::sign_init()`] and zero or more calls
+    /// to [`Self::sign_update()`]. This method will terminate the multi-part
+    /// signing operation. The resulting signature will be returned in a `Vec`.
+    ///
+    /// Use [`Self::sign_final_into()`] if (an upper bound for) the size of
+    /// the signature is known, to avoid the heap allocation of `Vec`. Use
+    /// [`Self::sign()`] if the input data is entirely available in a single
+    /// buffer (i.e. does not have to be streamed).
     pub fn sign_final(&self) -> Result<Vec<u8>> {
         let mut signature_len = 0;
 
-        // Get the output buffer length
+        // Get the output buffer length.
         unsafe {
             Rv::from(get_pkcs11!(self.client(), C_SignFinal)(
                 self.handle(),
@@ -102,8 +171,10 @@ impl Session {
             .into_result(Function::SignFinal)?;
         }
 
+        // Allocate the output buffer.
         let mut signature = vec![0; signature_len.try_into()?];
 
+        // Perform the actual signing.
         unsafe {
             Rv::from(get_pkcs11!(self.client(), C_SignFinal)(
                 self.handle(),
@@ -113,11 +184,56 @@ impl Session {
             .into_result(Function::SignFinal)?;
         }
 
+        // Limit the output buffer to the size of the generated signature.
         signature.truncate(signature_len.try_into()?);
 
         Ok(signature)
     }
 
+    /// Complete an ongoing multi-part signing operation, writing the signature
+    /// into the given buffer.
+    ///
+    /// This must be preceded by [`Self::sign_init()`] and zero or more calls
+    /// to [`Self::sign_update()`]. This method will terminate the multi-part
+    /// signing operation.
+    ///
+    /// The output buffer should be large enough to store the signature.
+    /// If it is large enough, the number of filled bytes is returned
+    /// (i.e. `sig[0..size]` contains the prepared signature). Otherwise,
+    /// [`RvError::BufferTooSmall`] will be returned.
+    ///
+    /// Use [`Self::sign_final()`] if (an upper bound for) the size of the
+    /// signature is not known. Use [`Self::sign_into()`] if the input data
+    /// is entirely available in a single buffer (i.e. does not have to be
+    /// streamed).
+    //
+    // TODO: Is it possible to re-do this call if an incorrect buffer size
+    // was passed in? I think so?
+    pub fn sign_final_into(&mut self, sig: &mut [u8]) -> Result<usize> {
+        // The size of the signature buffer, into which 'C_SignFinal' will write
+        // the size of the generated signature.
+        let sig_buf_len = sig.len().try_into()?;
+        let mut sig_len = sig_buf_len;
+
+        // Perform the underlying finalization.
+        unsafe {
+            Rv::from(get_pkcs11!(self.client(), C_SignFinal)(
+                self.handle(),
+                sig.as_mut_ptr(),
+                &mut sig_len,
+            ))
+            .into_result(Function::SignFinal)?;
+        }
+
+        assert!(
+            sig_len <= sig_buf_len,
+            "'C_SignFinal' succeeded but increased 'sig_len', possibly indicating out-of-bounds accesses"
+        );
+
+        // NOTE: As checked above, 'sig_len <= sig_buf_len <= usize::MAX'.
+        Ok(sig_len as usize)
+    }
+
     /// Verify data in single-part
     pub fn verify(
         &self,