From aabbd5518798f03fc4fc076a0ee834d71c5100db Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Wed, 21 Aug 2024 14:42:41 +0100
Subject: [PATCH] add ktls handling

---
 Cargo.lock                   |  6 ++++--
 proxy/src/config.rs          | 18 ++++++++++++------
 proxy/src/proxy/handshake.rs |  9 ++++++++-
 workspace_hack/Cargo.toml    |  2 ++
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 39a51d750a..38931143ff 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5894,9 +5894,9 @@ dependencies = [
 
 [[package]]
 name = "sha2-asm"
-version = "0.6.3"
+version = "0.6.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f27ba7066011e3fb30d808b51affff34f0a66d3a03a58edd787c6e420e40e44e"
+checksum = "b845214d6175804686b2bd482bcffe96651bb2d1200742b712003504a2dac1ab"
 dependencies = [
  "cc",
 ]
@@ -7863,6 +7863,8 @@ dependencies = [
  "reqwest 0.11.19",
  "reqwest 0.12.4",
  "rustls 0.21.11",
+ "rustls-pki-types",
+ "rustls-webpki 0.102.6",
  "scopeguard",
  "serde",
  "serde_json",
diff --git a/proxy/src/config.rs b/proxy/src/config.rs
index 1ab96f7cfe..b26ee733d9 100644
--- a/proxy/src/config.rs
+++ b/proxy/src/config.rs
@@ -110,13 +110,19 @@ pub fn configure_tls(
 
     let cert_resolver = Arc::new(cert_resolver);
 
+    let provider = rustls::crypto::aws_lc_rs::default_provider();
+    #[cfg(target_os = "linux")]
+    let provider = {
+        let mut provider = provider;
+        let compat = ktls::CompatibleCiphers::new()?;
+        provider.cipher_suites.retain(|s| compat.is_supported(s));
+    };
+
     // allow TLS 1.2 to be compatible with older client libraries
-    let mut config = rustls::ServerConfig::builder_with_protocol_versions(&[
-        &rustls::version::TLS13,
-        &rustls::version::TLS12,
-    ])
-    .with_no_client_auth()
-    .with_cert_resolver(cert_resolver.clone());
+    let mut config = rustls::ServerConfig::builder_with_provider(Arc::new(provider))
+        .with_protocol_versions(&[&rustls::version::TLS13, &rustls::version::TLS12])?
+        .with_no_client_auth()
+        .with_cert_resolver(cert_resolver.clone());
 
     config.alpn_protocols = vec![PG_ALPN_PROTOCOL.to_vec()];
 
diff --git a/proxy/src/proxy/handshake.rs b/proxy/src/proxy/handshake.rs
index d6b6464fb9..92c569bc1e 100644
--- a/proxy/src/proxy/handshake.rs
+++ b/proxy/src/proxy/handshake.rs
@@ -32,6 +32,10 @@ pub enum HandshakeError {
     #[error("{0}")]
     StreamUpgradeError(#[from] StreamUpgradeError),
 
+    #[cfg(target_os = "linux")]
+    #[error("{0}")]
+    KtlsUpgradeError(#[from] ktls::Error),
+
     #[error("{0}")]
     Io(#[from] std::io::Error),
 
@@ -119,6 +123,9 @@ pub async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
                             ));
                         };
 
+                        #[cfg(target_os = "linux")]
+                        let raw = ktls::CorkStream::new(raw);
+
                         let mut read_buf = read_buf.reader();
                         let mut res = Ok(());
                         let accept = tokio_rustls::TlsAcceptor::from(tls.to_server_config())
@@ -179,7 +186,7 @@ pub async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
                                     #[cfg(not(target_os = "linux"))]
                                     tls: Box::new(tls_stream),
                                     #[cfg(target_os = "linux")]
-                                    tls: {},
+                                    tls: ktls::config_ktls_server(tls_stream)?,
                                     tls_server_end_point,
                                 },
                                 read_buf,
diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml
index 2d9b372654..ef1dd0e9c5 100644
--- a/workspace_hack/Cargo.toml
+++ b/workspace_hack/Cargo.toml
@@ -66,6 +66,8 @@ regex-syntax = { version = "0.8" }
 reqwest-5ef9efb8ec2df382 = { package = "reqwest", version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "stream"] }
 reqwest-a6292c17cd707f01 = { package = "reqwest", version = "0.11", default-features = false, features = ["blocking", "rustls-tls", "stream"] }
 rustls = { version = "0.21", features = ["dangerous_configuration"] }
+rustls-pki-types = { version = "1", features = ["std"] }
+rustls-webpki = { version = "0.102", default-features = false, features = ["aws_lc_rs", "ring", "std"] }
 scopeguard = { version = "1" }
 serde = { version = "1", features = ["alloc", "derive"] }
 serde_json = { version = "1", features = ["raw_value"] }