Rust wrapper: check DH::shared_secret buffer size against prime size

wrapper/rust/wolfssl-wolfcrypt/CHANGELOG.md

@@ -1,5 +1,17 @@
 # wolfssl-wolfcrypt Change Log
 
+## Unreleased
+
+New features:
+
+- Add DH::prime_size() to query the DH prime size, which is the minimum output
+  buffer size DH::shared_secret() requires
+
+Fixes and improvements:
+
+- Fix DH::shared_secret() writing past the end of an output buffer smaller than
+  the DH prime; it now returns BUFFER_E instead
+
 ## v2.1.0
 
 Fixes:

wrapper/rust/wolfssl-wolfcrypt/src/dh.rs

@@ -1443,6 +1443,48 @@ impl DH {
         Ok(())
     }
 
+    /// Get the size in bytes of the DH `p` parameter (the prime).
+    ///
+    /// This is the size of the DH prime `p` and the maximum length that
+    /// `shared_secret()` may write, and therefore the minimum length of the
+    /// output buffer.
+    ///
+    /// # Returns
+    ///
+    /// Returns either Ok(size) containing the size in bytes of the DH prime or
+    /// Err(e) containing the wolfSSL library error code value.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// #[cfg(dh_ffdhe_2048)]
+    /// {
+    /// use wolfssl_wolfcrypt::dh::DH;
+    /// let mut dh = DH::new_named(DH::FFDHE_2048).expect("Error with new_named()");
+    /// let prime_size = dh.prime_size().expect("Error with prime_size()");
+    /// assert_eq!(prime_size, 256);
+    /// }
+    /// ```
+    pub fn prime_size(&mut self) -> Result<usize, i32> {
+        let mut p_size = 0u32;
+        let mut q_size = 0u32;
+        let mut g_size = 0u32;
+        /* Passing null buffers for p, q and g requests the parameter sizes
+         * only, which wc_DhExportParamsRaw() reports by returning
+         * LENGTH_ONLY_E. Any other return value, including success, means the
+         * sizes were not filled in. */
+        let rc = unsafe {
+            sys::wc_DhExportParamsRaw(&mut self.wc_dhkey,
+                core::ptr::null_mut(), &mut p_size,
+                core::ptr::null_mut(), &mut q_size,
+                core::ptr::null_mut(), &mut g_size)
+        };
+        if rc != sys::wolfCrypt_ErrorCodes_LENGTH_ONLY_E {
+            return Err(rc);
+        }
+        Ok(p_size as usize)
+    }
+
     /// Export Diffie-Hellman context parameters.
     ///
     /// # Parameters
@@ -1533,6 +1575,14 @@ impl DH {
     /// symmetric communication. On successfully generating a shared secret
     /// key, the size of the secret key written to `dout` will be returned.
     ///
+    /// # Parameters
+    ///
+    /// * `dout`: Output buffer containing the generated shared secret value.
+    ///   The buffer must be at least as large as `self.prime_size()` or
+    ///   a `BUFFER_E` error code will be returned.
+    /// * `private`: Private key buffer.
+    /// * `other_pub`: Other side public key used to generate shared secret.
+    ///
     /// # Returns
     ///
     /// Returns either Ok(size) containing the size of the key written to
@@ -1563,6 +1613,11 @@ impl DH {
     /// }
     /// ```
     pub fn shared_secret(&mut self, dout: &mut [u8], private: &[u8], other_pub: &[u8]) -> Result<usize, i32> {
+        /* wc_DhAgree() does not check the buffer size, so reject a short
+         * buffer before the call to avoid a buffer overrun. */
+        if dout.len() < self.prime_size()? {
+            return Err(sys::wolfCrypt_ErrorCodes_BUFFER_E);
+        }
         let mut dout_size = crate::buffer_len_to_u32(dout.len())?;
         let private_size = crate::buffer_len_to_u32(private.len())?;
         let other_pub_size = crate::buffer_len_to_u32(other_pub.len())?;

wrapper/rust/wolfssl-wolfcrypt/tests/test_dh.rs

@@ -2,7 +2,7 @@
 
 mod common;
 
-#[cfg(any(all(dh_keygen, dh_ffdhe_2048), random))]
+#[cfg(any(dh_ffdhe_2048, random))]
 use wolfssl_wolfcrypt::dh::DH;
 #[cfg(random)]
 use wolfssl_wolfcrypt::random::RNG;
@@ -210,3 +210,39 @@ fn test_dh_shared_secret() {
     assert_eq!(ss0_size, ss1_size);
     assert_eq!(*ss0, *ss1);
 }
+
+#[test]
+#[cfg(dh_ffdhe_2048)]
+fn test_dh_prime_size() {
+    common::setup();
+    let mut dh = DH::new_named(DH::FFDHE_2048).expect("Error with new_named()");
+    assert_eq!(dh.prime_size().expect("Error with prime_size()"), 256);
+}
+
+#[test]
+#[cfg(all(dh_ffdhe_2048, random))]
+fn test_dh_shared_secret_buffer_too_small() {
+    common::setup();
+    let mut rng = RNG::new().expect("Error with RNG::new()");
+    let mut dh = DH::new_named(DH::FFDHE_2048).expect("Error with new_named()");
+
+    let mut private = [0u8; 256];
+    let mut private_size = 0u32;
+    let mut public = [0u8; 256];
+    let mut public_size = 0u32;
+    dh.generate_key_pair(&mut rng, &mut private, &mut private_size, &mut public, &mut public_size).expect("Error with generate_key_pair()");
+    let private = &private[0..(private_size as usize)];
+    let public = &public[0..(public_size as usize)];
+
+    /* A buffer one byte short of the prime size must be reported rather than
+     * written past the end. The canary guards against the shared secret being
+     * written beyond the slice that was passed in. */
+    let prime_size = dh.prime_size().expect("Error with prime_size()");
+    let mut ss = [0u8; 256];
+    let (dout, canary) = ss.split_at_mut(prime_size - 1);
+    canary[0] = 0xA5;
+    let rc = dh.shared_secret(dout, private, public)
+        .expect_err("shared_secret() must reject an undersized buffer");
+    assert_eq!(rc, wolfssl_wolfcrypt::sys::wolfCrypt_ErrorCodes_BUFFER_E);
+    assert_eq!(canary[0], 0xA5);
+}