From 0c576da38ae20937bd5c99a479f96648787c5cdc Mon Sep 17 00:00:00 2001
From: Marco Neumann <marco@crepererum.net>
Date: Wed, 26 Feb 2025 11:17:54 +0100
Subject: [PATCH] feat: optional `SSLKEYLOGFILE` support

Add a `use_key_log` option to server and client TLS configs that -- when
set -- will enable rustls's `SSLKEYLOGFILE` handling.

This is helpful when you want to intercept TLS traffic for debugging and
is generally supported by many libraries and browsers. Also see:
https://wiki.wireshark.org/TLS#using-the-pre-master-secret
---
 tonic/src/transport/channel/service/tls.rs |  6 ++++++
 tonic/src/transport/channel/tls.rs         | 10 ++++++++++
 tonic/src/transport/server/service/tls.rs  |  5 +++++
 tonic/src/transport/server/tls.rs          | 10 ++++++++++
 4 files changed, 31 insertions(+)

diff --git a/tonic/src/transport/channel/service/tls.rs b/tonic/src/transport/channel/service/tls.rs
index 386cc8a6c..7510099a1 100644
--- a/tonic/src/transport/channel/service/tls.rs
+++ b/tonic/src/transport/channel/service/tls.rs
@@ -26,12 +26,14 @@ pub(crate) struct TlsConnector {
 }
 
 impl TlsConnector {
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
         ca_certs: Vec<Certificate>,
         trust_anchors: Vec<TrustAnchor<'static>>,
         identity: Option<Identity>,
         domain: &str,
         assume_http2: bool,
+        use_key_log: bool,
         #[cfg(feature = "tls-native-roots")] with_native_roots: bool,
         #[cfg(feature = "tls-webpki-roots")] with_webpki_roots: bool,
     ) -> Result<Self, crate::BoxError> {
@@ -87,6 +89,10 @@ impl TlsConnector {
             None => builder.with_no_client_auth(),
         };
 
+        if use_key_log {
+            config.key_log = Arc::new(tokio_rustls::rustls::KeyLogFile::new());
+        }
+
         config.alpn_protocols.push(ALPN_H2.into());
         Ok(Self {
             config: Arc::new(config),
diff --git a/tonic/src/transport/channel/tls.rs b/tonic/src/transport/channel/tls.rs
index 0c2eb37e0..945384fd2 100644
--- a/tonic/src/transport/channel/tls.rs
+++ b/tonic/src/transport/channel/tls.rs
@@ -18,6 +18,7 @@ pub struct ClientTlsConfig {
     with_native_roots: bool,
     #[cfg(feature = "tls-webpki-roots")]
     with_webpki_roots: bool,
+    use_key_log: bool,
 }
 
 impl ClientTlsConfig {
@@ -84,6 +85,14 @@ impl ClientTlsConfig {
         }
     }
 
+    /// Use key log as specified by the `SSLKEYLOGFILE` environment variable.
+    pub fn use_key_log(self) -> Self {
+        ClientTlsConfig {
+            use_key_log: true,
+            ..self
+        }
+    }
+
     /// Enables the platform's trusted certs.
     #[cfg(feature = "tls-native-roots")]
     pub fn with_native_roots(self) -> Self {
@@ -123,6 +132,7 @@ impl ClientTlsConfig {
             self.identity,
             domain,
             self.assume_http2,
+            self.use_key_log,
             #[cfg(feature = "tls-native-roots")]
             self.with_native_roots,
             #[cfg(feature = "tls-webpki-roots")]
diff --git a/tonic/src/transport/server/service/tls.rs b/tonic/src/transport/server/service/tls.rs
index 0a24d9482..42c376249 100644
--- a/tonic/src/transport/server/service/tls.rs
+++ b/tonic/src/transport/server/service/tls.rs
@@ -23,6 +23,7 @@ impl TlsAcceptor {
         client_ca_root: Option<&Certificate>,
         client_auth_optional: bool,
         ignore_client_order: bool,
+        use_key_log: bool,
     ) -> Result<Self, crate::BoxError> {
         let builder = ServerConfig::builder();
 
@@ -45,6 +46,10 @@ impl TlsAcceptor {
         let mut config = builder.with_single_cert(cert, key)?;
         config.ignore_client_order = ignore_client_order;
 
+        if use_key_log {
+            config.key_log = Arc::new(tokio_rustls::rustls::KeyLogFile::new());
+        }
+
         config.alpn_protocols.push(ALPN_H2.into());
         Ok(Self {
             inner: Arc::new(config),
diff --git a/tonic/src/transport/server/tls.rs b/tonic/src/transport/server/tls.rs
index 92da18347..2e1ef8213 100644
--- a/tonic/src/transport/server/tls.rs
+++ b/tonic/src/transport/server/tls.rs
@@ -10,6 +10,7 @@ pub struct ServerTlsConfig {
     client_ca_root: Option<Certificate>,
     client_auth_optional: bool,
     ignore_client_order: bool,
+    use_key_log: bool,
 }
 
 impl fmt::Debug for ServerTlsConfig {
@@ -64,12 +65,21 @@ impl ServerTlsConfig {
         }
     }
 
+    /// Use key log as specified by the `SSLKEYLOGFILE` environment variable.
+    pub fn use_key_log(self) -> Self {
+        ServerTlsConfig {
+            use_key_log: true,
+            ..self
+        }
+    }
+
     pub(crate) fn tls_acceptor(&self) -> Result<TlsAcceptor, crate::BoxError> {
         TlsAcceptor::new(
             self.identity.as_ref().unwrap(),
             self.client_ca_root.as_ref(),
             self.client_auth_optional,
             self.ignore_client_order,
+            self.use_key_log,
         )
     }
 }