use crate::conn::resolve_pg_config;
use crate::db::ConnectError;
use crate::types::SessionConfig;
use std::pin::Pin;
use std::process::Stdio;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, ReadBuf};
use tokio::process::{ChildStderr, ChildStdin, ChildStdout};
use tokio::sync::oneshot;
use tokio_postgres::config::Host;
use tokio_postgres::{Client, NoTls};
const DEFAULT_DRIVER: &str = "docker";
const DEFAULT_REMOTE_HOST: &str = "127.0.0.1";
const DEFAULT_REMOTE_PORT: u16 = 5432;
const STDERR_CAPTURE_LIMIT: usize = 8 * 1024;
const STDERR_HINT_BYTES: usize = 512;
const BRIDGE_READY_PREFIX: &str = "AFPSQL_BRIDGE_OK";
const BRIDGE_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(2);
static BRIDGE_NONCE_COUNTER: AtomicU64 = AtomicU64::new(0);
fn generate_bridge_nonce() -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos() as u64)
.unwrap_or(0);
let counter = BRIDGE_NONCE_COUNTER.fetch_add(1, Ordering::Relaxed);
let mixed = nanos
.wrapping_mul(0x9E37_79B9_7F4A_7C15)
.wrapping_add(counter);
format!("{mixed:016x}")
}
const PYTHON_BRIDGE: &str = r#"import os,select,socket,sys
mode=sys.argv[1]
if mode=="tcp":
s=socket.create_connection((sys.argv[2], int(sys.argv[3])))
elif mode=="unix":
s=socket.socket(socket.AF_UNIX)
s.connect(sys.argv[2])
else:
sys.stderr.write("unsupported bridge mode: "+mode+"\n")
sys.exit(2)
stdin_obj=getattr(sys.stdin,"buffer",sys.stdin)
stdout_obj=getattr(sys.stdout,"buffer",sys.stdout)
stdin_fd=stdin_obj.fileno()
stdout_fd=stdout_obj.fileno()
stdin_open=True
while True:
readers=[s]
if stdin_open:
readers.append(stdin_fd)
ready,_,_=select.select(readers,[],[])
if stdin_fd in ready:
data=os.read(stdin_fd,65536)
if data:
s.sendall(data)
else:
stdin_open=False
try:
s.shutdown(socket.SHUT_WR)
except OSError:
pass
if s in ready:
data=s.recv(65536)
if data:
os.write(stdout_fd,data)
else:
break
"#;
const PERL_BRIDGE: &str = r#"use strict; use warnings; use IO::Socket::INET; use IO::Socket::UNIX; use IO::Select; use Socket qw(SOCK_STREAM);
my $mode = shift @ARGV;
my $sock;
if ($mode eq "tcp") {
my ($host, $port) = @ARGV;
$sock = IO::Socket::INET->new(PeerHost => $host, PeerPort => $port, Proto => "tcp") or die "connect tcp failed: $!";
} elsif ($mode eq "unix") {
my ($path) = @ARGV;
$sock = IO::Socket::UNIX->new(Type => SOCK_STREAM, Peer => $path) or die "connect unix failed: $!";
} else {
die "unsupported bridge mode: $mode";
}
binmode STDIN; binmode STDOUT; binmode $sock;
my $select = IO::Select->new($sock, \*STDIN);
while (1) {
for my $fh ($select->can_read) {
if ($fh == \*STDIN) {
my $buf = "";
my $n = sysread(STDIN, $buf, 65536);
die "read stdin failed: $!" unless defined $n;
if ($n == 0) {
$select->remove(\*STDIN);
shutdown($sock, 1);
} else {
write_all($sock, $buf);
}
} else {
my $buf = "";
my $n = sysread($sock, $buf, 65536);
die "read socket failed: $!" unless defined $n;
exit 0 if $n == 0;
write_all(\*STDOUT, $buf);
}
}
}
sub write_all {
my ($fh, $buf) = @_;
my $off = 0;
my $len = length($buf);
while ($off < $len) {
my $n = syswrite($fh, $buf, $len - $off, $off);
die "write failed: $!" unless defined $n;
$off += $n;
}
}
"#;
const SHELL_BRIDGE_BODY: &str = r#"if command -v python3 >/dev/null 2>&1; then
echo "AFPSQL_BRIDGE_OK $AFPSQL_BRIDGE_NONCE" >&2
exec python3 -c "$AFPSQL_CONTAINER_PY_BRIDGE" "$@"
fi
if command -v python >/dev/null 2>&1; then
echo "AFPSQL_BRIDGE_OK $AFPSQL_BRIDGE_NONCE" >&2
exec python -c "$AFPSQL_CONTAINER_PY_BRIDGE" "$@"
fi
if command -v perl >/dev/null 2>&1; then
echo "AFPSQL_BRIDGE_OK $AFPSQL_BRIDGE_NONCE" >&2
exec perl -e "$AFPSQL_CONTAINER_PERL_BRIDGE" "$@"
fi
echo "afpsql container bridge requires python3, python, or perl in the container" >&2
exit 127
"#;
pub struct ContainerBridgeGuard {
child: Option<tokio::process::Child>,
connection_task: Option<tokio::task::JoinHandle<()>>,
stderr_task: Option<tokio::task::JoinHandle<()>>,
}
impl ContainerBridgeGuard {
pub fn is_finished(&self) -> bool {
self.connection_task
.as_ref()
.map(|task| task.is_finished())
.unwrap_or(true)
}
pub async fn shutdown(mut self, timeout: Duration) {
if let Some(task) = self.connection_task.take() {
let mut task = task;
if tokio::time::timeout(timeout, &mut task).await.is_err() {
task.abort();
}
}
if let Some(task) = self.stderr_task.take() {
let mut task = task;
if tokio::time::timeout(timeout, &mut task).await.is_err() {
task.abort();
}
}
if let Some(mut child) = self.child.take() {
let _ = child.start_kill();
let _ = tokio::time::timeout(timeout, child.wait()).await;
}
}
}
impl Drop for ContainerBridgeGuard {
fn drop(&mut self) {
if let Some(task) = self.connection_task.as_ref() {
task.abort();
}
if let Some(task) = self.stderr_task.as_ref() {
task.abort();
}
if let Some(child) = self.child.as_mut() {
let _ = child.start_kill();
}
}
}
struct ContainerStdioStream {
stdout: ChildStdout,
stdin: ChildStdin,
}
impl AsyncRead for ContainerStdioStream {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.stdout).poll_read(cx, buf)
}
}
impl AsyncWrite for ContainerStdioStream {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
Pin::new(&mut self.stdin).poll_write(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.stdin).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.stdin).poll_shutdown(cx)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ContainerDriver {
Docker,
Podman,
Nerdctl,
Compose,
Kubectl,
}
impl ContainerDriver {
fn parse(value: &str) -> Result<Self, String> {
match value {
"docker" => Ok(Self::Docker),
"podman" => Ok(Self::Podman),
"nerdctl" => Ok(Self::Nerdctl),
"compose" | "docker-compose" => Ok(Self::Compose),
"kubectl" | "kubernetes" | "k8s" => Ok(Self::Kubectl),
_ => Err(format!(
"unsupported container driver `{value}`; expected docker, podman, nerdctl, compose, or kubectl"
)),
}
}
fn default_runtime(self) -> &'static str {
match self {
Self::Docker | Self::Compose => "docker",
Self::Podman => "podman",
Self::Nerdctl => "nerdctl",
Self::Kubectl => "kubectl",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct ContainerSettings {
driver: ContainerDriver,
runtime: String,
target: String,
user: Option<String>,
namespace: Option<String>,
context: Option<String>,
compose_files: Vec<String>,
compose_project: Option<String>,
pod_container: Option<String>,
ssh_destination: Option<String>,
ssh_options: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum ContainerTarget {
Tcp { host: String, port: u16 },
UnixSocket { path: String },
}
pub async fn connect_stdio_bridge(
cfg: &SessionConfig,
) -> Result<(Client, ContainerBridgeGuard), ConnectError> {
let settings = resolve_container_settings(cfg)?;
let target = resolve_container_target(cfg)?;
let nonce = generate_bridge_nonce();
let (program, args) =
build_bridge_process(&settings, &target, &nonce).map_err(ConnectError::new)?;
let mut child = tokio::process::Command::new(&program)
.args(&args)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true)
.spawn()
.map_err(|e| ConnectError::new(format!("start container bridge failed: {e}")))?;
let stderr = child.stderr.take();
let stderr_capture = Arc::new(Mutex::new(Vec::new()));
let (stderr_task, handshake_rx) = stderr
.map(|stderr| spawn_stderr_capture(stderr, Arc::clone(&stderr_capture), nonce.clone()))
.map(|(task, rx)| (Some(task), Some(rx)))
.unwrap_or((None, None));
let stdin = child.stdin.take().ok_or_else(|| {
ConnectError::new("start container bridge failed: stdin pipe unavailable")
})?;
let stdout = child.stdout.take().ok_or_else(|| {
ConnectError::new("start container bridge failed: stdout pipe unavailable")
})?;
if let Err(mut err) = wait_bridge_handshake(&mut child, &stderr_capture, handshake_rx).await {
let stderr_text = captured_stderr(&stderr_capture);
if stderr_indicates_missing_target(&stderr_text)
&& let Some(list) = list_container_targets(&settings).await
{
err.hint = Some(match err.hint.take() {
Some(base) => format!("{base}; available containers: {list}"),
None => format!("available containers: {list}"),
});
}
return Err(err);
}
let stream = ContainerStdioStream { stdout, stdin };
let pg_cfg = resolve_pg_config(cfg)
.map_err(|e| ConnectError::new(format!("invalid container connection config: {e}")))?;
let (client, connection) = match pg_cfg.connect_raw(stream, NoTls).await {
Ok(connection) => connection,
Err(e) => {
let status = match tokio::time::timeout(Duration::from_millis(100), child.wait()).await
{
Ok(Ok(status)) => Some(status),
_ => child.try_wait().ok().flatten(),
};
tokio::time::sleep(Duration::from_millis(20)).await;
let stderr_text = captured_stderr(&stderr_capture);
return Err(enrich_container_connect_error(
ConnectError::from_pg_error("connect through container bridge failed", e),
status,
&stderr_text,
));
}
};
let connection_task = tokio::spawn(async move {
let _ = connection.await;
});
Ok((
client,
ContainerBridgeGuard {
child: Some(child),
connection_task: Some(connection_task),
stderr_task,
},
))
}
fn spawn_stderr_capture(
mut stderr: ChildStderr,
capture: Arc<Mutex<Vec<u8>>>,
nonce: String,
) -> (tokio::task::JoinHandle<()>, oneshot::Receiver<bool>) {
let (handshake_tx, handshake_rx) = oneshot::channel();
let task = tokio::spawn(async move {
let mut buf = [0u8; 1024];
let mut pending = Vec::new();
let mut handshake_tx = Some(handshake_tx);
while let Ok(n) = stderr.read(&mut buf).await {
if n == 0 {
break;
}
pending.extend_from_slice(&buf[..n]);
while let Some(pos) = pending.iter().position(|b| *b == b'\n') {
let line = pending.drain(..=pos).collect::<Vec<_>>();
handle_stderr_line(line, &capture, &mut handshake_tx, &nonce);
}
}
if !pending.is_empty() {
handle_stderr_line(pending, &capture, &mut handshake_tx, &nonce);
}
if let Some(tx) = handshake_tx.take() {
let _ = tx.send(false);
}
});
(task, handshake_rx)
}
fn handle_stderr_line(
line: Vec<u8>,
capture: &Arc<Mutex<Vec<u8>>>,
handshake_tx: &mut Option<oneshot::Sender<bool>>,
nonce: &str,
) {
if stderr_line_is_banner(&line, nonce) {
if let Some(tx) = handshake_tx.take() {
let _ = tx.send(true);
}
return;
}
append_stderr_capture(capture, &line);
}
fn stderr_line_is_banner(line: &[u8], nonce: &str) -> bool {
let line = line.strip_suffix(b"\n").unwrap_or(line);
let line = line.strip_suffix(b"\r").unwrap_or(line);
let expected = format!("{BRIDGE_READY_PREFIX} {nonce}");
line == expected.as_bytes()
}
fn append_stderr_capture(capture: &Arc<Mutex<Vec<u8>>>, bytes: &[u8]) {
if let Ok(mut captured) = capture.lock() {
let remaining = STDERR_CAPTURE_LIMIT.saturating_sub(captured.len());
if remaining > 0 {
captured.extend_from_slice(&bytes[..bytes.len().min(remaining)]);
}
}
}
async fn wait_bridge_handshake(
child: &mut tokio::process::Child,
stderr_capture: &Arc<Mutex<Vec<u8>>>,
handshake_rx: Option<oneshot::Receiver<bool>>,
) -> Result<(), ConnectError> {
let Some(handshake_rx) = handshake_rx else {
return Err(handshake_connect_error(
child,
stderr_capture,
"container bridge stderr pipe unavailable for startup handshake",
)
.await);
};
match tokio::time::timeout(BRIDGE_HANDSHAKE_TIMEOUT, handshake_rx).await {
Ok(Ok(true)) => Ok(()),
Ok(Ok(false)) => Err(handshake_connect_error(
child,
stderr_capture,
"container bridge exited before startup handshake",
)
.await),
Ok(Err(_)) => Err(handshake_connect_error(
child,
stderr_capture,
"container bridge startup handshake channel closed",
)
.await),
Err(_) => Err(handshake_connect_error(
child,
stderr_capture,
"container bridge did not emit startup handshake before timeout",
)
.await),
}
}
async fn handshake_connect_error(
child: &mut tokio::process::Child,
stderr_capture: &Arc<Mutex<Vec<u8>>>,
message: &str,
) -> ConnectError {
let status = child.try_wait().ok().flatten();
let stderr_text = captured_stderr(stderr_capture);
let mut err = ConnectError::new(message);
err.hint = Some(if stderr_text.is_empty() {
format!(
"the container bridge never reported {BRIDGE_READY_PREFIX}; check target name, runtime access, /bin/sh, and bridge interpreter prerequisites"
)
} else {
format!(
"the container bridge wrote diagnostics before {BRIDGE_READY_PREFIX}; container bridge stderr: {stderr_text}"
)
});
enrich_container_connect_error(err, status, &stderr_text)
}
fn captured_stderr(capture: &Arc<Mutex<Vec<u8>>>) -> String {
capture
.lock()
.ok()
.map(|captured| sanitize_diagnostic(&captured))
.filter(|text| !text.is_empty())
.unwrap_or_default()
}
fn sanitize_diagnostic(bytes: &[u8]) -> String {
let text = String::from_utf8_lossy(bytes);
let mut out = String::with_capacity(STDERR_HINT_BYTES);
for ch in text.chars() {
let mapped = if matches!(ch, '\n' | '\t') {
ch
} else if ch.is_control() {
' '
} else {
ch
};
if out.len() + mapped.len_utf8() > STDERR_HINT_BYTES {
break;
}
out.push(mapped);
}
out.trim().to_string()
}
fn enrich_container_connect_error(
mut err: ConnectError,
status: Option<std::process::ExitStatus>,
stderr: &str,
) -> ConnectError {
if let Some(hint) = container_bridge_hint(status, stderr) {
err.hint = Some(match err.hint.take() {
Some(base) => format!("{base}; {hint}"),
None => hint,
});
}
err
}
fn container_bridge_hint(status: Option<std::process::ExitStatus>, stderr: &str) -> Option<String> {
let trimmed = stderr.trim();
let lower = trimmed.to_ascii_lowercase();
let base = if lower.contains("requires python3, python, or perl") {
"the container bridge started but the container has no supported interpreter; install python3/python/perl, use a sidecar, or connect through the host instead"
} else if lower.contains("sh:") && lower.contains("not found") {
"the container bridge requires /bin/sh or compatible shell in the target container"
} else if lower.contains("no such container")
|| lower.contains("not found")
|| lower.contains("is not running")
{
"the container runtime could not exec into the target; check the container target name and running state"
} else if lower.contains("error from server") || lower.contains("pods") {
"kubectl could not exec into the target; check context, namespace, pod name, and cluster access"
} else if matches!(status.and_then(|s| s.code()), Some(125..=127)) {
"the container runtime or bridge command exited before PostgreSQL handshake; check runtime access, target name, shell, and bridge prerequisites"
} else if !trimmed.is_empty() {
"the container bridge wrote diagnostics before PostgreSQL handshake failed"
} else {
return None;
};
if trimmed.is_empty() {
Some(base.to_string())
} else {
Some(format!("{base}; container bridge stderr: {trimmed}"))
}
}
fn stderr_indicates_missing_target(stderr: &str) -> bool {
let lower = stderr.to_ascii_lowercase();
lower.contains("no such container")
|| lower.contains("is not running")
|| (lower.contains("not found") && lower.contains("container"))
}
const CONTAINER_LIST_TIMEOUT: Duration = Duration::from_secs(3);
async fn list_container_targets(settings: &ContainerSettings) -> Option<String> {
if settings.ssh_destination.is_some() {
return None;
}
let mut args = Vec::new();
match settings.driver {
ContainerDriver::Docker | ContainerDriver::Podman | ContainerDriver::Nerdctl => {
if settings.driver == ContainerDriver::Docker
&& let Some(context) = settings.context.as_ref()
{
args.push(format!("--context={context}"));
}
args.push("ps".to_string());
args.push("--format".to_string());
args.push("{{.Names}}".to_string());
}
ContainerDriver::Compose | ContainerDriver::Kubectl => return None,
}
let output = tokio::time::timeout(
CONTAINER_LIST_TIMEOUT,
tokio::process::Command::new(&settings.runtime)
.args(&args)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.kill_on_drop(true)
.output(),
)
.await
.ok()?
.ok()?;
let names: Vec<String> = String::from_utf8_lossy(&output.stdout)
.lines()
.map(|line| line.trim())
.filter(|line| !line.is_empty())
.map(|line| line.to_string())
.collect();
if names.is_empty() {
return None;
}
const MAX_LISTED: usize = 30;
if names.len() > MAX_LISTED {
let shown = names[..MAX_LISTED].join(", ");
let extra = names.len() - MAX_LISTED;
Some(format!("{shown} (+{extra} more)"))
} else {
Some(names.join(", "))
}
}
fn resolve_container_settings(cfg: &SessionConfig) -> Result<ContainerSettings, String> {
let target = cfg
.container
.target
.clone()
.or_else(|| env_nonempty("AFPSQL_CONTAINER"));
let ssh_destination = cfg
.ssh
.destination
.clone()
.or_else(|| env_nonempty("AFPSQL_SSH"));
let namespace = cfg
.container
.namespace
.clone()
.or_else(|| env_nonempty("AFPSQL_CONTAINER_NAMESPACE"));
let context = cfg
.container
.context
.clone()
.or_else(|| env_nonempty("AFPSQL_CONTAINER_CONTEXT"));
let compose_files = if cfg.container.compose_files.is_empty() {
env_colon_list("AFPSQL_CONTAINER_COMPOSE_FILE")
} else {
cfg.container.compose_files.clone()
};
let compose_project = cfg
.container
.compose_project
.clone()
.or_else(|| env_nonempty("AFPSQL_CONTAINER_COMPOSE_PROJECT"));
let pod_container = cfg
.container
.pod_container
.clone()
.or_else(|| env_nonempty("AFPSQL_CONTAINER_POD_CONTAINER"));
let has_container_fields = target.is_some()
|| cfg.container.driver.is_some()
|| cfg.container.runtime.is_some()
|| cfg.container.user.is_some()
|| namespace.is_some()
|| context.is_some()
|| !compose_files.is_empty()
|| compose_project.is_some()
|| pod_container.is_some();
let Some(target) = target else {
if has_container_fields {
return Err(
"--container is required when container transport options are set".to_string(),
);
}
return Err("--container is required for container transport".to_string());
};
if target.trim().is_empty() {
return Err("--container requires a non-empty target name".to_string());
}
if let Some(destination) = ssh_destination.as_ref() {
if destination.trim().is_empty() {
return Err("--ssh requires a non-empty USER@HOST destination".to_string());
}
} else if !cfg.ssh.options.is_empty() {
return Err(
"--ssh is required when --ssh-option is combined with container transport".to_string(),
);
}
if cfg.ssh.has_tunnel_or_bridge_options() {
return Err("container transport with --ssh supports only --ssh and --ssh-option; SSH tunnel and sudo bridge options are for non-container SSH transport".to_string());
}
let driver_name = cfg
.container
.driver
.clone()
.or_else(|| env_nonempty("AFPSQL_CONTAINER_DRIVER"))
.unwrap_or_else(|| DEFAULT_DRIVER.to_string());
let driver = ContainerDriver::parse(&driver_name)?;
let user = cfg
.container
.user
.clone()
.or_else(|| env_nonempty("AFPSQL_CONTAINER_USER"));
if driver == ContainerDriver::Kubectl && user.is_some() {
return Err(
"--container-user is not supported with --container-driver kubectl".to_string(),
);
}
validate_driver_scoped_options(
driver,
namespace.as_ref(),
context.as_ref(),
&compose_files,
compose_project.as_ref(),
pod_container.as_ref(),
)?;
Ok(ContainerSettings {
runtime: cfg
.container
.runtime
.clone()
.or_else(|| env_nonempty("AFPSQL_CONTAINER_RUNTIME"))
.unwrap_or_else(|| driver.default_runtime().to_string()),
driver,
target,
user,
namespace,
context,
compose_files,
compose_project,
pod_container,
ssh_destination,
ssh_options: cfg.ssh.options.clone(),
})
}
fn validate_driver_scoped_options(
driver: ContainerDriver,
namespace: Option<&String>,
context: Option<&String>,
compose_files: &[String],
compose_project: Option<&String>,
pod_container: Option<&String>,
) -> Result<(), String> {
match driver {
ContainerDriver::Docker => {
reject_present(
namespace,
"--container-namespace requires --container-driver kubectl",
)?;
reject_present(
pod_container,
"--container-pod-container requires --container-driver kubectl",
)?;
reject_nonempty(
compose_files,
"--container-compose-file requires --container-driver compose",
)?;
reject_present(
compose_project,
"--container-compose-project requires --container-driver compose",
)?;
}
ContainerDriver::Compose => {
reject_present(
namespace,
"--container-namespace requires --container-driver kubectl",
)?;
reject_present(
context,
"--container-context requires --container-driver docker or kubectl",
)?;
reject_present(
pod_container,
"--container-pod-container requires --container-driver kubectl",
)?;
}
ContainerDriver::Kubectl => {
reject_nonempty(
compose_files,
"--container-compose-file requires --container-driver compose",
)?;
reject_present(
compose_project,
"--container-compose-project requires --container-driver compose",
)?;
}
ContainerDriver::Podman | ContainerDriver::Nerdctl => {
reject_present(
namespace,
"--container-namespace requires --container-driver kubectl",
)?;
reject_present(
context,
"--container-context requires --container-driver docker or kubectl",
)?;
reject_nonempty(
compose_files,
"--container-compose-file requires --container-driver compose",
)?;
reject_present(
compose_project,
"--container-compose-project requires --container-driver compose",
)?;
reject_present(
pod_container,
"--container-pod-container requires --container-driver kubectl",
)?;
}
}
Ok(())
}
fn reject_present<T>(value: Option<T>, message: &str) -> Result<(), String> {
if value.is_some() {
Err(message.to_string())
} else {
Ok(())
}
}
fn reject_nonempty<T>(value: &[T], message: &str) -> Result<(), String> {
if value.is_empty() {
Ok(())
} else {
Err(message.to_string())
}
}
fn resolve_container_target(cfg: &SessionConfig) -> Result<ContainerTarget, String> {
let pg_cfg =
resolve_pg_config(cfg).map_err(|e| format!("invalid container connection config: {e}"))?;
let hosts = pg_cfg.get_hosts();
if hosts.len() > 1 {
return Err("container transport supports a single PostgreSQL host".to_string());
}
let ports = pg_cfg.get_ports();
if ports.len() > 1 {
return Err("container transport supports a single PostgreSQL port".to_string());
}
let port = ports.first().copied().unwrap_or(DEFAULT_REMOTE_PORT);
match hosts.first() {
Some(Host::Tcp(host)) => {
if host.starts_with('/') {
Ok(ContainerTarget::UnixSocket {
path: socket_file_from_dir(host, port),
})
} else {
Ok(ContainerTarget::Tcp {
host: host.clone(),
port,
})
}
}
#[cfg(unix)]
Some(Host::Unix(path)) => Ok(ContainerTarget::UnixSocket {
path: socket_file_from_dir(&path.to_string_lossy(), port),
}),
None => Ok(ContainerTarget::Tcp {
host: DEFAULT_REMOTE_HOST.to_string(),
port,
}),
}
}
fn socket_file_from_dir(dir: &str, port: u16) -> String {
format!("{}/.s.PGSQL.{port}", dir.trim_end_matches('/'))
}
fn build_bridge_process(
settings: &ContainerSettings,
target: &ContainerTarget,
nonce: &str,
) -> Result<(String, Vec<String>), String> {
if settings.ssh_destination.is_some() {
Ok((
"ssh".to_string(),
build_bridge_ssh_args(settings, target, nonce)?,
))
} else {
Ok((
settings.runtime.clone(),
build_container_exec_args(settings, target, nonce)?,
))
}
}
fn build_container_exec_args(
settings: &ContainerSettings,
target: &ContainerTarget,
nonce: &str,
) -> Result<Vec<String>, String> {
let command = bridge_command_args(target, nonce);
let mut args = Vec::new();
match settings.driver {
ContainerDriver::Docker | ContainerDriver::Podman | ContainerDriver::Nerdctl => {
if settings.driver == ContainerDriver::Docker
&& let Some(context) = settings.context.as_ref()
{
args.push(format!("--context={context}"));
}
args.push("exec".to_string());
args.push("-i".to_string());
if let Some(user) = settings.user.as_ref() {
args.push("--user".to_string());
args.push(user.clone());
}
args.push(settings.target.clone());
args.extend(command);
}
ContainerDriver::Compose => {
if !is_docker_compose_runtime(&settings.runtime) {
args.push("compose".to_string());
}
for file in &settings.compose_files {
args.push("-f".to_string());
args.push(file.clone());
}
if let Some(project) = settings.compose_project.as_ref() {
args.push("-p".to_string());
args.push(project.clone());
}
args.push("exec".to_string());
args.push("-T".to_string());
if let Some(user) = settings.user.as_ref() {
args.push("--user".to_string());
args.push(user.clone());
}
args.push(settings.target.clone());
args.extend(command);
}
ContainerDriver::Kubectl => {
if let Some(context) = settings.context.as_ref() {
args.push(format!("--context={context}"));
}
if let Some(namespace) = settings.namespace.as_ref() {
args.push(format!("--namespace={namespace}"));
}
args.push("exec".to_string());
args.push("-i".to_string());
args.push(settings.target.clone());
if let Some(container) = settings.pod_container.as_ref() {
args.push("-c".to_string());
args.push(container.clone());
}
args.push("--".to_string());
args.extend(command);
}
}
Ok(args)
}
fn is_docker_compose_runtime(runtime: &str) -> bool {
runtime
.rsplit(['/', '\\'])
.next()
.is_some_and(|name| name == "docker-compose")
}
fn bridge_command_args(target: &ContainerTarget, nonce: &str) -> Vec<String> {
let mut args = vec![
"sh".to_string(),
"-c".to_string(),
shell_bridge_script(nonce),
"afpsql-container-bridge".to_string(),
];
match target {
ContainerTarget::Tcp { host, port } => {
args.push("tcp".to_string());
args.push(host.clone());
args.push(port.to_string());
}
ContainerTarget::UnixSocket { path } => {
args.push("unix".to_string());
args.push(path.clone());
}
}
args
}
fn shell_bridge_script(nonce: &str) -> String {
format!(
"AFPSQL_BRIDGE_NONCE={}; AFPSQL_CONTAINER_PY_BRIDGE={}; AFPSQL_CONTAINER_PERL_BRIDGE={}; {}",
shell_quote(nonce),
shell_quote(PYTHON_BRIDGE),
shell_quote(PERL_BRIDGE),
SHELL_BRIDGE_BODY
)
}
fn build_bridge_ssh_args(
settings: &ContainerSettings,
target: &ContainerTarget,
nonce: &str,
) -> Result<Vec<String>, String> {
let mut args = vec![
"-T".to_string(),
"-o".to_string(),
"BatchMode=yes".to_string(),
];
for option in &settings.ssh_options {
args.push("-o".to_string());
args.push(option.clone());
}
if let Some(destination) = settings.ssh_destination.as_ref() {
args.push(destination.clone());
}
args.push(remote_container_command(settings, target, nonce)?);
Ok(args)
}
fn remote_container_command(
settings: &ContainerSettings,
target: &ContainerTarget,
nonce: &str,
) -> Result<String, String> {
Ok(std::iter::once(settings.runtime.clone())
.chain(build_container_exec_args(settings, target, nonce)?)
.map(|arg| shell_quote(&arg))
.collect::<Vec<_>>()
.join(" "))
}
fn shell_quote(value: &str) -> String {
if value.is_empty() {
return "''".to_string();
}
if value.bytes().all(|b| {
b.is_ascii_alphanumeric() || matches!(b, b'_' | b'-' | b'.' | b'/' | b':' | b'@' | b'=')
}) {
return value.to_string();
}
format!("'{}'", value.replace('\'', "'\\''"))
}
fn env_nonempty(name: &str) -> Option<String> {
std::env::var(name).ok().filter(|value| !value.is_empty())
}
fn env_colon_list(name: &str) -> Vec<String> {
std::env::var(name)
.ok()
.map(|value| {
value
.split(':')
.filter(|part| !part.is_empty())
.map(std::string::ToString::to_string)
.collect()
})
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use super::*;
const TEST_NONCE: &str = "deadbeefcafef00d";
fn settings(driver: ContainerDriver) -> ContainerSettings {
ContainerSettings {
driver,
runtime: driver.default_runtime().to_string(),
target: "pg".to_string(),
user: Some("postgres".to_string()),
namespace: None,
context: None,
compose_files: vec![],
compose_project: None,
pod_container: None,
ssh_destination: None,
ssh_options: vec![],
}
}
#[test]
fn typed_drivers_select_fixed_default_runtimes() {
for (driver, runtime) in [
(ContainerDriver::Docker, "docker"),
(ContainerDriver::Podman, "podman"),
(ContainerDriver::Nerdctl, "nerdctl"),
(ContainerDriver::Compose, "docker"),
(ContainerDriver::Kubectl, "kubectl"),
] {
assert_eq!(driver.default_runtime(), runtime);
}
}
#[test]
fn missing_target_detection() {
assert!(stderr_indicates_missing_target(
"Error response from daemon: No such container: pg-typo"
));
assert!(stderr_indicates_missing_target("container is not running"));
assert!(!stderr_indicates_missing_target(
"sh: python3: not found in PATH"
));
}
#[test]
fn bridge_args_target_tcp_for_container_cli_like_driver() -> Result<(), String> {
let args = build_container_exec_args(
&settings(ContainerDriver::Docker),
&ContainerTarget::Tcp {
host: "127.0.0.1".to_string(),
port: 5432,
},
TEST_NONCE,
)?;
assert_eq!(args[0], "exec");
assert!(args.contains(&"-i".to_string()));
assert!(args.contains(&"--user".to_string()));
assert!(args.contains(&"postgres".to_string()));
assert!(args.contains(&"pg".to_string()));
assert!(args.contains(&"tcp".to_string()));
assert!(args.contains(&"127.0.0.1".to_string()));
assert!(args.contains(&"5432".to_string()));
assert!(!args.contains(&"-e".to_string()));
Ok(())
}
#[test]
fn bridge_args_target_unix_socket() -> Result<(), String> {
let args = build_container_exec_args(
&ContainerSettings {
user: None,
..settings(ContainerDriver::Docker)
},
&ContainerTarget::UnixSocket {
path: "/var/run/postgresql/.s.PGSQL.5432".to_string(),
},
TEST_NONCE,
)?;
assert!(args.contains(&"unix".to_string()));
assert!(args.contains(&"/var/run/postgresql/.s.PGSQL.5432".to_string()));
Ok(())
}
#[test]
fn compose_driver_uses_compose_exec_without_tty() -> Result<(), String> {
let settings = ContainerSettings {
compose_files: vec!["compose.yml".to_string(), "compose.prod.yml".to_string()],
compose_project: Some("demo".to_string()),
..settings(ContainerDriver::Compose)
};
let (program, args) = build_bridge_process(
&settings,
&ContainerTarget::Tcp {
host: "db".to_string(),
port: 5432,
},
TEST_NONCE,
)?;
assert_eq!(program, "docker");
assert_eq!(
&args[0..8],
[
"compose",
"-f",
"compose.yml",
"-f",
"compose.prod.yml",
"-p",
"demo",
"exec"
]
);
assert_eq!(args[8], "-T");
assert!(args.contains(&"pg".to_string()));
Ok(())
}
#[test]
fn compose_runtime_docker_compose_skips_subcommand_prefix() -> Result<(), String> {
let settings = ContainerSettings {
runtime: "/usr/local/bin/docker-compose".to_string(),
..settings(ContainerDriver::Compose)
};
let (program, args) = build_bridge_process(
&settings,
&ContainerTarget::Tcp {
host: "db".to_string(),
port: 5432,
},
TEST_NONCE,
)?;
assert_eq!(program, "/usr/local/bin/docker-compose");
assert_eq!(&args[0..2], ["exec", "-T"]);
Ok(())
}
#[test]
fn kubectl_driver_uses_exec_separator() -> Result<(), String> {
let settings = ContainerSettings {
user: None,
namespace: Some("prod".to_string()),
context: Some("cluster-a".to_string()),
..settings(ContainerDriver::Kubectl)
};
let (program, args) = build_bridge_process(
&settings,
&ContainerTarget::Tcp {
host: "127.0.0.1".to_string(),
port: 5432,
},
TEST_NONCE,
)?;
assert_eq!(program, "kubectl");
assert_eq!(
&args[0..6],
[
"--context=cluster-a",
"--namespace=prod",
"exec",
"-i",
"pg",
"--"
]
);
assert!(args.contains(&"sh".to_string()));
Ok(())
}
#[test]
fn kubectl_driver_inserts_pod_container_before_separator() -> Result<(), String> {
let settings = ContainerSettings {
user: None,
pod_container: Some("postgres".to_string()),
..settings(ContainerDriver::Kubectl)
};
let (_, args) = build_bridge_process(
&settings,
&ContainerTarget::Tcp {
host: "127.0.0.1".to_string(),
port: 5432,
},
TEST_NONCE,
)?;
assert_eq!(
&args[0..7],
["exec", "-i", "pg", "-c", "postgres", "--", "sh"]
);
Ok(())
}
#[test]
fn bridge_banner_line_is_not_captured_as_diagnostic() {
let capture = Arc::new(Mutex::new(Vec::new()));
let (tx, mut rx) = oneshot::channel();
let mut tx = Some(tx);
let line = format!("AFPSQL_BRIDGE_OK {TEST_NONCE}\n").into_bytes();
handle_stderr_line(line, &capture, &mut tx, TEST_NONCE);
assert!(matches!(rx.try_recv(), Ok(true)));
assert!(captured_stderr(&capture).is_empty());
}
#[test]
fn bridge_banner_without_matching_nonce_is_treated_as_diagnostic() {
let capture = Arc::new(Mutex::new(Vec::new()));
let (tx, mut rx) = oneshot::channel();
let mut tx = Some(tx);
let line = b"AFPSQL_BRIDGE_OK 0000000000000000\n".to_vec();
handle_stderr_line(line, &capture, &mut tx, TEST_NONCE);
assert!(rx.try_recv().is_err());
assert!(captured_stderr(&capture).contains("AFPSQL_BRIDGE_OK"));
}
#[test]
fn bridge_banner_without_nonce_is_treated_as_diagnostic() {
let capture = Arc::new(Mutex::new(Vec::new()));
let (tx, mut rx) = oneshot::channel();
let mut tx = Some(tx);
handle_stderr_line(
b"AFPSQL_BRIDGE_OK\n".to_vec(),
&capture,
&mut tx,
TEST_NONCE,
);
assert!(rx.try_recv().is_err());
assert!(captured_stderr(&capture).contains("AFPSQL_BRIDGE_OK"));
}
#[test]
fn sanitize_diagnostic_strips_control_chars_and_truncates() {
let input = b"\x1b[31mboom\x1b[0m\nnext line\x00trailing";
let cleaned = sanitize_diagnostic(input);
assert!(!cleaned.contains('\x1b'));
assert!(!cleaned.contains('\x00'));
assert!(cleaned.contains("boom"));
assert!(cleaned.contains("next line"));
let big = vec![b'A'; 4096];
let cleaned = sanitize_diagnostic(&big);
assert!(cleaned.len() <= STDERR_HINT_BYTES);
}
#[test]
fn bridge_process_uses_remote_ssh_container_command() -> Result<(), String> {
let settings = ContainerSettings {
driver: ContainerDriver::Podman,
runtime: "podman".to_string(),
target: "pg remote".to_string(),
user: Some("postgres".to_string()),
namespace: None,
context: None,
compose_files: vec![],
compose_project: None,
pod_container: None,
ssh_destination: Some("root@example.com".to_string()),
ssh_options: vec!["ProxyJump=bastion".to_string()],
};
let (program, args) = build_bridge_process(
&settings,
&ContainerTarget::Tcp {
host: "host.containers.internal".to_string(),
port: 5432,
},
TEST_NONCE,
)?;
assert_eq!(program, "ssh");
assert!(args.contains(&"BatchMode=yes".to_string()));
assert!(args.contains(&"ProxyJump=bastion".to_string()));
assert_eq!(
args.iter().rev().nth(1),
Some(&"root@example.com".to_string())
);
let command = args.last().cloned().unwrap_or_default();
assert!(command.starts_with("podman exec -i "));
assert!(command.contains("'pg remote'"));
assert!(command.contains("host.containers.internal"));
Ok(())
}
#[test]
fn remote_container_command_quotes_shell_torture_values() -> Result<(), String> {
let target = "pg 'quoted' \"$USER\" `whoami`\nnext".to_string();
let user = "postgres '$HOME' `id`\nnext".to_string();
let context = "ctx '$VAR' `cmd`\nnext".to_string();
let settings = ContainerSettings {
driver: ContainerDriver::Docker,
runtime: "docker".to_string(),
target: target.clone(),
user: Some(user.clone()),
namespace: None,
context: Some(context.clone()),
compose_files: vec![],
compose_project: None,
pod_container: None,
ssh_destination: Some("root@example.com".to_string()),
ssh_options: vec![],
};
let command = remote_container_command(
&settings,
&ContainerTarget::Tcp {
host: "127.0.0.1".to_string(),
port: 5432,
},
TEST_NONCE,
)?;
assert!(command.contains(&shell_quote(&format!("--context={context}"))));
assert!(command.contains(&shell_quote(&target)));
assert!(command.contains(&shell_quote(&user)));
Ok(())
}
#[test]
fn target_from_dsn_tcp() -> Result<(), String> {
let cfg = SessionConfig {
container: crate::types::ContainerConfig {
target: Some("pg".to_string()),
..Default::default()
},
dsn_secret: Some("postgresql://u:p@127.0.0.1:6543/db".to_string()),
..Default::default()
};
assert_eq!(
resolve_container_target(&cfg)?,
ContainerTarget::Tcp {
host: "127.0.0.1".to_string(),
port: 6543,
}
);
Ok(())
}
#[test]
fn target_from_unix_socket_dir() -> Result<(), String> {
let cfg = SessionConfig {
container: crate::types::ContainerConfig {
target: Some("pg".to_string()),
..Default::default()
},
host: Some("/var/run/postgresql".to_string()),
port: Some(5433),
..Default::default()
};
assert_eq!(
resolve_container_target(&cfg)?,
ContainerTarget::UnixSocket {
path: "/var/run/postgresql/.s.PGSQL.5433".to_string(),
}
);
Ok(())
}
#[test]
fn settings_require_container_when_options_set() {
let cfg = SessionConfig {
container: crate::types::ContainerConfig {
user: Some("postgres".to_string()),
..Default::default()
},
..Default::default()
};
let err = resolve_container_settings(&cfg);
assert!(matches!(err, Err(message) if message.contains("--container is required")));
}
#[test]
fn settings_require_ssh_destination_when_ssh_options_set() {
let cfg = SessionConfig {
container: crate::types::ContainerConfig {
target: Some("pg".to_string()),
..Default::default()
},
ssh: crate::types::SshConfig {
options: vec!["ProxyJump=bastion".to_string()],
..Default::default()
},
..Default::default()
};
let err = resolve_container_settings(&cfg);
assert!(matches!(err, Err(message) if message.contains("--ssh is required")));
}
#[test]
fn settings_reject_tunnel_only_ssh_options_for_container_transport() {
let cfg = SessionConfig {
container: crate::types::ContainerConfig {
target: Some("pg".to_string()),
..Default::default()
},
ssh: crate::types::SshConfig {
destination: Some("user@example.com".to_string()),
local_port: Some(15432),
..Default::default()
},
..Default::default()
};
let err = resolve_container_settings(&cfg);
assert!(matches!(err, Err(message) if message.contains("supports only --ssh")));
}
#[test]
fn settings_reject_kubectl_user() {
let cfg = SessionConfig {
container: crate::types::ContainerConfig {
target: Some("pod/app".to_string()),
driver: Some("kubectl".to_string()),
user: Some("postgres".to_string()),
..Default::default()
},
..Default::default()
};
let err = resolve_container_settings(&cfg);
assert!(matches!(err, Err(message) if message.contains("not supported")));
}
#[test]
fn settings_reject_scoped_flags_for_wrong_driver() {
let cfg = SessionConfig {
container: crate::types::ContainerConfig {
target: Some("pg".to_string()),
driver: Some("podman".to_string()),
context: Some("prod".to_string()),
..Default::default()
},
..Default::default()
};
let err = resolve_container_settings(&cfg);
assert!(matches!(err, Err(message) if message.contains("--container-context requires")));
}
#[test]
fn settings_accept_compose_file_env_fallback() {
let _guard = crate::test_env::env_lock();
let old = std::env::var("AFPSQL_CONTAINER_COMPOSE_FILE").ok();
unsafe {
std::env::set_var("AFPSQL_CONTAINER_COMPOSE_FILE", "base.yml:prod.yml");
}
let cfg = SessionConfig {
container: crate::types::ContainerConfig {
target: Some("pg".to_string()),
driver: Some("compose".to_string()),
..Default::default()
},
..Default::default()
};
let settings = resolve_container_settings(&cfg);
match old {
Some(value) => unsafe { std::env::set_var("AFPSQL_CONTAINER_COMPOSE_FILE", value) },
None => unsafe { std::env::remove_var("AFPSQL_CONTAINER_COMPOSE_FILE") },
}
assert!(matches!(
settings,
Ok(ContainerSettings {
compose_files,
..
}) if compose_files == vec!["base.yml".to_string(), "prod.yml".to_string()]
));
}
#[test]
fn container_bridge_hint_classifies_interpreter_failure() {
let hint = container_bridge_hint(
None,
"afpsql container bridge requires python3, python, or perl in the container",
)
.unwrap_or_default();
assert!(hint.contains("no supported interpreter"));
assert!(hint.contains("container bridge stderr"));
}
}