Return SharedSecret from ActiveKeyExchange::complete

provider-example/src/kx.rs

@@ -10,15 +10,13 @@ impl crypto::ActiveKeyExchange for KeyExchange {
     fn complete(
         self: Box<KeyExchange>,
         peer: &[u8],
-        sink: &mut dyn FnMut(&[u8]),
-    ) -> Result<(), rustls::Error> {
+    ) -> Result<crypto::SharedSecret, rustls::Error> {
         let peer_array: [u8; 32] = peer
             .try_into()
             .map_err(|_| rustls::Error::from(rustls::PeerMisbehaved::InvalidKeyShare))?;
         let their_pub = x25519_dalek::PublicKey::from(peer_array);
         let shared_secret = self.priv_key.diffie_hellman(&their_pub);
-        sink(shared_secret.as_bytes());
-        Ok(())
+        Ok(crypto::SharedSecret::from(&shared_secret.as_bytes()[..]))
     }
 
     fn pub_key(&self) -> &[u8] {

rustls/src/client/tls13.rs

@@ -150,18 +150,8 @@ pub(super) fn handle_server_hello(
         KeySchedulePreHandshake::new(suite)
     };
 
-    let mut kx_input = Some(key_schedule_pre_handshake);
-    let mut kx_output = None;
-
-    our_key_share.complete(&their_key_share.payload.0, &mut |secret| {
-        kx_output = Some(
-            kx_input
-                .take()
-                .unwrap()
-                .into_handshake(secret),
-        );
-    })?;
-    let key_schedule = kx_output.unwrap();
+    let shared_secret = our_key_share.complete(&their_key_share.payload.0)?;
+    let key_schedule = key_schedule_pre_handshake.into_handshake(shared_secret.as_ref());
 
     // Remember what KX group the server liked for next time.
     config

rustls/src/crypto/mod.rs

@@ -59,18 +59,30 @@ pub trait SupportedKxGroup: Send + Sync + Debug {
 pub trait ActiveKeyExchange: Send + Sync {
     /// Completes the key exchange, given the peer's public key.
     ///
-    /// The shared secret is passed into the `sink` function for further processing.
+    /// The shared secret is returned as a [`SharedSecret`] which can be constructed
+    /// from a `&[u8]`.
     ///
     /// This consumes and so terminates the [`ActiveKeyExchange`].
-    fn complete(
-        self: Box<Self>,
-        peer_pub_key: &[u8],
-        sink: &mut dyn FnMut(&[u8]),
-    ) -> Result<(), Error>;
+    fn complete(self: Box<Self>, peer_pub_key: &[u8]) -> Result<SharedSecret, Error>;
 
     /// Return the public key being used.
     fn pub_key(&self) -> &[u8];
 
     /// Return the group being used.
     fn group(&self) -> NamedGroup;
 }
+
+/// The result from `ActiveKeyExchange::complete` as a value.
+pub struct SharedSecret(Vec<u8>);
+
+impl From<&[u8]> for SharedSecret {
+    fn from(source: &[u8]) -> Self {
+        Self(source.to_vec())
+    }
+}
+
+impl AsRef<[u8]> for SharedSecret {
+    fn as_ref(&self) -> &[u8] {
+        &self.0
+    }
+}

rustls/src/crypto/ring/mod.rs

@@ -1,4 +1,4 @@
-use crate::crypto::{ActiveKeyExchange, CryptoProvider, SupportedKxGroup};
+use crate::crypto::{ActiveKeyExchange, CryptoProvider, SharedSecret, SupportedKxGroup};
 use crate::error::{Error, PeerMisbehaved};
 use crate::msgs::enums::NamedGroup;
 use crate::rand::GetRandomFailed;
@@ -150,11 +150,10 @@ pub struct KeyExchange {
 
 impl ActiveKeyExchange for KeyExchange {
     /// Completes the key exchange, given the peer's public key.
-    fn complete(self: Box<Self>, peer: &[u8], sink: &mut dyn FnMut(&[u8])) -> Result<(), Error> {
+    fn complete(self: Box<Self>, peer: &[u8]) -> Result<SharedSecret, Error> {
         let peer_key = UnparsedPublicKey::new(self.agreement_algorithm, peer);
-        agree_ephemeral(self.priv_key, &peer_key, (), move |secret| {
-            sink(secret);
-            Ok(())
+        agree_ephemeral(self.priv_key, &peer_key, (), |secret| {
+            Ok(SharedSecret::from(secret))
         })
         .map_err(|()| PeerMisbehaved::InvalidKeyShare.into())
     }

rustls/src/server/tls13.rs

@@ -541,18 +541,8 @@ mod client_hello {
         };
 
         // Do key exchange
-        let mut kx_input = Some(key_schedule_pre_handshake);
-        let mut kx_output = None;
-
-        kx.complete(&share.payload.0, &mut |secret| {
-            kx_output = Some(
-                kx_input
-                    .take()
-                    .unwrap()
-                    .into_handshake(secret),
-            );
-        })?;
-        let key_schedule = kx_output.unwrap();
+        let shared_secret = kx.complete(&share.payload.0)?;
+        let key_schedule = key_schedule_pre_handshake.into_handshake(shared_secret.as_ref());
 
         let handshake_hash = transcript.get_current_hash();
         let key_schedule = key_schedule.derive_server_handshake_secrets(

rustls/src/tls12/mod.rs

@@ -96,14 +96,16 @@ impl ConnectionSecrets {
             ),
         };
 
-        kx.complete(peer_pub_key, &mut |secret| {
-            prf::prf(
-                &mut ret.master_secret,
-                &*ret.suite.hmac_provider.with_key(secret),
-                label.as_bytes(),
-                seed.as_ref(),
-            );
-        })?;
+        let shared_secret = kx.complete(peer_pub_key)?;
+        prf::prf(
+            &mut ret.master_secret,
+            &*ret
+                .suite
+                .hmac_provider
+                .with_key(shared_secret.as_ref()),
+            label.as_bytes(),
+            seed.as_ref(),
+        );
 
         Ok(ret)
     }