rust: bssl-tls: Introduce asynchronous private key delegate

This allows the authenticating party to sign or decrypt data
asynchronously.

Bug: 479599893

Signed-off-by: Xiangfei Ding <xfding@google.com>
Change-Id: I61172bf18bfce94c287413c436b881926a6a6964
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/91088
Presubmit-BoringSSL-Verified: boringssl-scoped@luci-project-accounts.iam.gserviceaccount.com <boringssl-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Adam Langley <agl@google.com>
diff --git a/rust/bssl-tls/src/connection.rs b/rust/bssl-tls/src/connection.rs
index 35b2c2d..03347b2 100644
--- a/rust/bssl-tls/src/connection.rs
+++ b/rust/bssl-tls/src/connection.rs
@@ -18,7 +18,10 @@
 use core::{
     ffi::c_int,
     marker::PhantomData,
-    mem::forget,
+    mem::{
+        forget,
+        transmute, //
+    },
     ptr::NonNull,
     task::Waker, //
 };
@@ -28,7 +31,10 @@
         ConnectionMode,
         ProtocolVersion, //
     },
-    connection::methods::waker_data_ref_from_ssl,
+    connection::{
+        lifecycle::TlsConnectionInHandshake,
+        methods::waker_data_ref_from_ssl, //
+    },
     context::TlsMode,
     errors::{
         Error,
@@ -113,6 +119,17 @@
         self
     }
 
+    fn in_handshake(&mut self) -> TlsConnectionInHandshake<'_, R, M> {
+        unsafe {
+            // Safety:
+            // - the connection is still technically in handshake phase, so it is safe for internal
+            //   use to configure the handshake through the associated methods.
+            // - `TlsConnection` is a transparent wrapper around `NonNull<bssl_sys::SSL>`.
+            // - the `Role` and `Mode` are matching.
+            TlsConnectionInHandshake(transmute(&mut self.ptr))
+        }
+    }
+
     /// Set the session for resumption.
     pub fn with_session(&mut self, session: &TlsSession) -> &mut Self {
         self.as_in_handshake().set_session(session);
diff --git a/rust/bssl-tls/src/connection/credentials.rs b/rust/bssl-tls/src/connection/credentials.rs
index 8cce65e..11b16e5 100644
--- a/rust/bssl-tls/src/connection/credentials.rs
+++ b/rust/bssl-tls/src/connection/credentials.rs
@@ -31,22 +31,22 @@
 
 use super::{
     Client,
+    TlsConnection,
+    TlsConnectionBuilder,
+    lifecycle::{
+        EstablishedTlsConnection,
+        TlsConnectionInHandshake, //
+    },
+    methods::HasPrivateKeyMethods,
     methods::HasTlsConnectionMethod, //
 };
 use crate::{
     check_lib_error,
     config::ConfigurationError,
-    connection::{
-        TlsConnection,
-        TlsConnectionBuilder,
-        lifecycle::{
-            EstablishedTlsConnection,
-            TlsConnectionInHandshake, //
-        }, //
-    }, //
     credentials::{
         CertificateType,
         CertificateVerificationMode,
+        PrivateKeyDelegate,
         SignatureAlgorithm,
         TlsCredential,
         VerifyCertificate,
@@ -57,6 +57,53 @@
     has_duplicates, //
 };
 
+/// # Asynchronous private key operations
+impl<R, M> TlsConnectionBuilder<R, M>
+where
+    M: HasPrivateKeyMethods + HasTlsConnectionMethod,
+{
+    /// Set the private key delegate.
+    ///
+    /// This will override the [`crate::context::TlsContext`] private key delegate.
+    pub fn with_private_key_delegate(
+        &mut self,
+        key_method: Option<Box<dyn PrivateKeyDelegate>>,
+    ) -> &mut Self {
+        self.in_handshake().set_private_key_delegate(key_method);
+        self
+    }
+}
+
+/// # Asynchronous private key operations
+impl<R, M> TlsConnectionInHandshake<'_, R, M>
+where
+    M: HasPrivateKeyMethods + HasTlsConnectionMethod,
+{
+    /// Set the private key delegate.
+    ///
+    /// This will override the [`crate::context::TlsContext`] private key delegate.
+    pub fn set_private_key_delegate(
+        &mut self,
+        key_method: Option<Box<dyn PrivateKeyDelegate>>,
+    ) -> &mut Self {
+        let ctx = self.ptr();
+        if key_method.is_some() {
+            unsafe {
+                // Safety: we only install our own vtable.
+                bssl_sys::SSL_set_private_key_method(ctx, <M as HasPrivateKeyMethods>::METHODS);
+            }
+        } else {
+            unsafe {
+                // Safety: we only uninstall the vtable.
+                bssl_sys::SSL_set_private_key_method(ctx, core::ptr::null());
+            }
+        }
+        self.get_connection_methods().private_key_delegate = key_method;
+        self
+    }
+}
+
+/// # Certificate verification
 impl<R, M> TlsConnectionBuilder<R, M>
 where
     M: HasTlsConnectionMethod,
@@ -70,7 +117,13 @@
             .set_certificate_verification_mode(mode);
         self
     }
+}
 
+/// # Custom certificate verification
+impl<R, M> TlsConnectionBuilder<R, M>
+where
+    M: HasPrivateKeyMethods + HasTlsConnectionMethod,
+{
     /// Configure the certificate verifier.
     ///
     /// See [`VerifyCertificate`] for how to implement a custom verifier.
diff --git a/rust/bssl-tls/src/connection/methods.rs b/rust/bssl-tls/src/connection/methods.rs
index 0d1745d..c22aae3 100644
--- a/rust/bssl-tls/src/connection/methods.rs
+++ b/rust/bssl-tls/src/connection/methods.rs
@@ -31,6 +31,7 @@
 
 use crate::{
     Methods,
+    PrivateKeyMethods,
     VerifyCertificateMethods,
     abort_on_panic,
     context::{
@@ -39,8 +40,15 @@
         TlsMode, //
     },
     credentials::{
+        PrivateKeyDelegate,
+        PrivateKeyOperation,
         VerifyCertificate,
-        VerifyCertificateTask, //
+        VerifyCertificateTask,
+        methods::{
+            complete,
+            decrypt,
+            sign, //
+        }, //
     },
     errors::TlsRetryReason,
     io::RustBioHandle,
@@ -51,6 +59,8 @@
 pub(super) struct RustConnectionMethods<Mode> {
     /// A handle to a `BIO` managed by this crate.
     pub bio: Option<RustBioHandle>,
+    /// Private key delegate.
+    pub private_key_delegate: Option<Box<dyn PrivateKeyDelegate>>,
     /// Certificate verifier handle.
     pub verify_certificate_methods: Option<Box<dyn VerifyCertificate>>,
     /// A mailbox to propagate IO retrying reasons.
@@ -62,6 +72,7 @@
     pub fn new() -> Self {
         Self {
             bio: None,
+            private_key_delegate: None,
             verify_certificate_methods: None,
             pending_reason: None,
             _p: PhantomData,
@@ -99,6 +110,12 @@
     }
 }
 
+impl<M: HasTlsConnectionMethod> PrivateKeyMethods for RustConnectionMethods<M> {
+    fn private_key_methods(&self) -> Option<&dyn PrivateKeyDelegate> {
+        self.private_key_delegate.as_deref()
+    }
+}
+
 impl<Mode: HasTlsConnectionMethod> VerifyCertificateMethods for RustConnectionMethods<Mode> {
     fn verify_certificate_methods(&self) -> Option<&dyn VerifyCertificate> {
         self.verify_certificate_methods.as_deref()
@@ -151,6 +168,18 @@
 /// Safety:
 /// - `ssl` must be constructed from `TlsConnection` and outlived by `'a`.
 /// - `ssl` must be exclusively owned.
+pub(crate) unsafe fn private_key_op_from_ssl<'a>(
+    ssl: NonNull<bssl_sys::SSL>,
+) -> &'a mut Option<Box<dyn PrivateKeyOperation>> {
+    unsafe {
+        // Safety: `ssl` outlives `'a` and is constructed by `TlsConnection`.
+        <ExDataRegistration as ExData<Option<Box<dyn PrivateKeyOperation>>>>::get_mut(ssl)
+    }
+}
+
+/// Safety:
+/// - `ssl` must be constructed from `TlsConnection` and outlived by `'a`.
+/// - `ssl` must be exclusively owned.
 pub(crate) unsafe fn verify_cert_task_from_ssl<'a>(
     ssl: NonNull<bssl_sys::SSL>,
 ) -> &'a mut Option<Box<dyn VerifyCertificateTask>> {
@@ -263,6 +292,7 @@
     };
 }
 
+register_ex_data!(Option<Box<dyn PrivateKeyOperation>>);
 register_ex_data!(Option<Waker>);
 register_ex_data!(Option<Box<dyn VerifyCertificateTask>>);
 
@@ -296,3 +326,37 @@
         *TLS_CONTEXT_METHOD
     }
 }
+
+pub(super) trait HasPrivateKeyMethods {
+    const METHODS: *const bssl_sys::SSL_PRIVATE_KEY_METHOD;
+}
+
+impl HasPrivateKeyMethods for TlsMode {
+    const METHODS: *const bssl_sys::SSL_PRIVATE_KEY_METHOD = {
+        &bssl_sys::SSL_PRIVATE_KEY_METHOD {
+            sign: Some(sign::<RustConnectionMethods<TlsMode>>),
+            decrypt: Some(decrypt::<RustConnectionMethods<TlsMode>>),
+            complete: Some(complete::<RustConnectionMethods<TlsMode>>),
+        } as _
+    };
+}
+
+impl HasPrivateKeyMethods for DtlsMode {
+    const METHODS: *const bssl_sys::SSL_PRIVATE_KEY_METHOD = {
+        &bssl_sys::SSL_PRIVATE_KEY_METHOD {
+            sign: Some(sign::<RustConnectionMethods<DtlsMode>>),
+            decrypt: Some(decrypt::<RustConnectionMethods<DtlsMode>>),
+            complete: Some(complete::<RustConnectionMethods<DtlsMode>>),
+        } as _
+    };
+}
+
+impl HasPrivateKeyMethods for QuicMode {
+    const METHODS: *const bssl_sys::SSL_PRIVATE_KEY_METHOD = {
+        &bssl_sys::SSL_PRIVATE_KEY_METHOD {
+            sign: Some(sign::<RustConnectionMethods<QuicMode>>),
+            decrypt: Some(decrypt::<RustConnectionMethods<QuicMode>>),
+            complete: Some(complete::<RustConnectionMethods<QuicMode>>),
+        } as _
+    };
+}
diff --git a/rust/bssl-tls/src/context.rs b/rust/bssl-tls/src/context.rs
index 0ddc865..fe9e475 100644
--- a/rust/bssl-tls/src/context.rs
+++ b/rust/bssl-tls/src/context.rs
@@ -36,7 +36,10 @@
         TlsConnectionBuilder,
         methods::HasTlsConnectionMethod, //
     },
-    context::methods::HasTlsContextMethod,
+    context::methods::{
+        HasPrivateKeyMethods,
+        HasTlsContextMethod, //
+    },
     errors::Error,
     has_duplicates, //
 };
@@ -62,7 +65,10 @@
 pub(crate) trait HasBasicIo {}
 
 /// A collection of supported mode of operations.
-pub trait SupportedMode: HasTlsContextMethod + HasTlsConnectionMethod {}
+pub trait SupportedMode:
+    HasTlsContextMethod + HasTlsConnectionMethod + HasPrivateKeyMethods
+{
+}
 
 impl SupportedMode for TlsMode {}
 impl SupportedMode for DtlsMode {}
diff --git a/rust/bssl-tls/src/context/credentials.rs b/rust/bssl-tls/src/context/credentials.rs
index caa24f6..8159b8a 100644
--- a/rust/bssl-tls/src/context/credentials.rs
+++ b/rust/bssl-tls/src/context/credentials.rs
@@ -25,11 +25,13 @@
     config::ConfigurationError,
     context::{
         CertificateCache,
-        SupportedMode, //
+        SupportedMode,
+        methods::HasPrivateKeyMethods, //
     },
     credentials::{
         CertificateType,
         CertificateVerificationMode,
+        PrivateKeyDelegate,
         SignatureAlgorithm,
         TlsCredential,
         VerifyCertificate,
@@ -71,6 +73,30 @@
         self
     }
 
+    /// Set the private key method.
+    ///
+    /// This private key method delegation may be replaced the next moment when
+    /// a new TLS private key is supplied.
+    pub fn with_private_key_delegate(
+        &mut self,
+        key_method: Option<Box<dyn PrivateKeyDelegate>>,
+    ) -> &mut Self {
+        let ctx = self.ptr();
+        if key_method.is_some() {
+            unsafe {
+                // Safety: we only install our own vtable.
+                bssl_sys::SSL_CTX_set_private_key_method(ctx, <M as HasPrivateKeyMethods>::METHODS);
+            }
+        } else {
+            unsafe {
+                // Safety: we only uninstall the vtable.
+                bssl_sys::SSL_CTX_set_private_key_method(ctx, core::ptr::null());
+            }
+        }
+        self.get_context_methods().private_key_methods = key_method;
+        self
+    }
+
     /// Set certificate verification mode.
     ///
     /// # Client certificate verification for servers, mutual TLS
@@ -145,7 +171,6 @@
         S: EarlyCallback<M> + 'static,
     {
         let ctx = self.ptr();
-        let methods = self.get_context_methods();
         unsafe {
             // Safety: we only install our own vtable.
             bssl_sys::SSL_CTX_set_select_certificate_cb(
@@ -153,19 +178,18 @@
                 Some(early_select_cert_cb::<M, super::methods::RustContextMethods<M>>),
             );
         }
-        methods.early_callback_handler = Some(Box::new(handler) as _);
+        self.get_context_methods().early_callback_handler = Some(Box::new(handler) as _);
         self
     }
 
     /// Remove custom certificate selection callback.
     pub fn without_early_callback(&mut self) -> &mut Self {
         let ctx = self.ptr();
-        let methods = self.get_context_methods();
         unsafe {
             // Safety: we only uninstall the vtable.
             bssl_sys::SSL_CTX_set_select_certificate_cb(ctx, None);
         }
-        methods.early_callback_handler = None;
+        self.get_context_methods().early_callback_handler = None;
         self
     }
 
diff --git a/rust/bssl-tls/src/context/methods.rs b/rust/bssl-tls/src/context/methods.rs
index fd51c52..49c4669 100644
--- a/rust/bssl-tls/src/context/methods.rs
+++ b/rust/bssl-tls/src/context/methods.rs
@@ -23,18 +23,28 @@
 use crate::{
     EarlyCallbackMethods,
     Methods,
+    PrivateKeyMethods,
     VerifyCertificateMethods,
     context::{
         DtlsMode,
         QuicMode,
         TlsMode, //
     },
-    credentials::VerifyCertificate,
-    credentials::early_callback::EarlyCallback,
+    credentials::{
+        PrivateKeyDelegate,
+        VerifyCertificate,
+        early_callback::EarlyCallback,
+        methods::{
+            complete,
+            decrypt,
+            sign, //
+        }, //
+    },
     methods::drop_box_rust_methods, //
 };
 
 pub(crate) struct RustContextMethods<M> {
+    pub(crate) private_key_methods: Option<Box<dyn PrivateKeyDelegate>>,
     pub(crate) verify_certificate_methods: Option<Box<dyn VerifyCertificate>>,
     pub(crate) early_callback_handler: Option<Box<dyn EarlyCallback<M>>>,
     _p: PhantomData<fn() -> M>,
@@ -45,6 +55,7 @@
 impl<M> RustContextMethods<M> {
     pub fn new() -> Self {
         Self {
+            private_key_methods: None,
             verify_certificate_methods: None,
             early_callback_handler: None,
             _p: PhantomData,
@@ -68,6 +79,12 @@
     }
 }
 
+impl<M: HasTlsContextMethod> PrivateKeyMethods for RustContextMethods<M> {
+    fn private_key_methods(&self) -> Option<&dyn PrivateKeyDelegate> {
+        self.private_key_methods.as_deref()
+    }
+}
+
 impl<M: HasTlsContextMethod> VerifyCertificateMethods for RustContextMethods<M> {
     fn verify_certificate_methods(&self) -> Option<&dyn VerifyCertificate> {
         self.verify_certificate_methods.as_deref()
@@ -125,3 +142,35 @@
         *TLS_CONTEXT_METHOD
     }
 }
+
+pub(super) trait HasPrivateKeyMethods {
+    const METHODS: *const bssl_sys::SSL_PRIVATE_KEY_METHOD;
+}
+
+impl HasPrivateKeyMethods for TlsMode {
+    const METHODS: *const bssl_sys::SSL_PRIVATE_KEY_METHOD = {
+        &bssl_sys::SSL_PRIVATE_KEY_METHOD {
+            sign: Some(sign::<RustContextMethods<TlsMode>>),
+            decrypt: Some(decrypt::<RustContextMethods<TlsMode>>),
+            complete: Some(complete::<RustContextMethods<TlsMode>>),
+        } as _
+    };
+}
+impl HasPrivateKeyMethods for DtlsMode {
+    const METHODS: *const bssl_sys::SSL_PRIVATE_KEY_METHOD = {
+        &bssl_sys::SSL_PRIVATE_KEY_METHOD {
+            sign: Some(sign::<RustContextMethods<DtlsMode>>),
+            decrypt: Some(decrypt::<RustContextMethods<DtlsMode>>),
+            complete: Some(complete::<RustContextMethods<DtlsMode>>),
+        } as _
+    };
+}
+impl HasPrivateKeyMethods for QuicMode {
+    const METHODS: *const bssl_sys::SSL_PRIVATE_KEY_METHOD = {
+        &bssl_sys::SSL_PRIVATE_KEY_METHOD {
+            sign: Some(sign::<RustContextMethods<QuicMode>>),
+            decrypt: Some(decrypt::<RustContextMethods<QuicMode>>),
+            complete: Some(complete::<RustContextMethods<QuicMode>>),
+        } as _
+    };
+}
diff --git a/rust/bssl-tls/src/credentials.rs b/rust/bssl-tls/src/credentials.rs
index 2a05af8..7fe822b 100644
--- a/rust/bssl-tls/src/credentials.rs
+++ b/rust/bssl-tls/src/credentials.rs
@@ -120,6 +120,20 @@
         assert_eq!(rc, 1);
         self
     }
+
+    fn get_credential_methods(&mut self) -> &mut methods::RustCredentialMethods {
+        let methods = unsafe {
+            // Safety: the validity of the handle `self.0` is witnessed by `self`.
+            bssl_sys::SSL_CREDENTIAL_get_ex_data(self.ptr(), *methods::TLS_CREDENTIAL_METHOD)
+        };
+        if methods.is_null() {
+            panic!("context method goes missing")
+        }
+        unsafe {
+            // Safety: `methods` must be constructed by `new_inner`
+            &mut *(methods as *mut methods::RustCredentialMethods)
+        }
+    }
 }
 
 impl TlsCredentialBuilder<X509Mode> {
@@ -186,6 +200,30 @@
         Ok(self)
     }
 
+    /// Set private key delegate.
+    ///
+    /// This will override the `TlsConnection` private key delegate.
+    pub fn with_private_key_delegate<T: 'static + PrivateKeyDelegate>(
+        &mut self,
+        key_method: Option<T>,
+    ) -> &mut Self {
+        let cred = self.ptr();
+        if let Some(key_method) = key_method {
+            unsafe {
+                // Safety: we only install our own vtable.
+                bssl_sys::SSL_CREDENTIAL_set_private_key_method(cred, methods::PRIVATE_KEY_METHODS);
+            }
+            self.get_credential_methods().private_key_methods = Some(Box::new(key_method) as _);
+        } else {
+            unsafe {
+                // Safety: we only uninstall the vtable.
+                bssl_sys::SSL_CREDENTIAL_set_private_key_method(cred, core::ptr::null());
+            }
+            self.get_credential_methods().private_key_methods.take();
+        }
+        self
+    }
+
     /// Set a private key.
     ///
     /// **NOTE**: Call this method after setting the certificates with
@@ -520,6 +558,142 @@
     pub struct DelegatedCredential
 }
 
+/// Private key operation result.
+pub enum PrivateKeyOperationResult {
+    /// Private key operation is still in-flight.
+    Pending,
+    /// Private key operation has successfully completed, with a result of a certain size.
+    Success(usize),
+    /// Private key operation has failed.
+    Error,
+}
+
+/// Signature operation parameters.
+#[non_exhaustive]
+pub struct SignatureOperation<'a> {
+    /// Output location of the signature.
+    pub output: &'a mut [u8],
+    /// Input message to be signed.
+    pub message: &'a [u8],
+    /// Signature algorithm to use.
+    pub algorithm: SignatureAlgorithm,
+}
+
+/// Decryption operation parameters.
+#[non_exhaustive]
+pub struct DecryptionOperation<'a> {
+    /// Output location of the plaintext.
+    pub output: &'a mut [u8],
+    /// Input ciphertext.
+    pub ciphertext: &'a [u8],
+}
+
+/// Private key operation delegate.
+///
+/// This protocol allows for asynchronous signing and decryption operations.
+/// BoringSSL will call one of [`Self::sign`] or [`Self::decrypt`] for operation initiation,
+/// and call [`Self::complete`] to poll for completion as long as
+/// [`PrivateKeyOperationResult::Pending`] is returned.
+pub trait PrivateKeyDelegate: Send + Sync {
+    /// Sign operation.
+    ///
+    /// The underlying task will be immediately polled once as optimisation,
+    /// in case the operation becomes ready instantly.
+    fn sign(&self, sign_op: SignatureOperation<'_>) -> Box<dyn PrivateKeyOperation>;
+    /// Decryption operation.
+    ///
+    /// The underlying task will be immediately polled once as optimisation,
+    /// in case the operation becomes ready instantly.
+    fn decrypt(&self, decrypt_op: DecryptionOperation<'_>) -> Box<dyn PrivateKeyOperation>;
+}
+
+// NOTE: only `Send` bound is necessary; BoringSSL will not drive an operation from
+// multiple threads.
+/// An outstanding private key operation.
+pub trait PrivateKeyOperation: Send {
+    /// Try to complete the operation.
+    ///
+    /// To complete the operation, the implementation should write the results into `output` and
+    /// signal success by returning [`PrivateKeyOperationResult::Success`] with the number of bytes
+    /// written to the `output`.
+    fn complete(
+        &mut self,
+        context: Option<&mut Context<'_>>,
+        output: &mut [u8],
+    ) -> PrivateKeyOperationResult;
+}
+
+// NOTE: we require `Send + Sync + Unpin` mostly due to practicality of working with
+// async futures.
+/// Asynchronous privatge key operation delegate.
+///
+/// This is the `async` analogue of [`PrivateKeyDelegate`].
+pub trait AsyncPrivateKeyDelegate: Send + Sync + Unpin {
+    /// The future to drive the signing operation.
+    type SignOp: 'static + Unpin + Send + Sync + Future<Output = Option<Vec<u8>>>;
+    /// The future to drive the decryption operation.
+    type DecryptOp: 'static + Unpin + Send + Sync + Future<Output = Option<Vec<u8>>>;
+    /// Sign operation.
+    ///
+    /// The underlying future will be immediately polled once as optimisation,
+    /// in case the operation becomes ready immediately.
+    fn sign(&self, message: &[u8], algorithm: SignatureAlgorithm) -> Self::SignOp;
+    /// Decryption operations.
+    ///
+    /// The underlying future will be immediately polled once as optimisation,
+    /// in case the operation becomes ready immediately.
+    fn decrypt(&self, ciphertext: &[u8]) -> Self::DecryptOp;
+}
+
+/// A convenient asynchronous private key decrypter adapter.
+pub struct AsyncPrivateKeyDelegateAdapter<Inner>(pub Inner);
+
+impl<Inner> PrivateKeyDelegate for AsyncPrivateKeyDelegateAdapter<Inner>
+where
+    Inner: AsyncPrivateKeyDelegate,
+{
+    fn sign(&self, sign_op: SignatureOperation<'_>) -> Box<dyn PrivateKeyOperation> {
+        Box::new(AsyncPrivateKeyOperationAdapter(
+            self.0.sign(sign_op.message, sign_op.algorithm),
+        ))
+    }
+
+    fn decrypt(&self, decrypt_op: DecryptionOperation<'_>) -> Box<dyn PrivateKeyOperation> {
+        Box::new(AsyncPrivateKeyOperationAdapter(
+            self.0.decrypt(decrypt_op.ciphertext),
+        ))
+    }
+}
+
+struct AsyncPrivateKeyOperationAdapter<Fut>(Fut);
+
+impl<Fut> PrivateKeyOperation for AsyncPrivateKeyOperationAdapter<Fut>
+where
+    Fut: Unpin + Send + Sync + Future<Output = Option<Vec<u8>>>,
+{
+    fn complete(
+        &mut self,
+        context: Option<&mut Context<'_>>,
+        output: &mut [u8],
+    ) -> PrivateKeyOperationResult {
+        let Some(cx) = context else {
+            return PrivateKeyOperationResult::Error;
+        };
+        match Pin::new(&mut self.0).poll(cx) {
+            Poll::Ready(None) => PrivateKeyOperationResult::Error,
+            Poll::Ready(Some(res)) => {
+                if res.len() > output.len() {
+                    PrivateKeyOperationResult::Error
+                } else {
+                    output[..res.len()].copy_from_slice(&res);
+                    PrivateKeyOperationResult::Success(res.len())
+                }
+            }
+            Poll::Pending => PrivateKeyOperationResult::Pending,
+        }
+    }
+}
+
 bssl_macros::bssl_enum! {
     /// [IANA] designation of TLS signature algorithms.
     ///
@@ -800,7 +974,7 @@
 
 impl ExactSizeIterator for CertificateChainIterator<'_> {
     fn len(&self) -> usize {
-        self.len
+        self.len - self.curr
     }
 }
 
diff --git a/rust/bssl-tls/src/credentials/methods.rs b/rust/bssl-tls/src/credentials/methods.rs
index 6ca7449..11bc865 100644
--- a/rust/bssl-tls/src/credentials/methods.rs
+++ b/rust/bssl-tls/src/credentials/methods.rs
@@ -12,15 +12,34 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use alloc::boxed::Box;
 use core::{
     ffi::c_int,
-    ptr::null_mut, //
+
+    ptr::{
+        NonNull,
+        null_mut, //
+    }, //
 };
 
 use once_cell::sync::Lazy;
 
 use crate::{
     Methods,
+    PrivateKeyMethods,
+    abort_on_panic,
+    connection::methods::private_key_op_from_ssl,
+    credentials::{
+        DecryptionOperation,
+        PrivateKeyDelegate,
+        PrivateKeyOperationResult,
+        SignatureOperation,
+        waker_data_from_ssl, //
+    },
+    ffi::{
+        sanitise_mut_byteslice,
+        sanitize_slice, //
+    },
     methods::drop_box_rust_methods, //
 };
 
@@ -41,7 +60,9 @@
 });
 
 #[derive(Default)]
-pub(crate) struct RustCredentialMethods {}
+pub(crate) struct RustCredentialMethods {
+    pub(crate) private_key_methods: Option<Box<dyn PrivateKeyDelegate>>,
+}
 
 impl Methods for RustCredentialMethods {
     unsafe extern "C" fn from_ssl<'a>(ssl: *mut bssl_sys::SSL) -> Option<&'a Self> {
@@ -61,3 +82,212 @@
         }
     }
 }
+
+impl PrivateKeyMethods for RustCredentialMethods {
+    fn private_key_methods(&self) -> Option<&dyn PrivateKeyDelegate> {
+        self.private_key_methods.as_deref()
+    }
+}
+
+macro_rules! private_key_method_prelude {
+    ($M:ty, $ssl:ident => $context:ident, $private_key_methods:ident, $task:ident) => {
+        let Some(ssl) = NonNull::new($ssl) else {
+            return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+        };
+        let Some(methods) = (unsafe {
+            // Safety: `ssl` outlives `methods` and must be valid by BoringSSL contract.
+            <$M>::from_ssl(ssl.as_ptr())
+        }) else {
+            return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+        };
+        let waker = unsafe {
+            // Safety: `ssl` must be constructed by `TlsConnection` as this is called from
+            waker_data_from_ssl(ssl)
+        };
+        let mut $context = if let Some(waker) = &waker {
+            Some(core::task::Context::from_waker(waker))
+        } else {
+            None
+        };
+        #[allow(unused_mut)]
+        let $task = unsafe {
+            // Safety: `ssl` must be constructed by `TlsConnection` as this is called from
+            private_key_op_from_ssl(ssl)
+        };
+        let Some($private_key_methods) = methods.private_key_methods() else {
+            return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+        };
+    };
+}
+
+pub(crate) unsafe extern "C" fn sign<Method: PrivateKeyMethods>(
+    ssl: *mut bssl_sys::SSL,
+    out: *mut u8,
+    out_len: *mut usize,
+    max_out: usize,
+    sig_alg: u16,
+    msg: *const u8,
+    msg_len: usize,
+) -> bssl_sys::ssl_private_key_result_t {
+    private_key_method_prelude!(Method, ssl => context, private_key_methods, task);
+    if task.is_some() {
+        abort_on_panic(|| {
+            let _ = task.take();
+        });
+        return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+    }
+
+    // Unwind-safety: later when panic happens, we detach the poisoned private key method
+    // without calling destructor.
+    let algorithm = match sig_alg.try_into() {
+        Ok(sig_alg) => sig_alg,
+        // TODO(@xfding) maybe we should log this error?
+        Err(_) => return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure,
+    };
+    let Some(output) = (unsafe {
+        // Safety: the slice will only be used within this callback.
+        sanitise_mut_byteslice(out, max_out)
+    }) else {
+        return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+    };
+    if output.is_empty() {
+        return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+    }
+    let Some(message) = (unsafe {
+        // Safety: `msg` outlives `message` because it is owned by BoringSSL.
+        sanitize_slice(msg, msg_len)
+    }) else {
+        return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+    };
+    // Unwind-safety: when panic happens, we will not inspect the `output` buffer.
+    abort_on_panic(move || {
+        let sign_op = SignatureOperation {
+            output,
+            message,
+            algorithm,
+        };
+        let mut outstanding_task = private_key_methods.sign(sign_op);
+        match outstanding_task.complete(context.as_mut(), output) {
+            PrivateKeyOperationResult::Success(len) => {
+                unsafe {
+                    // Safety: `out_len` is a valid pointer by BoringSSL invariant.
+                    *out_len = len;
+                }
+                bssl_sys::ssl_private_key_result_t_ssl_private_key_success
+            }
+            PrivateKeyOperationResult::Pending => {
+                *task = Some(outstanding_task);
+                bssl_sys::ssl_private_key_result_t_ssl_private_key_retry
+            }
+            PrivateKeyOperationResult::Error => {
+                bssl_sys::ssl_private_key_result_t_ssl_private_key_failure
+            }
+        }
+    })
+}
+
+pub(crate) unsafe extern "C" fn decrypt<Method: PrivateKeyMethods>(
+    ssl: *mut bssl_sys::SSL,
+    out: *mut u8,
+    out_len: *mut usize,
+    max_out: usize,
+    ciphertext: *const u8,
+    ciphertext_len: usize,
+) -> bssl_sys::ssl_private_key_result_t {
+    private_key_method_prelude!(Method, ssl => context, private_key_methods, task);
+    if task.is_some() {
+        abort_on_panic(|| {
+            let _ = task.take();
+        });
+        return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+    }
+
+    // Unwind-safety: later when panic happens, we detach the poisoned private key method
+    // without calling destructor.
+    let Some(output) = (unsafe {
+        // Safety: the slice will only be used within this callback.
+        sanitise_mut_byteslice(out, max_out)
+    }) else {
+        return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+    };
+    if output.is_empty() {
+        return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+    }
+    let Some(ciphertext) = (unsafe {
+        // Safety: `ciphertext` is now owned by BoringSSL and outlives the slice.
+        sanitize_slice(ciphertext, ciphertext_len)
+    }) else {
+        return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+    };
+    // Unwind-safety: when panic happens, we will not inspect the `output` buffer.
+    abort_on_panic(move || {
+        let decrypt_op = DecryptionOperation { output, ciphertext };
+        let mut outstanding_task = private_key_methods.decrypt(decrypt_op);
+        match outstanding_task.complete(context.as_mut(), output) {
+            PrivateKeyOperationResult::Success(len) => {
+                unsafe {
+                    // Safety: `out_len` is a valid pointer by BoringSSL invariant.
+                    *out_len = len;
+                }
+                bssl_sys::ssl_private_key_result_t_ssl_private_key_success
+            }
+            PrivateKeyOperationResult::Pending => {
+                *task = Some(outstanding_task);
+                bssl_sys::ssl_private_key_result_t_ssl_private_key_retry
+            }
+            PrivateKeyOperationResult::Error => {
+                bssl_sys::ssl_private_key_result_t_ssl_private_key_failure
+            }
+        }
+    })
+}
+
+pub(crate) unsafe extern "C" fn complete<Method: PrivateKeyMethods>(
+    ssl: *mut bssl_sys::SSL,
+    out: *mut u8,
+    out_len: *mut usize,
+    max_out: usize,
+) -> bssl_sys::ssl_private_key_result_t {
+    private_key_method_prelude!(Method, ssl => context, _private_key_methods, task);
+
+    // Unwind-safety: later when panic happens, we detach the poisoned private key method
+    // without calling destructor.
+    let Some(output) = (unsafe {
+        // Safety: the slice will only be used within this callback.
+        sanitise_mut_byteslice(out, max_out)
+    }) else {
+        return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+    };
+    if output.is_empty() {
+        return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+    }
+    // Unwind-safety: when panic happens, we will not inspect the `output` buffer.
+    abort_on_panic(move || {
+        let Some(outstanding_task) = task else {
+            return bssl_sys::ssl_private_key_result_t_ssl_private_key_failure;
+        };
+        match outstanding_task.complete(context.as_mut(), output) {
+            PrivateKeyOperationResult::Success(len) => {
+                unsafe {
+                    // Safety: `out_len` is a valid pointer by BoringSSL invariant.
+                    *out_len = len;
+                }
+                bssl_sys::ssl_private_key_result_t_ssl_private_key_success
+            }
+            PrivateKeyOperationResult::Pending => {
+                bssl_sys::ssl_private_key_result_t_ssl_private_key_retry
+            }
+            PrivateKeyOperationResult::Error => {
+                bssl_sys::ssl_private_key_result_t_ssl_private_key_failure
+            }
+        }
+    })
+}
+
+pub(super) const PRIVATE_KEY_METHODS: *const bssl_sys::SSL_PRIVATE_KEY_METHOD = {
+    &bssl_sys::SSL_PRIVATE_KEY_METHOD {
+        sign: Some(sign::<RustCredentialMethods>),
+        decrypt: Some(decrypt::<RustCredentialMethods>),
+        complete: Some(complete::<RustCredentialMethods>),
+    } as _
+};
diff --git a/rust/bssl-tls/src/lib.rs b/rust/bssl-tls/src/lib.rs
index ac98e1f..54fd6bd 100644
--- a/rust/bssl-tls/src/lib.rs
+++ b/rust/bssl-tls/src/lib.rs
@@ -65,6 +65,10 @@
     unsafe extern "C" fn from_ssl<'a>(ssl: *mut bssl_sys::SSL) -> Option<&'a Self>;
 }
 
+pub(crate) trait PrivateKeyMethods: Methods {
+    fn private_key_methods(&self) -> Option<&dyn credentials::PrivateKeyDelegate>;
+}
+
 pub(crate) trait VerifyCertificateMethods: Methods {
     fn verify_certificate_methods(&self) -> Option<&dyn credentials::VerifyCertificate>;
 }
diff --git a/rust/bssl-tls/src/tests.rs b/rust/bssl-tls/src/tests.rs
index f7e92ea..83c7f5f 100644
--- a/rust/bssl-tls/src/tests.rs
+++ b/rust/bssl-tls/src/tests.rs
@@ -58,7 +58,12 @@
     include_bytes!("../../test-data/BoringSSLServerTest-RSA.crt");
 pub(crate) const RSA_SERVER_KEY: &[u8] =
     include_bytes!("../../test-data/BoringSSLServerTest-RSA.key");
+pub(crate) const P256_SERVER_CERT: &[u8] =
+    include_bytes!("../../test-data/BoringSSLServerTest-ECDSA-P256.crt");
+pub(crate) const P256_SERVER_KEY_DER: &[u8] =
+    include_bytes!("../../test-data/BoringSSLServerTest-ECDSA-P256.der");
 
+mod credentials;
 mod datagram;
 mod handshake;
 mod transport;
diff --git a/rust/bssl-tls/src/tests/credentials.rs b/rust/bssl-tls/src/tests/credentials.rs
new file mode 100644
index 0000000..0dfe046
--- /dev/null
+++ b/rust/bssl-tls/src/tests/credentials.rs
@@ -0,0 +1,205 @@
+// Copyright 2026 The BoringSSL Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::{
+        Arc,
+        Mutex, //
+    },
+    task::{
+        Context,
+        Poll, //
+    },
+};
+
+use bssl_crypto::ecdsa::ParsedPrivateKey;
+use bssl_x509::{
+    certificates::X509Certificate,
+    params::Trust,
+    store::X509StoreBuilder, //
+};
+use futures::channel::oneshot;
+
+use super::{
+    CA,
+    P256_SERVER_CERT, //
+};
+use crate::{
+    context::TlsContextBuilder,
+    credentials::{
+        AsyncPrivateKeyDelegate,
+        Certificate,
+        CertificateVerificationMode,
+        SignatureAlgorithm,
+        TlsCredentialBuilder, //
+    },
+    errors::TlsRetryReason,
+    io::IoStatus,
+    tests::create_mock_pipe, //
+};
+
+#[test]
+fn test_private_key_methods() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    let ca = Certificate::parse_one_from_pem(CA, None)?;
+    let server_cert = Certificate::parse_one_from_pem(P256_SERVER_CERT, None)?;
+
+    let (client_to_server_tx, client_to_server_rx) = oneshot::channel::<()>();
+    let (server_to_client_tx, server_to_client_rx) = oneshot::channel::<()>();
+
+    let private_key_method = MyPrivateKeyMethod {
+        key: crate::tests::P256_SERVER_KEY_DER,
+        client_to_server_rx: Arc::new(Mutex::new(Some(client_to_server_rx))),
+        server_to_client_tx: Arc::new(Mutex::new(Some(server_to_client_tx))),
+    };
+
+    struct MyPrivateKeyMethod {
+        key: &'static [u8],
+        client_to_server_rx: Arc<Mutex<Option<oneshot::Receiver<()>>>>,
+        server_to_client_tx: Arc<Mutex<Option<oneshot::Sender<()>>>>,
+    }
+
+    impl AsyncPrivateKeyDelegate for MyPrivateKeyMethod {
+        type DecryptOp = Pin<Box<dyn Send + Sync + Future<Output = Option<Vec<u8>>>>>;
+        type SignOp = Pin<Box<dyn Send + Sync + Future<Output = Option<Vec<u8>>>>>;
+        fn sign(&self, message: &[u8], algorithm: SignatureAlgorithm) -> Self::SignOp {
+            let message = message.to_vec();
+            let Some(ParsedPrivateKey::P256(key)) = ParsedPrivateKey::from_der(self.key) else {
+                panic!()
+            };
+            let client_to_server_rx = self.client_to_server_rx.clone();
+            let server_to_client_tx = self.server_to_client_tx.clone();
+            Box::pin(async move {
+                let tx = server_to_client_tx.lock().unwrap().take();
+                if let Some(tx) = tx {
+                    tx.send(()).unwrap();
+                }
+
+                let rx = client_to_server_rx.lock().unwrap().take();
+                if let Some(rx) = rx {
+                    rx.await.unwrap();
+                }
+
+                assert!(matches!(
+                    algorithm,
+                    SignatureAlgorithm::EcdsaSecp256r1Sha256
+                ));
+                Some(key.sign(&message))
+            })
+        }
+
+        fn decrypt(&self, _: &[u8]) -> Self::DecryptOp {
+            unreachable!()
+        }
+    }
+
+    let mut server_ctx_builder = TlsContextBuilder::new_tls();
+    let server_cred = {
+        let mut builder = TlsCredentialBuilder::new();
+        builder
+            .with_certificate_chain(&[server_cert, ca])?
+            .with_private_key_delegate(Some(crate::credentials::AsyncPrivateKeyDelegateAdapter(
+                private_key_method,
+            )));
+        builder.build().unwrap()
+    };
+    server_ctx_builder.with_credential(server_cred)?;
+    let server_ctx = server_ctx_builder.build();
+    let mut server_conn = server_ctx.new_server_connection(None)?.build();
+
+    let mut client_ctx_builder = TlsContextBuilder::new_tls();
+    let mut cert_store = X509StoreBuilder::new();
+    cert_store
+        .set_trust(Trust::SslServer)?
+        .add_cert(X509Certificate::parse_one_from_pem(CA)?)?;
+    let cert_store = cert_store.build();
+    client_ctx_builder.with_certificate_store(&cert_store);
+    let client_ctx = client_ctx_builder.build();
+    let mut client_conn = client_ctx.new_client_connection(None)?;
+    client_conn.with_certificate_verification_mode(CertificateVerificationMode::PeerCertMandatory);
+    let mut client_conn = client_conn.build();
+    client_conn
+        .in_handshake()
+        .unwrap()
+        .set_host("www.google.com")?;
+
+    let (client_socket, server_socket, mut executor) = create_mock_pipe();
+    client_conn.set_io(client_socket)?;
+    server_conn.set_io(server_socket)?;
+
+    let mut server_task = async move || -> Result<(), crate::errors::Error> {
+        let mut in_handshake = server_conn.in_handshake().unwrap();
+
+        loop {
+            match in_handshake.async_handshake().await {
+                Ok(None) => break,
+                Ok(Some(TlsRetryReason::PendingPrivateKeyOperation)) => {
+                    struct Yield(bool);
+                    impl std::future::Future for Yield {
+                        type Output = ();
+                        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+                            if self.0 {
+                                return Poll::Ready(());
+                            }
+                            self.0 = true;
+                            cx.waker().wake_by_ref();
+                            Poll::Pending
+                        }
+                    }
+                    Yield(false).await;
+                }
+                res => panic!("Unexpected handshake result: {:?}", res),
+            }
+        }
+
+        let mut message = [0; 21];
+        assert!(matches!(
+            server_conn.as_pin_mut().async_read(&mut message).await?,
+            IoStatus::Ok(21)
+        ));
+        assert_eq!(message, *b"BoringSSL is awesome!");
+        server_conn.as_pin_mut().async_shutdown().await?;
+        Ok(())
+    };
+
+    let client_task = async move || -> Result<(), crate::errors::Error> {
+        let res = client_conn.in_handshake().unwrap().do_handshake();
+        assert!(
+            matches!(res, Ok(Some(TlsRetryReason::WantRead))),
+            "Expected WantRead, got {:?}",
+            res
+        );
+
+        server_to_client_rx.await.unwrap();
+        client_to_server_tx.send(()).unwrap();
+
+        client_conn
+            .as_pin_mut()
+            .async_write(b"BoringSSL is awesome!")
+            .await?;
+
+        client_conn.as_pin_mut().async_shutdown().await?;
+        Ok(())
+    };
+
+    let test_closure = async move || -> Result<(), crate::errors::Error> {
+        futures::future::try_join(server_task(), client_task()).await?;
+        Ok(())
+    };
+
+    executor.run(test_closure())?;
+
+    Ok(())
+}