From 7fa8abb8f925591ee9e26ed70813551ae8d34fed Mon Sep 17 00:00:00 2001
From: Marco Neumann <marco@crepererum.net>
Date: Thu, 28 Sep 2023 16:14:46 +0200
Subject: [PATCH] feat: `SSLKEYLOGFILE` support for flight CLI

Allows analysis of TLS traffic with an external tool like Wireshark.

See https://wiki.wireshark.org/TLS#using-the-pre-master-secret
---
 Cargo.toml                                |  3 +++
 arrow-flight/src/bin/flight_sql_client.rs | 15 ++++++++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 936935ec7e3d..dea1baa13dad 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -93,3 +93,6 @@ arrow-string = { version = "47.0.0", path = "./arrow-string" }
 parquet = { version = "47.0.0", path = "./parquet", default-features = false }
 
 chrono = { version = "0.4.31", default-features = false, features = ["clock"] }
+
+[patch.crates-io]
+tonic = { git = "https://github.com/crepererum/tonic.git", branch = "crepererum/keylog" }
diff --git a/arrow-flight/src/bin/flight_sql_client.rs b/arrow-flight/src/bin/flight_sql_client.rs
index c6aaccf376eb..b607fd9f9af1 100644
--- a/arrow-flight/src/bin/flight_sql_client.rs
+++ b/arrow-flight/src/bin/flight_sql_client.rs
@@ -85,6 +85,14 @@ struct ClientArgs {
     #[clap(long)]
     tls: bool,
 
+    /// Dump TLS key log.
+    ///
+    /// The target file is specified by the `SSLKEYLOGFILE` environment variable.
+    ///
+    /// Required `--tls`.
+    #[clap(long, requires = "tls")]
+    key_log: bool,
+
     /// Server host.
     #[clap(long)]
     host: String,
@@ -235,7 +243,12 @@ async fn setup_client(args: ClientArgs) -> Result<FlightSqlServiceClient<Channel
         .keep_alive_while_idle(true);
 
     if args.tls {
-        let tls_config = ClientTlsConfig::new();
+        let mut tls_config = ClientTlsConfig::new();
+
+        if args.key_log {
+            tls_config = tls_config.use_key_log();
+        }
+
         endpoint = endpoint
             .tls_config(tls_config)
             .context("create TLS endpoint")?;