1 files changed, 65 insertions, 1 deletions

diff --git a/vendor/openssl/src/dh.rs b/vendor/openssl/src/dh.rs
index e781543e2..7445e3408 100644
--- a/vendor/openssl/src/dh.rs
+++ b/vendor/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;
 
@@ -39,6 +39,16 @@ where
         params_to_der,
         ffi::i2d_DHparams
     }
+
+    /// Validates DH parameters for correctness
+    #[corresponds(DH_check_key)]
+    pub fn check_key(&self) -> Result<bool, ErrorStack> {
+        unsafe {
+            let mut codes = 0;
+            cvt(ffi::DH_check(self.as_ptr(), &mut codes))?;
+            Ok(codes == 0)
+        }
+    }
 }
 
 impl Dh<Params> {
@@ -66,6 +76,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 +99,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> {
@@ -368,6 +398,30 @@ mod tests {
     }
 
     #[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();
         let params = include_bytes!("../test/dhparams.pem");
@@ -413,4 +467,14 @@ mod tests {
 
         assert_eq!(shared_a, shared_b);
     }
+
+    #[test]
+    fn test_dh_check_key() {
+        let dh1 = Dh::generate_params(512, 2).unwrap();
+        let p = BigNum::from_hex_str("04").unwrap();
+        let g = BigNum::from_hex_str("02").unwrap();
+        let dh2 = Dh::from_pqg(p, None, g).unwrap();
+        assert!(dh1.check_key().unwrap());
+        assert!(!dh2.check_key().unwrap());
+    }
 }