1use std::ffi::OsString;
27use std::fs::File;
28use std::io::{self, Write};
29use std::path::{Path, PathBuf};
30use std::process::{Command, Stdio};
31
32#[derive(Debug, Clone, PartialEq, Eq)]
34pub struct LaunchConfig {
35 pub service_name: String,
37 pub console_mode: bool,
39 pub debug_log_path: Option<PathBuf>,
41 pub service_args: Vec<OsString>,
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
47pub struct ServiceArgs {
48 pub script_path: PathBuf,
50 pub stdout_path: PathBuf,
52 pub stderr_path: PathBuf,
54 pub exit_code_path: PathBuf,
56}
57
58#[derive(Debug, Clone, PartialEq, Eq)]
60pub enum PayloadRequest {
61 File(ServiceArgs),
63 Pipe(PipeServiceArgs),
65}
66
67#[derive(Debug, Clone, PartialEq, Eq)]
69pub struct PipeServiceArgs {
70 pub pipe_prefix: String,
72 pub command: Option<String>,
74 pub working_directory: Option<PathBuf>,
76 pub columns: Option<u16>,
78 pub rows: Option<u16>,
80}
81
82pub fn parse_launch_config(args: &[OsString]) -> Result<LaunchConfig, String> {
84 let mut service_name = None;
85 let mut console_mode = false;
86 let mut debug_log_path = None;
87 let mut service_args = Vec::new();
88 let mut index = 0;
89 while index < args.len() {
90 match args[index].to_string_lossy().as_ref() {
91 "--service-name" => {
92 index += 1;
93 let value = args
94 .get(index)
95 .ok_or_else(|| "missing value for --service-name".to_string())?;
96 service_name = Some(value.to_string_lossy().into_owned());
97 }
98 "--console" => {
99 console_mode = true;
100 }
101 "--debug-log" => {
102 index += 1;
103 let value = args
104 .get(index)
105 .ok_or_else(|| "missing value for --debug-log".to_string())?;
106 debug_log_path = Some(PathBuf::from(value));
107 }
108 _ => {
109 service_args.push(args[index].clone());
110 }
111 }
112 index += 1;
113 }
114
115 Ok(LaunchConfig {
116 service_name: service_name.unwrap_or_else(|| "smolder-psexecsvc".to_string()),
117 console_mode,
118 debug_log_path,
119 service_args,
120 })
121}
122
123pub fn parse_service_args(args: &[OsString]) -> Result<ServiceArgs, String> {
125 let mut script_path = None;
126 let mut stdout_path = None;
127 let mut stderr_path = None;
128 let mut exit_code_path = None;
129 let mut index = 0;
130 while index < args.len() {
131 let key = args[index].to_string_lossy();
132 let value = match key.as_ref() {
133 "--script" | "--stdout" | "--stderr" | "--exit-code" => {
134 index += 1;
135 args.get(index)
136 .ok_or_else(|| format!("missing value for {key}"))?
137 .clone()
138 }
139 _ => return Err(format!("unknown service argument: {key}")),
140 };
141
142 match key.as_ref() {
143 "--script" => script_path = Some(PathBuf::from(value)),
144 "--stdout" => stdout_path = Some(PathBuf::from(value)),
145 "--stderr" => stderr_path = Some(PathBuf::from(value)),
146 "--exit-code" => exit_code_path = Some(PathBuf::from(value)),
147 _ => unreachable!(),
148 }
149 index += 1;
150 }
151
152 Ok(ServiceArgs {
153 script_path: script_path.ok_or_else(|| "missing --script".to_string())?,
154 stdout_path: stdout_path.ok_or_else(|| "missing --stdout".to_string())?,
155 stderr_path: stderr_path.ok_or_else(|| "missing --stderr".to_string())?,
156 exit_code_path: exit_code_path.ok_or_else(|| "missing --exit-code".to_string())?,
157 })
158}
159
160pub fn parse_payload_request(args: &[OsString]) -> Result<PayloadRequest, String> {
162 if args
163 .iter()
164 .any(|arg| arg.to_string_lossy() == "--pipe-prefix")
165 {
166 return parse_pipe_service_args(args).map(PayloadRequest::Pipe);
167 }
168 parse_service_args(args).map(PayloadRequest::File)
169}
170
171pub fn parse_pipe_service_args(args: &[OsString]) -> Result<PipeServiceArgs, String> {
173 let mut pipe_prefix = None;
174 let mut command = None;
175 let mut working_directory = None;
176 let mut columns = None;
177 let mut rows = None;
178 let mut index = 0;
179 while index < args.len() {
180 let key = args[index].to_string_lossy();
181 let value = match key.as_ref() {
182 "--pipe-prefix" | "--command" | "--workdir" | "--cols" | "--rows" => {
183 index += 1;
184 args.get(index)
185 .ok_or_else(|| format!("missing value for {key}"))?
186 .clone()
187 }
188 _ => return Err(format!("unknown service argument: {key}")),
189 };
190
191 match key.as_ref() {
192 "--pipe-prefix" => pipe_prefix = Some(value.to_string_lossy().into_owned()),
193 "--command" => command = Some(value.to_string_lossy().into_owned()),
194 "--workdir" => working_directory = Some(PathBuf::from(value)),
195 "--cols" => {
196 columns = Some(
197 value
198 .to_string_lossy()
199 .parse::<u16>()
200 .map_err(|_| "invalid value for --cols".to_string())?,
201 );
202 }
203 "--rows" => {
204 rows = Some(
205 value
206 .to_string_lossy()
207 .parse::<u16>()
208 .map_err(|_| "invalid value for --rows".to_string())?,
209 );
210 }
211 _ => unreachable!(),
212 }
213 index += 1;
214 }
215
216 Ok(PipeServiceArgs {
217 pipe_prefix: pipe_prefix.ok_or_else(|| "missing --pipe-prefix".to_string())?,
218 command,
219 working_directory,
220 columns,
221 rows,
222 })
223}
224
225#[derive(Debug, Clone, PartialEq, Eq)]
227pub struct PipeNames {
228 pub stdin: String,
230 pub stdout: String,
232 pub stderr: String,
234 pub control: String,
236}
237
238impl PipeNames {
239 #[must_use]
241 pub fn new(prefix: &str) -> Self {
242 Self {
243 stdin: format!(r"\\.\pipe\{prefix}.stdin"),
244 stdout: format!(r"\\.\pipe\{prefix}.stdout"),
245 stderr: format!(r"\\.\pipe\{prefix}.stderr"),
246 control: format!(r"\\.\pipe\{prefix}.control"),
247 }
248 }
249}
250
251pub fn run_service_once(args: &ServiceArgs) -> io::Result<i32> {
253 let stdout = File::create(&args.stdout_path)?;
254 let stderr = File::create(&args.stderr_path)?;
255 let status = child_command(&args.script_path)
256 .stdout(Stdio::from(stdout))
257 .stderr(Stdio::from(stderr))
258 .status()?;
259 let exit_code = status.code().unwrap_or(1);
260 write_exit_code(&args.exit_code_path, exit_code)?;
261 Ok(exit_code)
262}
263
264fn child_command(script_path: &Path) -> Command {
265 #[cfg(windows)]
266 {
267 let mut command =
268 Command::new(current_comspec().unwrap_or_else(|| OsString::from("cmd.exe")));
269 command.arg("/Q").arg("/C").arg(script_path.as_os_str());
270 command
271 }
272
273 #[cfg(not(windows))]
274 {
275 let mut command = Command::new("sh");
276 command.arg(script_path.as_os_str());
277 command
278 }
279}
280
281#[cfg(windows)]
282fn current_comspec() -> Option<OsString> {
283 std::env::var_os("COMSPEC")
284}
285
286fn write_exit_code(path: &Path, exit_code: i32) -> io::Result<()> {
287 let mut file = File::create(path)?;
288 writeln!(file, "{exit_code}")?;
289 Ok(())
290}
291
292#[cfg(test)]
293mod tests {
294 use std::fs;
295 use std::path::PathBuf;
296 use std::time::{SystemTime, UNIX_EPOCH};
297
298 use super::{
299 parse_launch_config, parse_payload_request, parse_pipe_service_args, parse_service_args,
300 run_service_once, LaunchConfig, PayloadRequest, PipeNames, PipeServiceArgs, ServiceArgs,
301 };
302
303 #[test]
304 fn parse_launch_config_extracts_service_name_and_console_mode() {
305 let config = parse_launch_config(&[
306 "--service-name".into(),
307 "SMOLDERTEST".into(),
308 "--console".into(),
309 "--script".into(),
310 "run.cmd".into(),
311 ])
312 .expect("launch config should parse");
313 assert_eq!(
314 config,
315 LaunchConfig {
316 service_name: "SMOLDERTEST".to_string(),
317 console_mode: true,
318 debug_log_path: None,
319 service_args: vec!["--script".into(), "run.cmd".into()],
320 }
321 );
322 }
323
324 #[test]
325 fn parse_launch_config_extracts_optional_debug_log_path() {
326 let config = parse_launch_config(&[
327 "--service-name".into(),
328 "SMOLDERTEST".into(),
329 "--debug-log".into(),
330 "C:\\Temp\\svc.log".into(),
331 "--script".into(),
332 "run.cmd".into(),
333 ])
334 .expect("launch config should parse");
335 assert_eq!(
336 config,
337 LaunchConfig {
338 service_name: "SMOLDERTEST".to_string(),
339 console_mode: false,
340 debug_log_path: Some(PathBuf::from("C:\\Temp\\svc.log")),
341 service_args: vec!["--script".into(), "run.cmd".into()],
342 }
343 );
344 }
345
346 #[test]
347 fn parse_service_args_extracts_required_paths() {
348 let args = parse_service_args(&[
349 "--script".into(),
350 "run.cmd".into(),
351 "--stdout".into(),
352 "stdout.txt".into(),
353 "--stderr".into(),
354 "stderr.txt".into(),
355 "--exit-code".into(),
356 "exit.txt".into(),
357 ])
358 .expect("service args should parse");
359 assert_eq!(args.script_path, PathBuf::from("run.cmd"));
360 assert_eq!(args.stdout_path, PathBuf::from("stdout.txt"));
361 assert_eq!(args.stderr_path, PathBuf::from("stderr.txt"));
362 assert_eq!(args.exit_code_path, PathBuf::from("exit.txt"));
363 }
364
365 #[test]
366 fn parse_pipe_service_args_extracts_prefix_command_workdir_and_size() {
367 let args = parse_pipe_service_args(&[
368 "--pipe-prefix".into(),
369 "SMOLDER-ABC".into(),
370 "--command".into(),
371 "whoami".into(),
372 "--workdir".into(),
373 "C:\\Temp".into(),
374 "--cols".into(),
375 "132".into(),
376 "--rows".into(),
377 "43".into(),
378 ])
379 .expect("pipe service args should parse");
380 assert_eq!(
381 args,
382 PipeServiceArgs {
383 pipe_prefix: "SMOLDER-ABC".to_string(),
384 command: Some("whoami".to_string()),
385 working_directory: Some(PathBuf::from("C:\\Temp")),
386 columns: Some(132),
387 rows: Some(43),
388 }
389 );
390 }
391
392 #[test]
393 fn pipe_names_expand_from_prefix() {
394 let pipes = PipeNames::new("SMOLDER-ABC");
395 assert_eq!(pipes.stdin, r"\\.\pipe\SMOLDER-ABC.stdin");
396 assert_eq!(pipes.stdout, r"\\.\pipe\SMOLDER-ABC.stdout");
397 assert_eq!(pipes.stderr, r"\\.\pipe\SMOLDER-ABC.stderr");
398 assert_eq!(pipes.control, r"\\.\pipe\SMOLDER-ABC.control");
399 }
400
401 #[test]
402 fn parse_payload_request_detects_pipe_mode() {
403 let request = parse_payload_request(&[
404 "--pipe-prefix".into(),
405 "SMOLDER-ABC".into(),
406 "--command".into(),
407 "whoami".into(),
408 ])
409 .expect("payload request should parse");
410
411 assert_eq!(
412 request,
413 PayloadRequest::Pipe(PipeServiceArgs {
414 pipe_prefix: "SMOLDER-ABC".to_string(),
415 command: Some("whoami".to_string()),
416 working_directory: None,
417 columns: None,
418 rows: None,
419 })
420 );
421 }
422
423 #[test]
424 fn run_service_once_executes_script_and_writes_exit_code() {
425 let unique = SystemTime::now()
426 .duration_since(UNIX_EPOCH)
427 .expect("time should move forward")
428 .as_nanos();
429 let base = std::env::temp_dir().join(format!("smolder-psexecsvc-{unique}"));
430 fs::create_dir_all(&base).expect("temp dir should create");
431
432 #[cfg(windows)]
433 let script_path = {
434 let path = base.join("run.cmd");
435 fs::write(&path, "@echo hello\r\n@echo oops 1>&2\r\n@exit /b 7\r\n")
436 .expect("script should write");
437 path
438 };
439
440 #[cfg(not(windows))]
441 let script_path = {
442 let path = base.join("run.sh");
443 fs::write(&path, "#!/bin/sh\necho hello\necho oops >&2\nexit 7\n")
444 .expect("script should write");
445 #[cfg(unix)]
446 {
447 use std::os::unix::fs::PermissionsExt;
448 let mut perms = fs::metadata(&path)
449 .expect("metadata should load")
450 .permissions();
451 perms.set_mode(0o755);
452 fs::set_permissions(&path, perms).expect("permissions should set");
453 }
454 path
455 };
456
457 let stdout_path = base.join("stdout.txt");
458 let stderr_path = base.join("stderr.txt");
459 let exit_code_path = base.join("exit.txt");
460 let exit_code = run_service_once(&ServiceArgs {
461 script_path,
462 stdout_path: stdout_path.clone(),
463 stderr_path: stderr_path.clone(),
464 exit_code_path: exit_code_path.clone(),
465 })
466 .expect("service execution should succeed");
467
468 assert_eq!(exit_code, 7);
469 assert!(fs::read_to_string(&stdout_path)
470 .expect("stdout should read")
471 .contains("hello"));
472 assert!(fs::read_to_string(&stderr_path)
473 .expect("stderr should read")
474 .contains("oops"));
475 assert_eq!(
476 fs::read_to_string(&exit_code_path)
477 .expect("exit code should read")
478 .trim(),
479 "7"
480 );
481 let _ = fs::remove_dir_all(base);
482 }
483}