1use crate::conn::resolve_pg_config;
2use crate::types::SessionConfig;
3use std::io::Read as _;
4use std::net::{TcpListener, TcpStream};
5use std::pin::Pin;
6use std::process::{Child, Command, ExitStatus, Stdio};
7use std::sync::{Arc, Mutex};
8use std::task::{Context, Poll};
9use std::time::{Duration, Instant};
10use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, ReadBuf};
11use tokio::process::{ChildStderr, ChildStdin, ChildStdout};
12use tokio_postgres::{Client, NoTls};
13
14const DEFAULT_LOCAL_HOST: &str = "127.0.0.1";
15const DEFAULT_REMOTE_HOST: &str = "127.0.0.1";
16const DEFAULT_REMOTE_PORT: u16 = 5432;
17const TUNNEL_READY_TIMEOUT: Duration = Duration::from_secs(5);
18const TUNNEL_READY_POLL: Duration = Duration::from_millis(50);
19const TUNNEL_READY_SETTLE: Duration = Duration::from_millis(100);
20const STDERR_CAPTURE_LIMIT: usize = 8192;
21const STDERR_HINT_BYTES: usize = 2048;
22
23const PYTHON_STREAM_BRIDGE: &str = r#"import os,select,socket,sys
24mode=sys.argv[1] if len(sys.argv)>1 else ""
25last_error=""
26s=None
27timeout=float(os.environ.get("AFPSQL_SSH_BRIDGE_CONNECT_TIMEOUT","10"))
28if mode=="tcp":
29 host=sys.argv[2]
30 port=int(sys.argv[3])
31 s=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
32 s.settimeout(timeout)
33 try:
34 s.connect((host,port))
35 s.settimeout(None)
36 except OSError as e:
37 sys.stderr.write("could not connect PostgreSQL tcp "+host+":"+str(port)+": "+str(e)+"\n")
38 sys.exit(1)
39elif mode=="unix":
40 paths=sys.argv[2:]
41 for p in paths:
42 s=socket.socket(socket.AF_UNIX,socket.SOCK_STREAM)
43 s.settimeout(timeout)
44 try:
45 s.connect(p)
46 s.settimeout(None)
47 break
48 except OSError as e:
49 last_error=p+": "+str(e)
50 s.close()
51 s=None
52 if s is None:
53 sys.stderr.write("could not connect PostgreSQL socket; tried "+", ".join(paths)+"; "+last_error+"\n")
54 sys.exit(1)
55else:
56 sys.stderr.write("unsupported afpsql ssh bridge mode: "+mode+"\n")
57 sys.exit(1)
58stdin=getattr(sys.stdin,"buffer",sys.stdin)
59stdout=getattr(sys.stdout,"buffer",sys.stdout)
60stdin_fd=stdin.fileno()
61stdout_fd=stdout.fileno()
62stdin_open=True
63while True:
64 readers=[s]
65 if stdin_open:
66 readers.append(stdin_fd)
67 ready,_,_=select.select(readers,[],[])
68 if stdin_fd in ready:
69 data=os.read(stdin_fd,65536)
70 if data:
71 s.sendall(data)
72 else:
73 stdin_open=False
74 try:
75 s.shutdown(socket.SHUT_WR)
76 except OSError:
77 pass
78 if s in ready:
79 data=s.recv(65536)
80 if data:
81 os.write(stdout_fd,data)
82 else:
83 break
84"#;
85
86const PERL_STREAM_BRIDGE: &str = r#"my $mode=shift @ARGV||"";
87my $s;
88if ($mode eq "tcp") {
89 my ($host,$port)=@ARGV;
90 socket($s, PF_INET, SOCK_STREAM, getprotobyname("tcp")) or die "socket tcp: $!\n";
91 my $addr=inet_aton($host) or die "could not resolve PostgreSQL tcp $host:$port\n";
92 eval {
93 local $SIG{ALRM}=sub { die "connect timeout\n" };
94 alarm(10);
95 connect($s, sockaddr_in($port,$addr)) or die "could not connect PostgreSQL tcp $host:$port: $!\n";
96 alarm(0);
97 };
98 if ($@) { print STDERR $@; exit 1; }
99} elsif ($mode eq "unix") {
100 my $last="";
101 for my $path (@ARGV) {
102 socket($s, PF_UNIX, SOCK_STREAM, 0) or die "socket unix: $!\n";
103 if (connect($s, sockaddr_un($path))) { $last=""; last; }
104 $last="$path: $!";
105 close($s);
106 undef $s;
107 }
108 if (!$s) {
109 print STDERR "could not connect PostgreSQL socket; tried ".join(", ", @ARGV)."; $last\n";
110 exit 1;
111 }
112} else {
113 print STDERR "unsupported afpsql ssh bridge mode: $mode\n";
114 exit 1;
115}
116binmode STDIN;
117binmode STDOUT;
118binmode $s;
119sub write_all {
120 my ($fh,$buf)=@_;
121 my $off=0;
122 while ($off < length($buf)) {
123 my $n=syswrite($fh,$buf,length($buf)-$off,$off);
124 if (!defined $n) { print STDERR "bridge write failed: $!\n"; exit 1; }
125 $off += $n;
126 }
127}
128my $sel=IO::Select->new($s, \*STDIN);
129while (1) {
130 for my $fh ($sel->can_read) {
131 if (fileno($fh) == fileno(STDIN)) {
132 my $buf="";
133 my $n=sysread(STDIN,$buf,65536);
134 if ($n) {
135 write_all($s,$buf);
136 } else {
137 $sel->remove(\*STDIN);
138 shutdown($s,1);
139 }
140 } else {
141 my $buf="";
142 my $n=sysread($s,$buf,65536);
143 if ($n) {
144 write_all(\*STDOUT,$buf);
145 } else {
146 exit 0;
147 }
148 }
149 }
150}
151"#;
152
153pub struct SshTunnelGuard {
154 child: Mutex<Child>,
155 stderr_capture: Arc<Mutex<Vec<u8>>>,
156 pub local_host: String,
157 pub local_port: u16,
158}
159
160impl Drop for SshTunnelGuard {
161 fn drop(&mut self) {
162 if let Ok(mut child) = self.child.lock() {
163 let _ = child.kill();
164 let _ = child.wait();
165 }
166 }
167}
168
169pub struct SshBridgeGuard {
170 child: Option<tokio::process::Child>,
171 connection_task: Option<tokio::task::JoinHandle<()>>,
172 stderr_task: Option<tokio::task::JoinHandle<()>>,
173}
174
175impl SshBridgeGuard {
176 pub fn is_finished(&self) -> bool {
177 self.connection_task
178 .as_ref()
179 .map(|task| task.is_finished())
180 .unwrap_or(true)
181 }
182
183 pub async fn shutdown(mut self, timeout: Duration) {
184 if let Some(task) = self.connection_task.take() {
185 let mut task = task;
186 if tokio::time::timeout(timeout, &mut task).await.is_err() {
187 task.abort();
188 }
189 }
190 if let Some(mut child) = self.child.take() {
191 let _ = child.start_kill();
192 let _ = tokio::time::timeout(timeout, child.wait()).await;
193 }
194 if let Some(task) = self.stderr_task.take() {
195 task.abort();
196 }
197 }
198}
199
200impl Drop for SshBridgeGuard {
201 fn drop(&mut self) {
202 if let Some(task) = self.connection_task.as_ref() {
203 task.abort();
204 }
205 if let Some(child) = self.child.as_mut() {
206 let _ = child.start_kill();
207 }
208 if let Some(task) = self.stderr_task.as_ref() {
209 task.abort();
210 }
211 }
212}
213
214struct SshStdioStream {
215 stdout: ChildStdout,
216 stdin: ChildStdin,
217}
218
219impl AsyncRead for SshStdioStream {
220 fn poll_read(
221 mut self: Pin<&mut Self>,
222 cx: &mut Context<'_>,
223 buf: &mut ReadBuf<'_>,
224 ) -> Poll<std::io::Result<()>> {
225 Pin::new(&mut self.stdout).poll_read(cx, buf)
226 }
227}
228
229impl AsyncWrite for SshStdioStream {
230 fn poll_write(
231 mut self: Pin<&mut Self>,
232 cx: &mut Context<'_>,
233 buf: &[u8],
234 ) -> Poll<std::io::Result<usize>> {
235 Pin::new(&mut self.stdin).poll_write(cx, buf)
236 }
237
238 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
239 Pin::new(&mut self.stdin).poll_flush(cx)
240 }
241
242 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
243 Pin::new(&mut self.stdin).poll_shutdown(cx)
244 }
245}
246
247#[derive(Debug, Clone, PartialEq, Eq)]
248struct SshSettings {
249 destination: String,
250 via: Vec<String>,
251 options: Vec<String>,
252 local_host: String,
253 local_port: Option<u16>,
254 remote_socket: Option<String>,
255 sudo_user: Option<String>,
256}
257
258enum TunnelTarget {
259 Tcp { host: String, port: u16 },
260 UnixSocket { path: String },
261}
262
263pub fn needs_stdio_bridge(cfg: &SessionConfig) -> bool {
264 resolve_ssh_settings(cfg)
265 .ok()
266 .flatten()
267 .map(|settings| settings.sudo_user.is_some() || !settings.via.is_empty())
268 .unwrap_or(false)
269}
270
271pub async fn prepare_session(
272 cfg: &SessionConfig,
273) -> Result<(SessionConfig, Option<SshTunnelGuard>), String> {
274 let Some(settings) = resolve_ssh_settings(cfg)? else {
275 return Ok((cfg.clone(), None));
276 };
277
278 if settings.sudo_user.is_some() {
279 return Err("--ssh-sudo-user requires SSH Unix-socket bridge mode".to_string());
280 }
281 if !settings.via.is_empty() {
282 return Err("--ssh-via requires SSH stdio bridge mode".to_string());
283 }
284 reject_secret_conn_strings_with_ssh(cfg)?;
285
286 let target = if let Some(path) = settings.remote_socket.clone() {
287 TunnelTarget::UnixSocket { path }
288 } else {
289 let host = effective_remote_host(cfg);
290 let port = effective_remote_port(cfg);
291 if host.starts_with('/') {
292 TunnelTarget::UnixSocket {
293 path: socket_file_from_dir(&host, port),
294 }
295 } else {
296 TunnelTarget::Tcp { host, port }
297 }
298 };
299 let tunnel = start_tunnel(&settings, target).await?;
300
301 let mut local_cfg = cfg.clone();
302 local_cfg.dsn_secret = None;
303 local_cfg.conninfo_secret = None;
304 local_cfg.host = Some(tunnel.local_host.clone());
305 local_cfg.port = Some(tunnel.local_port);
306 Ok((local_cfg, Some(tunnel)))
307}
308
309pub async fn connect_stdio_bridge(cfg: &SessionConfig) -> Result<(Client, SshBridgeGuard), String> {
310 let Some(settings) = resolve_ssh_settings(cfg)? else {
311 return Err("--ssh is required for SSH stdio bridge mode".to_string());
312 };
313 if settings.sudo_user.is_none() && settings.via.is_empty() {
314 return Err(
315 "--ssh-via or --ssh-sudo-user is required for SSH stdio bridge mode".to_string(),
316 );
317 }
318 reject_secret_conn_strings_with_ssh(cfg)?;
319
320 let target = bridge_target(&settings, cfg)?;
321 let command = remote_stream_bridge_command(&settings, &target);
322 let args = build_bridge_ssh_args(&settings, &command);
323 let mut child = tokio::process::Command::new("ssh")
324 .args(&args)
325 .stdin(Stdio::piped())
326 .stdout(Stdio::piped())
327 .stderr(Stdio::piped())
328 .kill_on_drop(true)
329 .spawn()
330 .map_err(|e| format!("start ssh bridge failed: {e}"))?;
331 let stderr_capture = Arc::new(Mutex::new(Vec::new()));
332 let stderr_task = child
333 .stderr
334 .take()
335 .map(|stderr| spawn_async_stderr_capture(stderr, Arc::clone(&stderr_capture)));
336
337 let stdin = child
338 .stdin
339 .take()
340 .ok_or_else(|| "start ssh bridge failed: stdin pipe unavailable".to_string())?;
341 let stdout = child
342 .stdout
343 .take()
344 .ok_or_else(|| "start ssh bridge failed: stdout pipe unavailable".to_string())?;
345 let stream = SshStdioStream { stdout, stdin };
346 let pg_cfg =
347 resolve_pg_config(cfg).map_err(|e| format!("invalid bridge connection config: {e}"))?;
348 let (client, connection) = match pg_cfg.connect_raw(stream, NoTls).await {
349 Ok(connection) => connection,
350 Err(e) => {
351 let status = match tokio::time::timeout(Duration::from_millis(100), child.wait()).await
352 {
353 Ok(Ok(status)) => Some(status),
354 _ => child.try_wait().ok().flatten(),
355 };
356 tokio::time::sleep(Duration::from_millis(20)).await;
357 let stderr_text = captured_stderr(&stderr_capture);
358 return Err(ssh_bridge_error(
359 format!("connect through ssh bridge failed: {e}"),
360 status,
361 &stderr_text,
362 ));
363 }
364 };
365 let connection_task = tokio::spawn(async move {
366 let _ = connection.await;
367 });
368
369 Ok((
370 client,
371 SshBridgeGuard {
372 child: Some(child),
373 connection_task: Some(connection_task),
374 stderr_task,
375 },
376 ))
377}
378
379enum BridgeTarget {
380 Tcp { host: String, port: u16 },
381 UnixSocket { paths: Vec<String> },
382}
383
384fn bridge_target(settings: &SshSettings, cfg: &SessionConfig) -> Result<BridgeTarget, String> {
385 if settings.sudo_user.is_some() {
386 return bridge_socket_candidates(settings, cfg)
387 .map(|paths| BridgeTarget::UnixSocket { paths });
388 }
389 if let Some(remote_socket) = settings.remote_socket.clone() {
390 return Ok(BridgeTarget::UnixSocket {
391 paths: vec![remote_socket],
392 });
393 }
394 let host = effective_remote_host(cfg);
395 let port = effective_remote_port(cfg);
396 if host.starts_with('/') {
397 Ok(BridgeTarget::UnixSocket {
398 paths: vec![socket_file_from_dir(&host, port)],
399 })
400 } else {
401 Ok(BridgeTarget::Tcp { host, port })
402 }
403}
404
405fn socket_file_from_dir(dir: &str, port: u16) -> String {
406 format!("{}/.s.PGSQL.{port}", dir.trim_end_matches('/'))
407}
408
409fn bridge_socket_candidates(
410 settings: &SshSettings,
411 cfg: &SessionConfig,
412) -> Result<Vec<String>, String> {
413 if let Some(remote_socket) = settings.remote_socket.clone() {
414 return Ok(vec![remote_socket]);
415 }
416
417 let host = effective_remote_host(cfg);
418 let port = effective_remote_port(cfg);
419 if host.starts_with('/') {
420 return Ok(vec![socket_file_from_dir(&host, port)]);
421 }
422
423 Err("--ssh-sudo-user requires an explicit remote PostgreSQL Unix socket".to_string())
424}
425
426async fn start_tunnel(
427 settings: &SshSettings,
428 target: TunnelTarget,
429) -> Result<SshTunnelGuard, String> {
430 let local_port = match settings.local_port {
431 Some(port) => {
432 ensure_local_port_available(&settings.local_host, port)?;
433 port
434 }
435 None => allocate_local_port(&settings.local_host)?,
436 };
437 let args = build_tunnel_ssh_args(settings, &target, local_port);
438 let mut child = Command::new("ssh")
439 .args(&args)
440 .stdin(Stdio::null())
441 .stdout(Stdio::null())
442 .stderr(Stdio::piped())
443 .spawn()
444 .map_err(|e| format!("start ssh tunnel failed: {e}"))?;
445 let stderr_capture = Arc::new(Mutex::new(Vec::new()));
446 if let Some(stderr) = child.stderr.take() {
447 spawn_blocking_stderr_capture(stderr, Arc::clone(&stderr_capture));
448 }
449
450 let guard = SshTunnelGuard {
451 child: Mutex::new(child),
452 stderr_capture,
453 local_host: settings.local_host.clone(),
454 local_port,
455 };
456 wait_for_tunnel_ready(&guard).await?;
457 Ok(guard)
458}
459
460fn allocate_local_port(local_host: &str) -> Result<u16, String> {
461 TcpListener::bind((local_host, 0))
462 .map_err(|e| format!("allocate ssh local port on {local_host} failed: {e}"))?
463 .local_addr()
464 .map(|addr| addr.port())
465 .map_err(|e| format!("read allocated ssh local port failed: {e}"))
466}
467
468fn ensure_local_port_available(local_host: &str, local_port: u16) -> Result<(), String> {
469 TcpListener::bind((local_host, local_port))
470 .map(|_| ())
471 .map_err(|e| format!("ssh local port {local_host}:{local_port} is not available: {e}"))
472}
473
474async fn wait_for_tunnel_ready(guard: &SshTunnelGuard) -> Result<(), String> {
475 let start = Instant::now();
476 loop {
477 if let Some(status) = tunnel_child_status(guard)? {
478 return Err(tunnel_error(
479 guard,
480 format!("ssh tunnel exited before it became ready with status {status}"),
481 ));
482 }
483 if TcpStream::connect((guard.local_host.as_str(), guard.local_port)).is_ok() {
484 tokio::time::sleep(TUNNEL_READY_SETTLE).await;
485 if let Some(status) = tunnel_child_status(guard)? {
486 return Err(tunnel_error(
487 guard,
488 format!(
489 "ssh tunnel exited after local port became reachable with status {status}"
490 ),
491 ));
492 }
493 return Ok(());
494 }
495 if start.elapsed() >= TUNNEL_READY_TIMEOUT {
496 return Err(tunnel_error(
497 guard,
498 format!(
499 "ssh tunnel did not become ready on {}:{}",
500 guard.local_host, guard.local_port
501 ),
502 ));
503 }
504 tokio::time::sleep(TUNNEL_READY_POLL).await;
505 }
506}
507
508fn tunnel_error(guard: &SshTunnelGuard, base: String) -> String {
509 append_ssh_stderr(base, &captured_stderr(&guard.stderr_capture))
510}
511
512fn tunnel_child_status(guard: &SshTunnelGuard) -> Result<Option<ExitStatus>, String> {
513 let mut child = guard
514 .child
515 .lock()
516 .map_err(|_| "ssh tunnel child lock poisoned".to_string())?;
517 child
518 .try_wait()
519 .map_err(|e| format!("check ssh tunnel status failed: {e}"))
520}
521
522fn resolve_ssh_settings(cfg: &SessionConfig) -> Result<Option<SshSettings>, String> {
523 let destination = cfg.ssh.destination.clone();
524 let has_ssh_fields = cfg.ssh.has_transport_fields();
525
526 let Some(destination) = destination else {
527 if has_ssh_fields {
528 return Err("--ssh is required when SSH transport options are set".to_string());
529 }
530 return Ok(None);
531 };
532 if destination.trim().is_empty() {
533 return Err("--ssh requires a non-empty USER@HOST destination".to_string());
534 }
535 let via = cfg
536 .ssh
537 .via
538 .iter()
539 .map(|hop| hop.trim())
540 .filter(|hop| !hop.is_empty())
541 .map(std::string::ToString::to_string)
542 .collect::<Vec<_>>();
543 if !cfg.ssh.via.is_empty() && via.len() != cfg.ssh.via.len() {
544 return Err("--ssh-via requires non-empty USER@HOST hop values".to_string());
545 }
546 if !via.is_empty() && (cfg.ssh.local_host.is_some() || cfg.ssh.local_port.is_some()) {
547 return Err("--ssh-via uses SSH stdio bridge mode and cannot be combined with --ssh-local-host or --ssh-local-port".to_string());
548 }
549
550 Ok(Some(SshSettings {
551 destination,
552 via,
553 options: cfg.ssh.options.clone(),
554 local_host: cfg
555 .ssh
556 .local_host
557 .clone()
558 .unwrap_or_else(|| DEFAULT_LOCAL_HOST.to_string()),
559 local_port: cfg.ssh.local_port,
560 remote_socket: cfg.ssh.remote_socket.clone(),
561 sudo_user: cfg.ssh.sudo_user.clone(),
562 }))
563}
564
565fn reject_secret_conn_strings_with_ssh(cfg: &SessionConfig) -> Result<(), String> {
566 if cfg.dsn_secret.is_some()
567 || cfg.conninfo_secret.is_some()
568 || env_nonempty("AFPSQL_DSN_SECRET").is_some()
569 || env_nonempty("AFPSQL_CONNINFO_SECRET").is_some()
570 {
571 return Err("SSH transport currently supports discrete connection fields only; use --host/--port/--user/--dbname/--password-secret-env instead of --dsn-secret or --conninfo-secret".to_string());
572 }
573 Ok(())
574}
575
576fn effective_remote_host(cfg: &SessionConfig) -> String {
577 cfg.host
578 .clone()
579 .or_else(|| env_nonempty("AFPSQL_HOST"))
580 .or_else(|| env_nonempty("PGHOST"))
581 .unwrap_or_else(|| DEFAULT_REMOTE_HOST.to_string())
582}
583
584fn effective_remote_port(cfg: &SessionConfig) -> u16 {
585 cfg.port
586 .or_else(|| env_u16("AFPSQL_PORT").and_then(Result::ok))
587 .or_else(|| env_u16("PGPORT").and_then(Result::ok))
588 .unwrap_or(DEFAULT_REMOTE_PORT)
589}
590
591fn env_nonempty(name: &str) -> Option<String> {
592 std::env::var(name).ok().filter(|value| !value.is_empty())
593}
594
595fn env_u16(name: &str) -> Option<Result<u16, String>> {
596 std::env::var(name)
597 .ok()
598 .filter(|value| !value.is_empty())
599 .map(|value| {
600 value
601 .parse::<u16>()
602 .map_err(|_| format!("{name} must be a valid TCP port"))
603 })
604}
605
606fn build_tunnel_ssh_args(
607 settings: &SshSettings,
608 target: &TunnelTarget,
609 local_port: u16,
610) -> Vec<String> {
611 let mut args = base_ssh_args(settings);
612 args.push("-N".to_string());
613 args.push("-L".to_string());
614 args.push(match target {
615 TunnelTarget::Tcp { host, port } => {
616 format!("{}:{local_port}:{host}:{port}", settings.local_host)
617 }
618 TunnelTarget::UnixSocket { path } => {
619 format!("{}:{local_port}:{path}", settings.local_host)
620 }
621 });
622 args.push(settings.destination.clone());
623 args
624}
625
626fn build_bridge_ssh_args(settings: &SshSettings, final_command: &str) -> Vec<String> {
627 let mut args = base_ssh_args(settings);
628 let mut chain = settings.via.clone();
629 chain.push(settings.destination.clone());
630 let Some(first_hop) = chain.first() else {
631 return args;
632 };
633 let mut command = final_command.to_string();
634 for hop in chain.iter().skip(1).rev() {
635 command = remote_ssh_command(hop, &command);
636 }
637 args.push(first_hop.clone());
638 args.push(command);
639 args
640}
641
642fn base_ssh_args(settings: &SshSettings) -> Vec<String> {
643 let mut args = vec![
644 "-T".to_string(),
645 "-o".to_string(),
646 "BatchMode=yes".to_string(),
647 "-o".to_string(),
648 "ExitOnForwardFailure=yes".to_string(),
649 ];
650 for option in &settings.options {
651 args.push("-o".to_string());
652 args.push(option.clone());
653 }
654 args
655}
656
657fn remote_ssh_command(destination: &str, remote_command: &str) -> String {
658 shell_join(&[
659 "ssh".to_string(),
660 "-T".to_string(),
661 "-o".to_string(),
662 "BatchMode=yes".to_string(),
663 "-o".to_string(),
664 "ExitOnForwardFailure=yes".to_string(),
665 "-o".to_string(),
666 "ConnectTimeout=10".to_string(),
667 destination.to_string(),
668 remote_command.to_string(),
669 ])
670}
671
672fn remote_stream_bridge_command(settings: &SshSettings, target: &BridgeTarget) -> String {
673 let bridge_args = match target {
674 BridgeTarget::Tcp { host, port } => {
675 vec!["tcp".to_string(), host.clone(), port.to_string()]
676 }
677 BridgeTarget::UnixSocket { paths } => {
678 let mut args = vec!["unix".to_string()];
679 args.extend(paths.iter().cloned());
680 args
681 }
682 };
683 let launcher = bridge_launcher_command(&bridge_args);
684 if let Some(sudo_user) = settings.sudo_user.as_ref() {
685 return shell_join(&[
686 "sudo".to_string(),
687 "-n".to_string(),
688 "-u".to_string(),
689 sudo_user.clone(),
690 "sh".to_string(),
691 "-c".to_string(),
692 launcher,
693 ]);
694 }
695 launcher
696}
697
698fn bridge_launcher_command(args: &[String]) -> String {
699 let bridge_args = shell_join(args);
700 format!(
701 "if command -v python3 >/dev/null 2>&1; then exec python3 -c {} {}; fi; \
702 if command -v python >/dev/null 2>&1; then exec python -c {} {}; fi; \
703 if command -v perl >/dev/null 2>&1; then exec perl -MIO::Select -MSocket -e {} {}; fi; \
704 echo {} >&2; exit 127",
705 shell_quote(PYTHON_STREAM_BRIDGE),
706 bridge_args,
707 shell_quote(PYTHON_STREAM_BRIDGE),
708 bridge_args,
709 shell_quote(PERL_STREAM_BRIDGE),
710 bridge_args,
711 shell_quote("afpsql ssh bridge requires python3, python, or perl on the remote host")
712 )
713}
714
715fn shell_join(parts: &[String]) -> String {
716 parts
717 .iter()
718 .map(|part| shell_quote(part))
719 .collect::<Vec<_>>()
720 .join(" ")
721}
722
723fn shell_quote(value: &str) -> String {
724 if value.is_empty() {
725 return "''".to_string();
726 }
727 if value
728 .bytes()
729 .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'_' | b'-' | b'.' | b'/' | b':' | b'@'))
730 {
731 return value.to_string();
732 }
733 format!("'{}'", value.replace('\'', "'\\''"))
734}
735
736fn spawn_blocking_stderr_capture(
737 mut stderr: std::process::ChildStderr,
738 capture: Arc<Mutex<Vec<u8>>>,
739) {
740 let _handle = std::thread::spawn(move || {
741 let mut buf = [0u8; 1024];
742 loop {
743 match stderr.read(&mut buf) {
744 Ok(0) | Err(_) => break,
745 Ok(n) => append_stderr_capture(&capture, &buf[..n]),
746 }
747 }
748 });
749}
750
751fn spawn_async_stderr_capture(
752 mut stderr: ChildStderr,
753 capture: Arc<Mutex<Vec<u8>>>,
754) -> tokio::task::JoinHandle<()> {
755 tokio::spawn(async move {
756 let mut buf = [0u8; 1024];
757 while let Ok(n) = stderr.read(&mut buf).await {
758 if n == 0 {
759 break;
760 }
761 append_stderr_capture(&capture, &buf[..n]);
762 }
763 })
764}
765
766fn append_stderr_capture(capture: &Arc<Mutex<Vec<u8>>>, bytes: &[u8]) {
767 if let Ok(mut captured) = capture.lock() {
768 let remaining = STDERR_CAPTURE_LIMIT.saturating_sub(captured.len());
769 if remaining > 0 {
770 captured.extend_from_slice(&bytes[..bytes.len().min(remaining)]);
771 }
772 }
773}
774
775fn captured_stderr(capture: &Arc<Mutex<Vec<u8>>>) -> String {
776 capture
777 .lock()
778 .ok()
779 .map(|captured| sanitize_diagnostic(&captured))
780 .filter(|text| !text.is_empty())
781 .unwrap_or_default()
782}
783
784fn sanitize_diagnostic(bytes: &[u8]) -> String {
785 let text = String::from_utf8_lossy(bytes);
786 let mut out = String::with_capacity(STDERR_HINT_BYTES);
787 for ch in text.chars() {
788 let mapped = if matches!(ch, '\n' | '\t') {
789 ch
790 } else if ch.is_control() {
791 ' '
792 } else {
793 ch
794 };
795 if out.len() + mapped.len_utf8() > STDERR_HINT_BYTES {
796 break;
797 }
798 out.push(mapped);
799 }
800 out.trim().to_string()
801}
802
803fn ssh_bridge_error(base: String, status: Option<ExitStatus>, stderr: &str) -> String {
804 let with_status = match status {
805 Some(status) => format!("{base}; ssh status: {status}"),
806 None => base,
807 };
808 append_ssh_stderr(with_status, stderr)
809}
810
811fn append_ssh_stderr(base: String, stderr: &str) -> String {
812 if stderr.trim().is_empty() {
813 base
814 } else {
815 format!("{base}; ssh stderr: {stderr}")
816 }
817}
818
819#[cfg(test)]
820mod tests {
821 use super::*;
822
823 #[test]
824 fn tunnel_args_target_remote_tcp() {
825 let settings = SshSettings {
826 destination: "app@example.com".to_string(),
827 via: vec![],
828 options: vec!["ProxyJump=bastion".to_string()],
829 local_host: "127.0.0.1".to_string(),
830 local_port: Some(15432),
831 remote_socket: None,
832 sudo_user: None,
833 };
834 let args = build_tunnel_ssh_args(
835 &settings,
836 &TunnelTarget::Tcp {
837 host: "127.0.0.1".to_string(),
838 port: 5432,
839 },
840 15432,
841 );
842 assert!(args.contains(&"ProxyJump=bastion".to_string()));
843 assert!(args.contains(&"127.0.0.1:15432:127.0.0.1:5432".to_string()));
844 assert_eq!(args.last(), Some(&"app@example.com".to_string()));
845 }
846
847 #[test]
848 fn bridge_args_use_sudo_noninteractive_python_socket_bridge() {
849 let settings = SshSettings {
850 destination: "user@example.com".to_string(),
851 via: vec![],
852 options: vec![],
853 local_host: "127.0.0.1".to_string(),
854 local_port: None,
855 remote_socket: Some("/var/run/postgresql/.s.PGSQL.5432".to_string()),
856 sudo_user: Some("postgres".to_string()),
857 };
858 let command = remote_stream_bridge_command(
859 &settings,
860 &BridgeTarget::UnixSocket {
861 paths: vec!["/var/run/postgresql/.s.PGSQL.5432".to_string()],
862 },
863 );
864 let args = build_bridge_ssh_args(&settings, &command);
865 assert!(args.iter().any(|arg| arg == "BatchMode=yes"));
866 let command = args.last().cloned().unwrap_or_default();
867 assert!(command.contains("sudo -n -u postgres sh -c"));
868 assert!(command.contains("python3 -c"));
869 assert!(command.contains("/var/run/postgresql/.s.PGSQL.5432"));
870 }
871
872 #[test]
873 fn bridge_args_chain_via_hosts_and_final_tcp_bridge() {
874 let settings = SshSettings {
875 destination: "ubuntu@db.internal".to_string(),
876 via: vec!["ubuntu@bastion".to_string()],
877 options: vec!["ConnectTimeout=10".to_string()],
878 local_host: "127.0.0.1".to_string(),
879 local_port: None,
880 remote_socket: None,
881 sudo_user: None,
882 };
883 let command = remote_stream_bridge_command(
884 &settings,
885 &BridgeTarget::Tcp {
886 host: "localhost".to_string(),
887 port: 5432,
888 },
889 );
890 let args = build_bridge_ssh_args(&settings, &command);
891 assert!(args.contains(&"ubuntu@bastion".to_string()));
892 assert!(args.iter().any(|arg| arg == "ConnectTimeout=10"));
893 let remote = args.last().cloned().unwrap_or_default();
894 assert!(remote.contains("ssh -T -o 'BatchMode=yes'"));
895 assert!(remote.contains("ubuntu@db.internal"));
896 assert!(remote.contains("python3 -c"));
897 assert!(remote.contains("tcp localhost 5432"));
898 }
899
900 #[test]
901 fn bridge_candidates_require_explicit_socket_for_sudo_bridge() -> Result<(), String> {
902 let settings = SshSettings {
903 destination: "user@example.com".to_string(),
904 via: vec![],
905 options: vec![],
906 local_host: "127.0.0.1".to_string(),
907 local_port: None,
908 remote_socket: None,
909 sudo_user: Some("postgres".to_string()),
910 };
911 let cfg = SessionConfig {
912 ssh: crate::types::SshConfig {
913 destination: Some("user@example.com".to_string()),
914 sudo_user: Some("postgres".to_string()),
915 ..Default::default()
916 },
917 ..Default::default()
918 };
919
920 let Err(err) = bridge_socket_candidates(&settings, &cfg) else {
921 return Err("expected explicit socket error".to_string());
922 };
923 assert!(err.contains("explicit remote PostgreSQL Unix socket"));
924 Ok(())
925 }
926
927 #[test]
928 fn bridge_candidates_honor_explicit_socket_and_host_dir() -> Result<(), String> {
929 let explicit_settings = SshSettings {
930 destination: "user@example.com".to_string(),
931 via: vec![],
932 options: vec![],
933 local_host: "127.0.0.1".to_string(),
934 local_port: None,
935 remote_socket: Some("/custom/.s.PGSQL.6543".to_string()),
936 sudo_user: Some("postgres".to_string()),
937 };
938 let cfg = SessionConfig::default();
939 assert_eq!(
940 bridge_socket_candidates(&explicit_settings, &cfg)?,
941 vec!["/custom/.s.PGSQL.6543".to_string()]
942 );
943
944 let dir_settings = SshSettings {
945 destination: "user@example.com".to_string(),
946 via: vec![],
947 options: vec![],
948 local_host: "127.0.0.1".to_string(),
949 local_port: None,
950 remote_socket: None,
951 sudo_user: Some("postgres".to_string()),
952 };
953 let cfg = SessionConfig {
954 host: Some("/run/postgresql".to_string()),
955 port: Some(5433),
956 ..Default::default()
957 };
958 assert_eq!(
959 bridge_socket_candidates(&dir_settings, &cfg)?,
960 vec!["/run/postgresql/.s.PGSQL.5433".to_string()]
961 );
962 Ok(())
963 }
964
965 #[test]
966 fn shell_quote_handles_spaces_and_quotes() {
967 assert_eq!(shell_quote("postgres"), "postgres");
968 assert_eq!(shell_quote("/tmp/socket path"), "'/tmp/socket path'");
969 assert_eq!(shell_quote("a'b"), "'a'\\''b'");
970 }
971
972 #[test]
973 fn ssh_bridge_error_includes_sanitized_stderr() {
974 let capture = Arc::new(Mutex::new(Vec::new()));
975 append_stderr_capture(&capture, b"Permission denied (publickey).\x1b\n");
976 let err = ssh_bridge_error(
977 "connect through ssh bridge failed".to_string(),
978 None,
979 &captured_stderr(&capture),
980 );
981 assert!(err.contains("connect through ssh bridge failed"));
982 assert!(err.contains("ssh stderr: Permission denied (publickey)."));
983 assert!(!err.contains('\x1b'));
984 }
985
986 #[test]
987 fn socket_dir_maps_to_postgres_socket_file() {
988 assert_eq!(
989 socket_file_from_dir("/var/run/postgresql/", 5433),
990 "/var/run/postgresql/.s.PGSQL.5433"
991 );
992 }
993
994 #[test]
995 fn local_port_availability_rejects_bound_port() -> Result<(), String> {
996 let listener = TcpListener::bind(("127.0.0.1", 0)).map_err(|e| e.to_string())?;
997 let port = listener.local_addr().map_err(|e| e.to_string())?.port();
998
999 let unavailable = ensure_local_port_available("127.0.0.1", port);
1000 assert!(matches!(
1001 unavailable,
1002 Err(message) if message.contains("not available")
1003 ));
1004
1005 drop(listener);
1006 assert!(ensure_local_port_available("127.0.0.1", port).is_ok());
1007 Ok(())
1008 }
1009}