Reload client certificates

Author: behos

Cargo.toml

@@ -16,7 +16,7 @@ rust-version = "1.76"
 
 [features]
 default = []
-tls = ["rustls", "rustls-pemfile", "webpki-roots", "futures-rustls"]
+tls = ["rustls", "rustls-pemfile", "webpki-roots", "futures-rustls", "async-fs"]
 sasl = ["sasl-gssapi", "sasl-digest-md5"]
 sasl-digest-md5 = ["rsasl/unstable_custom_mechanism", "md5", "linkme", "hex"]
 sasl-gssapi = ["rsasl/gssapi"]
@@ -48,6 +48,7 @@ md5 = { version = "0.7.0", optional = true }
 hex = { version = "0.4.3", optional = true }
 linkme = { version = "0.3", optional = true }
 async-io = "2.3.2"
+async-fs = { version = "2.1.2", optional = true }
 futures = "0.3.30"
 async-net = "2.0.0"
 futures-rustls = { version = "0.26.0", optional = true }
@@ -78,3 +79,8 @@ all-features = true
 [profile.dev]
 # Need this for linkme crate to work for spawns in macOS
 lto = "thin"
+
+[[example]]
+name = "tls_file_based"
+path = "examples/tls_file_based.rs"
+required-features = ["tls", "tokio"]

examples/tls_file_based.rs

@@ -0,0 +1,104 @@
+use std::env;
+use std::io::{self, Write};
+use std::path::PathBuf;
+use std::time::Duration;
+
+use futures::executor::block_on;
+use zookeeper_client::Error::NodeExists;
+use zookeeper_client::{Acls, Client, CreateMode, TlsOptions};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    block_on(run()).unwrap_or_else(|e| {
+        eprintln!("Error: {}", e);
+        std::process::exit(1);
+    });
+    Ok(())
+}
+
+async fn run() -> Result<(), Box<dyn std::error::Error>> {
+    let connect_string = env::var("ZK_CONNECT_STRING").unwrap_or_else(|_| "tcp+tls://localhost:2281".to_string());
+    let ca_cert = PathBuf::from(env::var("ZK_CA_CERT").expect("ZK_CA_CERT environment variable is required"));
+    let client_cert =
+        PathBuf::from(env::var("ZK_CLIENT_CERT").expect("ZK_CLIENT_CERT environment variable is required"));
+    let client_key = PathBuf::from(env::var("ZK_CLIENT_KEY").expect("ZK_CLIENT_KEY environment variable is required"));
+
+    println!("Connecting to ZooKeeper with file-based TLS...");
+    println!("Server: {}", connect_string);
+    println!("CA cert: {}", ca_cert.display());
+    println!("Client cert: {}", client_cert.display());
+    println!("Client key: {}", client_key.display());
+
+    let tls_options = TlsOptions::default()
+        .with_pem_ca_certs(&ca_cert.to_str().unwrap())?
+        .with_pem_identity_files(&client_cert, &client_key)
+        .await?;
+
+    let tls_options = unsafe { tls_options.with_no_hostname_verification() };
+
+    println!("WARNING: Hostname verification disabled!");
+
+    let client = Client::connector()
+        .connection_timeout(Duration::from_secs(10))
+        .session_timeout(Duration::from_secs(30))
+        .tls(tls_options)
+        .secure_connect(&connect_string)
+        .await?;
+
+    println!("Connected to ZooKeeper successfully!");
+
+    let path = "/tls_example";
+
+    loop {
+        print!("\nOptions:\ne. Edit key\nq. Quit\nEnter choice (e/q): ");
+        io::stdout().flush()?;
+
+        let mut input = String::new();
+        io::stdin().read_line(&mut input)?;
+
+        match input.trim() {
+            "e" => {
+                print!("Enter new data for the key: ");
+                io::stdout().flush()?;
+
+                let mut data_input = String::new();
+                io::stdin().read_line(&mut data_input)?;
+                let data = data_input.trim().as_bytes();
+
+                println!("Setting data at path: {}", path);
+                match client.create(path, data, &CreateMode::Ephemeral.with_acls(Acls::anyone_all())).await {
+                    Ok(_) => println!("ZNode created successfully"),
+                    Err(NodeExists) => {
+                        println!("ZNode already exists, updating data...");
+                        client.set_data(path, data, None).await?;
+                        println!("ZNode data updated successfully");
+                    },
+                    Err(e) => {
+                        println!("Error creating/updating ZNode: {}", e);
+                        continue;
+                    },
+                }
+
+                match client.get_data(path).await {
+                    Ok((data, _stat)) => {
+                        println!("Current data: {}", String::from_utf8_lossy(&data));
+                    },
+                    Err(e) => println!("Error reading data: {}", e),
+                }
+            },
+            "q" => {
+                println!("Cleaning up and exiting...");
+                match client.delete(path, None).await {
+                    Ok(_) => println!("ZNode deleted successfully"),
+                    Err(_) => println!("ZNode may not exist or already deleted"),
+                }
+                break;
+            },
+            _ => {
+                println!("Invalid choice. Please enter 'e' or 'q'.");
+            },
+        }
+    }
+
+    println!("Example completed successfully!");
+    Ok(())
+}

src/session/connection.rs

@@ -1,5 +1,7 @@
 use std::io::{Error, ErrorKind, IoSlice, Result};
 use std::pin::Pin;
+#[cfg(feature = "tls")]
+use std::sync::Arc;
 use std::task::{Context, Poll};
 use std::time::Duration;
 
@@ -15,18 +17,18 @@ use tracing::{debug, trace};
 
 #[cfg(feature = "tls")]
 mod tls {
-    pub use std::sync::Arc;
-
     pub use futures_rustls::client::TlsStream;
     pub use futures_rustls::TlsConnector;
     pub use rustls::pki_types::ServerName;
-    pub use rustls::ClientConfig;
 }
+
 #[cfg(feature = "tls")]
 use tls::*;
 
 use crate::deadline::Deadline;
 use crate::endpoint::{EndpointRef, IterableEndpoints};
+#[cfg(feature = "tls")]
+use crate::TlsOptions;
 
 #[derive(Debug)]
 pub enum Connection {
@@ -170,15 +172,15 @@ impl Connection {
 #[derive(Clone)]
 pub struct Connector {
     #[cfg(feature = "tls")]
-    tls: Option<TlsConnector>,
+    tls_options: Option<TlsOptions>,
     timeout: Duration,
 }
 
 impl Connector {
     #[cfg(feature = "tls")]
     #[allow(dead_code)]
     pub fn new() -> Self {
-        Self { tls: None, timeout: Duration::from_secs(10) }
+        Self { tls_options: None, timeout: Duration::from_secs(10) }
     }
 
     #[cfg(not(feature = "tls"))]
@@ -187,14 +189,27 @@ impl Connector {
     }
 
     #[cfg(feature = "tls")]
-    pub fn with_tls(config: ClientConfig) -> Self {
-        Self { tls: Some(TlsConnector::from(Arc::new(config))), timeout: Duration::from_secs(10) }
+    pub fn with_tls_options(tls_options: TlsOptions) -> Self {
+        Self { tls_options: Some(tls_options), timeout: Duration::from_secs(10) }
+    }
+
+    #[cfg(feature = "tls")]
+    async fn get_current_tls_connector(&self) -> Result<TlsConnector> {
+        let Some(ref tls_opts) = self.tls_options else {
+            return Err(Error::new(ErrorKind::InvalidInput, "no TLS configuration"));
+        };
+        let config = tls_opts
+            .to_config()
+            .await
+            .map_err(|e| Error::new(ErrorKind::InvalidData, format!("TLS config creation failed: {}", e)))?;
+        Ok(TlsConnector::from(Arc::new(config)))
     }
 
     #[cfg(feature = "tls")]
     async fn connect_tls(&self, stream: TcpStream, host: &str) -> Result<Connection> {
+        let tls_connector = self.get_current_tls_connector().await?;
         let domain = ServerName::try_from(host).unwrap().to_owned();
-        let stream = self.tls.as_ref().unwrap().connect(domain, stream).await?;
+        let stream = tls_connector.connect(domain, stream).await?;
         Ok(Connection::new_tls(stream))
     }
 
@@ -209,7 +224,7 @@ impl Connector {
     pub async fn connect(&self, endpoint: EndpointRef<'_>, deadline: &mut Deadline) -> Result<Connection> {
         if endpoint.tls {
             #[cfg(feature = "tls")]
-            if self.tls.is_none() {
+            if self.tls_options.is_none() {
                 return Err(Error::new(ErrorKind::Unsupported, "tls not supported"));
             }
             #[cfg(not(feature = "tls"))]
@@ -288,4 +303,12 @@ mod tests {
         let err = connector.connect(endpoint, &mut Deadline::never()).await.unwrap_err();
         assert_eq!(err.kind(), ErrorKind::Unsupported);
     }
+
+    #[cfg(feature = "tls")]
+    #[test]
+    fn test_with_tls_options() {
+        let tls_options = crate::TlsOptions::default();
+        let connector = Connector::with_tls_options(tls_options);
+        assert!(connector.tls_options.is_some());
+    }
 }

src/session/mod.rs

@@ -130,7 +130,7 @@ impl Builder {
             return Err(Error::BadArguments(&"connection timeout must not be negative"));
         }
         #[cfg(feature = "tls")]
-        let connector = Connector::with_tls(self.tls.unwrap_or_default().into_config()?);
+        let connector = Connector::with_tls_options(self.tls.unwrap_or_default());
         #[cfg(not(feature = "tls"))]
         let connector = Connector::new();
         let (state_sender, state_receiver) = asyncs::sync::watch::channel(SessionState::Disconnected);