You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
161 lines
4.6 KiB
161 lines
4.6 KiB
use std::{time::SystemTime}; |
|
use u_lib::{ |
|
errors::UError, |
|
models::{ |
|
jobs::{JobMeta}, |
|
ExecResult, |
|
misc::JobType |
|
}, |
|
builder::{JobBuilder, NamedJobBuilder} |
|
}; |
|
|
|
type TestResult<R = ()> = Result<R, Box<dyn std::error::Error>>; |
|
|
|
#[tokio::test] |
|
async fn test_is_really_async() { |
|
const SLEEP_SECS: u64 = 1; |
|
let job = JobMeta::from_shell(format!("sleep {}", SLEEP_SECS)).unwrap(); |
|
let sleep_jobs: Vec<JobMeta> = (0..50).map(|_| job.clone()).collect(); |
|
let now = SystemTime::now(); |
|
JobBuilder::from_meta(sleep_jobs).unwrap_one().wait().await; |
|
assert!(now.elapsed().unwrap().as_secs() < SLEEP_SECS + 2) |
|
} |
|
|
|
#[test_case( |
|
"/bin/sh {}", |
|
Some(b"echo test01 > /tmp/asd; cat /tmp/asd"), |
|
"test01" |
|
;"sh payload" |
|
)] |
|
#[test_case( |
|
r#"/usr/bin/python -c 'print("test02")'"#, |
|
None, |
|
"test02" |
|
;"python cmd" |
|
)] |
|
#[test_case( |
|
"/{}", |
|
Some( |
|
br#"#!/bin/sh |
|
TMPPATH=/tmp/lol |
|
mkdir -p $TMPPATH |
|
echo test03 > $TMPPATH/t |
|
cat $TMPPATH/t"# |
|
), |
|
"test03" |
|
;"sh multiline payload" |
|
)] |
|
#[test_case( |
|
"/{} 'some msg as arg'", |
|
Some(include_bytes!("../fixtures/echoer")), |
|
"some msg as arg" |
|
;"standalone binary with args" |
|
)] |
|
#[tokio::test] |
|
async fn test_shell_job(cmd: &str, payload: Option<&[u8]>, expected_result: &str) -> TestResult { |
|
let mut job = JobMeta::builder().with_shell(cmd); |
|
if let Some(p) = payload { |
|
job = job.with_payload(p); |
|
} |
|
let job = job.build().unwrap(); |
|
let job_result = JobBuilder::from_meta(job).unwrap_one().wait_one().await; |
|
let result = unwrap_enum!(job_result, ExecResult::Assigned); |
|
let result = result.to_string_result().unwrap(); |
|
assert_eq!(result.trim(), expected_result); |
|
Ok(()) |
|
} |
|
|
|
#[tokio::test] |
|
async fn test_complex_load() -> TestResult { |
|
const SLEEP_SECS: u64 = 1; |
|
let now = SystemTime::now(); |
|
let longest_job = JobMeta::from_shell(format!("sleep {}", SLEEP_SECS)).unwrap(); |
|
let longest_job = JobBuilder::from_meta(longest_job).unwrap_one().spawn().await; |
|
let ls = JobBuilder::from_meta(JobMeta::from_shell("ls")?).unwrap_one() |
|
.wait_one() |
|
.await; |
|
let ls = unwrap_enum!(ls, ExecResult::Assigned); |
|
assert_eq!(ls.retcode.unwrap(), 0); |
|
let folders = ls.to_string_result().unwrap(); |
|
let subfolders_jobs: Vec<JobMeta> = folders |
|
.lines() |
|
.map(|f| JobMeta::from_shell(format!("ls {}", f)).unwrap()) |
|
.collect(); |
|
let ls_subfolders = JobBuilder::from_meta(subfolders_jobs) |
|
.unwrap_one() |
|
.wait() |
|
.await; |
|
for result in ls_subfolders { |
|
let result = unwrap_enum!(result, ExecResult::Assigned); |
|
assert_eq!(result.retcode.unwrap(), 0); |
|
} |
|
longest_job.wait().await; |
|
assert_eq!(now.elapsed().unwrap().as_secs(), SLEEP_SECS); |
|
Ok(()) |
|
} |
|
/* |
|
#[tokio::test] |
|
async fn test_exec_multiple_jobs_nowait() -> UResult<()> { |
|
const REPEATS: usize = 10; |
|
let job = JobMeta::from_shell("whoami"); |
|
let sleep_jobs: Vec<JobMeta> = (0..=REPEATS).map(|_| job.clone()).collect(); |
|
build_jobs(sleep_jobs).spawn().await; |
|
let mut completed = 0; |
|
while completed < REPEATS { |
|
let c = pop_completed().await.len(); |
|
if c > 0 { |
|
completed += c; |
|
println!("{}", c); |
|
} |
|
} |
|
Ok(()) |
|
} |
|
*/ |
|
#[tokio::test] |
|
async fn test_failing_shell_job() -> TestResult { |
|
let job = JobMeta::from_shell("lol_kek_puk")?; |
|
let job_result = JobBuilder::from_meta(job) |
|
.unwrap_one() |
|
.wait_one() |
|
.await; |
|
let job_result = unwrap_enum!(job_result, ExecResult::Assigned); |
|
let output = job_result.to_string_result().unwrap(); |
|
assert!(output.contains("No such file")); |
|
assert!(job_result.retcode.is_none()); |
|
Ok(()) |
|
} |
|
|
|
#[test_case( |
|
"/bin/bash {}", |
|
None, |
|
"contains executable" |
|
; "no binary" |
|
)] |
|
#[test_case( |
|
"/bin/bash", |
|
Some(b"whoami"), |
|
"contains no executable" |
|
; "no path to binary" |
|
)] |
|
#[tokio::test] |
|
async fn test_job_building_failed(cmd: &str, payload: Option<&[u8]>, err_str: &str) -> TestResult { |
|
let mut job = JobMeta::builder().with_shell(cmd); |
|
if let Some(p) = payload { |
|
job = job.with_payload(p); |
|
} |
|
let err = job.build().unwrap_err(); |
|
let err_msg = unwrap_enum!(err, UError::JobArgsError); |
|
assert!(err_msg.contains(err_str)); |
|
Ok(()) |
|
} |
|
|
|
#[tokio::test] |
|
async fn test_different_job_types() -> TestResult { |
|
let mut jobs = NamedJobBuilder::from_meta(vec![ |
|
("sleeper", JobMeta::from_shell("sleep 3")?), |
|
("gatherer", JobMeta::builder().with_type(JobType::Manage).build()?) |
|
]).wait().await; |
|
let gathered = jobs.pop("gatherer"); |
|
assert_eq!(unwrap_enum!(gathered, ExecResult::Agent).alias, None); |
|
Ok(()) |
|
}
|
|
|