use std::path::PathBuf;
use aetherdb::{DatabaseRef, OpenOptions, open_bundled};
use thiserror::Error;
use crate::aether_tunnel;
use crate::dns;
use super::kdl_build;
use super::topology::Topology;
use kdl_build::WIRE_REPLICA_PRESET;
const DEFAULT_EMBEDDED_PRESET: &str = "kv-local";
const DEFAULT_WIRE_PRESET: &str = "kv";
#[derive(Clone)]
pub enum ConnectMode {
Embedded {
preset: String,
},
File {
path: PathBuf,
preset: String,
},
Bundled {
preset: String,
},
Colocated {
db: aetherdb::DatabaseRef,
},
Wire {
endpoint: String,
preset: String,
tls: Option<aether_tunnel::WireTls>,
insecure: bool,
replicate_key: Option<Vec<u8>>,
},
}
#[derive(Debug, Error)]
pub enum ConnectError {
#[error("AetherDB open failed: {0}")]
Db(#[from] aetherdb::Error),
#[error("invalid file URL (expected file:/path/to.kdl#preset): {0}")]
BadFileUrl(String),
#[error("distributed AetherDB requires a tcp:// endpoint in KWT_REPLAY_AETHERDB_URL")]
SharedEndpointRequired,
#[error(
"distributed AetherDB requires mTLS tunnel (set NAUTALID_AETHER_TUNNEL_* or KWT_REPLAY_AETHERDB_TLS_* paths; NAUTALID_AETHER_TUNNEL_INSECURE=1 for dev only)"
)]
TlsRequired,
#[error("AetherDB wire tunnel: {0}")]
Tunnel(#[from] aether_tunnel::TunnelError),
#[error("DNS resolution failed: {0}")]
Dns(#[from] crate::dns::DnsError),
#[error("invalid KWT_REPLAY_AETHERDB_WIRE_KWT_KEY / master key hex: {0}")]
BadWireKwtKey(String),
#[error("distributed AetherDB requires KWT_REPLAY_AETHERDB_REPLICATE_KEY (64 hex chars)")]
ReplicateKeyRequired,
#[error("invalid KWT_REPLAY_AETHERDB_REPLICATE_KEY hex: {0}")]
BadReplicateKey(String),
}
pub fn spec_from_env() -> ConnectMode {
let url = std::env::var("KWT_REPLAY_AETHERDB_URL")
.ok()
.filter(|s| !s.trim().is_empty())
.map(|s| s.trim().to_string())
.unwrap_or_else(|| "local".into());
parse_url(&url)
}
pub fn preset_from_env() -> String {
std::env::var("KWT_REPLAY_AETHERDB_PRESET")
.ok()
.filter(|s| !s.trim().is_empty())
.unwrap_or_else(|| DEFAULT_EMBEDDED_PRESET.into())
}
fn preset_for_wire() -> String {
std::env::var("KWT_REPLAY_AETHERDB_PRESET")
.ok()
.filter(|s| !s.trim().is_empty())
.unwrap_or_else(|| DEFAULT_WIRE_PRESET.into())
}
pub fn is_wire_endpoint(url: &str) -> bool {
let connect = url.strip_prefix("wire:").map(str::trim).unwrap_or(url);
connect.starts_with("tcp://")
}
pub fn shared_replay_required() -> bool {
matches!(
std::env::var("KWT_REPLAY_REQUIRE_SHARED").ok().as_deref(),
Some("1") | Some("true") | Some("yes")
)
}
impl std::fmt::Debug for ConnectMode {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Embedded { preset } => write!(f, "Embedded({preset})"),
Self::File { path, preset } => write!(f, "File({}, {preset})", path.display()),
Self::Bundled { preset } => write!(f, "Bundled({preset})"),
Self::Colocated { .. } => write!(f, "Colocated"),
Self::Wire {
endpoint,
preset,
insecure,
..
} => write!(f, "Wire({endpoint}, {preset}, insecure={insecure})"),
}
}
}
impl PartialEq for ConnectMode {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
(Self::Embedded { preset: a }, Self::Embedded { preset: b }) => a == b,
(
Self::File {
path: a,
preset: pa,
},
Self::File {
path: b,
preset: pb,
},
) => a == b && pa == pb,
(Self::Bundled { preset: a }, Self::Bundled { preset: b }) => a == b,
(Self::Colocated { .. }, Self::Colocated { .. }) => true,
(
Self::Wire {
endpoint: a,
preset: pa,
tls: ta,
insecure: ia,
replicate_key: ra,
},
Self::Wire {
endpoint: b,
preset: pb,
tls: tb,
insecure: ib,
replicate_key: rb,
},
) => a == b && pa == pb && ta == tb && ia == ib && ra == rb,
_ => false,
}
}
}
impl Eq for ConnectMode {}
fn tls_insecure_allowed() -> bool {
aether_tunnel::insecure_dev_allowed()
}
pub fn parse_url(url: &str) -> ConnectMode {
let lower = url.to_ascii_lowercase();
if lower == "local" || lower == "embedded" {
return ConnectMode::Embedded {
preset: preset_from_env(),
};
}
if let Some(name) = url.strip_prefix("bundled:") {
let name = name.trim();
return ConnectMode::Bundled {
preset: if name.is_empty() {
preset_from_env()
} else {
name.to_string()
},
};
}
if let Some(rest) = url.strip_prefix("file:") {
if let Some((path, preset)) = split_file_spec(rest) {
return ConnectMode::File {
path: PathBuf::from(path),
preset: preset.to_string(),
};
}
return ConnectMode::File {
path: PathBuf::from(rest),
preset: preset_from_env(),
};
}
let connect = url
.strip_prefix("wire:")
.map(str::trim)
.filter(|s| !s.is_empty())
.unwrap_or(url);
ConnectMode::Wire {
endpoint: connect.to_string(),
preset: preset_for_wire(),
tls: None,
insecure: tls_insecure_allowed(),
replicate_key: None,
}
}
fn split_file_spec(spec: &str) -> Option<(&str, &str)> {
if let Some((path, preset)) = spec.split_once('#') {
return Some((path, preset));
}
spec.rsplit_once(':')
.filter(|(path, preset)| !path.is_empty() && !preset.is_empty())
}
fn wire_replicate_key_from_env() -> Result<Option<Vec<u8>>, ConnectError> {
let Some(hex) = std::env::var("KWT_REPLAY_AETHERDB_REPLICATE_KEY")
.ok()
.filter(|s| !s.trim().is_empty())
else {
return Ok(None);
};
aetherdb::decode_master_key_hex(hex.trim())
.map(|key| Some(key.to_vec()))
.map_err(|e| ConnectError::BadReplicateKey(e.to_string()))
}
fn wire_kwt_master_key() -> Result<Option<Vec<u8>>, ConnectError> {
let Some(hex) = std::env::var("KWT_REPLAY_AETHERDB_WIRE_KWT_KEY")
.ok()
.filter(|s| !s.trim().is_empty())
.or_else(|| {
std::env::var("IMAGE_TRUST_KWT_MASTER_KEY")
.ok()
.filter(|s| !s.trim().is_empty())
})
else {
return Ok(None);
};
aetherdb::decode_master_key_hex(hex.trim())
.map(|key| Some(key.to_vec()))
.map_err(|e| ConnectError::BadWireKwtKey(e.to_string()))
}
fn validate_wire(spec: &ConnectMode) -> Result<(), ConnectError> {
validate_wire_spec(spec, shared_replay_required())
}
fn validate_wire_spec(spec: &ConnectMode, require_shared: bool) -> Result<(), ConnectError> {
let ConnectMode::Wire {
endpoint,
tls,
insecure,
..
} = spec
else {
if require_shared {
return Err(ConnectError::SharedEndpointRequired);
}
return Ok(());
};
if !is_wire_endpoint(endpoint) {
return Err(ConnectError::SharedEndpointRequired);
}
if *insecure {
tracing::warn!(
"NAUTALID_AETHER_TUNNEL_INSECURE=1 — AetherDB wire without mTLS tunnel (dev only)"
);
return Ok(());
}
if tls.is_none() {
return Err(ConnectError::TlsRequired);
}
Ok(())
}
pub async fn resolve_wire_endpoint(endpoint: &str) -> Result<String, ConnectError> {
if !is_wire_endpoint(endpoint) {
return Ok(endpoint.to_string());
}
Ok(dns::resolve_tcp_connect(endpoint).await?)
}
pub fn resolve_wire_endpoint_blocking(endpoint: &str) -> Result<String, ConnectError> {
if !is_wire_endpoint(endpoint) {
return Ok(endpoint.to_string());
}
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|e| ConnectError::Dns(crate::dns::DnsError::Hickory(e.to_string())))?;
runtime.block_on(resolve_wire_endpoint(endpoint))
}
pub fn open(spec: &ConnectMode) -> Result<DatabaseRef, ConnectError> {
if let ConnectMode::Colocated { db } = spec {
return Ok(db.clone());
}
validate_wire(spec)?;
let kwt_master_key = wire_kwt_master_key()?;
let wire_tls = aether_tunnel::load_tunnel()?;
let topology = Topology::from_env();
match spec {
ConnectMode::Colocated { db } => Ok(db.clone()),
ConnectMode::Embedded { preset } | ConnectMode::Bundled { preset } => {
Ok(open_bundled(preset)?)
}
ConnectMode::File { path, preset } => {
if !path.exists() {
return Err(ConnectError::BadFileUrl(format!(
"KDL file not found: {}",
path.display()
)));
}
Ok(DatabaseRef::open(path, preset)?)
}
ConnectMode::Wire {
endpoint,
tls: spec_tls,
insecure,
replicate_key,
..
} => {
let tls = spec_tls.as_ref().or(wire_tls.as_ref());
let resolved = resolve_wire_endpoint_blocking(endpoint)?;
let replicate_key = replicate_key
.clone()
.or(wire_replicate_key_from_env()?)
.ok_or(ConnectError::ReplicateKeyRequired)?;
let replicate_hex = replicate_key
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
let options = OpenOptions {
replication_handler: true,
replicate_master_key: Some(replicate_key),
kwt_master_key,
..OpenOptions::default()
};
if *insecure || tls.is_none() {
let kdl = kdl_build::wire_client_kdl(&topology, &replicate_hex);
return Ok(DatabaseRef::open_str_with_options(
&kdl,
WIRE_REPLICA_PRESET,
OpenOptions {
transport_connect: Some(resolved),
..options
},
)?);
}
let tls = tls.expect("validated above");
let kdl = kdl_build::wire_client_kdl_tls(&topology, &resolved, tls, &replicate_hex);
Ok(DatabaseRef::open_str_with_options(
&kdl,
WIRE_REPLICA_PRESET,
options,
)?)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::aether_tunnel::WireTls;
use super::super::topology::DeployMode;
#[test]
fn detects_wire_urls() {
assert!(is_wire_endpoint("tcp://db.internal:5555"));
assert!(is_wire_endpoint("wire:tcp://10.0.0.1:5555"));
assert!(!is_wire_endpoint("local"));
}
#[test]
fn parse_local_defaults_embedded() {
assert_eq!(
parse_url("local"),
ConnectMode::Embedded {
preset: DEFAULT_EMBEDDED_PRESET.into()
}
);
}
#[test]
fn parse_wire_defaults_kv_preset() {
assert_eq!(
parse_url("tcp://127.0.0.1:5555"),
ConnectMode::Wire {
endpoint: "tcp://127.0.0.1:5555".into(),
preset: DEFAULT_WIRE_PRESET.into(),
tls: None,
insecure: false,
replicate_key: None,
}
);
}
#[test]
fn wire_replica_kdl_includes_replicate_auth() {
let topo = Topology {
mode: DeployMode::Replicate,
shards: 1,
replication: 3,
node_id: 0,
query: "kv".into(),
};
let kdl = kdl_build::wire_client_kdl(&topo, "aa".repeat(64).as_str());
assert!(kdl.contains(r#"route "replicate" handler="replication""#));
assert!(kdl.contains("replicate_auth"));
}
#[test]
fn wire_replica_tls_kdl_includes_connect_and_tls() {
let topo = Topology {
mode: DeployMode::Cluster,
shards: 16,
replication: 3,
node_id: 0,
query: "sql".into(),
};
let kdl = kdl_build::wire_client_kdl_tls(
&topo,
"tcp://127.0.0.1:5555",
&WireTls {
cert_file: "/c/client.pem".into(),
key_file: "/c/client-key.pem".into(),
ca_file: "/c/ca.pem".into(),
},
"bb".repeat(32).as_str(),
);
assert!(kdl.contains(r#"connect="tcp://127.0.0.1:5555""#));
assert!(kdl.contains("shards 16"));
}
#[test]
fn shared_requires_tcp_when_flag_set() {
let spec = ConnectMode::Embedded {
preset: "kv-local".into(),
};
let err = validate_wire_spec(&spec, true).unwrap_err();
assert!(matches!(err, ConnectError::SharedEndpointRequired));
}
}