This is an automated email from the ASF dual-hosted git repository. CritasWang pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/iotdb-client-rust.git
commit c7dcdc140e072b5b764aaf6e3cc2a4896fcd1e5e Author: CritasWang <[email protected]> AuthorDate: Mon Jul 13 14:16:59 2026 +0800 V5: TLS (feature-gated, native-tls shared-stream channel) + RPC compression (compact protocol; server requires dn_rpc_thrift_compression_enable — no auto-detection, spec corrected) --- .github/workflows/ci.yml | 10 ++ Cargo.toml | 4 + README.md | 30 ++++ README_ZH.md | 30 ++++ src/client/pool.rs | 11 +- src/client/session.rs | 251 +++++++++++++++++++++++++-- src/client/table_session.rs | 57 +++++++ src/connection/mod.rs | 392 ++++++++++++++++++++++++++++++++++++++++--- src/error.rs | 5 + src/lib.rs | 4 +- tests/fixtures/tls/README.md | 39 +++++ tests/fixtures/tls/cert.pem | 20 +++ tests/fixtures/tls/key.pem | 28 ++++ 13 files changed, 845 insertions(+), 36 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d7b001c..6e02e1b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -47,13 +47,23 @@ jobs: - name: Clippy run: cargo clippy --all-targets -- -D warnings + - name: Clippy (tls) + run: cargo clippy --all-targets --features tls -- -D warnings + - name: Build run: cargo build + - name: Build (tls) + run: cargo build --features tls + # Live-server tests skip themselves when nothing listens on 6667. - name: Unit tests run: cargo test + # Adds the loopback TLS handshake tests (self-signed fixture cert). + - name: Unit tests (tls) + run: cargo test --features tls + integration: runs-on: ubuntu-latest services: diff --git a/Cargo.toml b/Cargo.toml index 23195dd..1399830 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,9 +31,13 @@ thrift = "0.23" byteorder = "1.5" chrono = "0.4" log = "0.4" +native-tls = { version = "0.2", optional = true } [dev-dependencies] env_logger = "0.11" [features] default = [] +# TLS support via the platform-native TLS stack (SecureTransport / +# SChannel / OpenSSL). Adds `use_ssl` & friends to SessionConfig. +tls = ["dep:native-tls"] diff --git a/README.md b/README.md index a6bc084..823be0b 100644 --- a/README.md +++ b/README.md @@ -172,6 +172,36 @@ cargo run --example table_session cargo run --example session_pool ``` +## TLS & RPC compression + +**RPC compression** (IoTDB's term for the Thrift *compact protocol*) is a plain config flag: + +```rust +let config = SessionConfig { enable_rpc_compression: true, ..Default::default() }; +// or: TableSession::builder().enable_rpc_compression(true)... +``` + +It must match the **server** setting `dn_rpc_thrift_compression_enable` (default `false`). The server speaks exactly one protocol — there is no per-connection negotiation, so a mismatch in either direction fails at the first RPC with a transport error. + +**TLS** is behind the `tls` cargo feature (platform-native TLS via [`native-tls`](https://crates.io/crates/native-tls)): + +```toml +iotdb-client = { version = "0.1", features = ["tls"] } +``` + +```rust +let config = SessionConfig { + use_ssl: true, + ca_cert_path: Some("ca.pem".into()), // trust a private CA / self-signed cert + accept_invalid_certs: false, // true skips verification (tests only!) + domain_override: None, // SNI/validation hostname when connecting by IP + ..Default::default() +}; +// or: TableSession::builder().use_ssl(true).ca_cert_path("ca.pem")... +``` + +The server needs Thrift SSL enabled (`enable_thrift_ssl=true` + key store). Pool configs pass both options through their embedded `session` config. + ## Thrift codegen Generated stubs live in `src/protocol/` (`client.rs`, `common.rs`); never hand-edit them. The IDL sources in `thrift/` are synced from the IoTDB repo's `iotdb-protocol/` (`thrift-datanode/src/main/thrift/client.thrift`, `thrift-commons/src/main/thrift/common.thrift`). diff --git a/README_ZH.md b/README_ZH.md index 2a5cb36..6f40256 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -172,6 +172,36 @@ cargo run --example table_session cargo run --example session_pool ``` +## TLS 与 RPC 压缩 + +**RPC 压缩**（IoTDB 术语，实为 Thrift *compact 协议*）只是一个配置开关： + +```rust +let config = SessionConfig { enable_rpc_compression: true, ..Default::default() }; +// 或：TableSession::builder().enable_rpc_compression(true)... +``` + +必须与**服务端**配置 `dn_rpc_thrift_compression_enable`（默认 `false`）一致。服务端只讲一种协议——没有按连接协商的机制，任意方向的不匹配都会在第一个 RPC 上以传输错误失败。 + +**TLS** 位于 `tls` cargo feature 之后（基于 [`native-tls`](https://crates.io/crates/native-tls) 的平台原生 TLS）： + +```toml +iotdb-client = { version = "0.1", features = ["tls"] } +``` + +```rust +let config = SessionConfig { + use_ssl: true, + ca_cert_path: Some("ca.pem".into()), // 信任私有 CA / 自签名证书 + accept_invalid_certs: false, // true 跳过证书校验（仅限测试！） + domain_override: None, // 按 IP 连接时用于 SNI/校验的主机名 + ..Default::default() +}; +// 或：TableSession::builder().use_ssl(true).ca_cert_path("ca.pem")... +``` + +服务端需开启 Thrift SSL（`enable_thrift_ssl=true` + 密钥库）。连接池配置通过其内嵌的 `session` 配置透传这两组选项。 + ## Thrift 代码生成 生成的桩代码位于 `src/protocol/`（`client.rs`、`common.rs`），请勿手动编辑。`thrift/` 中的 IDL 源文件从 IoTDB 仓库的 `iotdb-protocol/` 同步（`thrift-datanode/src/main/thrift/client.thrift`、`thrift-commons/src/main/thrift/common.thrift`）。 diff --git a/src/client/pool.rs b/src/client/pool.rs index 64b3564..3df10aa 100644 --- a/src/client/pool.rs +++ b/src/client/pool.rs @@ -621,9 +621,14 @@ mod tests { fn injected_session(endpoint: &Endpoint) -> Session { let mut session = Session::new(dead_endpoint_config()); - let connection = - crate::connection::Connection::open(endpoint.clone(), Duration::from_millis(500)) - .expect("connect to test listener"); + let connection = crate::connection::Connection::open( + endpoint.clone(), + &crate::connection::ConnectionOptions { + connect_timeout: Duration::from_millis(500), + ..Default::default() + }, + ) + .expect("connect to test listener"); session.test_inject_connection(connection); session } diff --git a/src/client/session.rs b/src/client/session.rs index f98a308..5f4425c 100644 --- a/src/client/session.rs +++ b/src/client/session.rs @@ -23,7 +23,7 @@ use std::time::Duration; use crate::client::dataset::SessionDataSet; use crate::client::redirect::{self, RedirectCache, RedirectCacheStats}; -use crate::connection::{Connection, Endpoint}; +use crate::connection::{Connection, ConnectionOptions, Endpoint, RpcProtocol}; use crate::data::record::serialize_record_values; use crate::data::{Tablet, Value}; use crate::error::{Error, Result}; @@ -73,6 +73,28 @@ pub struct SessionConfig { /// Harvest status-400 `redirectNode` hints from insert responses into /// the per-session [`RedirectCache`]. Default `true`. pub enable_redirection: bool, + /// Speak TCompactProtocol instead of TBinaryProtocol ("RPC compression" + /// in IoTDB terms). Must match the **server** setting + /// `dn_rpc_thrift_compression_enable` (default `false`) — there is no + /// per-connection negotiation; a mismatch fails at the first RPC. + /// Default `false`. + pub enable_rpc_compression: bool, + /// Wrap connections in TLS (cargo feature `tls`). The server must have + /// `enable_thrift_ssl=true`. Default `false`. + #[cfg(feature = "tls")] + pub use_ssl: bool, + /// PEM certificate added as trusted root (private CA / self-signed + /// server cert). `None` → platform trust store only. + #[cfg(feature = "tls")] + pub ca_cert_path: Option<std::path::PathBuf>, + /// Skip certificate verification (self-signed test certs). + /// **Dangerous** outside tests. Default `false`. + #[cfg(feature = "tls")] + pub accept_invalid_certs: bool, + /// Hostname for SNI + certificate validation instead of the endpoint + /// host (e.g. when connecting by IP). + #[cfg(feature = "tls")] + pub domain_override: Option<String>, } impl Default for SessionConfig { @@ -91,6 +113,15 @@ impl Default for SessionConfig { max_reconnect_attempts: 3, retry_interval: Duration::from_secs(1), enable_redirection: true, + enable_rpc_compression: false, + #[cfg(feature = "tls")] + use_ssl: false, + #[cfg(feature = "tls")] + ca_cert_path: None, + #[cfg(feature = "tls")] + accept_invalid_certs: false, + #[cfg(feature = "tls")] + domain_override: None, } } } @@ -105,6 +136,25 @@ impl SessionConfig { .collect::<Result<Vec<_>>>()?; Ok(self) } + + /// Resolve the connection-level options (timeout, wire protocol, TLS) + /// this config implies. + pub fn connection_options(&self) -> ConnectionOptions { + ConnectionOptions { + connect_timeout: self.connect_timeout, + protocol: if self.enable_rpc_compression { + RpcProtocol::Compact + } else { + RpcProtocol::Binary + }, + #[cfg(feature = "tls")] + tls: self.use_ssl.then(|| crate::connection::TlsOptions { + ca_cert_path: self.ca_cert_path.clone(), + accept_invalid_certs: self.accept_invalid_certs, + domain_override: self.domain_override.clone(), + }), + } + } } /// Raw result of a query statement: response metadata plus undecoded TsBlocks. @@ -167,11 +217,12 @@ impl Session { } let start = ENDPOINT_START_INDEX.fetch_add(1, Ordering::Relaxed) % n; + let options = self.config.connection_options(); let mut connection = None; let mut last_err: Option<Error> = None; for i in 0..n { let endpoint = self.config.endpoints[(start + i) % n].clone(); - match Connection::open(endpoint, self.config.connect_timeout) { + match Connection::open(endpoint, &options) { Ok(c) => { connection = Some(c); break; @@ -242,6 +293,7 @@ impl Session { .and_then(|ep| self.config.endpoints.iter().position(|e| e == ep)) .unwrap_or(0); let attempts = self.config.max_reconnect_attempts.max(1); + let options = self.config.connection_options(); let mut last_err: Option<Error> = None; for attempt in 0..attempts { if attempt > 0 { @@ -249,12 +301,10 @@ impl Session { } for i in 0..n { let endpoint = self.config.endpoints[(start + i) % n].clone(); - let result = Connection::open(endpoint, self.config.connect_timeout).and_then( - |mut connection| { - let ids = self.authenticate(&mut connection)?; - Ok((connection, ids)) - }, - ); + let result = Connection::open(endpoint, &options).and_then(|mut connection| { + let ids = self.authenticate(&mut connection)?; + Ok((connection, ids)) + }); match result { Ok((connection, (session_id, statement_id))) => { log::info!("reconnected to {}", connection.endpoint()); @@ -970,6 +1020,49 @@ mod tests { assert_eq!(cfg.max_reconnect_attempts, 3); assert_eq!(cfg.retry_interval, Duration::from_secs(1)); assert!(cfg.enable_redirection); + assert!(!cfg.enable_rpc_compression); + #[cfg(feature = "tls")] + { + assert!(!cfg.use_ssl); + assert!(cfg.ca_cert_path.is_none()); + assert!(!cfg.accept_invalid_certs); + assert!(cfg.domain_override.is_none()); + } + } + + /// Config → connection options mapping: compression selects the compact + /// protocol, `use_ssl` (feature `tls`) carries the TLS fields through. + #[test] + fn connection_options_from_config() { + use crate::connection::RpcProtocol; + + let cfg = SessionConfig::default(); + let options = cfg.connection_options(); + assert_eq!(options.connect_timeout, cfg.connect_timeout); + assert_eq!(options.protocol, RpcProtocol::Binary); + #[cfg(feature = "tls")] + assert!(options.tls.is_none()); + + let cfg = SessionConfig { + enable_rpc_compression: true, + ..Default::default() + }; + assert_eq!(cfg.connection_options().protocol, RpcProtocol::Compact); + + #[cfg(feature = "tls")] + { + let cfg = SessionConfig { + use_ssl: true, + ca_cert_path: Some("/certs/ca.pem".into()), + accept_invalid_certs: true, + domain_override: Some("iotdb.internal".into()), + ..Default::default() + }; + let tls = cfg.connection_options().tls.expect("tls options"); + assert_eq!(tls.ca_cert_path.as_deref(), Some("/certs/ca.pem".as_ref())); + assert!(tls.accept_invalid_certs); + assert_eq!(tls.domain_override.as_deref(), Some("iotdb.internal")); + } } #[test] @@ -1302,8 +1395,14 @@ mod tests { retry_interval: Duration::from_millis(10), ..Default::default() }); - let connection = Connection::open(endpoint.clone(), Duration::from_millis(500)) // - .expect("connect to listener"); + let connection = Connection::open( + endpoint.clone(), + &ConnectionOptions { + connect_timeout: Duration::from_millis(500), + ..Default::default() + }, + ) + .expect("connect to listener"); session.test_inject_connection(connection); assert_eq!(session.current_endpoint(), Some(&endpoint)); @@ -1328,8 +1427,14 @@ mod tests { enable_auto_reconnect: false, ..Default::default() }); - let connection = Connection::open(endpoint, Duration::from_millis(500)) // - .expect("connect to listener"); + let connection = Connection::open( + endpoint, + &ConnectionOptions { + connect_timeout: Duration::from_millis(500), + ..Default::default() + }, + ) + .expect("connect to listener"); session.test_inject_connection(connection); let err = session.execute_non_query("SHOW DATABASES").unwrap_err(); @@ -1378,6 +1483,128 @@ mod tests { assert!(!session.is_open()); } + /// RPC compression (TCompactProtocol) against a live server. + /// + /// Verified 2026-07-13 against apache/iotdb:2.0.6-standalone: the server + /// speaks exactly **one** protocol, fixed by its config + /// `dn_rpc_thrift_compression_enable` (default `false` → binary) — there + /// is no per-connection auto-detection (server source: + /// `AbstractThriftServiceThread.getProtocolFactory(compress)` picks a + /// single factory). Matrix observed live: + /// compact↔compression-enabled server: OK; compact↔default server: EOF + /// at openSession; binary↔compression-enabled server: EOF. + /// + /// So this test adapts: if the compact open succeeds the server has + /// compression enabled and we assert a full insert+query roundtrip; + /// if it fails at the transport level (the expected outcome against the + /// default docker image) we skip with a message — the clean transport + /// failure is itself the verified mismatch behavior. + #[test] + fn live_rpc_compression_roundtrip() { + use crate::data::{tablet::Tablet, TSDataType, Value}; + use std::net::TcpStream; + // IOTDB_COMPACT_URL points at a compression-enabled server (e.g. + // docker run -e dn_rpc_thrift_compression_enable=true …) to force + // the positive roundtrip; default is the standard test server. + let url = std::env::var("IOTDB_COMPACT_URL").unwrap_or_else(|_| "127.0.0.1:6667".into()); + let endpoint = Endpoint::parse(&url).expect("IOTDB_COMPACT_URL"); + if TcpStream::connect_timeout( + &format!("{}:{}", endpoint.host, endpoint.port) + .parse() + .unwrap(), + Duration::from_millis(300), + ) + .is_err() + { + eprintln!("skipping live_rpc_compression_roundtrip: no IoTDB server on {url}"); + return; + } + + const DB: &str = "root.rusttest_compact"; + let _guard = LIVE_DB_LOCK.lock().unwrap_or_else(|p| p.into_inner()); + + let mut session = Session::new(SessionConfig { + endpoints: vec![endpoint], + enable_rpc_compression: true, + // Fail fast: a protocol mismatch dies at the first RPC, and + // reconnect passes cannot fix it. + enable_auto_reconnect: false, + max_reconnect_attempts: 1, + connect_timeout: Duration::from_secs(2), + ..Default::default() + }); + match session.open() { + Err(Error::Thrift(e)) => { + eprintln!( + "skipping live_rpc_compression_roundtrip: server rejected the compact \ + protocol ({e}) — expected when dn_rpc_thrift_compression_enable=false \ + (the default); IoTDB has no per-connection protocol auto-detection" + ); + return; + } + other => other.expect("open compact session"), + } + + // The server accepted the compact handshake ⇒ compression is + // enabled server-side; the whole roundtrip must work. + let _ = session.execute_non_query(&format!("DELETE DATABASE {DB}")); + session + .execute_non_query(&format!("CREATE DATABASE {DB}")) + .expect("create database"); + let mut tablet = Tablet::new( + format!("{DB}.d1"), + vec!["v".into()], + vec![TSDataType::Int32], + ) + .expect("tablet"); + tablet + .add_row(1, vec![Some(Value::Int32(42))]) + .expect("add_row"); + session.insert_tablet(&tablet).expect("insert_tablet"); + let mut rows = Vec::new(); + { + let mut dataset = session + .execute_query(&format!("SELECT v FROM {DB}.d1")) + .expect("query"); + while let Some(row) = dataset.next_row().expect("next_row") { + rows.push((row.timestamp.expect("timestamp"), row.values[0].clone())); + } + } + assert_eq!(rows, [(1, Value::Int32(42))]); + session + .execute_non_query(&format!("DELETE DATABASE {DB}")) + .expect("cleanup"); + session.close().expect("close session"); + } + + /// Live TLS against a real IoTDB requires a server with + /// `enable_thrift_ssl=true` plus its certificate — not part of the + /// standard docker test topology. Opt in by pointing + /// `IOTDB_TLS_URL` (`host:port`) at such a server; otherwise this + /// skips. The TLS handshake/transport path itself is covered by the + /// loopback tests in `connection::tls_tests`. + #[cfg(feature = "tls")] + #[test] + fn live_tls_roundtrip() { + let Ok(url) = std::env::var("IOTDB_TLS_URL") else { + eprintln!( + "skipping live_tls_roundtrip: set IOTDB_TLS_URL=host:port to a TLS-enabled \ + IoTDB (enable_thrift_ssl=true); the standard test server is plain TCP" + ); + return; + }; + let mut session = Session::new(SessionConfig { + endpoints: vec![Endpoint::parse(&url).expect("IOTDB_TLS_URL")], + use_ssl: true, + ca_cert_path: std::env::var_os("IOTDB_TLS_CA").map(Into::into), + accept_invalid_certs: std::env::var_os("IOTDB_TLS_INSECURE").is_some(), + ..Default::default() + }); + session.open().expect("open TLS session"); + let mut dataset = session.execute_query("SHOW DATABASES").expect("query"); + while dataset.next_row().expect("next_row").is_some() {} + } + /// DATE wire-format adjudication test (goal V1). Inserts one DATE row via /// the tablet binary path (i32 yyyyMMdd) and one via a SQL date literal /// (parsed server-side), then reads both back: if the tablet encoding is diff --git a/src/client/table_session.rs b/src/client/table_session.rs index 10f740e..a37b02f 100644 --- a/src/client/table_session.rs +++ b/src/client/table_session.rs @@ -107,6 +107,43 @@ impl TableSessionBuilder { self } + /// Speak TCompactProtocol ("RPC compression"). Must match the server's + /// `dn_rpc_thrift_compression_enable` setting — see + /// [`SessionConfig::enable_rpc_compression`]. + pub fn enable_rpc_compression(mut self, enable: bool) -> Self { + self.config.enable_rpc_compression = enable; + self + } + + /// Wrap connections in TLS (cargo feature `tls`). + #[cfg(feature = "tls")] + pub fn use_ssl(mut self, use_ssl: bool) -> Self { + self.config.use_ssl = use_ssl; + self + } + + /// PEM certificate added as trusted root (cargo feature `tls`). + #[cfg(feature = "tls")] + pub fn ca_cert_path(mut self, path: impl Into<std::path::PathBuf>) -> Self { + self.config.ca_cert_path = Some(path.into()); + self + } + + /// Skip certificate verification — **dangerous** outside tests + /// (cargo feature `tls`). + #[cfg(feature = "tls")] + pub fn accept_invalid_certs(mut self, accept: bool) -> Self { + self.config.accept_invalid_certs = accept; + self + } + + /// Hostname for SNI + certificate validation (cargo feature `tls`). + #[cfg(feature = "tls")] + pub fn domain_override(mut self, domain: impl Into<String>) -> Self { + self.config.domain_override = Some(domain.into()); + self + } + /// The [`SessionConfig`] this builder resolves to (dialect always /// `"table"`). pub fn config(&self) -> &SessionConfig { @@ -201,6 +238,26 @@ mod tests { assert_eq!(cfg.query_timeout_ms, 1_000); } + #[test] + fn builder_passes_transport_options_through() { + let b = TableSessionBuilder::new().enable_rpc_compression(true); + assert!(b.config().enable_rpc_compression); + + #[cfg(feature = "tls")] + { + let b = TableSessionBuilder::new() + .use_ssl(true) + .ca_cert_path("/certs/ca.pem") + .accept_invalid_certs(true) + .domain_override("iotdb.internal"); + let cfg = b.config(); + assert!(cfg.use_ssl); + assert_eq!(cfg.ca_cert_path.as_deref(), Some("/certs/ca.pem".as_ref())); + assert!(cfg.accept_invalid_certs); + assert_eq!(cfg.domain_override.as_deref(), Some("iotdb.internal")); + } + } + #[test] fn insert_rejects_tree_model_tablets() { let mut session = TableSession { diff --git a/src/connection/mod.rs b/src/connection/mod.rs index c6cdc89..5f77f6c 100644 --- a/src/connection/mod.rs +++ b/src/connection/mod.rs @@ -18,15 +18,19 @@ //! Low-level Thrift connection to an IoTDB node. //! //! Mirrors `src/connection/Connection.ts` (Node.js) and the Thrift client setup -//! in the C# SDK: TCP → TFramedTransport → TBinaryProtocol → IClientRPCService client. +//! in the C# SDK: TCP (optionally TLS) → TFramedTransport → TBinaryProtocol +//! (or TCompactProtocol when RPC compression is enabled) → IClientRPCService +//! client. +use std::io::{Read, Write}; use std::net::{TcpStream, ToSocketAddrs}; use std::time::Duration; -use thrift::protocol::{TBinaryInputProtocol, TBinaryOutputProtocol}; -use thrift::transport::{ - ReadHalf, TFramedReadTransport, TFramedWriteTransport, TIoChannel, TTcpChannel, WriteHalf, +use thrift::protocol::{ + TBinaryInputProtocol, TBinaryOutputProtocol, TCompactInputProtocol, TCompactOutputProtocol, + TInputProtocol, TOutputProtocol, }; +use thrift::transport::{TFramedReadTransport, TFramedWriteTransport, TIoChannel, TTcpChannel}; use crate::error::{Error, Result}; use crate::protocol::client::IClientRPCServiceSyncClient; @@ -34,11 +38,73 @@ use crate::protocol::client::IClientRPCServiceSyncClient; /// Default IoTDB DataNode RPC port. pub const DEFAULT_PORT: u16 = 6667; -/// The concrete generated RPC client over framed transport + strict binary protocol. -pub type RpcClient = IClientRPCServiceSyncClient< - TBinaryInputProtocol<TFramedReadTransport<ReadHalf<TTcpChannel>>>, - TBinaryOutputProtocol<TFramedWriteTransport<WriteHalf<TTcpChannel>>>, ->; +/// Wire protocol used on top of the framed transport. +/// +/// IoTDB's Thrift server speaks **one** protocol per server instance, +/// chosen by the server config `dn_rpc_thrift_compression_enable` +/// (default `false` → binary). There is **no** per-connection +/// auto-detection: a compact-protocol client against a binary-protocol +/// server fails at the first RPC, and vice versa — pick the protocol that +/// matches the server (the C# SDK's `enableRpcCompression` does the same +/// blind switch). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum RpcProtocol { + /// Strict TBinaryProtocol — matches the server default. + #[default] + Binary, + /// TCompactProtocol — matches `dn_rpc_thrift_compression_enable=true` + /// ("RPC compression" in IoTDB terms is the compact protocol, not a + /// compressed stream). + Compact, +} + +/// TLS settings for a [`Connection`] (cargo feature `tls`). +#[cfg(feature = "tls")] +#[derive(Debug, Clone, Default)] +pub struct TlsOptions { + /// PEM certificate added as a trusted root (e.g. a private CA or the + /// server's self-signed certificate). + pub ca_cert_path: Option<std::path::PathBuf>, + /// Skip certificate verification entirely (self-signed test certs). + /// **Dangerous** outside tests. Default `false`. + pub accept_invalid_certs: bool, + /// Hostname used for SNI + certificate validation instead of the + /// endpoint host (e.g. when connecting by IP). + pub domain_override: Option<String>, +} + +/// How to open a [`Connection`]: timeout, wire protocol, optional TLS. +#[derive(Debug, Clone)] +pub struct ConnectionOptions { + /// TCP connect timeout per endpoint attempt. Default 10 s. + pub connect_timeout: Duration, + /// Wire protocol; must match the server (see [`RpcProtocol`]). + pub protocol: RpcProtocol, + /// Wrap the TCP stream in TLS before the Thrift transports. + #[cfg(feature = "tls")] + pub tls: Option<TlsOptions>, +} + +impl Default for ConnectionOptions { + fn default() -> Self { + Self { + connect_timeout: Duration::from_secs(10), + protocol: RpcProtocol::Binary, + #[cfg(feature = "tls")] + tls: None, + } + } +} + +/// Type-erased Thrift input protocol (binary or compact, plain or TLS). +pub type BoxedInputProtocol = Box<dyn TInputProtocol + Send>; +/// Type-erased Thrift output protocol (binary or compact, plain or TLS). +pub type BoxedOutputProtocol = Box<dyn TOutputProtocol + Send>; + +/// The generated RPC client over framed transport and a type-erased +/// protocol pair (the `thrift` crate forwards the protocol traits through +/// `Box`, so the generated generic client works unchanged). +pub type RpcClient = IClientRPCServiceSyncClient<BoxedInputProtocol, BoxedOutputProtocol>; /// A single endpoint `host:port` of an IoTDB DataNode (default port 6667). #[derive(Debug, Clone, PartialEq, Eq)] @@ -93,25 +159,38 @@ impl std::fmt::Display for Endpoint { /// and the generated `IClientRPCService` client. pub struct Connection { endpoint: Endpoint, + protocol: RpcProtocol, client: RpcClient, } impl Connection { - /// Establish a TCP connection to `endpoint` (bounded by `connect_timeout`) - /// and wrap it in framed transport + strict binary protocol. - pub fn open(endpoint: Endpoint, connect_timeout: Duration) -> Result<Self> { - let stream = connect_stream(&endpoint, connect_timeout)?; + /// Establish a TCP connection to `endpoint` (bounded by + /// `options.connect_timeout`), optionally wrap it in TLS, and stack + /// framed transport + the selected protocol on top. + pub fn open(endpoint: Endpoint, options: &ConnectionOptions) -> Result<Self> { + let stream = connect_stream(&endpoint, options.connect_timeout)?; stream.set_nodelay(true).map_err(thrift::Error::from)?; + #[cfg(feature = "tls")] + if let Some(tls) = &options.tls { + let stream = tls_handshake(&endpoint, stream, tls)?; + let shared = SharedTlsStream::new(stream); + let (input, output) = build_protocols(shared.clone(), shared, options.protocol); + return Ok(Self { + endpoint, + protocol: options.protocol, + client: IClientRPCServiceSyncClient::new(input, output), + }); + } + let channel = TTcpChannel::with_stream(stream); let (read_half, write_half) = channel.split()?; - let read_transport = TFramedReadTransport::new(read_half); - let write_transport = TFramedWriteTransport::new(write_half); - let input_protocol = TBinaryInputProtocol::new(read_transport, true); - let output_protocol = TBinaryOutputProtocol::new(write_transport, true); - let client = IClientRPCServiceSyncClient::new(input_protocol, output_protocol); - - Ok(Self { endpoint, client }) + let (input, output) = build_protocols(read_half, write_half, options.protocol); + Ok(Self { + endpoint, + protocol: options.protocol, + client: IClientRPCServiceSyncClient::new(input, output), + }) } /// Mutable access to the generated RPC client for issuing calls. @@ -122,6 +201,36 @@ impl Connection { pub fn endpoint(&self) -> &Endpoint { &self.endpoint } + + /// The wire protocol this connection was opened with. + pub fn protocol(&self) -> RpcProtocol { + self.protocol + } +} + +/// Stack framed transport + the selected protocol over a read/write pair +/// and type-erase the result (see [`RpcClient`]). +fn build_protocols<R, W>( + read: R, + write: W, + protocol: RpcProtocol, +) -> (BoxedInputProtocol, BoxedOutputProtocol) +where + R: Read + Send + 'static, + W: Write + Send + 'static, +{ + let read_transport = TFramedReadTransport::new(read); + let write_transport = TFramedWriteTransport::new(write); + match protocol { + RpcProtocol::Binary => ( + Box::new(TBinaryInputProtocol::new(read_transport, true)), + Box::new(TBinaryOutputProtocol::new(write_transport, true)), + ), + RpcProtocol::Compact => ( + Box::new(TCompactInputProtocol::new(read_transport)), + Box::new(TCompactOutputProtocol::new(write_transport)), + ), + } } /// Resolve the endpoint and try each resolved address with the connect timeout. @@ -142,6 +251,78 @@ fn connect_stream(endpoint: &Endpoint, connect_timeout: Duration) -> Result<TcpS }) } +/// Run the TLS handshake over an established TCP stream. +#[cfg(feature = "tls")] +fn tls_handshake( + endpoint: &Endpoint, + stream: TcpStream, + tls: &TlsOptions, +) -> Result<native_tls::TlsStream<TcpStream>> { + let mut builder = native_tls::TlsConnector::builder(); + if let Some(path) = &tls.ca_cert_path { + let pem = std::fs::read(path).map_err(|e| { + Error::Client(format!( + "cannot read CA certificate {}: {e}", + path.display() + )) + })?; + let cert = native_tls::Certificate::from_pem(&pem).map_err(Error::Tls)?; + builder.add_root_certificate(cert); + } + if tls.accept_invalid_certs { + builder.danger_accept_invalid_certs(true); + } + let connector = builder.build().map_err(Error::Tls)?; + let domain = tls.domain_override.as_deref().unwrap_or(&endpoint.host); + connector.connect(domain, stream).map_err(|e| match e { + native_tls::HandshakeError::Failure(e) => Error::Tls(e), + // Blocking sockets never yield the mid-handshake variant. + native_tls::HandshakeError::WouldBlock(_) => { + Error::Client("TLS handshake interrupted".into()) + } + }) +} + +/// A `TlsStream` shared between the read and write transports. +/// +/// `TTcpChannel::split` clones the underlying OS socket, but a TLS stream +/// cannot be split that way (record layer state is shared), so both framed +/// transports hold the same stream behind a mutex. The generated sync +/// client fully writes + flushes a request before reading the response, so +/// read and write never contend. +#[cfg(feature = "tls")] +#[derive(Clone)] +struct SharedTlsStream(std::sync::Arc<std::sync::Mutex<native_tls::TlsStream<TcpStream>>>); + +#[cfg(feature = "tls")] +impl SharedTlsStream { + fn new(stream: native_tls::TlsStream<TcpStream>) -> Self { + Self(std::sync::Arc::new(std::sync::Mutex::new(stream))) + } + + fn lock(&self) -> std::sync::MutexGuard<'_, native_tls::TlsStream<TcpStream>> { + self.0.lock().unwrap_or_else(|p| p.into_inner()) + } +} + +#[cfg(feature = "tls")] +impl Read for SharedTlsStream { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> { + self.lock().read(buf) + } +} + +#[cfg(feature = "tls")] +impl Write for SharedTlsStream { + fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> { + self.lock().write(buf) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.lock().flush() + } +} + #[cfg(test)] mod tests { use super::*; @@ -203,4 +384,175 @@ mod tests { Endpoint::new("::1", 6667) ); } + + #[test] + fn default_options_are_binary_no_tls() { + let options = ConnectionOptions::default(); + assert_eq!(options.connect_timeout, Duration::from_secs(10)); + assert_eq!(options.protocol, RpcProtocol::Binary); + #[cfg(feature = "tls")] + assert!(options.tls.is_none()); + } + + /// A local listener that accepts and immediately drops connections, so + /// `Connection::open` (which issues no RPC) succeeds for any protocol. + fn accept_then_drop_listener() -> Endpoint { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let port = listener.local_addr().expect("local_addr").port(); + std::thread::spawn(move || { + for stream in listener.incoming() { + match stream { + Ok(s) => drop(s), + Err(_) => break, + } + } + }); + Endpoint::new("127.0.0.1", port) + } + + /// Both protocol variants construct their transport/protocol stack and + /// report the choice back via `Connection::protocol`. (Wire-level + /// verification against a live server lives in the session tests.) + #[test] + fn open_with_each_protocol() { + let endpoint = accept_then_drop_listener(); + for protocol in [RpcProtocol::Binary, RpcProtocol::Compact] { + // The struct update is only "needless" without the tls feature, + // which adds a field this literal doesn't name. + #[allow(clippy::needless_update)] + let options = ConnectionOptions { + connect_timeout: Duration::from_millis(500), + protocol, + ..Default::default() + }; + let connection = Connection::open(endpoint.clone(), &options).expect("open"); + assert_eq!(connection.protocol(), protocol); + assert_eq!(connection.endpoint(), &endpoint); + } + } +} + +#[cfg(all(test, feature = "tls"))] +mod tls_tests { + use super::*; + use crate::protocol::client::TIClientRPCServiceSyncClient; + use std::path::PathBuf; + + fn fixture(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures/tls") + .join(name) + } + + /// Spawn a TLS acceptor on a loopback port that completes one handshake + /// and then drops the connection. Uses the checked-in self-signed cert + /// (CN=localhost, SAN DNS:localhost + IP:127.0.0.1, 100-year validity). + fn tls_acceptor_once() -> Endpoint { + let cert = std::fs::read(fixture("cert.pem")).expect("read cert fixture"); + let key = std::fs::read(fixture("key.pem")).expect("read key fixture"); + let identity = native_tls::Identity::from_pkcs8(&cert, &key).expect("identity"); + let acceptor = native_tls::TlsAcceptor::new(identity).expect("acceptor"); + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let port = listener.local_addr().expect("local_addr").port(); + std::thread::spawn(move || { + for stream in listener.incoming() { + match stream { + // Handshake (which may itself fail when the client + // rejects our cert — fine, just move on), then drop. + Ok(s) => drop(acceptor.accept(s)), + Err(_) => break, + } + } + }); + Endpoint::new("127.0.0.1", port) + } + + /// Full client-side TLS path with the fixture cert as trusted root and + /// hostname pinned via `domain_override`: the handshake completes (so + /// `Connection::open` succeeds), and the first RPC then dies at the + /// Thrift layer because the acceptor closed the connection — proving + /// the failure is post-handshake. + #[test] + fn tls_handshake_with_trusted_root_then_thrift_failure() { + let endpoint = tls_acceptor_once(); + let options = ConnectionOptions { + connect_timeout: Duration::from_millis(500), + protocol: RpcProtocol::Binary, + tls: Some(TlsOptions { + ca_cert_path: Some(fixture("cert.pem")), + accept_invalid_certs: false, + domain_override: Some("localhost".into()), + }), + }; + let mut connection = Connection::open(endpoint, &options).expect("TLS handshake"); + assert_eq!(connection.protocol(), RpcProtocol::Binary); + let err = connection + .client_mut() + .request_statement_id(1) + .expect_err("RPC on a closed TLS connection must fail"); + // Post-handshake: the error is a Thrift transport error, not TLS. + let msg = err.to_string(); + assert!(!msg.to_lowercase().contains("certificate"), "got: {msg}"); + } + + /// Without the fixture as trusted root, certificate validation rejects + /// the self-signed server during the handshake: `Connection::open` + /// itself fails with a TLS error. + #[test] + fn tls_untrusted_cert_fails_handshake() { + let endpoint = tls_acceptor_once(); + let options = ConnectionOptions { + connect_timeout: Duration::from_millis(500), + protocol: RpcProtocol::Binary, + tls: Some(TlsOptions { + ca_cert_path: None, + accept_invalid_certs: false, + domain_override: Some("localhost".into()), + }), + }; + let err = match Connection::open(endpoint, &options) { + Ok(_) => panic!("untrusted self-signed cert must fail the handshake"), + Err(e) => e, + }; + assert!(matches!(err, Error::Tls(_)), "got {err:?}"); + } + + /// `accept_invalid_certs` bypasses validation for self-signed test + /// certs — handshake succeeds without any trusted root. + #[test] + fn tls_accept_invalid_certs_bypasses_validation() { + let endpoint = tls_acceptor_once(); + let options = ConnectionOptions { + connect_timeout: Duration::from_millis(500), + protocol: RpcProtocol::Compact, // also exercise compact-over-TLS + tls: Some(TlsOptions { + accept_invalid_certs: true, + ..Default::default() + }), + }; + let connection = Connection::open(endpoint, &options).expect("TLS handshake"); + assert_eq!(connection.protocol(), RpcProtocol::Compact); + } + + /// A missing CA file is a clear client error before any connect. + #[test] + fn tls_missing_ca_file_is_client_error() { + let endpoint = tls_acceptor_once(); + let options = ConnectionOptions { + connect_timeout: Duration::from_millis(500), + protocol: RpcProtocol::Binary, + tls: Some(TlsOptions { + ca_cert_path: Some(fixture("does-not-exist.pem")), + ..Default::default() + }), + }; + let err = match Connection::open(endpoint, &options) { + Ok(_) => panic!("missing CA file must fail"), + Err(e) => e, + }; + assert!( + matches!(&err, Error::Client(m) if m.contains("cannot read CA certificate")), + "got {err:?}" + ); + } } diff --git a/src/error.rs b/src/error.rs index d28dbc8..52b126c 100644 --- a/src/error.rs +++ b/src/error.rs @@ -31,6 +31,9 @@ pub enum Error { Client(String), /// Malformed binary payload received from the server (e.g. truncated TsBlock). Decode(String), + /// TLS setup or handshake failure (cargo feature `tls`). + #[cfg(feature = "tls")] + Tls(native_tls::Error), } impl fmt::Display for Error { @@ -40,6 +43,8 @@ impl fmt::Display for Error { Error::Server { code, message } => write!(f, "server error {code}: {message}"), Error::Client(msg) => write!(f, "client error: {msg}"), Error::Decode(msg) => write!(f, "decode error: {msg}"), + #[cfg(feature = "tls")] + Error::Tls(e) => write!(f, "tls error: {e}"), } } } diff --git a/src/lib.rs b/src/lib.rs index 35a71a2..bd26ccd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -35,6 +35,8 @@ pub use client::pool::{PooledSession, SessionPool, SessionPoolConfig, TableSessi pub use client::redirect::{RedirectCache, RedirectCacheStats}; pub use client::session::{QueryHandle, Session, SessionConfig}; pub use client::table_session::{TableSession, TableSessionBuilder}; -pub use connection::Endpoint; +#[cfg(feature = "tls")] +pub use connection::TlsOptions; +pub use connection::{ConnectionOptions, Endpoint, RpcProtocol}; pub use data::{ColumnCategory, TSDataType, Tablet, TsBlock, Value}; pub use error::{Error, Result}; diff --git a/tests/fixtures/tls/README.md b/tests/fixtures/tls/README.md new file mode 100644 index 0000000..4ff8a2a --- /dev/null +++ b/tests/fixtures/tls/README.md @@ -0,0 +1,39 @@ +<!-- +Licensed to the Apache Software Foundation (ASF) under one +or more contributor license agreements. See the NOTICE file +distributed with this work for additional information +regarding copyright ownership. The ASF licenses this file +to you under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance +with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +--> + +# TLS test fixtures + +Self-signed certificate + PKCS#8 key used only by the loopback TLS tests +(`connection::tls_tests`, cargo feature `tls`). Not a secret — the key +never protects anything. + +macOS SecureTransport imposes extra requirements even on explicitly +trusted roots: validity ≤ 825 days (error −67901) and an +`extendedKeyUsage=serverAuth` extension (error −67609). Current cert +expires **2028-10-10**; when the trusted-root test starts failing with a +validity/expiry error, regenerate: + +```sh +cd tests/fixtures/tls +openssl req -x509 -newkey rsa:2048 -keyout key.pem -out cert.pem \ + -days 820 -nodes -subj "/CN=localhost" \ + -addext "subjectAltName=DNS:localhost,IP:127.0.0.1" \ + -addext "extendedKeyUsage=serverAuth" \ + -addext "keyUsage=digitalSignature,keyEncipherment" +``` diff --git a/tests/fixtures/tls/cert.pem b/tests/fixtures/tls/cert.pem new file mode 100644 index 0000000..ca7e9d0 --- /dev/null +++ b/tests/fixtures/tls/cert.pem @@ -0,0 +1,20 @@ +-----BEGIN CERTIFICATE----- +MIIDSTCCAjGgAwIBAgIUXWTGi5P/NlLO5jb4yB/faKrLSmEwDQYJKoZIhvcNAQEL +BQAwFDESMBAGA1UEAwwJbG9jYWxob3N0MB4XDTI2MDcxMzA2MDcyMFoXDTI4MTAx +MDA2MDcyMFowFDESMBAGA1UEAwwJbG9jYWxob3N0MIIBIjANBgkqhkiG9w0BAQEF +AAOCAQ8AMIIBCgKCAQEAihtCghtrZyQRS08IclOEhez3pB5DUfQzfo8sZE4uUIsp +Rxb1gbbVpICi7k8rc3NGKmmIzOOPnLPMlpISunxfdzYWEO1bSJQKEk1DCSFYL60t +ZCBJzJ5XdEQypQ8AqGsC4QGsqEhzzvEF6ttOGVIHOsTzoRnClyZTVJxqU62RfXuk +MqPiXwSa+LTfr/Yz31DBossQyHgTgQvS2zHIwFPQfr/FncAjnAcRs9gWjDFnoIX7 +pxamN1BVa6ypSbPXO3W0PpNGsGTpnOzD+yJFeO2SHzHfVwgI4MC4eyZO1y5YHaF1 +73VGTCdEVCMzmT/o498xkK0Nsp9zLZfbGW82UHrWMQIDAQABo4GSMIGPMB0GA1Ud +DgQWBBSf+ZeI7VU6K87FdqUuvKDXkUpT2jAfBgNVHSMEGDAWgBSf+ZeI7VU6K87F +dqUuvKDXkUpT2jAPBgNVHRMBAf8EBTADAQH/MBoGA1UdEQQTMBGCCWxvY2FsaG9z +dIcEfwAAATATBgNVHSUEDDAKBggrBgEFBQcDATALBgNVHQ8EBAMCBaAwDQYJKoZI +hvcNAQELBQADggEBAGAGhMfPVg6sVzIky9vy0rcve1gRre7SltiBzA7rexjqJZfY +yL6DjxqbSHxbkM1HB02n118BawaVxATFdViYuSFv278AJMryMePKeQJx/rLVAvJn ++kmav9KuQRogxDsO0XwoE7cTa1er0oa9Y9uTLQ/5yoPSl/9dGyGObYQjbb1zavgz +RQs3Px4jrmLvv1iqFXiWSAxqhtV2b3pHOIrDaB4yPtBSNmQhjO+ngHoQ4W7MSv4X +n+5siHE0TSBKfmZx+ETVONRNnElmPGmog3l/A25AmZipq4Yr4VJZbxoBVUDYvVyu +QieLpNS4Owx3y8pCUiNBqpeaFgks4hOtBE9bs5k= +-----END CERTIFICATE----- diff --git a/tests/fixtures/tls/key.pem b/tests/fixtures/tls/key.pem new file mode 100644 index 0000000..88e8568 --- /dev/null +++ b/tests/fixtures/tls/key.pem @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCKG0KCG2tnJBFL +TwhyU4SF7PekHkNR9DN+jyxkTi5QiylHFvWBttWkgKLuTytzc0YqaYjM44+cs8yW +khK6fF93NhYQ7VtIlAoSTUMJIVgvrS1kIEnMnld0RDKlDwCoawLhAayoSHPO8QXq +204ZUgc6xPOhGcKXJlNUnGpTrZF9e6Qyo+JfBJr4tN+v9jPfUMGiyxDIeBOBC9Lb +McjAU9B+v8WdwCOcBxGz2BaMMWeghfunFqY3UFVrrKlJs9c7dbQ+k0awZOmc7MP7 +IkV47ZIfMd9XCAjgwLh7Jk7XLlgdoXXvdUZMJ0RUIzOZP+jj3zGQrQ2yn3Mtl9sZ +bzZQetYxAgMBAAECggEALpxFpZCofn5rQc3yLVenneWcrKy0DshKytd6ZX4HXpKh +A0ep4zX4Q9VQZ/qbURMiJrVIgNH6iaobDeRN4teQBrVf3CWnjca84XFnAwFYeHrL +m4PUNHVqUaikqRV3pN+88IC2q3MMdtbcpueOI7iODRUknoYJqSa7sA48SnKFbGJb +xqMwE8xPoSVnatCWaUVWC+YPwRUsEDFiWZbCA14n9zgXRwvK7vs0yImoi23z3Td2 +70XwxEsPr2Tz1B7Atu4FatL44ZnzQzqzY69N6vDV5VVw9E2ZILAYAehe4s9J0CV9 +M6XWoMb/7qiCHnKY+4q45PSBTAqOp6GTnPUnzJcxMwKBgQDDAIdKQ3x+mJKuotna +ukhOacgk0u5Rn8n+0WY69GZRwd/T2wnWojN05du4SvJ80HS5B4y96TZ6OVkKGNcW +Hx8WE1l4QLy2RDaZp2gM18xBQxWiHNnqBs+HQze029GlA6Si1BITykH7K/1XoD0s +9h6ljKmSCTRsOSlfweFx6sErywKBgQC1Tp5ySKeAG/81WTLDtVRKfbI50lqb31Sn +JhMzoi2iaoladqEkw7pfTxywDCZk7yDXyVQBuwI9XLcEV1hSDNCaoQWNl4Od3a+Z ++6Qeks3qVFHqiKm54FKaTgc7v/lJ9r86vTt7z2TRRu3/9POUJymmv+bSps9QVPk9 +a6wPKCs+cwKBgBrQk8oOhawS5vjExBhjzVWbDj6iEst+oZQ8z0YEHZ3YfyLu1+d/ +3nuQsCojhDzNnX9kHmJvE2KzSB8vU0Rjey7Z8k+q90hJEQkIEDLT6e5/fxYPyZd4 +4EjFYX+pSqbaXWVWrDW0dmZHokrOahsorQYu1ZKTWXYYViwoyQuVMIMJAoGACkN/ +GL0gLlJnah+4jfev49+lvTw5QOWtpyCyqZOevbkGOYbJrIkf/dE+sICfk8stssmE +5ewuPkcHXUmMiGiNTpOa+t7+5mrsS+1A5zIsUf2f/YTrBsi6JF4SbsF5XUSIosyf +l2ywKmC3jGvTdimZZAUtDfO/fK6yxVVZEiqV+B0CgYB2adpXgm38TWgIz7xwpFed +hQvNKn0SxMHfbk3TdXC8kVwl/49NWFBaEd6NHcxsubZICHCbo0BIUkkYCqe0r+Na +DDDRbwrTgcffQwiR+RT9qpxnOLw1/ff8q0hbdmd7Vutq8wPDW2Cx6tEngegKMlpd +5BsJT6Q5VbJ5nQHhz8tPMQ== +-----END PRIVATE KEY-----
