1use super::*;
2
3#[cfg(unix)]
4use std::fs;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::sync::{Condvar, Mutex, OnceLock};
7use std::time::Duration;
8
9pub fn ssh_connectivity_probe(ssh: &SshTarget) -> CommandSpec {
24 let mut args = vec![
25 "-o".to_owned(),
26 "BatchMode=yes".to_owned(),
27 "-o".to_owned(),
28 "StrictHostKeyChecking=yes".to_owned(),
29 ];
30 args.extend(ssh.ssh_args.iter().cloned());
31 args.push(ssh.destination.clone());
32 args.push(join_remote_command(&["true".to_owned()]));
33 CommandSpec::new("ssh", args)
36 .ssh_probe_session(ssh)
37 .purpose("verify SSH connectivity")
38}
39
40pub fn ssh_command(
41 ssh: &SshTarget,
42 args: impl IntoIterator<Item = impl AsRef<str>>,
43) -> CommandSpec {
44 ssh_command_owned(
45 ssh,
46 args.into_iter()
47 .map(|arg| arg.as_ref().to_owned())
48 .collect(),
49 )
50}
51
52pub fn ssh_command_owned(ssh: &SshTarget, remote_args: Vec<String>) -> CommandSpec {
56 let mut args = ssh.ssh_args.clone();
57 args.push(ssh.destination.clone());
58 args.push(join_remote_command(&remote_args));
59 CommandSpec::new("ssh", args).ssh_session(ssh)
60}
61
62pub const REMOTE_UPLOAD_STAGING: &str = ".cache/mjolnir/uploads";
66
67pub fn scp_upload(ssh: &SshTarget, source: &Path, remote: &str, recursive: bool) -> CommandSpec {
69 let mut args = scp_args(ssh);
70 if recursive {
71 args.push("-r".into());
72 }
73 args.push(source.to_string_lossy().into_owned());
74 args.push(format!("{}:{remote}", ssh.destination));
75 scp_command(ssh, args)
76}
77
78pub fn scp_download(ssh: &SshTarget, remote: &str, local: &str) -> CommandSpec {
80 let mut args = scp_args(ssh);
81 args.push(format!("{}:{remote}", ssh.destination));
82 args.push(local.into());
83 scp_command(ssh, args)
84}
85
86fn scp_args(ssh: &SshTarget) -> Vec<String> {
89 ssh.ssh_args
90 .iter()
91 .map(|argument| {
92 if argument == "-p" {
93 "-P".to_owned()
94 } else {
95 argument.clone()
96 }
97 })
98 .collect()
99}
100
101fn scp_command(ssh: &SshTarget, args: Vec<String>) -> CommandSpec {
102 CommandSpec::new("scp", args).ssh_session(ssh)
105}
106
107#[cfg(unix)]
112const CONTROL_PERSIST: &str = "60";
113
114pub const CONTROL_MASTER_ENV: &str = "MJ_SSH_CONTROL_MASTER";
117
118#[cfg(unix)]
121const MAX_CONTROL_PATH: usize = 103;
122
123#[cfg(unix)]
126const CONNECTION_HASH_HEX: usize = 16;
127
128#[cfg(unix)]
133const CONTROL_SOCKET_NAME_RESERVE: usize = 1 + CONNECTION_HASH_HEX + 1 + 4 + 17;
134
135#[cfg(unix)]
136fn sharing_disabled(value: Option<&std::ffi::OsStr>) -> bool {
137 let Some(value) = value else {
138 return false;
139 };
140 matches!(
141 value.to_string_lossy().trim().to_ascii_lowercase().as_str(),
142 "0" | "off" | "false" | "no"
143 )
144}
145
146#[doc(hidden)]
149#[derive(Debug, Clone)]
150pub enum SshSharingForTest {
151 Disabled,
153 Directory(PathBuf),
155}
156
157static SHARING_OVERRIDE: Mutex<Option<SshSharingForTest>> = Mutex::new(None);
158
159#[cfg(all(test, unix))]
162pub(super) static SHARING_TEST_LOCK: Mutex<()> = Mutex::new(());
163
164#[doc(hidden)]
169pub fn set_ssh_connection_sharing_for_test(setting: Option<SshSharingForTest>) {
170 *SHARING_OVERRIDE
171 .lock()
172 .unwrap_or_else(std::sync::PoisonError::into_inner) = setting;
173}
174
175#[cfg(unix)]
176fn sharing_override() -> Option<SshSharingForTest> {
177 SHARING_OVERRIDE
178 .lock()
179 .unwrap_or_else(std::sync::PoisonError::into_inner)
180 .clone()
181}
182
183#[cfg(unix)]
194fn control_socket_dir() -> Option<PathBuf> {
195 match sharing_override() {
196 Some(SshSharingForTest::Disabled) => return None,
197 Some(SshSharingForTest::Directory(dir)) => return prepare_control_dir(dir),
198 None => {}
199 }
200 static DIR: OnceLock<Option<PathBuf>> = OnceLock::new();
201 DIR.get_or_init(|| {
202 if sharing_disabled(std::env::var_os(CONTROL_MASTER_ENV).as_deref()) {
203 return None;
204 }
205 let data_dir_override = crate::config::env_override_os("DATA_DIR").map(PathBuf::from);
206 let identity = control_dir_identity(
207 crate::config::instance_name().as_deref(),
208 data_dir_override.as_deref(),
209 );
210 prepare_control_dir(default_control_dir(
211 std::env::var_os("XDG_RUNTIME_DIR"),
212 &identity,
213 ))
214 })
215 .clone()
216}
217
218#[cfg(unix)]
220fn default_control_dir(runtime: Option<std::ffi::OsString>, identity: &str) -> PathBuf {
221 match runtime {
222 Some(runtime) if !runtime.is_empty() => {
223 PathBuf::from(runtime).join("mjolnir").join(identity)
224 }
225 _ => crate::config::data_dir().join("ssh"),
227 }
228}
229
230#[cfg(unix)]
238fn control_dir_identity(instance: Option<&str>, data_dir_override: Option<&Path>) -> String {
239 if instance.is_none() && data_dir_override.is_none() {
240 return "default".to_owned();
241 }
242 let data_dir = data_dir_override.map_or_else(crate::config::data_dir, Path::to_path_buf);
243 crate::config::instance_identity_for(instance, &data_dir)
244}
245
246#[cfg(unix)]
249fn prepare_control_dir(dir: PathBuf) -> Option<PathBuf> {
250 if dir.as_os_str().len() + CONTROL_SOCKET_NAME_RESERVE > MAX_CONTROL_PATH {
251 tracing::debug!(
252 directory = %dir.display(),
253 "skipping SSH connection sharing: control socket path would be too long"
254 );
255 return None;
256 }
257 if let Err(error) = fs::create_dir_all(&dir) {
258 tracing::debug!(
259 directory = %dir.display(),
260 %error,
261 "skipping SSH connection sharing: control directory is unavailable"
262 );
263 return None;
264 }
265 use std::os::unix::fs::PermissionsExt;
266 if let Err(error) = fs::set_permissions(&dir, fs::Permissions::from_mode(0o700)) {
267 tracing::debug!(
268 directory = %dir.display(),
269 %error,
270 "skipping SSH connection sharing: cannot restrict control directory"
271 );
272 return None;
273 }
274 Some(dir)
275}
276
277#[cfg(unix)]
281fn connection_key(ssh: &SshTarget) -> String {
282 let mut key = ssh.destination.clone();
283 for argument in &ssh.ssh_args {
284 key.push('\0');
285 key.push_str(argument);
286 }
287 key
288}
289
290#[cfg(unix)]
296fn control_socket_name(ssh: &SshTarget, shard: usize) -> String {
297 use sha2::{Digest, Sha256};
298 let digest = Sha256::digest(connection_key(ssh).as_bytes());
299 let mut name = String::with_capacity(CONNECTION_HASH_HEX + 5);
300 for byte in digest.iter().take(CONNECTION_HASH_HEX / 2) {
301 name.push_str(&format!("{byte:02x}"));
302 }
303 name.push_str(&format!("-{shard}"));
304 name
305}
306
307#[cfg(unix)]
312fn user_configures_sharing(ssh_args: &[String]) -> bool {
313 ssh_args.iter().any(|argument| {
314 if argument.starts_with("-S") {
315 return true;
316 }
317 let option = argument.strip_prefix("-o").unwrap_or(argument).trim_start();
318 let option = option.to_ascii_lowercase();
319 ["controlmaster", "controlpath"].iter().any(|name| {
320 option
321 .strip_prefix(name)
322 .is_some_and(|rest| rest.starts_with(['=', ' ', '\t']))
323 })
324 })
325}
326
327pub fn push_connection_reuse_args(args: &mut Vec<String>, ssh: &SshTarget) {
345 #[cfg(unix)]
346 if !user_configures_sharing(&ssh.ssh_args)
347 && let Some(dir) = control_socket_dir()
348 {
349 let socket = dir.join(control_socket_name(ssh, 0));
350 args.extend([
351 "-o".to_owned(),
352 "ControlMaster=no".to_owned(),
353 "-o".to_owned(),
354 format!("ControlPath={}", socket.display()),
355 ]);
356 }
357 #[cfg(not(unix))]
358 let _ = (args, ssh);
359}
360
361pub fn join_remote_command(args: &[String]) -> String {
362 args.iter()
363 .map(|arg| posix_quote(arg))
364 .collect::<Vec<_>>()
365 .join(" ")
366}
367
368pub fn ssh_directory_exists(
370 ssh: &SshTarget,
371 path: &Path,
372 executor: &impl CommandExecutor,
373) -> Result<bool> {
374 let command = ssh_validation_command(
375 ssh,
376 vec![
377 "test".into(),
378 "-d".into(),
379 path.to_string_lossy().into_owned(),
380 ],
381 "validate remote directory",
382 );
383 let output = executor.execute(&command)?;
384 match output.status {
385 0 => Ok(true),
386 1 => Ok(false),
387 status => {
388 let stderr = String::from_utf8_lossy(&output.stderr);
389 let error = anyhow::anyhow!(
390 "remote directory check failed with status {status}: {}",
391 stderr.trim()
392 );
393 Err(match host_key_refusal(&stderr, &ssh.ssh_args) {
394 Some(refusal) => error.context(refusal),
395 None => error,
396 })
397 }
398 }
399}
400
401fn host_key_refusal(stderr: &str, ssh_args: &[String]) -> Option<crate::refusal::Refusal> {
410 if !stderr.contains("Host key verification failed") {
411 return None;
412 }
413 let known_hosts = KnownHostsFile::from_ssh_args(ssh_args);
414 let file = known_hosts.phrase();
415 Some(crate::refusal::Refusal::precondition(
416 if stderr.contains("REMOTE HOST IDENTIFICATION HAS CHANGED") {
417 let keygen = match known_hosts.first_named() {
418 Some(path) => format!("`ssh-keygen -f {path} -R`"),
419 None => "`ssh-keygen -R`".to_owned(),
420 };
421 format!(
422 "ssh reported \"Host key verification failed\": the machine's host key is not the one saved in {file}. If you expected the change, remove the old entry with {keygen} and the host name, add the new key, and try again."
423 )
424 } else {
425 format!(
426 "ssh reported \"Host key verification failed\": the machine's host key is not in {file}, and its ssh options require a known key. Add the host key (for example with `ssh-keyscan`, after checking the fingerprint), or put `-o StrictHostKeyChecking=accept-new` in the machine's extra_args, and try again."
427 )
428 },
429 ))
430}
431
432#[derive(Debug, Clone, PartialEq, Eq)]
435enum KnownHostsFile {
436 Default,
438 Named(String),
441 Unknown,
444}
445
446impl KnownHostsFile {
447 fn from_ssh_args(args: &[String]) -> Self {
451 let mut config_file = false;
452 let mut args = args.iter();
453 while let Some(arg) = args.next() {
454 let option = match arg.as_str() {
455 "-o" => args.next().map(String::as_str),
456 other => other.strip_prefix("-o"),
457 };
458 if arg.starts_with("-F") {
459 config_file = true;
460 }
461 let Some(value) = option.and_then(|option| option_value(option, "UserKnownHostsFile"))
462 else {
463 continue;
464 };
465 let value = value.trim_matches('"').trim();
466 return if value.is_empty() || value.eq_ignore_ascii_case("none") {
467 Self::Unknown
468 } else {
469 Self::Named(value.to_owned())
470 };
471 }
472 if config_file {
473 Self::Unknown
474 } else {
475 Self::Default
476 }
477 }
478
479 fn phrase(&self) -> String {
481 match self {
482 Self::Default => "~/.ssh/known_hosts".to_owned(),
483 Self::Named(paths) => paths.split_whitespace().collect::<Vec<_>>().join(" or "),
484 Self::Unknown => "the known_hosts file ssh uses".to_owned(),
485 }
486 }
487
488 fn first_named(&self) -> Option<&str> {
490 match self {
491 Self::Named(paths) => paths.split_whitespace().next(),
492 Self::Default | Self::Unknown => None,
493 }
494 }
495}
496
497fn option_value<'a>(option: &'a str, keyword: &str) -> Option<&'a str> {
500 let option = option.trim_start();
501 let end = option
502 .find(|character: char| character == '=' || character.is_whitespace())
503 .unwrap_or(option.len());
504 let (name, rest) = option.split_at(end);
505 if !name.eq_ignore_ascii_case(keyword) {
506 return None;
507 }
508 let rest = rest.trim_start();
509 Some(rest.strip_prefix('=').unwrap_or(rest).trim())
510}
511
512pub fn validate_bare_project_directory(
514 ssh: &SshTarget,
515 path: &Path,
516 executor: &impl CommandExecutor,
517) -> Result<()> {
518 validate_bare_project_path(path)?;
519 if !ssh_directory_exists(ssh, path, executor)? {
520 bail!(
521 "remote project directory {} does not exist or is not a directory",
522 path.display()
523 );
524 }
525 let output = executor.execute(&ssh_validation_command(
526 ssh,
527 vec![
528 "git".into(),
529 "-C".into(),
530 path.to_string_lossy().into_owned(),
531 "rev-parse".into(),
532 "--verify".into(),
533 "HEAD".into(),
534 ],
535 "validate bare SSH Git project",
536 ))?;
537 if output.status != 0 {
538 let detail = String::from_utf8_lossy(&output.stderr);
539 let detail = detail.trim();
540 if detail.is_empty() {
541 bail!(
542 "remote project directory {} has no valid Git HEAD",
543 path.display()
544 );
545 }
546 bail!(
547 "remote project directory {} has no valid Git HEAD: {detail}",
548 path.display()
549 );
550 }
551 Ok(())
552}
553
554pub fn validate_bare_project_path(path: &Path) -> Result<()> {
555 if !path.is_absolute()
556 || path
557 .components()
558 .any(|part| part == std::path::Component::ParentDir)
559 {
560 bail!("bare project directory must be an absolute safe path");
561 }
562 Ok(())
563}
564
565pub fn ssh_validation_command(
566 ssh: &SshTarget,
567 remote_args: Vec<String>,
568 purpose: &'static str,
569) -> CommandSpec {
570 let mut args = ssh.ssh_args.clone();
571 args.extend([
572 "-o".into(),
573 "BatchMode=yes".into(),
574 "-o".into(),
575 "ConnectTimeout=3".into(),
576 "-o".into(),
577 "ServerAliveInterval=2".into(),
578 "-o".into(),
579 "ServerAliveCountMax=1".into(),
580 ]);
581 args.extend([ssh.destination.clone(), join_remote_command(&remote_args)]);
582 CommandSpec::new("ssh", args)
583 .ssh_probe_session(ssh)
584 .purpose(purpose)
585}
586
587pub fn posix_quote(value: &str) -> String {
591 format!("'{}'", value.replace('\'', "'\\''"))
592}
593
594pub fn verify_locator(locator: &TargetLocator, session_id: &str) -> Result<()> {
595 let expected_name = resource_name(session_id)?;
596 match locator {
597 TargetLocator::LocalBare { worker_root } => {
598 let path = Path::new(worker_root);
599 if !path.is_absolute()
600 || path
601 .components()
602 .any(|part| part == std::path::Component::ParentDir)
603 || !path.ends_with(session_id)
604 {
605 bail!("refusing cleanup: invalid local bare worker root");
606 }
607 }
608 TargetLocator::LocalPodman {
609 container_id,
610 borrowed_from,
611 ..
612 }
613 | TargetLocator::LocalDocker {
614 container_id,
615 borrowed_from,
616 }
617 | TargetLocator::AppleContainer {
618 container_id,
619 borrowed_from,
620 }
621 | TargetLocator::SshPodman {
622 container_id,
623 borrowed_from,
624 ..
625 }
626 | TargetLocator::SshDocker {
627 container_id,
628 borrowed_from,
629 ..
630 } => match borrowed_from {
631 Some(owner) => {
632 validate_session_id(owner)?;
633 if owner == session_id {
634 bail!(
635 "refusing cleanup: a borrowed container cannot be owned by the borrowing session"
636 );
637 }
638 let owner_name = resource_name(owner)?;
639 if container_id != &owner_name && !is_runtime_container_id(container_id) {
640 bail!(
641 "refusing cleanup: borrowed container locator is neither the owning session's generated name nor an immutable runtime ID"
642 );
643 }
644 }
645 None => {
646 if container_id != &expected_name && !is_runtime_container_id(container_id) {
647 bail!(
648 "refusing cleanup: container locator is neither the generated name nor an immutable runtime ID"
649 );
650 }
651 }
652 },
653 TargetLocator::AwsEc2 {
654 instance_id,
655 workspace,
656 ..
657 } => {
658 if !valid_ec2_instance_id(instance_id) {
659 bail!("refusing cleanup: invalid EC2 instance ID");
660 }
661 verify_session_workspace(workspace, session_id)?;
662 }
663 TargetLocator::SshBare {
664 workspace,
665 worker_id,
666 ..
667 } => match worker_id {
668 Some(worker_id) => {
669 validate_session_id(worker_id)?;
670 if worker_id != session_id {
671 bail!("refusing cleanup: SSH worker identity does not match session ID");
672 }
673 validate_workspace_prefix(workspace)?;
674 }
675 None => verify_session_workspace(workspace, session_id)?,
676 },
677 }
678 Ok(())
679}
680
681pub fn is_borrowed(locator: &TargetLocator) -> bool {
685 match locator {
686 TargetLocator::LocalPodman { borrowed_from, .. }
687 | TargetLocator::LocalDocker { borrowed_from, .. }
688 | TargetLocator::AppleContainer { borrowed_from, .. }
689 | TargetLocator::SshPodman { borrowed_from, .. }
690 | TargetLocator::SshDocker { borrowed_from, .. } => borrowed_from.is_some(),
691 TargetLocator::SshBare { worker_id, .. } => worker_id.is_some(),
692 TargetLocator::LocalBare { .. } | TargetLocator::AwsEc2 { .. } => false,
693 }
694}
695
696pub fn verify_session_workspace(workspace: &str, session_id: &str) -> Result<()> {
697 validate_workspace_prefix(workspace)?;
698 let final_component = workspace.trim_end_matches('/').rsplit('/').next();
699 if final_component != Some(session_id) {
700 bail!("refusing cleanup: workspace does not end in the exact session ID");
701 }
702 Ok(())
703}
704
705pub fn validate_session_id(value: &str) -> Result<()> {
706 if value.len() < 8
707 || value.len() > 128
708 || !value
709 .chars()
710 .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_'))
711 {
712 bail!("session ID must be 8-128 ASCII letters, digits, '-' or '_'");
713 }
714 Ok(())
715}
716
717pub fn validate_relative_path(value: &str) -> Result<()> {
718 let path = std::path::Path::new(value);
719 if value.is_empty()
720 || path.is_absolute()
721 || path
722 .components()
723 .any(|part| !matches!(part, std::path::Component::Normal(_)))
724 {
725 bail!("unsafe relative bundle path {value:?}");
726 }
727 Ok(())
728}
729
730pub fn validate_workspace_prefix(value: &str) -> Result<()> {
731 if value.is_empty()
732 || value == "/"
733 || value == "~"
734 || value == "~/"
735 || value.contains('\0')
736 || value.split('/').any(|part| part == "..")
737 {
738 bail!("unsafe workspace path");
739 }
740 Ok(())
741}
742
743pub fn validate_container_template(template: &ContainerTemplate) -> Result<()> {
744 if template.image.trim().is_empty() || template.image.starts_with('-') {
745 bail!("invalid container image");
746 }
747 if template
748 .extra_run_args
749 .iter()
750 .any(|arg| arg == "--name" || arg.starts_with("--name="))
751 {
752 bail!("container template may not override the generated name");
753 }
754 if template.extra_run_args.iter().any(|arg| {
755 arg == "--label"
756 || [SESSION_LABEL, MANAGED_LABEL, INSTANCE_LABEL]
757 .iter()
758 .any(|label| arg.starts_with(&format!("--label={label}=")))
759 }) {
760 bail!("container template may not override Mjolnir ownership labels");
761 }
762 Ok(())
763}
764
765pub fn validate_ssh(ssh: &SshTarget) -> Result<()> {
766 if ssh.destination.trim().is_empty()
767 || ssh.destination.starts_with('-')
768 || ssh.destination.chars().any(char::is_whitespace)
769 {
770 bail!("invalid SSH destination");
771 }
772 Ok(())
773}
774
775pub fn validate_aws(aws: &AwsTemplate) -> Result<()> {
776 validate_ssh(&aws.ssh)?;
777 for (name, value) in [
778 ("AWS profile", &aws.profile),
779 ("AWS region", &aws.region),
780 ("launch template", &aws.launch_template),
781 ] {
782 if value.is_empty()
783 || value.starts_with('-')
784 || !value
785 .chars()
786 .all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.' | '/'))
787 {
788 bail!("invalid {name}");
789 }
790 }
791 Ok(())
792}
793
794pub fn validate_executable(value: &str) -> Result<()> {
795 if value.is_empty() || value.starts_with('-') || value.chars().any(char::is_whitespace) {
796 bail!("invalid executable name");
797 }
798 Ok(())
799}
800
801pub fn valid_ec2_instance_id(value: &str) -> bool {
802 value
803 .strip_prefix("i-")
804 .is_some_and(|rest| rest.len() >= 8 && rest.chars().all(|c| c.is_ascii_hexdigit()))
805}
806
807pub fn is_runtime_container_id(value: &str) -> bool {
808 value.len() >= 12 && value.len() <= 128 && value.chars().all(|c| c.is_ascii_hexdigit())
809}
810
811pub const SSH_TRANSPORT_EXIT_STATUS: i32 = 255;
814
815const TRANSPORT_REJECTION_MARKERS: [&str; 4] = [
819 "Connection closed by",
820 "Connection reset by",
821 "kex_exchange_identification",
822 "Connection timed out during banner exchange",
823];
824
825pub fn is_transport_rejection(status: i32, stderr: &str) -> bool {
831 ssh_refusal(status, stderr).is_some()
832}
833
834const SESSION_REFUSAL_MARKER: &str = "Session open refused by peer";
838
839#[derive(Debug, Clone, Copy, PartialEq, Eq)]
842pub enum SshRefusal {
843 BeforeAuthentication,
847 SessionLimit,
850}
851
852impl SshRefusal {
853 pub fn retry_message(self) -> &'static str {
855 match self {
856 Self::BeforeAuthentication => {
857 "the SSH server closed the connection before authentication; retrying"
858 }
859 Self::SessionLimit => {
860 "the SSH server refused another session on a shared connection (MaxSessions); retrying"
861 }
862 }
863 }
864
865 pub fn log_retry(
874 self,
875 destination: &str,
876 purpose: &str,
877 attempt: usize,
878 delay: Duration,
879 stderr: &str,
880 ) {
881 let delay_ms = delay.as_millis() as u64;
882 match self {
883 Self::SessionLimit => tracing::debug!(
884 destination,
885 purpose,
886 attempt,
887 attempts = SSH_RETRY_ATTEMPTS,
888 delay_ms,
889 stderr,
890 "{}",
891 self.retry_message()
892 ),
893 Self::BeforeAuthentication => tracing::warn!(
894 destination,
895 purpose,
896 attempt,
897 attempts = SSH_RETRY_ATTEMPTS,
898 delay_ms,
899 stderr,
900 "{}",
901 self.retry_message()
902 ),
903 }
904 }
905
906 pub fn log_exhausted(self, destination: &str, purpose: &str, stderr: &str) {
908 tracing::warn!(
909 destination,
910 purpose,
911 attempts = SSH_RETRY_ATTEMPTS,
912 stderr,
913 "{}",
914 match self {
915 Self::BeforeAuthentication =>
916 "the SSH server closed the connection before authentication on every attempt",
917 Self::SessionLimit =>
918 "the SSH server refused another session on a shared connection (MaxSessions) on every attempt",
919 }
920 );
921 }
922}
923
924pub fn ssh_refusal(status: i32, stderr: &str) -> Option<SshRefusal> {
928 if status != SSH_TRANSPORT_EXIT_STATUS {
929 return None;
930 }
931 if stderr.contains(SESSION_REFUSAL_MARKER) {
932 return Some(SshRefusal::SessionLimit);
933 }
934 TRANSPORT_REJECTION_MARKERS
935 .iter()
936 .any(|marker| stderr.contains(marker))
937 .then_some(SshRefusal::BeforeAuthentication)
938}
939
940const DEFAULT_MAX_CONCURRENT_SSH: usize = 6;
949
950pub const MAX_CONCURRENT_SSH_ENV: &str = "MJ_SSH_MAX_CONCURRENT";
952
953fn max_concurrent_ssh() -> usize {
954 static LIMIT: OnceLock<usize> = OnceLock::new();
955 *LIMIT.get_or_init(|| positive_env_limit(MAX_CONCURRENT_SSH_ENV, DEFAULT_MAX_CONCURRENT_SSH))
956}
957
958fn positive_env_limit(name: &str, default: usize) -> usize {
961 let Some(raw) = std::env::var_os(name) else {
962 return default;
963 };
964 match raw
965 .to_str()
966 .and_then(|value| value.trim().parse::<usize>().ok())
967 {
968 Some(limit) if limit > 0 => limit,
969 _ => {
970 tracing::warn!(
971 variable = name,
972 value = %raw.to_string_lossy(),
973 default,
974 "ignoring invalid SSH limit"
975 );
976 default
977 }
978 }
979}
980
981struct DestinationGate {
987 limit: usize,
988 in_flight: Mutex<usize>,
989 released: Condvar,
990}
991
992impl DestinationGate {
993 fn new(limit: usize) -> Arc<Self> {
994 Arc::new(Self {
995 limit,
996 in_flight: Mutex::new(0),
997 released: Condvar::new(),
998 })
999 }
1000
1001 fn acquire(self: &Arc<Self>) -> SshPermit {
1002 let mut in_flight = self
1003 .in_flight
1004 .lock()
1005 .unwrap_or_else(std::sync::PoisonError::into_inner);
1006 while *in_flight >= self.limit {
1007 in_flight = self
1008 .released
1009 .wait(in_flight)
1010 .unwrap_or_else(std::sync::PoisonError::into_inner);
1011 }
1012 *in_flight += 1;
1013 drop(in_flight);
1014 SshPermit {
1015 gate: Arc::clone(self),
1016 }
1017 }
1018}
1019
1020pub struct SshPermit {
1022 gate: Arc<DestinationGate>,
1023}
1024
1025impl std::fmt::Debug for SshPermit {
1026 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1027 formatter.write_str("SshPermit")
1028 }
1029}
1030
1031impl Drop for SshPermit {
1032 fn drop(&mut self) {
1033 let mut in_flight = self
1034 .gate
1035 .in_flight
1036 .lock()
1037 .unwrap_or_else(std::sync::PoisonError::into_inner);
1038 *in_flight = in_flight.saturating_sub(1);
1039 drop(in_flight);
1040 self.gate.released.notify_one();
1041 }
1042}
1043
1044pub struct SshAdmission;
1046
1047impl SshAdmission {
1048 pub fn acquire(destination: &str) -> SshPermit {
1051 Self::gate(destination).acquire()
1052 }
1053
1054 fn gate(destination: &str) -> Arc<DestinationGate> {
1055 static GATES: OnceLock<Mutex<BTreeMap<String, Arc<DestinationGate>>>> = OnceLock::new();
1056 let mut gates = GATES
1057 .get_or_init(|| Mutex::new(BTreeMap::new()))
1058 .lock()
1059 .unwrap_or_else(std::sync::PoisonError::into_inner);
1060 Arc::clone(
1061 gates
1062 .entry(destination.to_owned())
1063 .or_insert_with(|| DestinationGate::new(max_concurrent_ssh())),
1064 )
1065 }
1066}
1067
1068#[cfg(unix)]
1075const DEFAULT_SESSIONS_PER_CONNECTION: usize = 8;
1076
1077pub const SESSIONS_PER_CONNECTION_ENV: &str = "MJ_SSH_SESSIONS_PER_CONNECTION";
1079
1080#[cfg(unix)]
1083const MASTER_CHECK_INTERVAL: Duration = Duration::from_secs(5);
1084
1085pub const SSH_MASTER_OPEN_TIMEOUT: Duration = Duration::from_secs(60);
1088
1089#[cfg(unix)]
1090fn sessions_per_connection() -> usize {
1091 static LIMIT: OnceLock<usize> = OnceLock::new();
1092 *LIMIT.get_or_init(|| {
1093 positive_env_limit(SESSIONS_PER_CONNECTION_ENV, DEFAULT_SESSIONS_PER_CONNECTION)
1094 })
1095}
1096
1097#[cfg(unix)]
1099struct Shard {
1100 leased: usize,
1101 verified_at: Option<Instant>,
1103 opening: Arc<Mutex<()>>,
1106}
1107
1108#[cfg(unix)]
1110struct SessionLedger {
1111 per_connection: usize,
1112 connections: Mutex<BTreeMap<String, Vec<Shard>>>,
1113}
1114
1115#[cfg(unix)]
1116impl SessionLedger {
1117 fn new(per_connection: usize) -> Arc<Self> {
1118 Arc::new(Self {
1119 per_connection: per_connection.max(1),
1120 connections: Mutex::new(BTreeMap::new()),
1121 })
1122 }
1123
1124 fn global() -> Arc<Self> {
1125 static LEDGER: OnceLock<Arc<SessionLedger>> = OnceLock::new();
1126 Arc::clone(LEDGER.get_or_init(|| Self::new(sessions_per_connection())))
1127 }
1128
1129 fn connections(&self) -> std::sync::MutexGuard<'_, BTreeMap<String, Vec<Shard>>> {
1130 self.connections
1131 .lock()
1132 .unwrap_or_else(std::sync::PoisonError::into_inner)
1133 }
1134
1135 fn reserve(&self, key: &str) -> (usize, Arc<Mutex<()>>) {
1138 let mut connections = self.connections();
1139 let shards = connections.entry(key.to_owned()).or_default();
1140 let index = match shards
1141 .iter()
1142 .position(|shard| shard.leased < self.per_connection)
1143 {
1144 Some(index) => index,
1145 None => {
1146 shards.push(Shard {
1147 leased: 0,
1148 verified_at: None,
1149 opening: Arc::new(Mutex::new(())),
1150 });
1151 shards.len() - 1
1152 }
1153 };
1154 shards[index].leased += 1;
1155 (index, Arc::clone(&shards[index].opening))
1156 }
1157
1158 fn lease(
1161 self: &Arc<Self>,
1162 ssh: &SshTarget,
1163 dir: &Path,
1164 executor: &dyn CommandExecutor,
1165 ) -> Result<SshSessionLease> {
1166 let key = connection_key(ssh);
1167 let (shard, opening) = self.reserve(&key);
1168 let slot = LeasedSlot {
1170 ledger: Arc::clone(self),
1171 key,
1172 shard,
1173 socket: dir.join(control_socket_name(ssh, shard)),
1174 };
1175 if slot.needs_check() {
1176 let _opening = opening
1177 .lock()
1178 .unwrap_or_else(std::sync::PoisonError::into_inner);
1179 if slot.needs_check() {
1182 ensure_master(ssh, &slot.socket, executor)?;
1183 slot.set_verified(Some(Instant::now()));
1184 }
1185 }
1186 Ok(SshSessionLease {
1187 slot: Some(slot),
1188 probe: false,
1189 })
1190 }
1191
1192 fn lease_probe(self: &Arc<Self>, ssh: &SshTarget, dir: &Path) -> SshSessionLease {
1201 let key = connection_key(ssh);
1202 let (shard, _) = self.reserve(&key);
1203 SshSessionLease {
1204 slot: Some(LeasedSlot {
1205 ledger: Arc::clone(self),
1206 key,
1207 shard,
1208 socket: dir.join(control_socket_name(ssh, shard)),
1209 }),
1210 probe: true,
1211 }
1212 }
1213}
1214
1215#[cfg(unix)]
1221fn ensure_master(ssh: &SshTarget, socket: &Path, executor: &dyn CommandExecutor) -> Result<()> {
1222 let _opening = lock_master_opening(socket)?;
1230 if master_running(ssh, socket, executor)? {
1231 return Ok(());
1232 }
1233 match fs::remove_file(socket) {
1237 Ok(()) => tracing::debug!(
1238 socket = %socket.display(),
1239 "removed a stale SSH control socket"
1240 ),
1241 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
1242 Err(error) => {
1243 return Err(error)
1244 .with_context(|| format!("remove stale SSH control socket {}", socket.display()));
1245 }
1246 }
1247 let opened = executor.execute(&master_open_command(ssh, socket))?;
1248 if master_running(ssh, socket, executor)? {
1249 tracing::info!(
1250 destination = ssh.destination.as_str(),
1251 socket = %socket.display(),
1252 "opened a shared SSH connection"
1253 );
1254 return Ok(());
1255 }
1256 let stderr = String::from_utf8_lossy(&opened.stderr);
1257 let detail = match stderr.trim() {
1258 "" => format!("ssh exited with status {}", opened.status),
1259 stderr => stderr.to_owned(),
1260 };
1261 bail!(
1262 "could not open a shared SSH connection to {}: {detail}",
1263 ssh.destination
1264 )
1265}
1266
1267#[cfg(unix)]
1272fn lock_master_opening(socket: &Path) -> Result<fs::File> {
1273 let mut path = socket.as_os_str().to_owned();
1274 path.push(".lock");
1275 let path = PathBuf::from(path);
1276 loop {
1277 let file = fs::OpenOptions::new()
1278 .create(true)
1279 .truncate(false)
1280 .write(true)
1281 .open(&path)
1282 .with_context(|| format!("open SSH master lock {}", path.display()))?;
1283 file.lock()
1284 .with_context(|| format!("lock SSH master lock {}", path.display()))?;
1285 if is_file_at(&file, &path) {
1289 return Ok(file);
1290 }
1291 }
1292}
1293
1294#[cfg(unix)]
1303fn remove_stale_master_locks_in(dir: &Path) {
1304 let entries = match fs::read_dir(dir) {
1305 Ok(entries) => entries,
1306 Err(error) => {
1307 tracing::debug!(directory = %dir.display(), %error, "cannot list SSH master locks");
1308 return;
1309 }
1310 };
1311 for entry in entries.flatten() {
1312 let lock = entry.path();
1313 let Some(socket) = lock
1314 .file_name()
1315 .and_then(std::ffi::OsStr::to_str)
1316 .and_then(|name| name.strip_suffix(".lock"))
1317 .map(|name| dir.join(name))
1318 else {
1319 continue;
1320 };
1321 if fs::symlink_metadata(&socket).is_ok() {
1322 continue;
1323 }
1324 let Ok(file) = fs::OpenOptions::new().write(true).open(&lock) else {
1325 continue;
1326 };
1327 if file.try_lock().is_err() {
1329 continue;
1330 }
1331 if fs::symlink_metadata(&socket).is_ok() || !is_file_at(&file, &lock) {
1332 continue;
1333 }
1334 match fs::remove_file(&lock) {
1335 Ok(()) => tracing::debug!(lock = %lock.display(), "removed a stale SSH master lock"),
1336 Err(error) => {
1337 tracing::debug!(lock = %lock.display(), %error, "cannot remove a stale SSH master lock")
1338 }
1339 }
1340 }
1341}
1342
1343#[cfg(unix)]
1345fn is_file_at(file: &fs::File, path: &Path) -> bool {
1346 use std::os::unix::fs::MetadataExt;
1347 match (file.metadata(), fs::metadata(path)) {
1348 (Ok(open), Ok(named)) => open.dev() == named.dev() && open.ino() == named.ino(),
1349 _ => false,
1350 }
1351}
1352
1353#[cfg(unix)]
1354fn master_running(ssh: &SshTarget, socket: &Path, executor: &dyn CommandExecutor) -> Result<bool> {
1355 Ok(executor.execute(&master_check_command(ssh, socket))?.status == 0)
1356}
1357
1358#[cfg(unix)]
1361fn master_check_command(ssh: &SshTarget, socket: &Path) -> CommandSpec {
1362 let mut args = ssh.ssh_args.clone();
1363 args.extend([
1364 "-o".to_owned(),
1365 format!("ControlPath={}", socket.display()),
1366 "-O".to_owned(),
1367 "check".to_owned(),
1368 ssh.destination.clone(),
1369 ]);
1370 CommandSpec::new("ssh", args).purpose("check a shared SSH connection")
1371}
1372
1373#[cfg(unix)]
1381fn master_open_command(ssh: &SshTarget, socket: &Path) -> CommandSpec {
1382 let mut args = ssh.ssh_args.clone();
1383 args.extend([
1384 "-o".to_owned(),
1385 "BatchMode=yes".to_owned(),
1386 "-f".to_owned(),
1387 "-N".to_owned(),
1388 "-o".to_owned(),
1389 "ControlMaster=yes".to_owned(),
1390 "-o".to_owned(),
1391 format!("ControlPath={}", socket.display()),
1392 "-o".to_owned(),
1393 format!("ControlPersist={CONTROL_PERSIST}"),
1394 ssh.destination.clone(),
1395 ]);
1396 CommandSpec::new("ssh", args)
1397 .ssh_destination(ssh.destination.clone())
1398 .purpose("open a shared SSH connection")
1399}
1400
1401#[cfg(unix)]
1403struct LeasedSlot {
1404 ledger: Arc<SessionLedger>,
1405 key: String,
1406 shard: usize,
1407 socket: PathBuf,
1408}
1409
1410#[cfg(unix)]
1411impl LeasedSlot {
1412 fn needs_check(&self) -> bool {
1413 let connections = self.ledger.connections();
1414 connections
1415 .get(&self.key)
1416 .and_then(|shards| shards.get(self.shard))
1417 .is_none_or(|shard| {
1418 shard
1419 .verified_at
1420 .is_none_or(|verified| verified.elapsed() >= MASTER_CHECK_INTERVAL)
1421 })
1422 }
1423
1424 fn set_verified(&self, verified_at: Option<Instant>) {
1425 let mut connections = self.ledger.connections();
1426 if let Some(shard) = connections
1427 .get_mut(&self.key)
1428 .and_then(|shards| shards.get_mut(self.shard))
1429 {
1430 shard.verified_at = verified_at;
1431 }
1432 }
1433}
1434
1435#[cfg(unix)]
1436impl Drop for LeasedSlot {
1437 fn drop(&mut self) {
1438 let mut connections = self.ledger.connections();
1439 if let Some(shard) = connections
1440 .get_mut(&self.key)
1441 .and_then(|shards| shards.get_mut(self.shard))
1442 {
1443 shard.leased = shard.leased.saturating_sub(1);
1444 }
1445 }
1446}
1447
1448pub struct SshSessionLease {
1456 #[cfg(unix)]
1457 slot: Option<LeasedSlot>,
1458 #[cfg(unix)]
1461 probe: bool,
1462}
1463
1464impl std::fmt::Debug for SshSessionLease {
1465 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1466 formatter
1467 .debug_struct("SshSessionLease")
1468 .field("control_path", &self.control_path())
1469 .finish()
1470 }
1471}
1472
1473impl SshSessionLease {
1474 fn unshared() -> Self {
1475 Self {
1476 #[cfg(unix)]
1477 slot: None,
1478 #[cfg(unix)]
1479 probe: false,
1480 }
1481 }
1482
1483 pub fn control_path(&self) -> Option<&Path> {
1486 #[cfg(unix)]
1487 return self.slot.as_ref().map(|slot| slot.socket.as_path());
1488 #[cfg(not(unix))]
1489 None
1490 }
1491
1492 pub fn invalidate(&self) {
1496 #[cfg(unix)]
1497 if let Some(slot) = &self.slot {
1498 slot.set_verified(None);
1499 }
1500 }
1501}
1502
1503pub struct SshSessions;
1510
1511impl SshSessions {
1512 pub fn lease(ssh: &SshTarget, executor: &dyn CommandExecutor) -> Result<SshSessionLease> {
1517 #[cfg(unix)]
1518 {
1519 if user_configures_sharing(&ssh.ssh_args) {
1520 return Ok(SshSessionLease::unshared());
1521 }
1522 let Some(dir) = control_socket_dir() else {
1523 return Ok(SshSessionLease::unshared());
1524 };
1525 SessionLedger::global().lease(ssh, &dir, executor)
1526 }
1527 #[cfg(not(unix))]
1528 {
1529 let _ = (ssh, executor);
1530 Ok(SshSessionLease::unshared())
1531 }
1532 }
1533
1534 pub fn remove_stale_master_locks() {
1537 #[cfg(unix)]
1538 if let Some(dir) = control_socket_dir() {
1539 remove_stale_master_locks_in(&dir);
1540 }
1541 }
1542
1543 pub fn lease_probe(ssh: &SshTarget) -> SshSessionLease {
1551 #[cfg(unix)]
1552 {
1553 if user_configures_sharing(&ssh.ssh_args) {
1554 return SshSessionLease::unshared();
1555 }
1556 let Some(dir) = control_socket_dir() else {
1557 return SshSessionLease::unshared();
1558 };
1559 SessionLedger::global().lease_probe(ssh, &dir)
1560 }
1561 #[cfg(not(unix))]
1562 {
1563 let _ = ssh;
1564 SshSessionLease::unshared()
1565 }
1566 }
1567}
1568
1569pub fn push_session_args(args: &mut Vec<String>, lease: &SshSessionLease) {
1578 if let Some(socket) = lease.control_path() {
1579 args.extend([
1580 "-o".to_owned(),
1581 "ControlMaster=no".to_owned(),
1582 "-o".to_owned(),
1583 format!("ControlPath={}", socket.display()),
1584 ]);
1585 #[cfg(unix)]
1586 let probe = lease.probe;
1587 #[cfg(not(unix))]
1588 let probe = false;
1589 if !probe {
1590 args.extend(["-o".to_owned(), "ProxyCommand=false".to_owned()]);
1591 }
1592 }
1593}
1594
1595pub fn session_command_args(
1608 program: &str,
1609 args: &[String],
1610 ssh: &SshTarget,
1611 lease: &SshSessionLease,
1612) -> Vec<String> {
1613 let mut session = Vec::with_capacity(args.len() + 6);
1614 push_session_args(&mut session, lease);
1615 if session.is_empty() {
1616 return args.to_vec();
1617 }
1618 if program == "ssh" && args.starts_with(&ssh.ssh_args) {
1619 let (user, rest) = args.split_at(ssh.ssh_args.len());
1620 let mut user = user.iter();
1621 while let Some(argument) = user.next() {
1622 match argument.strip_prefix("-J") {
1623 Some("") => match user.next() {
1624 Some(jump) => session.extend(["-o".to_owned(), format!("ProxyJump={jump}")]),
1625 None => session.push(argument.clone()),
1626 },
1627 Some(jump) => session.extend(["-o".to_owned(), format!("ProxyJump={jump}")]),
1628 None => session.push(argument.clone()),
1629 }
1630 }
1631 session.extend(rest.iter().cloned());
1632 } else {
1633 session.extend(args.iter().cloned());
1634 }
1635 session
1636}
1637
1638pub const SSH_RETRY_ATTEMPTS: usize = 3;
1640
1641const SSH_RETRY_BACKOFF_MS: [(u64, u64); SSH_RETRY_ATTEMPTS - 1] = [(500, 2_000), (2_000, 4_000)];
1645
1646static SSH_RETRY_BACKOFF_OVERRIDE_MS: AtomicU64 = AtomicU64::new(u64::MAX);
1649
1650#[doc(hidden)]
1653pub fn set_ssh_retry_backoff_for_test(delay: Option<Duration>) {
1654 SSH_RETRY_BACKOFF_OVERRIDE_MS.store(
1655 delay.map_or(u64::MAX, |delay| delay.as_millis() as u64),
1656 Ordering::Relaxed,
1657 );
1658}
1659
1660pub fn ssh_retry_delay(attempts_made: usize) -> Duration {
1666 let override_ms = SSH_RETRY_BACKOFF_OVERRIDE_MS.load(Ordering::Relaxed);
1667 if override_ms != u64::MAX {
1668 return Duration::from_millis(override_ms);
1669 }
1670 let (low, high) = SSH_RETRY_BACKOFF_MS
1671 .get(attempts_made.saturating_sub(1))
1672 .copied()
1673 .unwrap_or(*SSH_RETRY_BACKOFF_MS.last().expect("non-empty schedule"));
1674 let mut bytes = [0_u8; 8];
1675 let spread = if getrandom::fill(&mut bytes).is_ok() {
1677 u64::from_le_bytes(bytes) % (high - low + 1)
1678 } else {
1679 0
1680 };
1681 Duration::from_millis(low + spread)
1682}
1683
1684#[cfg(test)]
1685mod tests {
1686 use super::*;
1687
1688 #[cfg(unix)]
1693 #[test]
1694 fn stale_master_locks_are_removed_but_live_or_held_ones_stay() {
1695 let dir = tempfile::tempdir().unwrap();
1696 let stale = dir.path().join("aaaaaaaaaaaaaaaa-0.lock");
1697 fs::write(&stale, b"").unwrap();
1698 fs::write(dir.path().join("bbbbbbbbbbbbbbbb-0"), b"").unwrap();
1699 let live = dir.path().join("bbbbbbbbbbbbbbbb-0.lock");
1700 fs::write(&live, b"").unwrap();
1701 let held = dir.path().join("cccccccccccccccc-0.lock");
1702 let holder = lock_master_opening(&dir.path().join("cccccccccccccccc-0")).unwrap();
1703
1704 remove_stale_master_locks_in(dir.path());
1705 assert!(!stale.exists(), "a lock whose master is gone is removed");
1706 assert!(live.exists(), "a lock beside a live socket stays");
1707 assert!(held.exists(), "a lock another opener holds stays");
1708
1709 drop(holder);
1710 for _ in 0..200 {
1715 remove_stale_master_locks_in(dir.path());
1716 if !held.exists() {
1717 break;
1718 }
1719 std::thread::sleep(Duration::from_millis(10));
1720 }
1721 assert!(!held.exists(), "a lock nobody holds any more is removed");
1722 }
1723
1724 const BORROW_PARENT: &str = "0123456789abcdef0123456789abcdef";
1725 const BORROW_CHILD: &str = "fedcba9876543210fedcba9876543210";
1726
1727 fn borrowed_podman(owner: &str) -> TargetLocator {
1728 TargetLocator::LocalPodman {
1729 container_id: crate::targets::resource_name(owner).unwrap(),
1730 workspace_storage: PodmanWorkspaceLocator::default(),
1731 borrowed_from: Some(owner.to_owned()),
1732 }
1733 }
1734
1735 #[test]
1736 fn verify_locator_accepts_a_container_borrowed_from_its_owner() {
1737 verify_locator(&borrowed_podman(BORROW_PARENT), BORROW_CHILD)
1738 .expect("a child may borrow its parent's container");
1739 }
1740
1741 #[test]
1742 fn verify_locator_rejects_a_container_borrowed_from_the_checking_session() {
1743 let error = verify_locator(&borrowed_podman(BORROW_PARENT), BORROW_PARENT)
1744 .expect_err("a session cannot borrow from itself");
1745 assert!(
1746 format!("{error:#}").contains("cannot be owned by the borrowing session"),
1747 "unexpected error: {error:#}"
1748 );
1749 }
1750
1751 #[test]
1752 fn verify_locator_rejects_a_borrowed_container_naming_another_session() {
1753 let locator = TargetLocator::LocalPodman {
1754 container_id: crate::targets::resource_name(BORROW_CHILD).unwrap(),
1755 workspace_storage: PodmanWorkspaceLocator::default(),
1756 borrowed_from: Some(BORROW_PARENT.to_owned()),
1757 };
1758 let error = verify_locator(&locator, BORROW_CHILD)
1759 .expect_err("the container must belong to the recorded owner");
1760 assert!(
1761 format!("{error:#}").contains("borrowed container locator"),
1762 "unexpected error: {error:#}"
1763 );
1764 }
1765
1766 #[test]
1767 fn worker_root_of_a_borrowed_container_is_the_childs_own_directory() {
1768 assert_eq!(
1769 crate::targets::worker_root(&borrowed_podman(BORROW_PARENT), BORROW_CHILD).unwrap(),
1770 format!("/var/lib/hel/workers/{BORROW_CHILD}")
1771 );
1772 }
1773
1774 #[test]
1775 fn is_borrowed_distinguishes_borrowed_targets_from_owned_ones() {
1776 assert!(is_borrowed(&borrowed_podman(BORROW_PARENT)));
1777 assert!(is_borrowed(&TargetLocator::SshBare {
1778 ssh: SshTarget {
1779 destination: "host".to_owned(),
1780 ssh_args: Vec::new(),
1781 },
1782 workspace: format!(".local/share/hel/workspaces/{BORROW_PARENT}"),
1783 worker_id: Some(BORROW_CHILD.to_owned()),
1784 }));
1785 assert!(!is_borrowed(&TargetLocator::LocalPodman {
1786 container_id: crate::targets::resource_name(BORROW_CHILD).unwrap(),
1787 workspace_storage: PodmanWorkspaceLocator::default(),
1788 borrowed_from: None,
1789 }));
1790 }
1791
1792 #[test]
1793 fn an_owned_container_locator_serializes_without_a_borrowed_from_key() {
1794 let owned = TargetLocator::LocalDocker {
1795 container_id: crate::targets::resource_name(BORROW_CHILD).unwrap(),
1796 borrowed_from: None,
1797 };
1798 let serialized = serde_json::to_string(&owned).unwrap();
1799 assert!(
1800 !serialized.contains("borrowed_from"),
1801 "owned locators must stay byte-identical for older readers: {serialized}"
1802 );
1803 assert_eq!(
1804 serde_json::from_str::<TargetLocator>(&serialized).unwrap(),
1805 owned
1806 );
1807
1808 let borrowed = borrowed_podman(BORROW_PARENT);
1809 let serialized = serde_json::to_string(&borrowed).unwrap();
1810 assert!(serialized.contains("borrowed_from"));
1811 assert_eq!(
1812 serde_json::from_str::<TargetLocator>(&serialized).unwrap(),
1813 borrowed
1814 );
1815 }
1816 use std::sync::atomic::{AtomicUsize, Ordering};
1817
1818 #[cfg(unix)]
1820 #[derive(Default)]
1821 struct RecordingExecutor {
1822 seen: std::cell::RefCell<Vec<CommandSpec>>,
1823 }
1824
1825 #[cfg(unix)]
1826 impl CommandExecutor for RecordingExecutor {
1827 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
1828 self.seen.borrow_mut().push(command.clone());
1829 Ok(CommandOutput {
1830 status: 0,
1831 stdout: Vec::new(),
1832 stderr: Vec::new(),
1833 })
1834 }
1835 }
1836
1837 #[cfg(unix)]
1838 fn sharing_socket_dir() -> tempfile::TempDir {
1839 tempfile::tempdir_in("/tmp").expect("short control socket directory")
1841 }
1842
1843 #[cfg(unix)]
1847 fn spawned_args(
1848 command: &CommandSpec,
1849 dir: Option<&Path>,
1850 masters: &FakeMasters,
1851 ) -> Vec<String> {
1852 set_ssh_connection_sharing_for_test(Some(match dir {
1853 Some(dir) => SshSharingForTest::Directory(dir.to_path_buf()),
1854 None => SshSharingForTest::Disabled,
1855 }));
1856 let session = command.open_ssh_session(masters);
1857 set_ssh_connection_sharing_for_test(None);
1858 session.expect("session").command().args.clone()
1859 }
1860
1861 #[test]
1865 #[cfg(unix)]
1866 fn session_options_lead_the_spawned_command() {
1867 let _guard = SHARING_TEST_LOCK
1868 .lock()
1869 .unwrap_or_else(std::sync::PoisonError::into_inner);
1870 let socket_dir = sharing_socket_dir();
1871 let ssh = SshTarget {
1872 destination: "session-options-host".to_owned(),
1873 ssh_args: vec![
1874 "-p".to_owned(),
1875 "2222".to_owned(),
1876 "-o".to_owned(),
1877 "ProxyCommand=nc %h %p".to_owned(),
1878 ],
1879 };
1880 let command = ssh_command(&ssh, ["true"]);
1881 assert_eq!(
1882 command.args,
1883 [
1884 "-p",
1885 "2222",
1886 "-o",
1887 "ProxyCommand=nc %h %p",
1888 "session-options-host",
1889 "'true'"
1890 ],
1891 "stored arguments never contain sharing options"
1892 );
1893 assert_eq!(command.ssh_session.as_ref(), Some(&ssh));
1894 assert_eq!(
1895 command.ssh_destination.as_deref(),
1896 Some("session-options-host")
1897 );
1898
1899 let masters = FakeMasters::default();
1900 let args = spawned_args(&command, Some(socket_dir.path()), &masters);
1901 let socket = socket_dir.path().join(control_socket_name(&ssh, 0));
1902 assert_eq!(
1903 args,
1904 [
1905 "-o".to_owned(),
1906 "ControlMaster=no".to_owned(),
1907 "-o".to_owned(),
1908 format!("ControlPath={}", socket.display()),
1909 "-o".to_owned(),
1910 "ProxyCommand=false".to_owned(),
1911 "-p".to_owned(),
1912 "2222".to_owned(),
1913 "-o".to_owned(),
1914 "ProxyCommand=nc %h %p".to_owned(),
1915 "session-options-host".to_owned(),
1916 "'true'".to_owned(),
1917 ]
1918 );
1919 assert_eq!(masters.openers(), 1);
1920 assert_eq!(
1921 std::os::unix::fs::MetadataExt::mode(
1922 &fs::metadata(socket_dir.path()).expect("socket directory")
1923 ) & 0o777,
1924 0o700
1925 );
1926 }
1927
1928 #[test]
1932 #[cfg(unix)]
1933 fn a_jump_host_flag_becomes_an_option_the_session_guard_overrides() {
1934 let _guard = SHARING_TEST_LOCK
1935 .lock()
1936 .unwrap_or_else(std::sync::PoisonError::into_inner);
1937 let socket_dir = sharing_socket_dir();
1938 let ssh = SshTarget {
1939 destination: "jump-rewrite-host".to_owned(),
1940 ssh_args: vec!["-J".to_owned(), "bastion".to_owned(), "-Jother".to_owned()],
1941 };
1942 let masters = FakeMasters::default();
1943 let args = spawned_args(
1944 &ssh_command(&ssh, ["-J"]),
1945 Some(socket_dir.path()),
1946 &masters,
1947 );
1948 assert_eq!(
1949 args[6..],
1950 [
1951 "-o",
1952 "ProxyJump=bastion",
1953 "-o",
1954 "ProxyJump=other",
1955 "jump-rewrite-host",
1956 "'-J'",
1957 ]
1958 );
1959 let upload = spawned_args(
1960 &scp_upload(&ssh, Path::new("/tmp/file"), "file", false),
1961 Some(socket_dir.path()),
1962 &masters,
1963 );
1964 assert_eq!(
1965 upload[6..],
1966 [
1967 "-J",
1968 "bastion",
1969 "-Jother",
1970 "/tmp/file",
1971 "jump-rewrite-host:file"
1972 ]
1973 );
1974 }
1975
1976 #[test]
1980 #[cfg(unix)]
1981 fn user_configured_sharing_suppresses_mjolnir_sharing() {
1982 let _guard = SHARING_TEST_LOCK
1983 .lock()
1984 .unwrap_or_else(std::sync::PoisonError::into_inner);
1985 let socket_dir = sharing_socket_dir();
1986 let spellings: [&[&str]; 5] = [
1987 &["-o", "ControlMaster=no"],
1988 &["-o", "controlpath /tmp/mine"],
1989 &["-oControlPath=/tmp/mine"],
1990 &["-S", "/tmp/mine"],
1991 &["-S/tmp/mine"],
1992 ];
1993 let masters = FakeMasters::default();
1994 for user in spellings {
1995 let ssh = SshTarget {
1996 destination: "user-sharing-host".to_owned(),
1997 ssh_args: user.iter().map(|arg| (*arg).to_owned()).collect(),
1998 };
1999 set_ssh_connection_sharing_for_test(Some(SshSharingForTest::Directory(
2000 socket_dir.path().to_path_buf(),
2001 )));
2002 let validation = ssh_validation_command(&ssh, vec!["true".to_owned()], "test");
2003 let validation = spawned_args(&validation, Some(socket_dir.path()), &masters);
2004 let command = ssh_command(&ssh, ["true"]);
2005 let args = spawned_args(&command, Some(socket_dir.path()), &masters);
2006 assert_eq!(args, command.args, "user args {user:?}");
2007 let socket_dir_text = socket_dir.path().display().to_string();
2008 assert!(
2009 !validation.iter().any(|arg| arg.contains(&socket_dir_text)),
2010 "user args {user:?}: {validation:?}"
2011 );
2012 }
2013 assert_eq!(masters.commands(), 0);
2014 }
2015
2016 #[test]
2019 #[cfg(unix)]
2020 fn control_sockets_live_in_a_directory_per_instance() {
2021 let runtime = Some(std::ffi::OsString::from("/run/user/1000"));
2022 assert_eq!(
2023 default_control_dir(runtime.clone(), &control_dir_identity(Some("hel2"), None)),
2024 PathBuf::from("/run/user/1000/mjolnir/hel2")
2025 );
2026 assert_eq!(
2027 default_control_dir(runtime, &control_dir_identity(None, None)),
2028 PathBuf::from("/run/user/1000/mjolnir/default")
2029 );
2030 }
2031
2032 #[test]
2037 #[cfg(unix)]
2038 fn a_data_directory_override_gets_its_own_socket_directory() {
2039 let runtime = Some(std::ffi::OsString::from("/run/user/1000"));
2040 let lab = Path::new("/tmp/lab-a/data");
2041 let other_lab = Path::new("/tmp/lab-b/data");
2042 let lab_dir = default_control_dir(runtime.clone(), &control_dir_identity(None, Some(lab)));
2043 assert_ne!(lab_dir, PathBuf::from("/run/user/1000/mjolnir/default"));
2044 assert_eq!(
2045 lab_dir,
2046 PathBuf::from("/run/user/1000/mjolnir")
2047 .join(crate::config::instance_identity_for(None, lab))
2048 );
2049 assert_ne!(
2050 lab_dir,
2051 default_control_dir(
2052 runtime.clone(),
2053 &control_dir_identity(None, Some(other_lab))
2054 )
2055 );
2056 assert_eq!(
2059 default_control_dir(runtime, &control_dir_identity(Some("hel2"), Some(lab))),
2060 PathBuf::from("/run/user/1000/mjolnir/hel2")
2061 );
2062 }
2063
2064 #[test]
2067 #[cfg(unix)]
2068 fn socket_names_identify_the_connection_and_the_shard() {
2069 let plain = SshTarget {
2070 destination: "host".to_owned(),
2071 ssh_args: Vec::new(),
2072 };
2073 let other_port = SshTarget {
2074 destination: "host".to_owned(),
2075 ssh_args: vec!["-p".to_owned(), "2222".to_owned()],
2076 };
2077 let first = control_socket_name(&plain, 0);
2078 let second = control_socket_name(&plain, 1);
2079 assert_eq!(first.len(), CONNECTION_HASH_HEX + 2, "{first}");
2080 assert!(first.ends_with("-0") && second.ends_with("-1"));
2081 assert_eq!(first[..CONNECTION_HASH_HEX], second[..CONNECTION_HASH_HEX]);
2082 assert_ne!(
2083 first[..CONNECTION_HASH_HEX],
2084 control_socket_name(&other_port, 0)[..CONNECTION_HASH_HEX]
2085 );
2086 }
2087
2088 #[test]
2092 #[cfg(unix)]
2093 fn fail_fast_commands_reuse_a_master_without_becoming_one() {
2094 let _guard = SHARING_TEST_LOCK
2095 .lock()
2096 .unwrap_or_else(std::sync::PoisonError::into_inner);
2097 let socket_dir = sharing_socket_dir();
2098 set_ssh_connection_sharing_for_test(Some(SshSharingForTest::Directory(
2099 socket_dir.path().to_path_buf(),
2100 )));
2101 let ssh = SshTarget {
2102 destination: "host".to_owned(),
2103 ssh_args: Vec::new(),
2104 };
2105 let masters = FakeMasters::default();
2106 let validation = spawned_args(
2107 &ssh_validation_command(&ssh, vec!["true".to_owned()], "test"),
2108 Some(socket_dir.path()),
2109 &masters,
2110 );
2111 assert_eq!(
2112 masters.commands(),
2113 0,
2114 "a probe never checks or opens a master"
2115 );
2116 assert!(!validation.contains(&"ProxyCommand=false".to_owned()));
2117 set_ssh_connection_sharing_for_test(Some(SshSharingForTest::Directory(
2118 socket_dir.path().to_path_buf(),
2119 )));
2120 let executor = RecordingExecutor::default();
2121 crate::path_completion::ssh_completions(
2122 &ssh,
2123 "/srv/pr",
2124 crate::path_completion::CompletionKind::Directories,
2125 &executor,
2126 )
2127 .expect("completion runs");
2128 let completion = executor.seen.borrow()[0].args.clone();
2129 set_ssh_connection_sharing_for_test(None);
2130
2131 let control_path = format!(
2132 "ControlPath={}/{}",
2133 socket_dir.path().display(),
2134 control_socket_name(&ssh, 0)
2135 );
2136 for args in [&validation, &completion] {
2137 assert!(args.contains(&"ControlMaster=no".to_owned()), "{args:?}");
2138 assert!(args.contains(&control_path), "{args:?}");
2139 assert!(
2140 !args.iter().any(|arg| arg.starts_with("ControlPersist")),
2141 "a fail-fast command must not set how long a master lingers: {args:?}"
2142 );
2143 let master = args
2144 .iter()
2145 .position(|arg| arg == "ControlMaster=no")
2146 .expect("sharing options");
2147 assert!(
2148 args.contains(&"ServerAliveCountMax=1".to_owned()),
2149 "its own keepalive: {args:?}"
2150 );
2151 assert!(
2152 master
2153 < args
2154 .iter()
2155 .position(|arg| arg == "host")
2156 .expect("destination"),
2157 "{args:?}"
2158 );
2159 }
2160 }
2161
2162 #[test]
2167 #[cfg(unix)]
2168 fn connectivity_probe_joins_a_master_without_becoming_one() {
2169 let _guard = SHARING_TEST_LOCK
2170 .lock()
2171 .unwrap_or_else(std::sync::PoisonError::into_inner);
2172 let socket_dir = sharing_socket_dir();
2173 set_ssh_connection_sharing_for_test(Some(SshSharingForTest::Directory(
2174 socket_dir.path().to_path_buf(),
2175 )));
2176 let ssh = SshTarget {
2177 destination: "host".to_owned(),
2178 ssh_args: Vec::new(),
2179 };
2180 let masters = FakeMasters::default();
2181 let args = spawned_args(
2182 &ssh_connectivity_probe(&ssh),
2183 Some(socket_dir.path()),
2184 &masters,
2185 );
2186
2187 assert_eq!(masters.commands(), 0, "a probe never opens a master");
2188 assert!(args.contains(&"ControlMaster=no".to_owned()), "{args:?}");
2189 assert!(
2190 args.contains(&format!(
2191 "ControlPath={}/{}",
2192 socket_dir.path().display(),
2193 control_socket_name(&ssh, 0)
2194 )),
2195 "the probe must still join an existing master: {args:?}"
2196 );
2197 assert!(
2198 !args.iter().any(|arg| arg.starts_with("ControlPersist")),
2199 "a doctor probe must not set how long a master lingers: {args:?}"
2200 );
2201 let master = args
2202 .iter()
2203 .position(|arg| arg == "ControlMaster=no")
2204 .expect("sharing options");
2205 let strict = args
2206 .iter()
2207 .position(|arg| arg == "StrictHostKeyChecking=yes")
2208 .expect("its own host key policy");
2209 assert!(master < strict, "{args:?}");
2210 }
2211
2212 #[test]
2213 #[cfg(unix)]
2214 fn connection_sharing_is_absent_when_turned_off() {
2215 let _guard = SHARING_TEST_LOCK
2216 .lock()
2217 .unwrap_or_else(std::sync::PoisonError::into_inner);
2218 let ssh = SshTarget {
2219 destination: "sharing-off-host".to_owned(),
2220 ssh_args: Vec::new(),
2221 };
2222 let masters = FakeMasters::default();
2223 let args = spawned_args(&ssh_command(&ssh, ["true"]), None, &masters);
2224 let validation = spawned_args(
2225 &ssh_validation_command(&ssh, vec!["true".to_owned()], "test"),
2226 None,
2227 &masters,
2228 );
2229 assert_eq!(args, ["sharing-off-host", "'true'"]);
2230 assert!(!validation.iter().any(|arg| arg.starts_with("Control")));
2231 assert_eq!(masters.commands(), 0);
2232 }
2233
2234 #[test]
2235 #[cfg(unix)]
2236 fn a_control_path_that_cannot_fit_a_socket_address_is_skipped() {
2237 let _guard = SHARING_TEST_LOCK
2238 .lock()
2239 .unwrap_or_else(std::sync::PoisonError::into_inner);
2240 let root = tempfile::tempdir().expect("temp dir");
2241 let long = root.path().join("a".repeat(MAX_CONTROL_PATH));
2242 let ssh = SshTarget {
2243 destination: "long-path-host".to_owned(),
2244 ssh_args: Vec::new(),
2245 };
2246 let masters = FakeMasters::default();
2247 let args = spawned_args(&ssh_command(&ssh, ["true"]), Some(&long), &masters);
2248 assert_eq!(args, ["long-path-host", "'true'"]);
2249 assert_eq!(masters.commands(), 0);
2250 assert!(!long.exists(), "an unusable directory must not be created");
2251 }
2252
2253 #[test]
2254 #[cfg(not(unix))]
2255 fn connection_sharing_is_unix_only() {
2256 let mut args = vec!["-o".to_owned(), "BatchMode=yes".to_owned()];
2257 let ssh = SshTarget {
2258 destination: "host".to_owned(),
2259 ssh_args: Vec::new(),
2260 };
2261 push_connection_reuse_args(&mut args, &ssh);
2262 assert_eq!(args, vec!["-o".to_owned(), "BatchMode=yes".to_owned()]);
2263 }
2264
2265 #[test]
2266 #[cfg(unix)]
2267 fn the_escape_hatch_accepts_the_usual_off_spellings() {
2268 for value in ["0", "off", "FALSE", " no "] {
2269 assert!(
2270 sharing_disabled(Some(std::ffi::OsStr::new(value))),
2271 "{value:?} must disable connection sharing"
2272 );
2273 }
2274 for value in ["1", "auto", "", "yes"] {
2275 assert!(
2276 !sharing_disabled(Some(std::ffi::OsStr::new(value))),
2277 "{value:?} must leave connection sharing on"
2278 );
2279 }
2280 assert!(!sharing_disabled(None));
2281 }
2282
2283 #[test]
2287 #[cfg(unix)]
2288 fn a_leased_session_runs_through_an_opened_master_on_a_real_host() {
2289 let _guard = SHARING_TEST_LOCK
2290 .lock()
2291 .unwrap_or_else(std::sync::PoisonError::into_inner);
2292 let Some(host) = std::env::var_os("MJ_E2E_SSH_HOST") else {
2293 return;
2294 };
2295 let host = host.to_string_lossy().into_owned();
2296 let socket_dir = sharing_socket_dir();
2297 let ssh = SshTarget {
2298 destination: host.clone(),
2299 ssh_args: vec!["-o".to_owned(), "BatchMode=yes".to_owned()],
2300 };
2301 set_ssh_connection_sharing_for_test(Some(SshSharingForTest::Directory(
2302 socket_dir.path().to_path_buf(),
2303 )));
2304 let output = ProcessExecutor.execute(&ssh_command(&ssh, ["true"]));
2305 set_ssh_connection_sharing_for_test(None);
2306 let socket = socket_dir.path().join(control_socket_name(&ssh, 0));
2307 let check = ProcessExecutor
2308 .execute(&master_check_command(&ssh, &socket))
2309 .expect("ssh -O check must run");
2310 let exit = std::process::Command::new("ssh")
2311 .args([
2312 "-O",
2313 "exit",
2314 "-o",
2315 &format!("ControlPath={}", socket.display()),
2316 &host,
2317 ])
2318 .output();
2319 let output = output.expect("ssh must run");
2320 assert_eq!(
2321 output.status,
2322 0,
2323 "ssh {host} true failed: {}",
2324 String::from_utf8_lossy(&output.stderr)
2325 );
2326 assert_eq!(
2327 check.status,
2328 0,
2329 "no master is running: {}",
2330 String::from_utf8_lossy(&check.stderr)
2331 );
2332 drop(exit);
2333 }
2334
2335 #[test]
2340 #[cfg(unix)]
2341 fn sessions_shard_across_masters_on_a_real_host() {
2342 let _guard = SHARING_TEST_LOCK
2343 .lock()
2344 .unwrap_or_else(std::sync::PoisonError::into_inner);
2345 let Some(host) = std::env::var_os("MJ_E2E_SSH_HOST") else {
2346 return;
2347 };
2348 let host = host.to_string_lossy().into_owned();
2349 let socket_dir = sharing_socket_dir();
2350 let ssh = SshTarget {
2351 destination: host.clone(),
2352 ssh_args: vec!["-o".to_owned(), "BatchMode=yes".to_owned()],
2353 };
2354 let ledger = SessionLedger::new(3);
2355 let leases: Vec<SshSessionLease> = (0..7)
2356 .map(|_| {
2357 ledger
2358 .lease(&ssh, socket_dir.path(), &ProcessExecutor)
2359 .expect("lease a session on a real host")
2360 })
2361 .collect();
2362 let sockets: Vec<PathBuf> = (0..4)
2363 .map(|shard| socket_dir.path().join(control_socket_name(&ssh, shard)))
2364 .collect();
2365 let exit_all = || {
2366 for socket in &sockets {
2367 let _ = std::process::Command::new("ssh")
2368 .args([
2369 "-O",
2370 "exit",
2371 "-o",
2372 &format!("ControlPath={}", socket.display()),
2373 &host,
2374 ])
2375 .output();
2376 }
2377 };
2378
2379 let base = ssh_command(&ssh, ["sleep", "2"]);
2381 let children: Vec<std::io::Result<std::process::Output>> = std::thread::scope(|scope| {
2382 let handles: Vec<_> = leases
2383 .iter()
2384 .map(|lease| {
2385 let args = session_command_args(&base.program, &base.args, &ssh, lease);
2386 scope.spawn(move || {
2387 std::process::Command::new("ssh")
2388 .args(args)
2389 .stdin(std::process::Stdio::null())
2390 .output()
2391 })
2392 })
2393 .collect();
2394 handles
2395 .into_iter()
2396 .map(|handle| handle.join().expect("session thread"))
2397 .collect()
2398 });
2399 let running: Vec<bool> = sockets
2400 .iter()
2401 .map(|socket| {
2402 ProcessExecutor
2403 .execute(&master_check_command(&ssh, socket))
2404 .map(|output| output.status == 0)
2405 .unwrap_or(false)
2406 })
2407 .collect();
2408 let orphan = std::process::Command::new("ssh")
2409 .args(session_command_args(
2410 &base.program,
2411 &base.args,
2412 &ssh,
2413 &SshSessionLease {
2414 slot: Some(LeasedSlot {
2415 ledger: Arc::clone(&ledger),
2416 key: connection_key(&ssh),
2417 shard: 9,
2418 socket: socket_dir.path().join(control_socket_name(&ssh, 9)),
2419 }),
2420 probe: false,
2421 },
2422 ))
2423 .stdin(std::process::Stdio::null())
2424 .output();
2425 drop(leases);
2426 exit_all();
2427
2428 let shards: Vec<usize> = leases_per_shard(&ledger, &ssh);
2429 assert_eq!(shards, [0, 0, 0], "every slot is freed on drop");
2430 for (index, output) in children.iter().enumerate() {
2431 let output = output.as_ref().expect("ssh must run");
2432 assert_eq!(
2433 output.status.code(),
2434 Some(0),
2435 "session {index} failed: {}",
2436 String::from_utf8_lossy(&output.stderr)
2437 );
2438 }
2439 assert_eq!(
2440 running,
2441 [true, true, true, false],
2442 "seven sessions at three per master"
2443 );
2444 let orphan = orphan.expect("ssh must run");
2445 assert_eq!(
2446 orphan.status.code(),
2447 Some(255),
2448 "a guarded session with no master must not connect: {}",
2449 String::from_utf8_lossy(&orphan.stderr)
2450 );
2451 }
2452
2453 #[cfg(unix)]
2454 fn leases_per_shard(ledger: &SessionLedger, ssh: &SshTarget) -> Vec<usize> {
2455 ledger
2456 .connections()
2457 .get(&connection_key(ssh))
2458 .map(|shards| shards.iter().map(|shard| shard.leased).collect())
2459 .unwrap_or_default()
2460 }
2461
2462 #[test]
2466 #[cfg(unix)]
2467 fn scp_translates_the_ssh_port_option_and_is_tagged_with_its_destination() {
2468 let _guard = SHARING_TEST_LOCK
2469 .lock()
2470 .unwrap_or_else(std::sync::PoisonError::into_inner);
2471 set_ssh_connection_sharing_for_test(Some(SshSharingForTest::Disabled));
2472 let ssh = SshTarget {
2473 destination: "build@10.0.0.1".into(),
2474 ssh_args: vec!["-p".into(), "2222".into()],
2475 };
2476
2477 let upload = scp_upload(&ssh, Path::new("/tmp/local"), "remote/path", true);
2478 let download = scp_download(&ssh, "remote/archive.zip", "/tmp/local.zip");
2479 set_ssh_connection_sharing_for_test(None);
2480
2481 assert_eq!(
2482 upload.args,
2483 [
2484 "-P",
2485 "2222",
2486 "-r",
2487 "/tmp/local",
2488 "build@10.0.0.1:remote/path"
2489 ]
2490 );
2491 assert_eq!(
2492 download.args,
2493 [
2494 "-P",
2495 "2222",
2496 "build@10.0.0.1:remote/archive.zip",
2497 "/tmp/local.zip"
2498 ]
2499 );
2500 for command in [upload, download] {
2501 assert_eq!(command.program, "scp");
2502 assert_eq!(command.ssh_destination.as_deref(), Some("build@10.0.0.1"));
2503 }
2504 }
2505
2506 #[cfg(unix)]
2510 #[derive(Default)]
2511 struct FakeMasters {
2512 running: std::cell::RefCell<BTreeSet<String>>,
2513 refuse_open: std::cell::Cell<Option<&'static str>>,
2514 socket_existed_at_open: std::cell::RefCell<Vec<bool>>,
2515 seen: std::cell::RefCell<Vec<CommandSpec>>,
2516 }
2517
2518 #[cfg(unix)]
2519 impl FakeMasters {
2520 fn socket(command: &CommandSpec) -> String {
2521 command
2522 .args
2523 .iter()
2524 .find_map(|arg| arg.strip_prefix("ControlPath="))
2525 .expect("every master command names its socket")
2526 .to_owned()
2527 }
2528
2529 fn kill(&self, socket: &Path) {
2530 self.running
2531 .borrow_mut()
2532 .remove(&socket.display().to_string());
2533 }
2534
2535 fn commands(&self) -> usize {
2536 self.seen.borrow().len()
2537 }
2538
2539 fn openers(&self) -> usize {
2540 self.seen
2541 .borrow()
2542 .iter()
2543 .filter(|command| command.args.contains(&"ControlMaster=yes".to_owned()))
2544 .count()
2545 }
2546 }
2547
2548 #[cfg(unix)]
2549 impl CommandExecutor for FakeMasters {
2550 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
2551 self.seen.borrow_mut().push(command.clone());
2552 assert_eq!(command.program, "ssh");
2553 let socket = Self::socket(command);
2554 let (status, stderr) = if command.args.windows(2).any(|pair| pair == ["-O", "check"]) {
2555 if self.running.borrow().contains(&socket) {
2556 (0, "")
2557 } else {
2558 (255, "Control socket connect: No such file or directory")
2559 }
2560 } else if command.args.contains(&"ControlMaster=yes".to_owned()) {
2561 self.socket_existed_at_open
2562 .borrow_mut()
2563 .push(Path::new(&socket).exists());
2564 match self.refuse_open.get() {
2565 Some(stderr) => (255, stderr),
2566 None => {
2567 self.running.borrow_mut().insert(socket);
2568 (0, "")
2569 }
2570 }
2571 } else {
2572 panic!("the ledger ran an unexpected command: {command:?}");
2573 };
2574 Ok(CommandOutput {
2575 status,
2576 stdout: Vec::new(),
2577 stderr: stderr.as_bytes().to_vec(),
2578 })
2579 }
2580 }
2581
2582 #[cfg(unix)]
2583 fn shard_of(lease: &SshSessionLease) -> String {
2584 let path = lease.control_path().expect("a shared lease has a socket");
2585 let name = path.file_name().unwrap().to_string_lossy().into_owned();
2586 name.rsplit('-').next().unwrap().to_owned()
2587 }
2588
2589 #[cfg(unix)]
2590 fn plain_target(destination: &str) -> SshTarget {
2591 SshTarget {
2592 destination: destination.to_owned(),
2593 ssh_args: Vec::new(),
2594 }
2595 }
2596
2597 #[test]
2598 #[cfg(unix)]
2599 fn leases_fill_the_lowest_shard_and_open_another_at_the_cap() {
2600 let dir = sharing_socket_dir();
2601 let ledger = SessionLedger::new(2);
2602 let ssh = plain_target("host");
2603 let masters = FakeMasters::default();
2604
2605 let first = ledger.lease(&ssh, dir.path(), &masters).expect("first");
2606 assert_eq!(masters.commands(), 3);
2608 let second = ledger.lease(&ssh, dir.path(), &masters).expect("second");
2609 assert_eq!(
2610 masters.commands(),
2611 3,
2612 "a master verified moments ago is not checked again"
2613 );
2614 let third = ledger.lease(&ssh, dir.path(), &masters).expect("third");
2615 assert_eq!(
2616 [&first, &second, &third].map(shard_of),
2617 ["0", "0", "1"].map(str::to_owned)
2618 );
2619 assert_eq!(masters.openers(), 2, "one master per shard");
2620 assert_eq!(
2621 third.control_path().unwrap(),
2622 dir.path().join(control_socket_name(&ssh, 1))
2623 );
2624
2625 drop(first);
2626 let fourth = ledger.lease(&ssh, dir.path(), &masters).expect("fourth");
2627 assert_eq!(shard_of(&fourth), "0", "a freed slot is reused first");
2628 assert_eq!(masters.openers(), 2);
2629 }
2630
2631 #[test]
2636 #[cfg(unix)]
2637 fn probes_are_counted_on_the_shard_they_join_without_opening_it() {
2638 let dir = sharing_socket_dir();
2639 let ledger = SessionLedger::new(2);
2640 let ssh = plain_target("probe-host");
2641 let masters = FakeMasters::default();
2642
2643 let session = ledger.lease(&ssh, dir.path(), &masters).expect("session");
2644 let probe = ledger.lease_probe(&ssh, dir.path());
2645 assert_eq!(leases_per_shard(&ledger, &ssh), [2]);
2646 let second_probe = ledger.lease_probe(&ssh, dir.path());
2647 assert_eq!(
2648 [&session, &probe, &second_probe].map(shard_of),
2649 ["0", "0", "1"].map(str::to_owned)
2650 );
2651 assert_eq!(masters.openers(), 1, "a probe never opens a master");
2652
2653 let mut args = Vec::new();
2654 push_session_args(&mut args, &probe);
2655 assert!(
2656 !args.contains(&"ProxyCommand=false".to_owned()),
2657 "a probe may connect directly when its master is down: {args:?}"
2658 );
2659 drop(probe);
2660 drop(second_probe);
2661 assert_eq!(leases_per_shard(&ledger, &ssh), [1, 0]);
2662 let next = ledger.lease(&ssh, dir.path(), &masters).expect("next");
2663 assert_eq!(shard_of(&next), "0");
2664 }
2665
2666 #[test]
2667 #[cfg(unix)]
2668 fn separate_connections_are_counted_separately() {
2669 let dir = sharing_socket_dir();
2670 let ledger = SessionLedger::new(1);
2671 let masters = FakeMasters::default();
2672 let first = ledger
2673 .lease(&plain_target("one"), dir.path(), &masters)
2674 .expect("one");
2675 let second = ledger
2676 .lease(&plain_target("two"), dir.path(), &masters)
2677 .expect("two");
2678 assert_eq!(
2679 [&first, &second].map(shard_of),
2680 ["0", "0"].map(str::to_owned)
2681 );
2682 assert_ne!(first.control_path(), second.control_path());
2683 }
2684
2685 #[test]
2686 #[cfg(unix)]
2687 fn an_invalidated_lease_makes_the_next_lease_reopen_a_dead_master() {
2688 let dir = sharing_socket_dir();
2689 let ledger = SessionLedger::new(8);
2690 let ssh = plain_target("host");
2691 let masters = FakeMasters::default();
2692 let first = ledger.lease(&ssh, dir.path(), &masters).expect("first");
2693 masters.kill(first.control_path().unwrap());
2694
2695 drop(ledger.lease(&ssh, dir.path(), &masters).expect("trusted"));
2697 assert_eq!(masters.openers(), 1);
2698
2699 first.invalidate();
2700 let second = ledger.lease(&ssh, dir.path(), &masters).expect("reopened");
2701 assert_eq!(masters.openers(), 2);
2702 assert_eq!(first.control_path(), second.control_path());
2703 }
2704
2705 #[test]
2706 #[cfg(unix)]
2707 fn a_master_that_cannot_be_opened_is_an_error_naming_the_destination() {
2708 let dir = sharing_socket_dir();
2709 let ledger = SessionLedger::new(1);
2710 let ssh = plain_target("build@10.0.0.1");
2711 let masters = FakeMasters::default();
2712 masters
2713 .refuse_open
2714 .set(Some("Permission denied (publickey)."));
2715
2716 let error = ledger
2717 .lease(&ssh, dir.path(), &masters)
2718 .expect_err("no master means no session");
2719 let message = format!("{error:#}");
2720 assert!(message.contains("build@10.0.0.1"), "{message}");
2721 assert!(message.contains("Permission denied"), "{message}");
2722 assert_eq!(
2723 masters.openers(),
2724 1,
2725 "the opener is not retried by the ledger"
2726 );
2727
2728 masters.refuse_open.set(None);
2731 let lease = ledger.lease(&ssh, dir.path(), &masters).expect("opens");
2732 assert_eq!(shard_of(&lease), "0");
2733 }
2734
2735 #[cfg(unix)]
2741 #[derive(Default)]
2742 struct RacingMasters {
2743 bound: Mutex<BTreeSet<String>>,
2744 masters: AtomicUsize,
2745 orphans: AtomicUsize,
2746 }
2747
2748 #[cfg(unix)]
2749 impl CommandExecutor for RacingMasters {
2750 fn execute(&self, command: &CommandSpec) -> Result<CommandOutput> {
2751 let socket = FakeMasters::socket(command);
2752 let status = if command.args.windows(2).any(|pair| pair == ["-O", "check"]) {
2753 if self.bound.lock().unwrap().contains(&socket) {
2754 0
2755 } else {
2756 255
2757 }
2758 } else {
2759 std::thread::sleep(Duration::from_millis(100));
2760 if self.bound.lock().unwrap().insert(socket) {
2761 self.masters.fetch_add(1, Ordering::SeqCst);
2762 } else {
2763 self.orphans.fetch_add(1, Ordering::SeqCst);
2764 }
2765 0
2766 };
2767 Ok(CommandOutput {
2768 status,
2769 stdout: Vec::new(),
2770 stderr: Vec::new(),
2771 })
2772 }
2773 }
2774
2775 #[test]
2779 #[cfg(unix)]
2780 fn two_processes_opening_one_socket_open_one_master() {
2781 let dir = sharing_socket_dir();
2782 let ssh = plain_target("racing-host");
2783 let fake = RacingMasters::default();
2784 std::thread::scope(|scope| {
2785 for _ in 0..2 {
2786 scope.spawn(|| {
2787 let ledger = SessionLedger::new(8);
2789 ledger.lease(&ssh, dir.path(), &fake).expect("lease");
2790 });
2791 }
2792 });
2793 assert_eq!(fake.masters.load(Ordering::SeqCst), 1);
2794 assert_eq!(fake.orphans.load(Ordering::SeqCst), 0);
2795 }
2796
2797 #[test]
2798 #[cfg(unix)]
2799 fn a_stale_socket_is_removed_before_the_master_is_opened() {
2800 let dir = sharing_socket_dir();
2801 let ledger = SessionLedger::new(8);
2802 let ssh = plain_target("host");
2803 let socket = dir.path().join(control_socket_name(&ssh, 0));
2804 fs::write(&socket, b"").expect("stale socket stand-in");
2805 let masters = FakeMasters::default();
2806
2807 ledger.lease(&ssh, dir.path(), &masters).expect("opens");
2808
2809 assert_eq!(*masters.socket_existed_at_open.borrow(), [false]);
2810 }
2811
2812 #[test]
2813 #[cfg(unix)]
2814 fn the_opener_is_an_admitted_batch_master_and_the_check_is_local() {
2815 let ssh = SshTarget {
2816 destination: "host".to_owned(),
2817 ssh_args: vec!["-J".to_owned(), "jump".to_owned()],
2818 };
2819 let socket = Path::new("/run/mj/abc-0");
2820 let open = master_open_command(&ssh, socket);
2821 assert_eq!(
2822 open.args,
2823 [
2824 "-J",
2825 "jump",
2826 "-o",
2827 "BatchMode=yes",
2828 "-f",
2829 "-N",
2830 "-o",
2831 "ControlMaster=yes",
2832 "-o",
2833 "ControlPath=/run/mj/abc-0",
2834 "-o",
2835 &format!("ControlPersist={CONTROL_PERSIST}"),
2836 "host",
2837 ]
2838 );
2839 assert_eq!(open.ssh_destination.as_deref(), Some("host"));
2840
2841 let check = master_check_command(&ssh, socket);
2842 assert_eq!(
2843 check.args,
2844 [
2845 "-J",
2846 "jump",
2847 "-o",
2848 "ControlPath=/run/mj/abc-0",
2849 "-O",
2850 "check",
2851 "host"
2852 ]
2853 );
2854 assert_eq!(
2855 check.ssh_destination, None,
2856 "a check opens no connection and takes no admission permit"
2857 );
2858 }
2859
2860 #[test]
2861 #[cfg(unix)]
2862 fn session_args_forbid_a_direct_connection() {
2863 let dir = sharing_socket_dir();
2864 let ledger = SessionLedger::new(8);
2865 let masters = FakeMasters::default();
2866 let lease = ledger
2867 .lease(&plain_target("host"), dir.path(), &masters)
2868 .expect("lease");
2869 let mut args = Vec::new();
2870 push_session_args(&mut args, &lease);
2871 assert_eq!(
2872 args,
2873 [
2874 "-o".to_owned(),
2875 "ControlMaster=no".to_owned(),
2876 "-o".to_owned(),
2877 format!("ControlPath={}", lease.control_path().unwrap().display()),
2878 "-o".to_owned(),
2879 "ProxyCommand=false".to_owned(),
2880 ]
2881 );
2882 }
2883
2884 #[test]
2887 #[cfg(unix)]
2888 fn unshared_connections_lease_without_a_socket() {
2889 let _guard = SHARING_TEST_LOCK
2890 .lock()
2891 .unwrap_or_else(std::sync::PoisonError::into_inner);
2892 let masters = FakeMasters::default();
2893 let dir = sharing_socket_dir();
2894 set_ssh_connection_sharing_for_test(Some(SshSharingForTest::Directory(
2895 dir.path().to_path_buf(),
2896 )));
2897 let user_owned = SshSessions::lease(
2898 &SshTarget {
2899 destination: "unshared-user-host".to_owned(),
2900 ssh_args: vec!["-S".to_owned(), "/tmp/mine".to_owned()],
2901 },
2902 &masters,
2903 );
2904 set_ssh_connection_sharing_for_test(Some(SshSharingForTest::Disabled));
2905 let disabled = SshSessions::lease(&plain_target("unshared-disabled-host"), &masters);
2906 set_ssh_connection_sharing_for_test(None);
2907
2908 for lease in [user_owned, disabled] {
2909 let lease = lease.expect("an unshared lease never fails");
2910 assert_eq!(lease.control_path(), None);
2911 let mut args = Vec::new();
2912 push_session_args(&mut args, &lease);
2913 assert!(args.is_empty());
2914 }
2915 assert_eq!(masters.commands(), 0);
2916 }
2917
2918 #[test]
2922 fn a_refused_session_is_named_apart_from_a_pre_authentication_hangup() {
2923 assert_eq!(
2924 ssh_refusal(
2925 255,
2926 "mux_client_request_session: session request failed: Session open refused by peer\n\
2927 kex_exchange_identification: Connection closed by remote host\n\
2928 Connection closed by UNKNOWN port 65535"
2929 ),
2930 Some(SshRefusal::SessionLimit)
2931 );
2932 assert_eq!(
2933 ssh_refusal(255, "Connection closed by 192.168.1.77 port 22"),
2934 Some(SshRefusal::BeforeAuthentication)
2935 );
2936 assert_eq!(ssh_refusal(1, "Session open refused by peer"), None);
2937 assert!(
2938 !SshRefusal::SessionLimit
2939 .retry_message()
2940 .contains("before authentication")
2941 );
2942 }
2943
2944 #[test]
2945 fn transport_rejection_matches_only_sshd_hangups() {
2946 let cases: [(i32, &str, bool); 7] = [
2947 (255, "Connection closed by 192.168.1.77 port 22", true),
2948 (
2949 255,
2950 "kex_exchange_identification: read: Connection reset by peer",
2951 true,
2952 ),
2953 (255, "ssh: Connection reset by 10.0.0.1 port 22", true),
2954 (255, "Connection timed out during banner exchange", true),
2955 (255, "Permission denied (publickey).", false),
2956 (
2957 255,
2958 "ssh: connect to host h port 22: Connection refused",
2959 false,
2960 ),
2961 (1, "Connection closed by 192.168.1.77 port 22", false),
2962 ];
2963 for (status, stderr, expected) in cases {
2964 assert_eq!(
2965 is_transport_rejection(status, stderr),
2966 expected,
2967 "status {status} stderr {stderr:?}"
2968 );
2969 }
2970 }
2971
2972 #[test]
2973 fn admission_never_admits_more_than_the_limit() {
2974 let gate = DestinationGate::new(2);
2975 let in_flight = Arc::new(AtomicUsize::new(0));
2976 let peak = Arc::new(AtomicUsize::new(0));
2977 let threads: Vec<_> = (0..12)
2978 .map(|_| {
2979 let gate = Arc::clone(&gate);
2980 let in_flight = Arc::clone(&in_flight);
2981 let peak = Arc::clone(&peak);
2982 std::thread::spawn(move || {
2983 for _ in 0..25 {
2984 let permit = gate.acquire();
2985 let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1;
2986 peak.fetch_max(now, Ordering::SeqCst);
2987 std::thread::yield_now();
2988 in_flight.fetch_sub(1, Ordering::SeqCst);
2989 drop(permit);
2990 }
2991 })
2992 })
2993 .collect();
2994 for thread in threads {
2995 thread.join().expect("admission worker must not panic");
2996 }
2997 assert!(
2998 peak.load(Ordering::SeqCst) <= 2,
2999 "admission let {} connections run against a 2-permit gate",
3000 peak.load(Ordering::SeqCst)
3001 );
3002 assert_eq!(in_flight.load(Ordering::SeqCst), 0);
3003 }
3004
3005 #[test]
3006 fn admission_blocks_once_every_permit_is_held() {
3007 let gate = DestinationGate::new(2);
3008 let first = gate.acquire();
3009 let second = gate.acquire();
3010 let waiter = {
3011 let gate = Arc::clone(&gate);
3012 std::thread::spawn(move || {
3013 let permit = gate.acquire();
3014 drop(permit);
3015 })
3016 };
3017 std::thread::sleep(std::time::Duration::from_millis(50));
3019 assert!(!waiter.is_finished());
3020 drop(first);
3021 waiter
3022 .join()
3023 .expect("waiter must be admitted once a permit frees");
3024 drop(second);
3025 }
3026}