rust: bssl-tls: RPK with private key methods

From Android feedback, their RPK use case requires off-application key
signing, so we should expose this API to them.

Update-Note: RPK credential must install its private key delegate at
construction time. The previous API made the installation a no-op.

Signed-off-by: Xiangfei Ding <xfding@google.com>
Change-Id: I60cfc6bb458872320a54c8df8258fcfd6a6a6964
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/99188
Reviewed-by: David Benjamin <davidben@google.com>
Reviewed-by: Adam Langley <agl@google.com>
diff --git a/rust/bssl-tls/src/credentials.rs b/rust/bssl-tls/src/credentials.rs
index b2b4db1..15a1bb7 100644
--- a/rust/bssl-tls/src/credentials.rs
+++ b/rust/bssl-tls/src/credentials.rs
@@ -42,7 +42,7 @@
 
 use bssl_x509::{
     errors::PemReason,
-    keys::PrivateKey, //
+    keys::{PrivateKey, PublicKey},
 };
 
 use crate::{
@@ -86,10 +86,10 @@
 /// Raw Public Key credential
 pub enum RawPublicKeyMode {}
 
-pub(crate) trait NeedsPrivateKey {}
+pub(crate) trait HasPrivateKey {}
 
-impl NeedsPrivateKey for X509Mode {}
-impl NeedsPrivateKey for RawPublicKeyMode {}
+impl HasPrivateKey for X509Mode {}
+impl HasPrivateKey for RawPublicKeyMode {}
 
 // Safety: At this type state, the credential handle is exclusively owned.
 unsafe impl<M> Send for TlsCredentialBuilder<M> {}
@@ -152,56 +152,6 @@
         );
         this.set_ex_data()
     }
-}
-
-impl TlsCredentialBuilder<RawPublicKeyMode> {
-    /// Construct raw public key credential instance.
-    ///
-    /// TLS connection may use this credential to perform authentication per [RFC 7250].
-    ///
-    /// [RFC 7250]: <https://tools.ietf.org/html/rfc7250>
-    pub fn new_raw_public_key(mut key: PrivateKey) -> Self {
-        let this = Self(
-            NonNull::new(unsafe {
-                // Safety:
-                // - the `PrivateKey` type already contains both public and private key parts.
-                // - the constructor call also claims the ownership of the key.
-                bssl_sys::SSL_CREDENTIAL_new_raw_public_key(key.as_mut_ptr())
-            })
-            .expect("allocation failure"),
-            PhantomData,
-        );
-        this.set_ex_data()
-    }
-}
-
-impl<M> TlsCredentialBuilder<M>
-where
-    M: NeedsPrivateKey,
-{
-    /// Set [`SignatureAlgorithm`] preferences.
-    ///
-    /// This controls which signature algorithms will be used with this credential.
-    pub fn with_signing_algorithm_preferences(
-        &mut self,
-        algs: &[SignatureAlgorithm],
-    ) -> Result<&mut Self, Error> {
-        let algs: &[u16] = unsafe {
-            // Safety: `SignatureAlgorithm` has a `repr(u16)`
-            core::mem::transmute(algs)
-        };
-        if has_duplicates(algs) {
-            return Err(Error::Configuration(
-                ConfigurationError::DuplicatedParameters,
-            ));
-        }
-        let (ptr, len) = slice_into_ffi_raw_parts(algs);
-        check_lib_error!(unsafe {
-            // Safety
-            bssl_sys::SSL_CREDENTIAL_set1_signing_algorithm_prefs(self.ptr(), ptr, len)
-        });
-        Ok(self)
-    }
 
     /// Set private key delegate.
     ///
@@ -248,9 +198,6 @@
             Ok(self)
         }
     }
-}
-
-impl TlsCredentialBuilder<X509Mode> {
     /// Set certificate chain.
     ///
     /// The leaf, also known as end-entity, certificate **must come first** in `certs`.
@@ -306,6 +253,85 @@
     }
 }
 
+impl TlsCredentialBuilder<RawPublicKeyMode> {
+    /// Construct raw public key credential instance.
+    ///
+    /// TLS connection may use this credential to perform authentication per [RFC 7250].
+    ///
+    /// [RFC 7250]: <https://tools.ietf.org/html/rfc7250>
+    pub fn new_raw_public_key(mut key: PrivateKey) -> Self {
+        let this = Self(
+            NonNull::new(unsafe {
+                // Safety:
+                // - the `PrivateKey` type already contains both public and private key parts.
+                // - the constructor call also claims the ownership of the key.
+                bssl_sys::SSL_CREDENTIAL_new_raw_public_key(key.as_mut_ptr())
+            })
+            .expect("allocation failure"),
+            PhantomData,
+        );
+        this.set_ex_data()
+    }
+
+    /// Construct raw public key credential, with a private key delegate.
+    ///
+    /// This credential is suitable for cases where the private key is stored off the device or in a
+    /// credential store.
+    /// TLS connection may use this credential to perform authentication per [RFC 7250].
+    ///
+    /// [RFC 7250]: <https://tools.ietf.org/html/rfc7250>
+    pub fn new_raw_public_key_with_delegate<T: PrivateKeyDelegate + 'static>(
+        mut key: PublicKey,
+        delegate: T,
+    ) -> Self {
+        let mut this = Self(
+            NonNull::new(unsafe {
+                // Safety:
+                // - the `PublicKey` type contains the public key part.
+                // - the constructor bumps the ref-count on the public key.
+                bssl_sys::SSL_CREDENTIAL_new_raw_public_key_custom(
+                    key.as_mut_ptr(),
+                    methods::PRIVATE_KEY_METHODS,
+                )
+            })
+            .expect("allocation failure"),
+            PhantomData,
+        )
+        .set_ex_data();
+        this.get_credential_methods().private_key_methods = Some(Box::new(delegate) as _);
+        this
+    }
+}
+
+impl<M> TlsCredentialBuilder<M>
+where
+    M: HasPrivateKey,
+{
+    /// Set [`SignatureAlgorithm`] preferences.
+    ///
+    /// This controls which signature algorithms will be used with this credential.
+    pub fn with_signing_algorithm_preferences(
+        &mut self,
+        algs: &[SignatureAlgorithm],
+    ) -> Result<&mut Self, Error> {
+        let algs: &[u16] = unsafe {
+            // Safety: `SignatureAlgorithm` has a `repr(u16)`
+            core::mem::transmute(algs)
+        };
+        if has_duplicates(algs) {
+            return Err(Error::Configuration(
+                ConfigurationError::DuplicatedParameters,
+            ));
+        }
+        let (ptr, len) = slice_into_ffi_raw_parts(algs);
+        check_lib_error!(unsafe {
+            // Safety
+            bssl_sys::SSL_CREDENTIAL_set1_signing_algorithm_prefs(self.ptr(), ptr, len)
+        });
+        Ok(self)
+    }
+}
+
 impl<M> TlsCredentialBuilder<M> {
     /// Finalise the credential.
     pub fn build(mut self) -> Option<TlsCredential> {
diff --git a/rust/bssl-tls/src/credentials/tests.rs b/rust/bssl-tls/src/credentials/tests.rs
index 36e92d6..33eae43 100644
--- a/rust/bssl-tls/src/credentials/tests.rs
+++ b/rust/bssl-tls/src/credentials/tests.rs
@@ -12,13 +12,30 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use bssl_x509::keys::PrivateKeyAlgorithm;
+use core::future::{
+    Ready,
+    ready, //
+};
+
+use bssl_crypto::ecdsa::ParsedPrivateKey;
+use bssl_x509::keys::{
+    PrivateKey,
+    PrivateKeyAlgorithm, //
+};
 use futures::future::try_join;
 
 use super::*;
 use crate::{
-    context::TlsMode,
-    tests::create_mock_pipe, //
+    context::{
+        TlsContextBuilder,
+        TlsMode, //
+    },
+    errors::Error,
+    tests::{
+        P256_SERVER_KEY,
+        P256_SERVER_KEY_DER,
+        create_mock_pipe, //
+    }, //
 };
 
 #[test]
@@ -148,10 +165,10 @@
 
     let cred = TlsCredential::new_pre_shared_key(key, identity, PskHash::Sha256, context)?;
 
-    let mut server_ctx = crate::context::TlsContextBuilder::new_tls();
+    let mut server_ctx = TlsContextBuilder::new_tls();
     server_ctx.with_credential(cred.clone())?;
 
-    let mut client_ctx = crate::context::TlsContextBuilder::new_tls();
+    let mut client_ctx = TlsContextBuilder::new_tls();
     client_ctx.with_credential(cred)?;
 
     let server_ctx = server_ctx.build();
@@ -168,16 +185,16 @@
     let test_future = async {
         let client_handshake = async {
             assert!(client_conn.async_handshake().await?.is_none());
-            Ok::<(), crate::errors::Error>(())
+            Ok::<(), Error>(())
         };
 
         let server_handshake = async {
             assert!(server_conn.async_handshake().await?.is_none());
-            Ok::<(), crate::errors::Error>(())
+            Ok::<(), Error>(())
         };
 
         try_join(client_handshake, server_handshake).await?;
-        Ok::<(), crate::errors::Error>(())
+        Ok::<(), Error>(())
     };
 
     executor.run(test_future)?;
@@ -205,11 +222,11 @@
 
     let cred = TlsCredential::new_pre_shared_key(key, identity, PskHash::Sha256, context)?;
 
-    let mut server_ctx = crate::context::TlsContextBuilder::new_tls();
+    let mut server_ctx = TlsContextBuilder::new_tls();
     server_ctx.with_credential(cred.clone())?;
     let server_ctx = server_ctx.build();
 
-    let mut client_ctx = crate::context::TlsContextBuilder::new_tls();
+    let mut client_ctx = TlsContextBuilder::new_tls();
     client_ctx.with_credential(cred)?;
     let client_ctx = client_ctx.build();
 
@@ -231,44 +248,22 @@
     Ok(())
 }
 
-unsafe extern "C" fn accept_any_verify(
-    _ssl: *mut bssl_sys::SSL,
-    _out_alert: *mut u8,
-) -> bssl_sys::ssl_verify_result_t {
-    bssl_sys::ssl_verify_result_t_ssl_verify_ok
-}
-
-#[test]
-fn rpk_tls13_handshake() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
-    use bssl_x509::keys::PrivateKey;
-
-    let priv_key = PrivateKey::from_pem(crate::tests::RSA_SERVER_KEY, || unreachable!()).unwrap();
-
-    let cred = TlsCredentialBuilder::<RawPublicKeyMode>::new_raw_public_key(priv_key)
-        .build()
-        .unwrap();
-
-    let mut server_ctx = crate::context::TlsContextBuilder::new_tls();
-    server_ctx.with_credential(cred.clone())?;
+/// Helper: perform an RPK handshake with the given server credential and client verifier.
+fn rpk_handshake_test(
+    cred: TlsCredential,
+    verifier: impl VerifyCertificate + 'static,
+) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    let mut server_ctx = TlsContextBuilder::new_tls();
+    server_ctx.with_credential(cred)?;
     let server_ctx = server_ctx.build();
 
-    let mut client_ctx = crate::context::TlsContextBuilder::new_tls();
-    client_ctx.with_accepted_peer_cert_types(&[CertificateType::Rpk])?;
+    let mut client_ctx = TlsContextBuilder::new_tls();
+    client_ctx
+        .with_accepted_peer_cert_types(&[CertificateType::Rpk])?
+        .with_certificate_verifier(CertificateVerificationMode::PeerCertMandatory, verifier);
     let client_ctx = client_ctx.build();
 
-    let mut client_conn_builder = client_ctx.new_client_connection();
-    client_conn_builder.with_certificate_verification_mode(CertificateVerificationMode::None);
-    let mut client_conn = client_conn_builder.build();
-    unsafe {
-        // Safety:
-        // - `client_conn.ptr()` is a valid `SSL` handle.
-        // - `accept_any_verify` is a valid callback.
-        bssl_sys::SSL_set_custom_verify(
-            client_conn.ptr(),
-            bssl_sys::SSL_VERIFY_PEER as _,
-            Some(accept_any_verify),
-        );
-    }
+    let mut client_conn = client_ctx.new_client_connection().build();
     let mut server_conn = server_ctx.new_server_connection().build();
 
     let (sock_client, sock_server, mut executor) = create_mock_pipe();
@@ -289,6 +284,61 @@
     Ok(())
 }
 
+#[test]
+fn rpk_tls13_handshake() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    let priv_key = PrivateKey::from_pem(crate::tests::RSA_SERVER_KEY, || unreachable!()).unwrap();
+
+    let cred = TlsCredentialBuilder::new_raw_public_key(priv_key)
+        .build()
+        .unwrap();
+
+    rpk_handshake_test(cred, AcceptAnyVerifier)
+}
+
+#[test]
+fn rpk_tls13_handshake_with_delegate() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    use crate::credentials::{
+        AsyncPrivateKeyDelegate, AsyncPrivateKeyDelegateAdapter, SignatureAlgorithm,
+    };
+
+    let priv_key = PrivateKey::from_pem(P256_SERVER_KEY, || unreachable!()).unwrap();
+    let pub_key = priv_key.to_public_key().clone();
+    let expected_spki_der = pub_key.to_der();
+
+    struct P256KeyDelegate {
+        key_der: &'static [u8],
+    }
+
+    impl AsyncPrivateKeyDelegate for P256KeyDelegate {
+        type SignOp = Ready<Option<Vec<u8>>>;
+        type DecryptOp = Ready<Option<Vec<u8>>>;
+
+        fn sign(&self, message: &[u8], algorithm: SignatureAlgorithm) -> Self::SignOp {
+            assert!(matches!(
+                algorithm,
+                SignatureAlgorithm::EcdsaSecp256r1Sha256
+            ));
+            let Some(ParsedPrivateKey::P256(key)) = ParsedPrivateKey::from_der(self.key_der) else {
+                panic!("failed to parse P256 key")
+            };
+            ready(Some(key.sign(message)))
+        }
+
+        fn decrypt(&self, _ciphertext: &[u8]) -> Self::DecryptOp {
+            unreachable!()
+        }
+    }
+
+    let delegate = AsyncPrivateKeyDelegateAdapter(P256KeyDelegate {
+        key_der: P256_SERVER_KEY_DER,
+    });
+    let cred = TlsCredentialBuilder::new_raw_public_key_with_delegate(pub_key, delegate)
+        .build()
+        .unwrap();
+
+    rpk_handshake_test(cred, RpkVerifier { expected_spki_der })
+}
+
 struct AcceptAnyVerifier;
 
 impl VerifyCertificate for AcceptAnyVerifier {
@@ -389,6 +439,7 @@
                 .is_some();
 
             if offered_rpk {
+                // TODO: switch to `try` block on stabilisation
                 if (|| {
                     client_hello
                         .connection_mut()
@@ -396,7 +447,7 @@
                         .set_certificate_verification_mode(
                             CertificateVerificationMode::PeerCertMandatory,
                         );
-                    Ok::<_, crate::errors::Error>(())
+                    Ok::<_, Error>(())
                 })()
                 .is_err()
                 {
@@ -478,7 +529,7 @@
 
             if !offer_psk && !offer_rpk {
                 assert!(handshake_result.is_err(), "Handshake should have failed");
-                return Ok::<(), crate::errors::Error>(());
+                return Ok::<(), Error>(());
             }
 
             handshake_result?;
@@ -497,7 +548,7 @@
             assert!(matches!(read_len, IoStatus::Ok(5)));
             assert_eq!(&client_buf, b"world");
 
-            Ok::<(), crate::errors::Error>(())
+            Ok::<(), Error>(())
         };
 
         executor.run(test_future)?;
diff --git a/rust/bssl-tls/src/tests.rs b/rust/bssl-tls/src/tests.rs
index 17997a0..f8413dc 100644
--- a/rust/bssl-tls/src/tests.rs
+++ b/rust/bssl-tls/src/tests.rs
@@ -60,6 +60,8 @@
     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: &[u8] =
+    include_bytes!("../../test-data/BoringSSLServerTest-ECDSA-P256.key");
 pub(crate) const P256_SERVER_KEY_DER: &[u8] =
     include_bytes!("../../test-data/BoringSSLServerTest-ECDSA-P256.der");