|
|
@ -1,12 +1,10 @@ |
|
|
|
use crate::{ |
|
|
|
use crate::{ |
|
|
|
UError,
|
|
|
|
|
|
|
|
UErrorBt,
|
|
|
|
|
|
|
|
UResult,
|
|
|
|
|
|
|
|
cache::JobCache, |
|
|
|
cache::JobCache, |
|
|
|
executor::{Waiter, DynFut},
|
|
|
|
executor::{DynFut, Waiter}, |
|
|
|
messaging::Reportable, |
|
|
|
messaging::Reportable, |
|
|
|
models::{Agent, AssignedJob, JobMeta, JobType}, |
|
|
|
models::{Agent, AssignedJob, JobMeta, JobType}, |
|
|
|
utils::{CombinedResult, OneOrVec} |
|
|
|
utils::{CombinedResult, OneOrVec}, |
|
|
|
|
|
|
|
UError, UErrorBt, UResult, |
|
|
|
}; |
|
|
|
}; |
|
|
|
use guess_host_triple::guess_host_triple; |
|
|
|
use guess_host_triple::guess_host_triple; |
|
|
|
use std::collections::HashMap; |
|
|
|
use std::collections::HashMap; |
|
|
@ -38,7 +36,8 @@ impl JobBuilder { |
|
|
|
return Err(UError::InsuitablePlatform( |
|
|
|
return Err(UError::InsuitablePlatform( |
|
|
|
meta.platform.clone(), |
|
|
|
meta.platform.clone(), |
|
|
|
curr_platform, |
|
|
|
curr_platform, |
|
|
|
).into()); |
|
|
|
) |
|
|
|
|
|
|
|
.into()); |
|
|
|
} |
|
|
|
} |
|
|
|
let job = AssignedJob::new(req.job_id, Some(&req)); |
|
|
|
let job = AssignedJob::new(req.job_id, Some(&req)); |
|
|
|
prepared.push(Box::pin(job.run())) |
|
|
|
prepared.push(Box::pin(job.run())) |
|
|
@ -152,19 +151,13 @@ impl NamedJobBuilder { |
|
|
|
|
|
|
|
|
|
|
|
#[cfg(test)] |
|
|
|
#[cfg(test)] |
|
|
|
mod tests { |
|
|
|
mod tests { |
|
|
|
|
|
|
|
|
|
|
|
use super::*; |
|
|
|
use super::*; |
|
|
|
use test_case::test_case; |
|
|
|
|
|
|
|
use std::{time::SystemTime}; |
|
|
|
|
|
|
|
use crate::{ |
|
|
|
use crate::{ |
|
|
|
errors::UError, |
|
|
|
|
|
|
|
models::{ |
|
|
|
|
|
|
|
JobMeta, |
|
|
|
|
|
|
|
misc::JobType |
|
|
|
|
|
|
|
}, |
|
|
|
|
|
|
|
builder::{JobBuilder, NamedJobBuilder}, |
|
|
|
builder::{JobBuilder, NamedJobBuilder}, |
|
|
|
|
|
|
|
models::{misc::JobType, JobMeta}, |
|
|
|
unwrap_enum, |
|
|
|
unwrap_enum, |
|
|
|
}; |
|
|
|
}; |
|
|
|
|
|
|
|
use std::time::SystemTime; |
|
|
|
|
|
|
|
|
|
|
|
type TestResult<R = ()> = Result<R, Box<dyn std::error::Error>>; |
|
|
|
type TestResult<R = ()> = Result<R, Box<dyn std::error::Error>>; |
|
|
|
|
|
|
|
|
|
|
@ -178,19 +171,14 @@ mod tests { |
|
|
|
assert!(now.elapsed().unwrap().as_secs() < SLEEP_SECS + 2) |
|
|
|
assert!(now.elapsed().unwrap().as_secs() < SLEEP_SECS + 2) |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
#[test_case(
|
|
|
|
#[rstest] |
|
|
|
|
|
|
|
#[case::sh_payload(
|
|
|
|
"/bin/sh {}", |
|
|
|
"/bin/sh {}", |
|
|
|
Some(b"echo test01 > /tmp/asd; cat /tmp/asd"),
|
|
|
|
Some(b"echo test01 > /tmp/asd; cat /tmp/asd".as_slice()),
|
|
|
|
"test01" |
|
|
|
"test01" |
|
|
|
;"sh payload" |
|
|
|
|
|
|
|
)] |
|
|
|
|
|
|
|
#[test_case(
|
|
|
|
|
|
|
|
r#"/usr/bin/python3 -c 'print("test02")'"#, |
|
|
|
|
|
|
|
None,
|
|
|
|
|
|
|
|
"test02" |
|
|
|
|
|
|
|
;"python cmd" |
|
|
|
|
|
|
|
)] |
|
|
|
)] |
|
|
|
#[test_case(
|
|
|
|
#[case::python_cmd(r#"/usr/bin/python3 -c 'print("test02")'"#, None, "test02")] |
|
|
|
|
|
|
|
#[case::sh_multiline_payload(
|
|
|
|
"/{}",
|
|
|
|
"/{}",
|
|
|
|
Some( |
|
|
|
Some( |
|
|
|
br#"#!/bin/sh |
|
|
|
br#"#!/bin/sh |
|
|
@ -198,19 +186,21 @@ mod tests { |
|
|
|
mkdir -p $TMPPATH |
|
|
|
mkdir -p $TMPPATH |
|
|
|
echo test03 > $TMPPATH/t |
|
|
|
echo test03 > $TMPPATH/t |
|
|
|
cat $TMPPATH/t |
|
|
|
cat $TMPPATH/t |
|
|
|
rm -rf $TMPPATH"# |
|
|
|
rm -rf $TMPPATH"#.as_slice() |
|
|
|
), |
|
|
|
), |
|
|
|
"test03" |
|
|
|
"test03" |
|
|
|
;"sh multiline payload" |
|
|
|
|
|
|
|
)] |
|
|
|
)] |
|
|
|
#[test_case(
|
|
|
|
#[case::standalone_binary_with_args(
|
|
|
|
"/{} 'some msg as arg'", |
|
|
|
"/{} 'some msg as arg'", |
|
|
|
Some(include_bytes!("../tests/fixtures/echoer")), |
|
|
|
Some(include_bytes!("../tests/fixtures/echoer").as_slice()),
|
|
|
|
"some msg as arg" |
|
|
|
"some msg as arg" |
|
|
|
;"standalone binary with args" |
|
|
|
|
|
|
|
)] |
|
|
|
)] |
|
|
|
#[tokio::test] |
|
|
|
#[tokio::test] |
|
|
|
async fn test_shell_job(cmd: &str, payload: Option<&[u8]>, expected_result: &str) -> TestResult { |
|
|
|
async fn test_shell_job( |
|
|
|
|
|
|
|
#[case] cmd: &str, |
|
|
|
|
|
|
|
#[case] payload: Option<&[u8]>, |
|
|
|
|
|
|
|
#[case] expected_result: &str, |
|
|
|
|
|
|
|
) -> TestResult { |
|
|
|
let mut job = JobMeta::builder().with_shell(cmd); |
|
|
|
let mut job = JobMeta::builder().with_shell(cmd); |
|
|
|
if let Some(p) = payload { |
|
|
|
if let Some(p) = payload { |
|
|
|
job = job.with_payload(p); |
|
|
|
job = job.with_payload(p); |
|
|
@ -228,8 +218,12 @@ mod tests { |
|
|
|
const SLEEP_SECS: u64 = 1; |
|
|
|
const SLEEP_SECS: u64 = 1; |
|
|
|
let now = SystemTime::now(); |
|
|
|
let now = SystemTime::now(); |
|
|
|
let longest_job = JobMeta::from_shell(format!("sleep {}", SLEEP_SECS)).unwrap(); |
|
|
|
let longest_job = JobMeta::from_shell(format!("sleep {}", SLEEP_SECS)).unwrap(); |
|
|
|
let longest_job = JobBuilder::from_meta(longest_job).unwrap_one().spawn().await;
|
|
|
|
let longest_job = JobBuilder::from_meta(longest_job) |
|
|
|
let ls = JobBuilder::from_meta(JobMeta::from_shell("ls")?).unwrap_one() |
|
|
|
.unwrap_one() |
|
|
|
|
|
|
|
.spawn() |
|
|
|
|
|
|
|
.await; |
|
|
|
|
|
|
|
let ls = JobBuilder::from_meta(JobMeta::from_shell("ls")?) |
|
|
|
|
|
|
|
.unwrap_one() |
|
|
|
.wait_one() |
|
|
|
.wait_one() |
|
|
|
.await; |
|
|
|
.await; |
|
|
|
let ls = unwrap_enum!(ls, Reportable::Assigned); |
|
|
|
let ls = unwrap_enum!(ls, Reportable::Assigned); |
|
|
@ -272,10 +266,7 @@ mod tests { |
|
|
|
#[tokio::test] |
|
|
|
#[tokio::test] |
|
|
|
async fn test_failing_shell_job() -> TestResult { |
|
|
|
async fn test_failing_shell_job() -> TestResult { |
|
|
|
let job = JobMeta::from_shell("lol_kek_puk")?; |
|
|
|
let job = JobMeta::from_shell("lol_kek_puk")?; |
|
|
|
let job_result = JobBuilder::from_meta(job) |
|
|
|
let job_result = JobBuilder::from_meta(job).unwrap_one().wait_one().await; |
|
|
|
.unwrap_one() |
|
|
|
|
|
|
|
.wait_one() |
|
|
|
|
|
|
|
.await; |
|
|
|
|
|
|
|
let job_result = unwrap_enum!(job_result, Reportable::Assigned); |
|
|
|
let job_result = unwrap_enum!(job_result, Reportable::Assigned); |
|
|
|
let output = job_result.to_string_result(); |
|
|
|
let output = job_result.to_string_result(); |
|
|
|
assert!(output.contains("No such file")); |
|
|
|
assert!(output.contains("No such file")); |
|
|
@ -283,20 +274,15 @@ mod tests { |
|
|
|
Ok(()) |
|
|
|
Ok(()) |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
#[test_case(
|
|
|
|
#[rstest] |
|
|
|
"/bin/bash {}",
|
|
|
|
#[case::no_binary("/bin/bash {}", None, "contains executable")] |
|
|
|
None,
|
|
|
|
#[case::no_path_to_binary("/bin/bash", Some(b"whoami".as_slice()), "contains no executable")] |
|
|
|
"contains executable" |
|
|
|
|
|
|
|
; "no binary" |
|
|
|
|
|
|
|
)] |
|
|
|
|
|
|
|
#[test_case(
|
|
|
|
|
|
|
|
"/bin/bash",
|
|
|
|
|
|
|
|
Some(b"whoami"),
|
|
|
|
|
|
|
|
"contains no executable" |
|
|
|
|
|
|
|
; "no path to binary" |
|
|
|
|
|
|
|
)] |
|
|
|
|
|
|
|
#[tokio::test] |
|
|
|
#[tokio::test] |
|
|
|
async fn test_job_building_failed(cmd: &str, payload: Option<&[u8]>, err_str: &str) -> TestResult { |
|
|
|
async fn test_job_building_failed( |
|
|
|
|
|
|
|
#[case] cmd: &str, |
|
|
|
|
|
|
|
#[case] payload: Option<&[u8]>, |
|
|
|
|
|
|
|
#[case] err_str: &str, |
|
|
|
|
|
|
|
) -> TestResult { |
|
|
|
let mut job = JobMeta::builder().with_shell(cmd); |
|
|
|
let mut job = JobMeta::builder().with_shell(cmd); |
|
|
|
if let Some(p) = payload { |
|
|
|
if let Some(p) = payload { |
|
|
|
job = job.with_payload(p); |
|
|
|
job = job.with_payload(p); |
|
|
@ -311,11 +297,15 @@ mod tests { |
|
|
|
async fn test_different_job_types() -> TestResult { |
|
|
|
async fn test_different_job_types() -> TestResult { |
|
|
|
let mut jobs = NamedJobBuilder::from_meta(vec![ |
|
|
|
let mut jobs = NamedJobBuilder::from_meta(vec![ |
|
|
|
("sleeper", JobMeta::from_shell("sleep 3")?), |
|
|
|
("sleeper", JobMeta::from_shell("sleep 3")?), |
|
|
|
("gatherer", JobMeta::builder().with_type(JobType::Manage).build()?) |
|
|
|
( |
|
|
|
]).wait().await; |
|
|
|
"gatherer", |
|
|
|
|
|
|
|
JobMeta::builder().with_type(JobType::Manage).build()?, |
|
|
|
|
|
|
|
), |
|
|
|
|
|
|
|
]) |
|
|
|
|
|
|
|
.wait() |
|
|
|
|
|
|
|
.await; |
|
|
|
let gathered = jobs.pop("gatherer"); |
|
|
|
let gathered = jobs.pop("gatherer"); |
|
|
|
assert_eq!(unwrap_enum!(gathered, Reportable::Agent).alias, None); |
|
|
|
assert_eq!(unwrap_enum!(gathered, Reportable::Agent).alias, None); |
|
|
|
Ok(()) |
|
|
|
Ok(()) |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
} |
|
|
|
} |