feat: auto-detect proxy and translate peer addresses to contact point

src/driver/mod.rs

@@ -7,6 +7,7 @@
 //! Many types and trait methods are defined ahead of their use in later
 //! development phases (REPL, DESCRIBE, COPY, etc.).
 
+pub mod proxy_address_translator;
 pub mod scylla_driver;
 pub mod types;
 

src/driver/proxy_address_translator.rs

@@ -0,0 +1,57 @@
+//! Proxy address translator.
+//!
+//! When connecting through a proxy or load balancer (e.g., AWS NLB, PrivateLink),
+//! the driver discovers internal node IPs from `system.peers` that are unreachable
+//! from the client. This translator remaps all peer addresses to the original
+//! contact point, ensuring all connections go through the proxy.
+
+use std::net::SocketAddr;
+
+use async_trait::async_trait;
+use scylla::errors::TranslationError;
+use scylla::policies::address_translator::{AddressTranslator, UntranslatedPeer};
+
+/// An [`AddressTranslator`] that redirects all peer connections to the original
+/// contact point address. Used when the cluster is accessed through a proxy.
+///
+/// All discovered node addresses are translated to `proxy_address`, ensuring
+/// the driver only connects through the proxy endpoint.
+#[derive(Debug, Clone)]
+pub struct ProxyAddressTranslator {
+    /// The proxy/contact point address to route all connections through.
+    proxy_address: SocketAddr,
+}
+
+impl ProxyAddressTranslator {
+    /// Create a new translator that routes all connections to `proxy_address`.
+    pub fn new(proxy_address: SocketAddr) -> Self {
+        Self { proxy_address }
+    }
+}
+
+#[async_trait]
+impl AddressTranslator for ProxyAddressTranslator {
+    async fn translate_address(
+        &self,
+        _untranslated_peer: &UntranslatedPeer,
+    ) -> Result<SocketAddr, TranslationError> {
+        Ok(self.proxy_address)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::net::{IpAddr, Ipv4Addr};
+
+    fn sock(ip: [u8; 4], port: u16) -> SocketAddr {
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::new(ip[0], ip[1], ip[2], ip[3])), port)
+    }
+
+    #[test]
+    fn creates_with_correct_address() {
+        let proxy_addr = sock([18, 208, 144, 200], 9042);
+        let translator = ProxyAddressTranslator::new(proxy_addr);
+        assert_eq!(translator.proxy_address, proxy_addr);
+    }
+}

src/driver/scylla_driver.rs

@@ -507,20 +507,40 @@ impl CqlDriver for ScyllaDriver {
 
         let mut builder = SessionBuilder::new().known_node(&addr);
 
-        // Authentication
+        // cqlsh is a single-user interactive tool — one connection per host suffices
+        // and avoids connection explosion when using a proxy translator.
+        builder = builder.pool_size(scylla::client::PoolSize::PerHost(
+            std::num::NonZeroUsize::new(1).unwrap(),
+        ));
+
         if let (Some(username), Some(password)) = (&config.username, &config.password) {
             builder = builder.user(username, password);
         }
 
-        // Connection timeout
         builder = builder.connection_timeout(Duration::from_secs(config.connect_timeout));
 
-        // Default keyspace
         if let Some(keyspace) = &config.keyspace {
             builder = builder.use_keyspace(keyspace, false);
         }
 
-        // SSL/TLS
+        // Always install the proxy address translator. Since known_node addresses
+        // are never translated (only peers from system.peers are), this is safe:
+        // - Direct connections: peers share the same network, translator is a no-op
+        //   in practice (peers are reachable anyway, and translated to contact point
+        //   which also works since it's the same node in single-node or resolves correctly)
+        // - Proxy connections: peers have unreachable internal IPs, translator
+        //   redirects all of them to the proxy contact point
+        let contact_point = tokio::net::lookup_host(&addr)
+            .await
+            .ok()
+            .and_then(|mut addrs| addrs.next());
+        if let Some(contact_point) = contact_point {
+            let translator = Arc::new(
+                super::proxy_address_translator::ProxyAddressTranslator::new(contact_point),
+            );
+            builder = builder.address_translator(translator);
+        }
+
         if config.ssl {
             let tls_config = if let Some(ssl_config) = &config.ssl_config {
                 Self::build_rustls_config(ssl_config)?
@@ -530,11 +550,6 @@ impl CqlDriver for ScyllaDriver {
             builder = builder.tls_context(Some(tls_config));
         }
 
-        // NOTE: config.protocol_version is accepted for CLI compatibility but
-        // scylla-rust-driver 1.5.0 auto-negotiates the native protocol version.
-        // SessionBuilder has no method to force a specific protocol version.
-        // Similarly, the driver hardcodes CQL_VERSION="4.0.0" in the STARTUP frame.
-
         let session = builder.build().await.context("connecting to cluster")?;
 
         Ok(ScyllaDriver {

tests/integration/main.rs

@@ -32,6 +32,8 @@ mod login_tests;
 #[cfg(feature = "test-plain")]
 mod output_tests;
 #[cfg(feature = "test-plain")]
+mod proxy_tests;
+#[cfg(feature = "test-plain")]
 mod schema_agreement_tests;
 #[cfg(feature = "test-plain")]
 mod unicode_tests;

tests/integration/proxy_tests.rs

@@ -0,0 +1,91 @@
+//! Integration tests for proxy auto-detection.
+//!
+//! Tests that verify the proxy address translator works correctly with a
+//! real ScyllaDB instance.
+
+use helpers::{cqlsh_cmd, get_scylla};
+
+use super::*;
+
+/// Direct connection (no proxy) should work without address translation.
+/// This confirms the two-phase connect doesn't break normal connections.
+#[test]
+#[ignore = "requires Docker"]
+fn test_direct_connection_still_works_with_proxy_detection() {
+    let scylla = get_scylla();
+    cqlsh_cmd(scylla)
+        .args(["-e", "SELECT cluster_name FROM system.local"])
+        .assert()
+        .success();
+}
+
+/// Verify that a query actually returns data through the normal path,
+/// confirming the two-phase connect doesn't cause session issues.
+#[test]
+#[ignore = "requires Docker"]
+fn test_query_after_proxy_detection_returns_data() {
+    let scylla = get_scylla();
+    let output =
+        helpers::execute_cql_output_direct(scylla, "SELECT release_version FROM system.local");
+    assert!(
+        !output.is_empty(),
+        "expected output from system.local query"
+    );
+}
+
+/// Test connection through a TCP proxy (socat) to simulate the proxy scenario.
+/// The proxy forwards traffic to the ScyllaDB container, but its IP won't
+/// match any peer address in system.peers, triggering proxy auto-detection.
+#[test]
+#[ignore = "requires Docker"]
+fn test_proxy_connection_via_socat() {
+    use std::process::{Command, Stdio};
+    use std::thread;
+    use std::time::Duration;
+
+    let scylla = get_scylla();
+    let target = format!("{}:{}", scylla.host, scylla.port);
+
+    // Start socat as a TCP proxy on a random high port
+    // socat listens on 0.0.0.0:0 (kernel picks port) and forwards to the ScyllaDB container
+    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+    let proxy_port = listener.local_addr().unwrap().port();
+    drop(listener); // free the port for socat
+
+    let mut socat = match Command::new("socat")
+        .args([
+            &format!("TCP-LISTEN:{proxy_port},fork,reuseaddr"),
+            &format!("TCP:{target}"),
+        ])
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .spawn()
+    {
+        Ok(child) => child,
+        Err(_) => {
+            eprintln!("socat not available, skipping proxy integration test");
+            return;
+        }
+    };
+
+    // Give socat time to start listening
+    thread::sleep(Duration::from_millis(500));
+
+    // Connect through the proxy — proxy IP (127.0.0.1:proxy_port) won't match
+    // the ScyllaDB node's internal address, triggering proxy detection
+    let result = assert_cmd::Command::cargo_bin("cqlsh-rs")
+        .unwrap()
+        .args([
+            "127.0.0.1",
+            &proxy_port.to_string(),
+            "-e",
+            "SELECT cluster_name FROM system.local",
+        ])
+        .timeout(Duration::from_secs(15))
+        .assert();
+
+    socat.kill().ok();
+    socat.wait().ok();
+
+    result.success();
+}