added local machine and no-root supprot for as_chain_der()

MS-megliu · Dec 11, 2024 · 3ffcf85 · 3ffcf85
1 parent ea50342
commit 3ffcf85
Show file tree

Hide file tree

Showing 4 changed files with 89 additions and 9 deletions.
diff --git a/examples/client.rs b/examples/client.rs
@@ -13,6 +13,7 @@ use rustls::{
 use rustls_pki_types::{CertificateDer, ServerName};
 
 use rustls_cng::{
+    cert::ChainEngineType,
     signer::CngSigningKey,
     store::{CertStore, CertStoreType},
 };
@@ -37,7 +38,7 @@ fn get_chain(
     let key = context.acquire_key()?;
     let signing_key = CngSigningKey::new(key)?;
     let chain = context
-        .as_chain_der()?
+        .as_chain_der(true, ChainEngineType::HkeyLocalMachine)?
         .into_iter()
         .map(Into::into)
         .collect();

diff --git a/examples/server.rs b/examples/server.rs
@@ -13,6 +13,7 @@ use rustls::{
 };
 
 use rustls_cng::{
+    cert::ChainEngineType,
     signer::CngSigningKey,
     store::{CertStore, CertStoreType},
 };
@@ -74,7 +75,9 @@ impl ResolvesServerCert for ServerCertResolver {
         println!("Key alg: {:?}", key.key().algorithm());
 
         // attempt to acquire a full certificate chain
-        let chain = context.as_chain_der().ok()?;
+        let chain = context
+            .as_chain_der(true, ChainEngineType::HkeyLocalMachine)
+            .ok()?;
         let certs = chain.into_iter().map(Into::into).collect();
 
         // return CertifiedKey instance

diff --git a/src/cert.rs b/src/cert.rs
@@ -6,12 +6,19 @@ use windows_sys::Win32::Security::Cryptography::*;
 
 use crate::{error::CngError, key::NCryptKey, Result};
 
+const HCCE_LOCAL_MACHINE: HCERTCHAINENGINE = 0x1 as HCERTCHAINENGINE;
+
 #[derive(Debug)]
 enum InnerContext {
     Owned(*const CERT_CONTEXT),
     Borrowed(*const CERT_CONTEXT),
 }
 
+pub enum ChainEngineType {
+    HkeyCurrentUser,
+    HkeyLocalMachine,
+}
+
 unsafe impl Send for InnerContext {}
 unsafe impl Sync for InnerContext {}
 
@@ -86,9 +93,59 @@ impl CertContext {
             )
         }
     }
+    /*
+        /// Return DER-encoded X.509 certificate chain
+        pub fn as_chain_der(&self) -> Result<Vec<Vec<u8>>> {
+            unsafe {
+                let param = CERT_CHAIN_PARA {
+                    cbSize: mem::size_of::<CERT_CHAIN_PARA>() as u32,
+                    RequestedUsage: std::mem::zeroed(),
+                };
+
+                let mut context: *mut CERT_CHAIN_CONTEXT = ptr::null_mut();
+
+                let result = CertGetCertificateChain(
+                    HCERTCHAINENGINE::default(),
+                    self.inner(),
+                    ptr::null(),
+                    ptr::null_mut(),
+                    &param,
+                    0,
+                    ptr::null(),
+                    &mut context,
+                ) != 0;
+
+                if result {
+                    let mut chain = vec![];
+
+                    if (*context).cChain > 0 {
+                        let chain_ptr = *(*context).rgpChain;
+                        let elements = slice::from_raw_parts(
+                            (*chain_ptr).rgpElement,
+                            (*chain_ptr).cElement as usize,
+                        );
+                        for element in elements {
+                            let context = (**element).pCertContext;
+                            chain.push(Self::new_borrowed(context).as_der().to_vec());
+                        }
+                    }
+
+                    CertFreeCertificateChain(&*context);
 
-    /// Return DER-encoded X.509 certificate chain
-    pub fn as_chain_der(&self) -> Result<Vec<Vec<u8>>> {
+                    Ok(chain)
+                } else {
+                    Err(CngError::from_win32_error())
+                }
+            }
+        }
+    */
+    /// Return DER-encoded X.509 certificate chain.
+    /// Giving user options to (1) not to include the root. (2) to use the HKLM engine instead of default HKCU engine
+    pub fn as_chain_der(
+        &self,
+        include_root: bool,
+        chain_engine: ChainEngineType,
+    ) -> Result<Vec<Vec<u8>>> {
         unsafe {
             let param = CERT_CHAIN_PARA {
                 cbSize: mem::size_of::<CERT_CHAIN_PARA>() as u32,
@@ -97,8 +154,13 @@ impl CertContext {
 
             let mut context: *mut CERT_CHAIN_CONTEXT = ptr::null_mut();
 
+            let chain_engine = match chain_engine {
+                ChainEngineType::HkeyLocalMachine => HCCE_LOCAL_MACHINE as isize,
+                ChainEngineType::HkeyCurrentUser => HCERTCHAINENGINE::default() as isize,
+            };
+
             let result = CertGetCertificateChain(
-                HCERTCHAINENGINE::default(),
+                chain_engine,
                 self.inner(),
                 ptr::null(),
                 ptr::null_mut(),
@@ -117,7 +179,19 @@ impl CertContext {
                         (*chain_ptr).rgpElement,
                         (*chain_ptr).cElement as usize,
                     );
+
+                    let mut first = true;
                     for element in elements {
+                        if first {
+                            first = false;
+                        } else if !include_root {
+                            if 0 != ((**element).TrustStatus.dwInfoStatus
+                                & CERT_TRUST_IS_SELF_SIGNED)
+                            {
+                                break;
+                            }
+                        }
+
                         let context = (**element).pCertContext;
                         chain.push(Self::new_borrowed(context).as_der().to_vec());
                     }

diff --git a/tests/test_client_server.rs b/tests/test_client_server.rs
@@ -16,7 +16,7 @@ mod client {
     };
     use rustls_pki_types::CertificateDer;
 
-    use rustls_cng::{signer::CngSigningKey, store::CertStore};
+    use rustls_cng::{cert::ChainEngineType, signer::CngSigningKey, store::CertStore};
 
     #[derive(Debug)]
     pub struct ClientCertResolver(CertStore, String);
@@ -32,7 +32,7 @@ mod client {
         let key = context.acquire_key()?;
         let signing_key = CngSigningKey::new(key)?;
         let chain = context
-            .as_chain_der()?
+            .as_chain_der(true, ChainEngineType::HkeyLocalMachine)?
             .into_iter()
             .map(Into::into)
             .collect();
@@ -111,7 +111,7 @@ mod server {
         RootCertStore, ServerConfig, ServerConnection, Stream,
     };
 
-    use rustls_cng::{signer::CngSigningKey, store::CertStore};
+    use rustls_cng::{cert::ChainEngineType, signer::CngSigningKey, store::CertStore};
 
     #[derive(Debug)]
     pub struct ServerCertResolver(CertStore);
@@ -127,7 +127,9 @@ mod server {
                 CngSigningKey::new(key).ok().map(|key| (ctx, key))
             })?;
 
-            let chain = context.as_chain_der().ok()?;
+            let chain = context
+                .as_chain_der(true, ChainEngineType::HkeyLocalMachine)
+                .ok()?;
             let certs = chain.into_iter().map(Into::into).collect();
 
             Some(Arc::new(CertifiedKey {