Merge pull request #36 from flashbots/peg/protocol-negotiation

Add ALPN protocol negotiation

Author: ameba23

src/lib.rs

@@ -30,6 +30,11 @@ use tokio_rustls::{
 
 use crate::attestation::{AttesationPayload, AttestationVerifier};
 
+/// This makes it possible to add breaking protocol changes and provide backwards compatibility.
+/// When adding more supported versions, note that ordering is important. ALPN will pick the first
+/// protocol which both parties support - so newer supported versions should come first.
+pub const SUPPORTED_ALPN_PROTOCOL_VERSIONS: [&[u8]; 1] = [b"flashbots-ratls/1"];
+
 /// The label used when exporting key material from a TLS session
 const EXPORTER_LABEL: &[u8; 24] = b"EXPORTER-Channel-Binding";
 
@@ -85,7 +90,7 @@ impl ProxyServer {
             return Err(ProxyError::NoClientAuth);
         }
 
-        let server_config = if client_auth {
+        let mut server_config = if client_auth {
             let root_store =
                 RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
             let verifier = WebPkiClientVerifier::builder(Arc::new(root_store)).build()?;
@@ -99,6 +104,11 @@ impl ProxyServer {
                 .with_single_cert(cert_and_key.cert_chain.clone(), cert_and_key.key)?
         };
 
+        server_config.alpn_protocols = SUPPORTED_ALPN_PROTOCOL_VERSIONS
+            .into_iter()
+            .map(|p| p.to_vec())
+            .collect();
+
         Self::new_with_tls_config(
             cert_and_key.cert_chain,
             server_config.into(),
@@ -179,6 +189,9 @@ impl ProxyServer {
         let mut tls_stream = acceptor.accept(inbound).await?;
         let (_io, connection) = tls_stream.get_ref();
 
+        // Ensure that we agreed a protocol
+        let _negotiated_protocol = connection.alpn_protocol().ok_or(ProxyError::AlpnFailed)?;
+
         // Compute an exporter unique to the session
         let mut exporter = [0u8; 32];
         connection.export_keying_material(
@@ -352,7 +365,7 @@ impl ProxyClient {
         };
 
         // Setup TLS client configuration, with or without client auth
-        let client_config = if let Some(ref cert_and_key) = cert_and_key {
+        let mut client_config = if let Some(ref cert_and_key) = cert_and_key {
             ClientConfig::builder()
                 .with_root_certificates(root_store)
                 .with_client_auth_cert(
@@ -365,6 +378,11 @@ impl ProxyClient {
                 .with_no_client_auth()
         };
 
+        client_config.alpn_protocols = SUPPORTED_ALPN_PROTOCOL_VERSIONS
+            .into_iter()
+            .map(|p| p.to_vec())
+            .collect();
+
         Self::new_with_tls_config(
             client_config.into(),
             address,
@@ -578,6 +596,11 @@ impl ProxyClient {
 
         let (_io, server_connection) = tls_stream.get_ref();
 
+        // Ensure that we agreed a protocol
+        let _negotiated_protocol = server_connection
+            .alpn_protocol()
+            .ok_or(ProxyError::AlpnFailed)?;
+
         // Compute an exporter unique to the channel
         let mut exporter = [0u8; 32];
         server_connection.export_keying_material(
@@ -737,6 +760,8 @@ pub enum ProxyError {
     MpscSend,
     #[error("Serialization: {0}")]
     Serialization(#[from] parity_scale_codec::Error),
+    #[error("Protocol negotiation failed - remote peer does not support this protocol")]
+    AlpnFailed,
 }
 
 impl From<mpsc::error::SendError<RequestWithResponseSender>> for ProxyError {

src/test_helpers.rs

@@ -13,7 +13,7 @@ use tokio_rustls::rustls::{
 
 use crate::{
     attestation::measurements::{CvmImageMeasurements, Measurements, PlatformMeasurements},
-    MEASUREMENT_HEADER,
+    MEASUREMENT_HEADER, SUPPORTED_ALPN_PROTOCOL_VERSIONS,
 };
 
 /// Helper to generate a self-signed certificate for testing
@@ -42,18 +42,27 @@ pub fn generate_tls_config(
     certificate_chain: Vec<CertificateDer<'static>>,
     key: PrivateKeyDer<'static>,
 ) -> (Arc<ServerConfig>, Arc<ClientConfig>) {
-    let server_config = ServerConfig::builder()
+    let supported_protocols: Vec<_> = SUPPORTED_ALPN_PROTOCOL_VERSIONS
+        .into_iter()
+        .map(|p| p.to_vec())
+        .collect();
+
+    let mut server_config = ServerConfig::builder()
         .with_no_client_auth()
         .with_single_cert(certificate_chain.clone(), key)
         .expect("Failed to create rustls server config");
 
+    server_config.alpn_protocols = supported_protocols.clone();
+
     let mut root_store = RootCertStore::empty();
     root_store.add(certificate_chain[0].clone()).unwrap();
 
-    let client_config = ClientConfig::builder()
+    let mut client_config = ClientConfig::builder()
         .with_root_certificates(root_store)
         .with_no_client_auth();
 
+    client_config.alpn_protocols = supported_protocols;
+
     (Arc::new(server_config), Arc::new(client_config))
 }
 
@@ -67,32 +76,44 @@ pub fn generate_tls_config_with_client_auth(
     (Arc<ServerConfig>, Arc<ClientConfig>),
     (Arc<ServerConfig>, Arc<ClientConfig>),
 ) {
+    let supported_protocols: Vec<_> = SUPPORTED_ALPN_PROTOCOL_VERSIONS
+        .into_iter()
+        .map(|p| p.to_vec())
+        .collect();
+
     let (alice_client_verifier, alice_root_store) =
         client_verifier_from_remote_cert(bob_certificate_chain[0].clone());
 
-    let alice_server_config = ServerConfig::builder()
+    let mut alice_server_config = ServerConfig::builder()
         .with_client_cert_verifier(alice_client_verifier)
         .with_single_cert(alice_certificate_chain.clone(), alice_key.clone_key())
         .expect("Failed to create rustls server config");
 
-    let alice_client_config = ClientConfig::builder()
+    alice_server_config.alpn_protocols = supported_protocols.clone();
+
+    let mut alice_client_config = ClientConfig::builder()
         .with_root_certificates(alice_root_store)
         .with_client_auth_cert(alice_certificate_chain.clone(), alice_key)
         .unwrap();
 
+    alice_client_config.alpn_protocols = supported_protocols.clone();
+
     let (bob_client_verifier, bob_root_store) =
         client_verifier_from_remote_cert(alice_certificate_chain[0].clone());
 
-    let bob_server_config = ServerConfig::builder()
+    let mut bob_server_config = ServerConfig::builder()
         .with_client_cert_verifier(bob_client_verifier)
         .with_single_cert(bob_certificate_chain.clone(), bob_key.clone_key())
         .expect("Failed to create rustls server config");
 
-    let bob_client_config = ClientConfig::builder()
+    bob_server_config.alpn_protocols = supported_protocols.clone();
+
+    let mut bob_client_config = ClientConfig::builder()
         .with_root_certificates(bob_root_store)
         .with_client_auth_cert(bob_certificate_chain, bob_key)
         .unwrap();
 
+    bob_client_config.alpn_protocols = supported_protocols;
     (
         (Arc::new(alice_server_config), Arc::new(alice_client_config)),
         (Arc::new(bob_server_config), Arc::new(bob_client_config)),