use std::{io, net::TcpStream, path::Path, sync::Arc, time::Duration};
use rustls_pki_types::ServerName;
use weir_core::Durability;
use crate::{ClientError, WeirClient};
#[derive(Debug, Clone)]
pub struct ClientTlsConfig<'a> {
pub client_cert: &'a Path,
pub client_key: &'a Path,
pub ca_cert: &'a Path,
pub server_name: &'a str,
pub default_durability: Option<Durability>,
}
pub type TlsStream = rustls::StreamOwned<rustls::ClientConnection, TcpStream>;
impl WeirClient<TlsStream> {
pub fn connect_tls(
addr: impl std::net::ToSocketAddrs,
cfg: ClientTlsConfig<'_>,
) -> Result<Self, ClientError> {
let mut roots = rustls::RootCertStore::empty();
for c in load_certs(cfg.ca_cert)? {
roots
.add(c)
.map_err(|e| ClientError::Protocol(format!("ca cert: {e}")))?;
}
let certs = load_certs(cfg.client_cert)?;
let key = load_key(cfg.client_key)?;
let tls_config = rustls::ClientConfig::builder_with_provider(Arc::new(
rustls::crypto::ring::default_provider(),
))
.with_safe_default_protocol_versions()
.map_err(|e| ClientError::Protocol(format!("tls versions: {e}")))?
.with_root_certificates(roots)
.with_client_auth_cert(certs, key)
.map_err(|e| ClientError::Protocol(format!("client auth cert: {e}")))?;
let server_name = ServerName::try_from(cfg.server_name.to_string())
.map_err(|e| ClientError::Protocol(format!("server name: {e}")))?;
let conn = rustls::ClientConnection::new(Arc::new(tls_config), server_name)
.map_err(|e| ClientError::Protocol(format!("tls connect: {e}")))?;
let sock = TcpStream::connect(addr)?;
let stream = rustls::StreamOwned::new(conn, sock);
Ok(Self::from_parts(stream, cfg.default_durability))
}
pub fn set_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
self.stream.sock.set_read_timeout(timeout)
}
pub fn set_write_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
self.stream.sock.set_write_timeout(timeout)
}
}
fn load_certs(p: &Path) -> Result<Vec<rustls_pki_types::CertificateDer<'static>>, ClientError> {
use rustls_pki_types::pem::PemObject;
let pem = std::fs::read(p)?;
rustls_pki_types::CertificateDer::pem_slice_iter(&pem)
.collect::<Result<_, _>>()
.map_err(|e| {
ClientError::Protocol(format!("invalid certificate PEM in {}: {e}", p.display()))
})
}
fn load_key(p: &Path) -> Result<rustls_pki_types::PrivateKeyDer<'static>, ClientError> {
use rustls_pki_types::pem::PemObject;
let pem = std::fs::read(p)?;
rustls_pki_types::PrivateKeyDer::from_pem_slice(&pem)
.map_err(|e| ClientError::Protocol(format!("no private key in file {}: {e}", p.display())))
}
#[cfg(test)]
mod tests {
use super::*;
fn tmp_file(label: &str, contents: &[u8]) -> std::path::PathBuf {
let p = std::env::temp_dir().join(format!("weir_tls_{label}_{}.pem", std::process::id()));
std::fs::write(&p, contents).unwrap();
p
}
#[test]
fn load_key_missing_file_is_io_error() {
let err = load_key(Path::new("/weir_no_such_key_xyzzy.pem")).unwrap_err();
assert!(matches!(err, ClientError::Io(_)), "{err:?}");
}
#[test]
fn load_certs_missing_file_is_io_error() {
let err = load_certs(Path::new("/weir_no_such_cert_xyzzy.pem")).unwrap_err();
assert!(matches!(err, ClientError::Io(_)), "{err:?}");
}
#[test]
fn load_key_file_without_a_key_is_protocol_error() {
let p = tmp_file("nokey", b"not a pem private key\n");
let err = load_key(&p).unwrap_err();
match err {
ClientError::Protocol(msg) => assert!(msg.contains("no private key"), "{msg}"),
other => panic!("expected Protocol, got {other:?}"),
}
std::fs::remove_file(&p).ok();
}
}