Merge pull request #1883 from alex/dh-set-keys

Fixes #1882 -- added APIs for setting public keys on Dh
sfackler · Apr 16, 2023 · 62189c1 · 62189c1
2 parents bdb165b + c2fbe9a
commit 62189c1
Showing 1 changed file with 45 additions and 1 deletion.
diff --git a/openssl/src/dh.rs b/openssl/src/dh.rs
@@ -7,7 +7,7 @@ use std::ptr;
 
 use crate::bn::{BigNum, BigNumRef};
 use crate::error::ErrorStack;
-use crate::pkey::{HasParams, HasPrivate, HasPublic, Params, Private};
+use crate::pkey::{HasParams, HasPrivate, HasPublic, Params, Private, Public};
 use crate::{cvt, cvt_p};
 use openssl_macros::corresponds;
 
@@ -66,6 +66,16 @@ impl Dh<Params> {
         }
     }
 
+    /// Sets the public key on the DH object.
+    pub fn set_public_key(self, pub_key: BigNum) -> Result<Dh<Public>, ErrorStack> {
+        unsafe {
+            let dh_ptr = self.0;
+            cvt(DH_set0_key(dh_ptr, pub_key.as_ptr(), ptr::null_mut()))?;
+            mem::forget((self, pub_key));
+            Ok(Dh::from_ptr(dh_ptr))
+        }
+    }
+
     /// Sets the private key on the DH object and recomputes the public key.
     pub fn set_private_key(self, priv_key: BigNum) -> Result<Dh<Private>, ErrorStack> {
         unsafe {
@@ -79,6 +89,16 @@ impl Dh<Params> {
         }
     }
 
+    /// Sets the public and private keys on the DH object.
+    pub fn set_key(self, pub_key: BigNum, priv_key: BigNum) -> Result<Dh<Private>, ErrorStack> {
+        unsafe {
+            let dh_ptr = self.0;
+            cvt(DH_set0_key(dh_ptr, pub_key.as_ptr(), priv_key.as_ptr()))?;
+            mem::forget((self, pub_key, priv_key));
+            Ok(Dh::from_ptr(dh_ptr))
+        }
+    }
+
     /// Generates DH params based on the given `prime_len` and a fixed `generator` value.
     #[corresponds(DH_generate_parameters_ex)]
     pub fn generate_params(prime_len: u32, generator: u32) -> Result<Dh<Params>, ErrorStack> {
@@ -367,6 +387,30 @@ mod tests {
         assert_eq!(key1.private_key(), key2.private_key());
     }
 
+    #[test]
+    #[cfg(ossl102)]
+    fn test_set_keys() {
+        let dh1 = Dh::get_2048_256().unwrap();
+        let key1 = dh1.generate_key().unwrap();
+
+        let dh2 = Dh::get_2048_256().unwrap();
+        let key2 = dh2
+            .set_public_key(key1.public_key().to_owned().unwrap())
+            .unwrap();
+
+        assert_eq!(key1.public_key(), key2.public_key());
+
+        let dh3 = Dh::get_2048_256().unwrap();
+        let key3 = dh3
+            .set_key(
+                key1.public_key().to_owned().unwrap(),
+                key1.private_key().to_owned().unwrap(),
+            )
+            .unwrap();
+        assert_eq!(key1.public_key(), key3.public_key());
+        assert_eq!(key1.private_key(), key3.private_key());
+    }
+
     #[test]
     fn test_dh_from_pem() {
         let mut ctx = SslContext::builder(SslMethod::tls()).unwrap();