consortium/worker/
ssh.rs

1//! SSH Worker implementation.
2//!
3//! Provides SshWorker, which wraps ExecWorker to execute commands on remote
4//! nodes via SSH. Mirrors ClusterShell's WorkerSsh / SshClient.
5
6use std::collections::HashMap;
7use std::os::unix::io::RawFd;
8use std::time::Duration;
9
10use crate::worker::exec::ExecWorker;
11use crate::worker::{EventHandler, Worker, WorkerState};
12
13/// SSH connection options.
14#[derive(Debug, Clone)]
15pub struct SshOptions {
16    /// Path to the SSH binary.
17    pub ssh_path: String,
18    /// Path to the SCP binary.
19    pub scp_path: String,
20    /// SSH user (if different from current user).
21    pub user: Option<String>,
22    /// Custom SSH options (passed as -o key=value).
23    pub options: Vec<String>,
24    /// Connect timeout in seconds.
25    pub connect_timeout: Option<u32>,
26    /// Whether to use strict host key checking.
27    pub strict_host_key_checking: bool,
28    /// Whether to enable password authentication.
29    pub password_auth: bool,
30    /// Path to SSH identity file (passed as -i).
31    pub identity_file: Option<String>,
32    /// Path to SSH config file (passed as -F).
33    pub config_file: Option<String>,
34    /// Enable SSH ControlMaster multiplexing.
35    pub control_master: Option<ControlMaster>,
36    /// Path for ControlMaster socket (passed as -oControlPath=...).
37    pub control_path: Option<String>,
38    /// ControlPersist timeout (e.g. "10m", "yes").
39    pub control_persist: Option<String>,
40    /// SSH ProxyJump host (passed as -J).
41    pub proxy_jump: Option<String>,
42    /// SSH port (passed as -p).
43    pub port: Option<u16>,
44}
45
46/// SSH ControlMaster modes.
47#[derive(Debug, Clone, PartialEq)]
48pub enum ControlMaster {
49    /// Automatically start a master if one isn't running.
50    Auto,
51    /// Automatically start and automatically close.
52    AutoAsk,
53    /// Always be the master.
54    Yes,
55    /// Never be the master.
56    No,
57}
58
59impl Default for SshOptions {
60    fn default() -> Self {
61        Self {
62            ssh_path: "ssh".to_string(),
63            scp_path: "scp".to_string(),
64            user: None,
65            options: Vec::new(),
66            connect_timeout: Some(10),
67            strict_host_key_checking: false,
68            password_auth: false,
69            identity_file: None,
70            config_file: None,
71            control_master: None,
72            control_path: None,
73            control_persist: None,
74            proxy_jump: None,
75            port: None,
76        }
77    }
78}
79
80impl SshOptions {
81    /// Build the SSH command prefix for a given node.
82    ///
83    /// Returns the full SSH command string that will execute `command` on `node`.
84    pub fn build_ssh_cmd(&self, node: &str, command: &str) -> String {
85        let mut parts = Vec::new();
86        parts.push(self.ssh_path.clone());
87
88        // Config file must come early (before other options)
89        if let Some(ref config_file) = self.config_file {
90            parts.push("-F".to_string());
91            parts.push(config_file.clone());
92        }
93
94        // Standard hardened options (mirrors Python SshClient._build_cmd)
95        parts.push("-oForwardX11=no".to_string());
96
97        if !self.strict_host_key_checking {
98            parts.push("-oStrictHostKeyChecking=no".to_string());
99        }
100
101        if !self.password_auth {
102            parts.push("-oPasswordAuthentication=no".to_string());
103        }
104
105        if let Some(timeout) = self.connect_timeout {
106            parts.push(format!("-oConnectTimeout={}", timeout));
107        }
108
109        // ControlMaster multiplexing
110        if let Some(ref cm) = self.control_master {
111            let val = match cm {
112                ControlMaster::Auto => "auto",
113                ControlMaster::AutoAsk => "autoask",
114                ControlMaster::Yes => "yes",
115                ControlMaster::No => "no",
116            };
117            parts.push(format!("-oControlMaster={}", val));
118        }
119        if let Some(ref cp) = self.control_path {
120            parts.push(format!("-oControlPath={}", cp));
121        }
122        if let Some(ref persist) = self.control_persist {
123            parts.push(format!("-oControlPersist={}", persist));
124        }
125
126        // Custom options
127        for opt in &self.options {
128            parts.push(format!("-o{}", opt));
129        }
130
131        // Identity file
132        if let Some(ref identity) = self.identity_file {
133            parts.push("-i".to_string());
134            parts.push(identity.clone());
135        }
136
137        // ProxyJump
138        if let Some(ref jump) = self.proxy_jump {
139            parts.push("-J".to_string());
140            parts.push(jump.clone());
141        }
142
143        // Port
144        if let Some(port) = self.port {
145            parts.push("-p".to_string());
146            parts.push(port.to_string());
147        }
148
149        // User
150        if let Some(ref user) = self.user {
151            parts.push("-l".to_string());
152            parts.push(user.clone());
153        }
154
155        // Target node and command
156        parts.push(node.to_string());
157        parts.push(command.to_string());
158
159        parts.join(" ")
160    }
161
162    /// Build an SCP command for copying files to/from a node.
163    ///
164    /// If `reverse` is true, copies from remote to local.
165    pub fn build_scp_cmd(
166        &self,
167        node: &str,
168        source: &str,
169        dest: &str,
170        reverse: bool,
171        preserve: bool,
172        directory: bool,
173    ) -> String {
174        let mut parts = Vec::new();
175        parts.push(self.scp_path.clone());
176
177        // Config file
178        if let Some(ref config_file) = self.config_file {
179            parts.push("-F".to_string());
180            parts.push(config_file.clone());
181        }
182
183        if !self.strict_host_key_checking {
184            parts.push("-oStrictHostKeyChecking=no".to_string());
185        }
186
187        if !self.password_auth {
188            parts.push("-oPasswordAuthentication=no".to_string());
189        }
190
191        if let Some(timeout) = self.connect_timeout {
192            parts.push(format!("-oConnectTimeout={}", timeout));
193        }
194
195        // ControlMaster multiplexing
196        if let Some(ref cm) = self.control_master {
197            let val = match cm {
198                ControlMaster::Auto => "auto",
199                ControlMaster::AutoAsk => "autoask",
200                ControlMaster::Yes => "yes",
201                ControlMaster::No => "no",
202            };
203            parts.push(format!("-oControlMaster={}", val));
204        }
205        if let Some(ref cp) = self.control_path {
206            parts.push(format!("-oControlPath={}", cp));
207        }
208        if let Some(ref persist) = self.control_persist {
209            parts.push(format!("-oControlPersist={}", persist));
210        }
211
212        if preserve {
213            parts.push("-p".to_string());
214        }
215
216        if directory {
217            parts.push("-r".to_string());
218        }
219
220        // Identity file
221        if let Some(ref identity) = self.identity_file {
222            parts.push("-i".to_string());
223            parts.push(identity.clone());
224        }
225
226        // ProxyJump
227        if let Some(ref jump) = self.proxy_jump {
228            parts.push("-J".to_string());
229            parts.push(jump.clone());
230        }
231
232        // Port
233        if let Some(port) = self.port {
234            parts.push("-P".to_string()); // SCP uses -P (uppercase) for port
235            parts.push(port.to_string());
236        }
237
238        // Custom options
239        for opt in &self.options {
240            parts.push(format!("-o{}", opt));
241        }
242
243        // User prefix
244        let user_prefix = if let Some(ref user) = self.user {
245            format!("{}@", user)
246        } else {
247            String::new()
248        };
249
250        if reverse {
251            // Remote -> Local
252            parts.push(format!("{}{}:{}", user_prefix, node, source));
253            parts.push(dest.to_string());
254        } else {
255            // Local -> Remote
256            parts.push(source.to_string());
257            parts.push(format!("{}{}:{}", user_prefix, node, dest));
258        }
259
260        parts.join(" ")
261    }
262}
263
264/// Worker that executes commands on remote nodes via SSH.
265///
266/// SshWorker is a thin wrapper around ExecWorker that builds SSH
267/// commands using configurable SshOptions. It mirrors ClusterShell's
268/// WorkerSsh class.
269pub struct SshWorker {
270    /// Underlying ExecWorker that does the actual process management.
271    inner: ExecWorker,
272    /// SSH options used for command building.
273    ssh_options: SshOptions,
274    /// Original command (before SSH wrapping).
275    original_command: String,
276}
277
278impl SshWorker {
279    /// Create a new SshWorker.
280    ///
281    /// # Arguments
282    /// * `nodes` - List of remote node names to execute on.
283    /// * `command` - Command to run on each remote node.
284    /// * `fanout` - Max concurrent SSH connections.
285    /// * `timeout` - Optional per-node timeout.
286    /// * `ssh_options` - SSH connection options.
287    pub fn new(
288        nodes: Vec<String>,
289        command: String,
290        fanout: usize,
291        timeout: Option<Duration>,
292        ssh_options: SshOptions,
293    ) -> Self {
294        // Build the SSH-wrapped command using %h for node substitution.
295        // ExecWorker will replace %h with the actual node name.
296        let ssh_cmd = ssh_options.build_ssh_cmd("%h", &command);
297
298        SshWorker {
299            inner: ExecWorker::new(nodes, ssh_cmd, fanout, timeout),
300            ssh_options,
301            original_command: command,
302        }
303    }
304
305    /// Create a new SshWorker with default SSH options.
306    pub fn with_defaults(
307        nodes: Vec<String>,
308        command: String,
309        fanout: usize,
310        timeout: Option<Duration>,
311    ) -> Self {
312        Self::new(nodes, command, fanout, timeout, SshOptions::default())
313    }
314
315    /// Enable stderr capture.
316    pub fn with_stderr(mut self, stderr: bool) -> Self {
317        self.inner = self.inner.with_stderr(stderr);
318        self
319    }
320
321    /// Get the SSH options.
322    pub fn ssh_options(&self) -> &SshOptions {
323        &self.ssh_options
324    }
325
326    /// Get the original (unwrapped) command.
327    pub fn original_command(&self) -> &str {
328        &self.original_command
329    }
330}
331
332impl Worker for SshWorker {
333    fn start(&mut self) -> crate::worker::Result<()> {
334        self.inner.start()
335    }
336
337    fn abort(&mut self, kill: bool) {
338        self.inner.abort(kill);
339    }
340
341    fn state(&self) -> WorkerState {
342        self.inner.state()
343    }
344
345    fn set_handler(&mut self, handler: Box<dyn EventHandler>) {
346        self.inner.set_handler(handler);
347    }
348
349    fn read_fds(&self) -> Vec<RawFd> {
350        self.inner.read_fds()
351    }
352
353    fn write_fds(&self) -> Vec<RawFd> {
354        self.inner.write_fds()
355    }
356
357    fn handle_read(&mut self, fd: RawFd) -> crate::worker::Result<()> {
358        self.inner.handle_read(fd)
359    }
360
361    fn handle_write(&mut self, fd: RawFd) -> crate::worker::Result<()> {
362        self.inner.handle_write(fd)
363    }
364
365    fn is_done(&self) -> bool {
366        self.inner.is_done()
367    }
368
369    fn retcodes(&self) -> &HashMap<String, i32> {
370        self.inner.retcodes()
371    }
372
373    fn num_nodes(&self) -> usize {
374        self.inner.num_nodes()
375    }
376
377    fn take_handler(&mut self) -> Option<Box<dyn EventHandler>> {
378        self.inner.take_handler()
379    }
380}
381
382/// Worker that copies files to/from remote nodes via SCP.
383///
384/// Mirrors ClusterShell's ScpClient / WorkerSsh copy mode.
385pub struct ScpWorker {
386    /// Underlying ExecWorker for running SCP commands.
387    inner: ExecWorker,
388    /// SSH/SCP options.
389    #[allow(dead_code)]
390    ssh_options: SshOptions,
391    /// Source path.
392    source: String,
393    /// Destination path.
394    dest: String,
395    /// Whether this is a reverse copy (remote -> local).
396    reverse: bool,
397}
398
399impl ScpWorker {
400    /// Create a new ScpWorker for file copy.
401    ///
402    /// # Arguments
403    /// * `nodes` - Remote nodes to copy to/from.
404    /// * `source` - Source file/directory path.
405    /// * `dest` - Destination file/directory path.
406    /// * `fanout` - Max concurrent SCP operations.
407    /// * `timeout` - Optional per-node timeout.
408    /// * `ssh_options` - SSH/SCP options.
409    /// * `reverse` - If true, copy from remote to local.
410    /// * `preserve` - If true, preserve file attributes.
411    /// * `directory` - If true, copy recursively.
412    #[allow(clippy::too_many_arguments)]
413    pub fn new(
414        nodes: Vec<String>,
415        source: String,
416        dest: String,
417        fanout: usize,
418        timeout: Option<Duration>,
419        ssh_options: SshOptions,
420        reverse: bool,
421        preserve: bool,
422        directory: bool,
423    ) -> Self {
424        // Build SCP command with %h placeholder for node substitution
425        let scp_cmd = ssh_options.build_scp_cmd("%h", &source, &dest, reverse, preserve, directory);
426
427        ScpWorker {
428            inner: ExecWorker::new(nodes, scp_cmd, fanout, timeout).with_stderr(true),
429            ssh_options,
430            source,
431            dest,
432            reverse,
433        }
434    }
435
436    /// Get the source path.
437    pub fn source(&self) -> &str {
438        &self.source
439    }
440
441    /// Get the destination path.
442    pub fn dest(&self) -> &str {
443        &self.dest
444    }
445
446    /// Whether this is a reverse copy.
447    pub fn is_reverse(&self) -> bool {
448        self.reverse
449    }
450}
451
452impl Worker for ScpWorker {
453    fn start(&mut self) -> crate::worker::Result<()> {
454        self.inner.start()
455    }
456
457    fn abort(&mut self, kill: bool) {
458        self.inner.abort(kill);
459    }
460
461    fn state(&self) -> WorkerState {
462        self.inner.state()
463    }
464
465    fn set_handler(&mut self, handler: Box<dyn EventHandler>) {
466        self.inner.set_handler(handler);
467    }
468
469    fn read_fds(&self) -> Vec<RawFd> {
470        self.inner.read_fds()
471    }
472
473    fn write_fds(&self) -> Vec<RawFd> {
474        self.inner.write_fds()
475    }
476
477    fn handle_read(&mut self, fd: RawFd) -> crate::worker::Result<()> {
478        self.inner.handle_read(fd)
479    }
480
481    fn handle_write(&mut self, fd: RawFd) -> crate::worker::Result<()> {
482        self.inner.handle_write(fd)
483    }
484
485    fn is_done(&self) -> bool {
486        self.inner.is_done()
487    }
488
489    fn retcodes(&self) -> &HashMap<String, i32> {
490        self.inner.retcodes()
491    }
492
493    fn num_nodes(&self) -> usize {
494        self.inner.num_nodes()
495    }
496
497    fn take_handler(&mut self) -> Option<Box<dyn EventHandler>> {
498        self.inner.take_handler()
499    }
500}
501
502#[cfg(test)]
503mod tests {
504    use super::*;
505
506    #[test]
507    fn test_ssh_options_default() {
508        let opts = SshOptions::default();
509        assert_eq!(opts.ssh_path, "ssh");
510        assert_eq!(opts.scp_path, "scp");
511        assert!(opts.user.is_none());
512        assert_eq!(opts.connect_timeout, Some(10));
513        assert!(!opts.strict_host_key_checking);
514        assert!(!opts.password_auth);
515    }
516
517    #[test]
518    fn test_ssh_options_build_ssh_cmd() {
519        let opts = SshOptions::default();
520        let cmd = opts.build_ssh_cmd("node1", "uname -a");
521
522        assert!(cmd.starts_with("ssh "));
523        assert!(cmd.contains("-oForwardX11=no"));
524        assert!(cmd.contains("-oStrictHostKeyChecking=no"));
525        assert!(cmd.contains("-oPasswordAuthentication=no"));
526        assert!(cmd.contains("-oConnectTimeout=10"));
527        assert!(cmd.contains("node1"));
528        assert!(cmd.ends_with("uname -a"));
529    }
530
531    #[test]
532    fn test_ssh_options_build_ssh_cmd_with_user() {
533        let opts = SshOptions {
534            user: Some("admin".to_string()),
535            ..SshOptions::default()
536        };
537        let cmd = opts.build_ssh_cmd("node1", "hostname");
538
539        assert!(cmd.contains("-l admin"));
540        assert!(cmd.contains("node1"));
541    }
542
543    #[test]
544    fn test_ssh_options_build_ssh_cmd_custom_options() {
545        let opts = SshOptions {
546            options: vec!["BatchMode=yes".to_string()],
547            ..SshOptions::default()
548        };
549        let cmd = opts.build_ssh_cmd("node1", "ls");
550
551        assert!(cmd.contains("-oBatchMode=yes"));
552    }
553
554    #[test]
555    fn test_ssh_options_build_scp_cmd_forward() {
556        let opts = SshOptions::default();
557        let cmd = opts.build_scp_cmd("node1", "/tmp/file.txt", "/home/user/", false, true, false);
558
559        assert!(cmd.starts_with("scp "));
560        assert!(cmd.contains("-p")); // preserve
561        assert!(cmd.contains("/tmp/file.txt"));
562        assert!(cmd.contains("node1:/home/user/"));
563    }
564
565    #[test]
566    fn test_ssh_options_build_scp_cmd_reverse() {
567        let opts = SshOptions::default();
568        let cmd = opts.build_scp_cmd(
569            "node1",
570            "/remote/file.txt",
571            "/local/dir/",
572            true,
573            false,
574            false,
575        );
576
577        assert!(cmd.contains("node1:/remote/file.txt"));
578        assert!(cmd.ends_with("/local/dir/"));
579    }
580
581    #[test]
582    fn test_ssh_options_build_scp_cmd_recursive() {
583        let opts = SshOptions::default();
584        let cmd = opts.build_scp_cmd("node1", "/src/", "/dst/", false, false, true);
585
586        assert!(cmd.contains("-r")); // recursive
587    }
588
589    #[test]
590    fn test_ssh_options_build_scp_cmd_with_user() {
591        let opts = SshOptions {
592            user: Some("deploy".to_string()),
593            ..SshOptions::default()
594        };
595        let cmd = opts.build_scp_cmd("node1", "/f", "/d", false, false, false);
596
597        assert!(cmd.contains("deploy@node1:/d"));
598    }
599
600    #[test]
601    fn test_ssh_worker_creation() {
602        let nodes = vec!["node1".to_string(), "node2".to_string()];
603        let worker = SshWorker::with_defaults(nodes, "hostname".to_string(), 64, None);
604
605        assert_eq!(worker.state(), WorkerState::Pending);
606        assert_eq!(worker.num_nodes(), 2);
607        assert_eq!(worker.original_command(), "hostname");
608    }
609
610    #[test]
611    fn test_ssh_worker_custom_options() {
612        let opts = SshOptions {
613            user: Some("root".to_string()),
614            connect_timeout: Some(5),
615            ..SshOptions::default()
616        };
617        let nodes = vec!["node1".to_string()];
618        let worker = SshWorker::new(nodes, "uptime".to_string(), 64, None, opts);
619
620        assert_eq!(worker.ssh_options().user.as_deref(), Some("root"));
621        assert_eq!(worker.ssh_options().connect_timeout, Some(5));
622    }
623
624    #[test]
625    fn test_ssh_worker_with_stderr() {
626        let nodes = vec!["node1".to_string()];
627        let worker = SshWorker::with_defaults(nodes, "ls".to_string(), 64, None).with_stderr(true);
628
629        assert_eq!(worker.state(), WorkerState::Pending);
630    }
631
632    #[test]
633    fn test_scp_worker_creation() {
634        let nodes = vec!["node1".to_string()];
635        let opts = SshOptions::default();
636        let worker = ScpWorker::new(
637            nodes,
638            "/local/file".to_string(),
639            "/remote/dir".to_string(),
640            64,
641            None,
642            opts,
643            false,
644            true,
645            false,
646        );
647
648        assert_eq!(worker.state(), WorkerState::Pending);
649        assert_eq!(worker.source(), "/local/file");
650        assert_eq!(worker.dest(), "/remote/dir");
651        assert!(!worker.is_reverse());
652    }
653
654    #[test]
655    fn test_scp_worker_reverse() {
656        let nodes = vec!["node1".to_string()];
657        let opts = SshOptions::default();
658        let worker = ScpWorker::new(
659            nodes,
660            "/remote/file".to_string(),
661            "/local/dir".to_string(),
662            64,
663            None,
664            opts,
665            true,
666            false,
667            false,
668        );
669
670        assert!(worker.is_reverse());
671        assert_eq!(worker.num_nodes(), 1);
672    }
673
674    #[test]
675    fn test_ssh_options_with_identity_file() {
676        let opts = SshOptions {
677            identity_file: Some("/path/to/key".to_string()),
678            ..SshOptions::default()
679        };
680        let cmd = opts.build_ssh_cmd("node1", "hostname");
681        assert!(cmd.contains("-i /path/to/key"));
682    }
683
684    #[test]
685    fn test_ssh_options_with_config_file() {
686        let opts = SshOptions {
687            config_file: Some("/etc/ssh/config".to_string()),
688            ..SshOptions::default()
689        };
690        let cmd = opts.build_ssh_cmd("node1", "hostname");
691        assert!(cmd.contains("-F /etc/ssh/config"));
692    }
693
694    #[test]
695    fn test_ssh_options_with_control_master() {
696        let opts = SshOptions {
697            control_master: Some(ControlMaster::Auto),
698            control_path: Some("/tmp/ssh-%r@%h:%p".to_string()),
699            control_persist: Some("10m".to_string()),
700            ..SshOptions::default()
701        };
702        let cmd = opts.build_ssh_cmd("node1", "hostname");
703        assert!(cmd.contains("-oControlMaster=auto"));
704        assert!(cmd.contains("-oControlPath=/tmp/ssh-%r@%h:%p"));
705        assert!(cmd.contains("-oControlPersist=10m"));
706    }
707
708    #[test]
709    fn test_ssh_options_with_proxy_jump() {
710        let opts = SshOptions {
711            proxy_jump: Some("bastion.example.com".to_string()),
712            ..SshOptions::default()
713        };
714        let cmd = opts.build_ssh_cmd("node1", "hostname");
715        assert!(cmd.contains("-J bastion.example.com"));
716    }
717
718    #[test]
719    fn test_ssh_options_with_port() {
720        let opts = SshOptions {
721            port: Some(2222),
722            ..SshOptions::default()
723        };
724        let ssh_cmd = opts.build_ssh_cmd("node1", "hostname");
725        assert!(ssh_cmd.contains("-p 2222"));
726
727        let scp_cmd = opts.build_scp_cmd("node1", "/src", "/dst", false, false, false);
728        assert!(scp_cmd.contains("-P 2222")); // SCP uses uppercase -P
729    }
730
731    #[test]
732    fn test_scp_options_with_control_master() {
733        let opts = SshOptions {
734            control_master: Some(ControlMaster::Auto),
735            control_path: Some("/tmp/ssh-%r@%h:%p".to_string()),
736            identity_file: Some("/key".to_string()),
737            ..SshOptions::default()
738        };
739        let cmd = opts.build_scp_cmd("node1", "/src", "/dst", false, false, false);
740        assert!(cmd.contains("-oControlMaster=auto"));
741        assert!(cmd.contains("-oControlPath=/tmp/ssh-%r@%h:%p"));
742        assert!(cmd.contains("-i /key"));
743    }
744
745    #[test]
746    fn test_ssh_worker_abort() {
747        let nodes = vec!["node1".to_string()];
748        let mut worker = SshWorker::with_defaults(nodes, "sleep 100".to_string(), 64, None);
749
750        worker.abort(true);
751        assert_eq!(worker.state(), WorkerState::Aborted);
752    }
753}
consortium/worker/ssh.rs

consortium/worker/
ssh.rs
