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    pub fn new(
413        nodes: Vec<String>,
414        source: String,
415        dest: String,
416        fanout: usize,
417        timeout: Option<Duration>,
418        ssh_options: SshOptions,
419        reverse: bool,
420        preserve: bool,
421        directory: bool,
422    ) -> Self {
423        // Build SCP command with %h placeholder for node substitution
424        let scp_cmd = ssh_options.build_scp_cmd("%h", &source, &dest, reverse, preserve, directory);
425
426        ScpWorker {
427            inner: ExecWorker::new(nodes, scp_cmd, fanout, timeout),
428            ssh_options,
429            source,
430            dest,
431            reverse,
432        }
433    }
434
435    /// Get the source path.
436    pub fn source(&self) -> &str {
437        &self.source
438    }
439
440    /// Get the destination path.
441    pub fn dest(&self) -> &str {
442        &self.dest
443    }
444
445    /// Whether this is a reverse copy.
446    pub fn is_reverse(&self) -> bool {
447        self.reverse
448    }
449}
450
451impl Worker for ScpWorker {
452    fn start(&mut self) -> crate::worker::Result<()> {
453        self.inner.start()
454    }
455
456    fn abort(&mut self, kill: bool) {
457        self.inner.abort(kill);
458    }
459
460    fn state(&self) -> WorkerState {
461        self.inner.state()
462    }
463
464    fn set_handler(&mut self, handler: Box<dyn EventHandler>) {
465        self.inner.set_handler(handler);
466    }
467
468    fn read_fds(&self) -> Vec<RawFd> {
469        self.inner.read_fds()
470    }
471
472    fn write_fds(&self) -> Vec<RawFd> {
473        self.inner.write_fds()
474    }
475
476    fn handle_read(&mut self, fd: RawFd) -> crate::worker::Result<()> {
477        self.inner.handle_read(fd)
478    }
479
480    fn handle_write(&mut self, fd: RawFd) -> crate::worker::Result<()> {
481        self.inner.handle_write(fd)
482    }
483
484    fn is_done(&self) -> bool {
485        self.inner.is_done()
486    }
487
488    fn retcodes(&self) -> &HashMap<String, i32> {
489        self.inner.retcodes()
490    }
491
492    fn num_nodes(&self) -> usize {
493        self.inner.num_nodes()
494    }
495
496    fn take_handler(&mut self) -> Option<Box<dyn EventHandler>> {
497        self.inner.take_handler()
498    }
499}
500
501#[cfg(test)]
502mod tests {
503    use super::*;
504
505    #[test]
506    fn test_ssh_options_default() {
507        let opts = SshOptions::default();
508        assert_eq!(opts.ssh_path, "ssh");
509        assert_eq!(opts.scp_path, "scp");
510        assert!(opts.user.is_none());
511        assert_eq!(opts.connect_timeout, Some(10));
512        assert!(!opts.strict_host_key_checking);
513        assert!(!opts.password_auth);
514    }
515
516    #[test]
517    fn test_ssh_options_build_ssh_cmd() {
518        let opts = SshOptions::default();
519        let cmd = opts.build_ssh_cmd("node1", "uname -a");
520
521        assert!(cmd.starts_with("ssh "));
522        assert!(cmd.contains("-oForwardX11=no"));
523        assert!(cmd.contains("-oStrictHostKeyChecking=no"));
524        assert!(cmd.contains("-oPasswordAuthentication=no"));
525        assert!(cmd.contains("-oConnectTimeout=10"));
526        assert!(cmd.contains("node1"));
527        assert!(cmd.ends_with("uname -a"));
528    }
529
530    #[test]
531    fn test_ssh_options_build_ssh_cmd_with_user() {
532        let opts = SshOptions {
533            user: Some("admin".to_string()),
534            ..SshOptions::default()
535        };
536        let cmd = opts.build_ssh_cmd("node1", "hostname");
537
538        assert!(cmd.contains("-l admin"));
539        assert!(cmd.contains("node1"));
540    }
541
542    #[test]
543    fn test_ssh_options_build_ssh_cmd_custom_options() {
544        let opts = SshOptions {
545            options: vec!["BatchMode=yes".to_string()],
546            ..SshOptions::default()
547        };
548        let cmd = opts.build_ssh_cmd("node1", "ls");
549
550        assert!(cmd.contains("-oBatchMode=yes"));
551    }
552
553    #[test]
554    fn test_ssh_options_build_scp_cmd_forward() {
555        let opts = SshOptions::default();
556        let cmd = opts.build_scp_cmd("node1", "/tmp/file.txt", "/home/user/", false, true, false);
557
558        assert!(cmd.starts_with("scp "));
559        assert!(cmd.contains("-p")); // preserve
560        assert!(cmd.contains("/tmp/file.txt"));
561        assert!(cmd.contains("node1:/home/user/"));
562    }
563
564    #[test]
565    fn test_ssh_options_build_scp_cmd_reverse() {
566        let opts = SshOptions::default();
567        let cmd = opts.build_scp_cmd(
568            "node1",
569            "/remote/file.txt",
570            "/local/dir/",
571            true,
572            false,
573            false,
574        );
575
576        assert!(cmd.contains("node1:/remote/file.txt"));
577        assert!(cmd.ends_with("/local/dir/"));
578    }
579
580    #[test]
581    fn test_ssh_options_build_scp_cmd_recursive() {
582        let opts = SshOptions::default();
583        let cmd = opts.build_scp_cmd("node1", "/src/", "/dst/", false, false, true);
584
585        assert!(cmd.contains("-r")); // recursive
586    }
587
588    #[test]
589    fn test_ssh_options_build_scp_cmd_with_user() {
590        let opts = SshOptions {
591            user: Some("deploy".to_string()),
592            ..SshOptions::default()
593        };
594        let cmd = opts.build_scp_cmd("node1", "/f", "/d", false, false, false);
595
596        assert!(cmd.contains("deploy@node1:/d"));
597    }
598
599    #[test]
600    fn test_ssh_worker_creation() {
601        let nodes = vec!["node1".to_string(), "node2".to_string()];
602        let worker = SshWorker::with_defaults(nodes, "hostname".to_string(), 64, None);
603
604        assert_eq!(worker.state(), WorkerState::Pending);
605        assert_eq!(worker.num_nodes(), 2);
606        assert_eq!(worker.original_command(), "hostname");
607    }
608
609    #[test]
610    fn test_ssh_worker_custom_options() {
611        let opts = SshOptions {
612            user: Some("root".to_string()),
613            connect_timeout: Some(5),
614            ..SshOptions::default()
615        };
616        let nodes = vec!["node1".to_string()];
617        let worker = SshWorker::new(nodes, "uptime".to_string(), 64, None, opts);
618
619        assert_eq!(worker.ssh_options().user.as_deref(), Some("root"));
620        assert_eq!(worker.ssh_options().connect_timeout, Some(5));
621    }
622
623    #[test]
624    fn test_ssh_worker_with_stderr() {
625        let nodes = vec!["node1".to_string()];
626        let worker = SshWorker::with_defaults(nodes, "ls".to_string(), 64, None).with_stderr(true);
627
628        assert_eq!(worker.state(), WorkerState::Pending);
629    }
630
631    #[test]
632    fn test_scp_worker_creation() {
633        let nodes = vec!["node1".to_string()];
634        let opts = SshOptions::default();
635        let worker = ScpWorker::new(
636            nodes,
637            "/local/file".to_string(),
638            "/remote/dir".to_string(),
639            64,
640            None,
641            opts,
642            false,
643            true,
644            false,
645        );
646
647        assert_eq!(worker.state(), WorkerState::Pending);
648        assert_eq!(worker.source(), "/local/file");
649        assert_eq!(worker.dest(), "/remote/dir");
650        assert!(!worker.is_reverse());
651    }
652
653    #[test]
654    fn test_scp_worker_reverse() {
655        let nodes = vec!["node1".to_string()];
656        let opts = SshOptions::default();
657        let worker = ScpWorker::new(
658            nodes,
659            "/remote/file".to_string(),
660            "/local/dir".to_string(),
661            64,
662            None,
663            opts,
664            true,
665            false,
666            false,
667        );
668
669        assert!(worker.is_reverse());
670        assert_eq!(worker.num_nodes(), 1);
671    }
672
673    #[test]
674    fn test_ssh_options_with_identity_file() {
675        let opts = SshOptions {
676            identity_file: Some("/path/to/key".to_string()),
677            ..SshOptions::default()
678        };
679        let cmd = opts.build_ssh_cmd("node1", "hostname");
680        assert!(cmd.contains("-i /path/to/key"));
681    }
682
683    #[test]
684    fn test_ssh_options_with_config_file() {
685        let opts = SshOptions {
686            config_file: Some("/etc/ssh/config".to_string()),
687            ..SshOptions::default()
688        };
689        let cmd = opts.build_ssh_cmd("node1", "hostname");
690        assert!(cmd.contains("-F /etc/ssh/config"));
691    }
692
693    #[test]
694    fn test_ssh_options_with_control_master() {
695        let opts = SshOptions {
696            control_master: Some(ControlMaster::Auto),
697            control_path: Some("/tmp/ssh-%r@%h:%p".to_string()),
698            control_persist: Some("10m".to_string()),
699            ..SshOptions::default()
700        };
701        let cmd = opts.build_ssh_cmd("node1", "hostname");
702        assert!(cmd.contains("-oControlMaster=auto"));
703        assert!(cmd.contains("-oControlPath=/tmp/ssh-%r@%h:%p"));
704        assert!(cmd.contains("-oControlPersist=10m"));
705    }
706
707    #[test]
708    fn test_ssh_options_with_proxy_jump() {
709        let opts = SshOptions {
710            proxy_jump: Some("bastion.example.com".to_string()),
711            ..SshOptions::default()
712        };
713        let cmd = opts.build_ssh_cmd("node1", "hostname");
714        assert!(cmd.contains("-J bastion.example.com"));
715    }
716
717    #[test]
718    fn test_ssh_options_with_port() {
719        let opts = SshOptions {
720            port: Some(2222),
721            ..SshOptions::default()
722        };
723        let ssh_cmd = opts.build_ssh_cmd("node1", "hostname");
724        assert!(ssh_cmd.contains("-p 2222"));
725
726        let scp_cmd = opts.build_scp_cmd("node1", "/src", "/dst", false, false, false);
727        assert!(scp_cmd.contains("-P 2222")); // SCP uses uppercase -P
728    }
729
730    #[test]
731    fn test_scp_options_with_control_master() {
732        let opts = SshOptions {
733            control_master: Some(ControlMaster::Auto),
734            control_path: Some("/tmp/ssh-%r@%h:%p".to_string()),
735            identity_file: Some("/key".to_string()),
736            ..SshOptions::default()
737        };
738        let cmd = opts.build_scp_cmd("node1", "/src", "/dst", false, false, false);
739        assert!(cmd.contains("-oControlMaster=auto"));
740        assert!(cmd.contains("-oControlPath=/tmp/ssh-%r@%h:%p"));
741        assert!(cmd.contains("-i /key"));
742    }
743
744    #[test]
745    fn test_ssh_worker_abort() {
746        let nodes = vec!["node1".to_string()];
747        let mut worker = SshWorker::with_defaults(nodes, "sleep 100".to_string(), 64, None);
748
749        worker.abort(true);
750        assert_eq!(worker.state(), WorkerState::Aborted);
751    }
752}
consortium/worker/ssh.rs

consortium/worker/
ssh.rs
