Resolve "Agent states" Closes #7, #12, #6, and #13 See merge request root/unki!34-update-check
ortem
4 years ago
commit
5bd1760d10
42 changed files with 1489 additions and 885 deletions
@ -1,2 +1,4 @@ |
||||
/target |
||||
**/*.rs.bk |
||||
/.idea |
||||
/data |
||||
@ -1,111 +1,142 @@ |
||||
use crate::db::{lock_db, UDB}; |
||||
use diesel::SaveChangesDsl; |
||||
use std::fmt::Display; |
||||
use hyper::Body; |
||||
use serde::Serialize; |
||||
use u_lib::{ |
||||
messaging::{BaseMessage, ToMsg}, |
||||
models::{ExactJob, IAgent, JobMeta, JobState}, |
||||
messaging::{AsMsg, BaseMessage}, |
||||
models::{Agent, AssignedJob, ExecResult, JobMeta, JobState}, |
||||
ULocalError, |
||||
}; |
||||
use uuid::Uuid; |
||||
use warp::{ |
||||
http::{Response, StatusCode}, |
||||
reply::with_status, |
||||
Rejection, Reply, |
||||
}; |
||||
|
||||
fn build_response<S: Display>(code: StatusCode, body: S) -> Result<Response<String>, Rejection> { |
||||
Ok(Response::builder() |
||||
.status(code) |
||||
.body(format!("{}", body)) |
||||
.unwrap()) |
||||
fn build_response<S: Into<Body>>(code: StatusCode, body: S) -> Response<Body> { |
||||
Response::builder().status(code).body(body.into()).unwrap() |
||||
} |
||||
|
||||
fn build_empty_200() -> Result<Response<String>, Rejection> { |
||||
build_response(StatusCode::OK, "") |
||||
fn build_ok<S: Into<Body>>(body: S) -> Response<Body> { |
||||
build_response(StatusCode::OK, body) |
||||
} |
||||
|
||||
pub async fn add_agent(msg: BaseMessage<'_, IAgent>) -> Result<impl Reply, Rejection> { |
||||
match lock_db().insert_agents(&msg.into_inner()) { |
||||
Ok(_) => build_empty_200(), |
||||
Err(e) => build_response(StatusCode::BAD_REQUEST, e), |
||||
} |
||||
fn build_err<S: ToString>(body: S) -> Response<Body> { |
||||
build_response(StatusCode::BAD_REQUEST, body.to_string()) |
||||
} |
||||
|
||||
fn build_message<M: AsMsg + Serialize>(m: M) -> Response<Body> { |
||||
warp::reply::json(&m.as_message()).into_response() |
||||
} |
||||
|
||||
pub async fn add_agent(msg: Agent) -> Result<impl Reply, Rejection> { |
||||
debug!("hnd: add_agent"); |
||||
lock_db() |
||||
.insert_agent(&msg) |
||||
.map(|_| build_ok("")) |
||||
.or_else(|e| Ok(build_err(e))) |
||||
} |
||||
|
||||
pub async fn get_agents(uid: Option<Uuid>) -> Result<impl Reply, Rejection> { |
||||
match lock_db().get_agents(uid) { |
||||
Ok(r) => Ok(warp::reply::json(&r.as_message())), |
||||
Err(e) => Err(warp::reject()), |
||||
} |
||||
debug!("hnd: get_agents"); |
||||
lock_db() |
||||
.get_agents(uid) |
||||
.map(|m| build_message(m)) |
||||
.or_else(|e| Ok(build_err(e))) |
||||
} |
||||
|
||||
pub async fn get_jobs(uid: Option<Uuid>) -> Result<impl Reply, Rejection> { |
||||
match lock_db().get_jobs(uid) { |
||||
Ok(r) => Ok(warp::reply::json(&r.as_message())), |
||||
Err(e) => Err(warp::reject()), |
||||
} |
||||
debug!("hnd: get_jobs"); |
||||
lock_db() |
||||
.get_jobs(uid) |
||||
.map(|m| build_message(m)) |
||||
.or_else(|e| Ok(build_err(e))) |
||||
} |
||||
|
||||
pub async fn get_agent_jobs(uid: Option<Uuid>, personal: bool) -> Result<impl Reply, Rejection> { |
||||
info!("hnd: get_agent_jobs {}", personal); |
||||
if personal { |
||||
let agents = lock_db().get_agents(uid).unwrap(); |
||||
if agents.len() == 0 { |
||||
let db = lock_db(); |
||||
db.insert_agent(&Agent::with_id(uid.unwrap())).unwrap(); |
||||
let job = db.find_job_by_alias("agent_hello").unwrap(); |
||||
if let Err(e) = db.set_jobs_for_agent(&uid.unwrap(), &[job.id]) { |
||||
return Ok(build_err(e)); |
||||
} |
||||
} |
||||
} |
||||
let result = lock_db().get_exact_jobs(uid, personal); |
||||
match result { |
||||
Ok(r) => { |
||||
let _db = lock_db(); |
||||
let db = lock_db(); |
||||
for j in r.iter() { |
||||
_db.update_job_status(j.id, JobState::Running).ok(); |
||||
db.update_job_status(j.id, JobState::Running).ok(); |
||||
} |
||||
Ok(warp::reply::json(&r.as_message())) |
||||
Ok(build_message(r)) |
||||
} |
||||
Err(e) => Err(warp::reject()), |
||||
Err(e) => Ok(build_err(e)), |
||||
} |
||||
} |
||||
|
||||
pub async fn upload_jobs(msg: BaseMessage<'_, Vec<JobMeta>>) -> Result<impl Reply, Rejection> { |
||||
match lock_db().insert_jobs(&msg.into_inner()) { |
||||
Ok(_) => build_empty_200(), |
||||
Err(e) => build_response(StatusCode::BAD_REQUEST, e), |
||||
} |
||||
debug!("hnd: upload_jobs"); |
||||
lock_db() |
||||
.insert_jobs(&msg.into_inner()) |
||||
.map(|_| build_ok("")) |
||||
.or_else(|e| Ok(build_err(e))) |
||||
} |
||||
|
||||
pub async fn del(uid: Uuid) -> Result<impl Reply, Rejection> { |
||||
debug!("hnd: del"); |
||||
let db = lock_db(); |
||||
let del_fns = &[UDB::del_agents, UDB::del_jobs, UDB::del_results]; |
||||
for del_fn in del_fns { |
||||
let affected = del_fn(&db, &vec![uid]).unwrap(); |
||||
if affected > 0 { |
||||
return build_response(StatusCode::OK, affected); |
||||
return Ok(build_ok(affected.to_string())); |
||||
} |
||||
} |
||||
build_response(StatusCode::BAD_REQUEST, 0) |
||||
Ok(build_err("0")) |
||||
} |
||||
|
||||
pub async fn set_jobs( |
||||
agent_uid: Uuid, |
||||
msg: BaseMessage<'_, Vec<Uuid>>, |
||||
) -> Result<impl Reply, Rejection> { |
||||
match lock_db().set_jobs_for_agent(&agent_uid, &msg.into_inner()) { |
||||
Ok(_) => build_empty_200(), |
||||
Err(e) => build_response(StatusCode::BAD_REQUEST, e), |
||||
} |
||||
debug!("hnd: set_jobs"); |
||||
lock_db() |
||||
.set_jobs_for_agent(&agent_uid, &msg.into_inner()) |
||||
.map(|_| build_ok("")) |
||||
.or_else(|e| Ok(build_err(e))) |
||||
} |
||||
|
||||
pub async fn report(msg: BaseMessage<'_, Vec<ExactJob>>) -> Result<impl Reply, Rejection> { |
||||
let db = lock_db(); |
||||
pub async fn report(msg: BaseMessage<'_, Vec<ExecResult>>) -> Result<impl Reply, Rejection> { |
||||
debug!("hnd: report"); |
||||
let id = msg.id; |
||||
let mut failed = vec![]; |
||||
for res in msg.into_inner() { |
||||
if id != res.agent_id { |
||||
continue; |
||||
} |
||||
if let Err(e) = res |
||||
.save_changes::<ExactJob>(&db.conn) |
||||
.map_err(ULocalError::from) |
||||
{ |
||||
failed.push(e.to_string()) |
||||
for entry in msg.into_inner() { |
||||
match entry { |
||||
ExecResult::Assigned(res) => { |
||||
if id != res.agent_id { |
||||
continue; |
||||
} |
||||
let db = lock_db(); |
||||
if let Err(e) = res |
||||
.save_changes::<AssignedJob>(&db.conn) |
||||
.map_err(ULocalError::from) |
||||
{ |
||||
failed.push(e.to_string()) |
||||
} |
||||
} |
||||
ExecResult::Agent(a) => { |
||||
add_agent(a).await?; |
||||
} |
||||
} |
||||
} |
||||
if failed.len() > 0 { |
||||
let err_msg = ULocalError::ProcessingError(failed.join(", ")); |
||||
return build_response(StatusCode::BAD_REQUEST, err_msg); |
||||
return Ok(build_err(err_msg)); |
||||
} |
||||
build_empty_200() |
||||
Ok(build_ok("")) |
||||
} |
||||
|
||||
@ -0,0 +1,16 @@ |
||||
[package] |
||||
name = "u_api_proc_macro" |
||||
version = "0.1.0" |
||||
authors = ["plazmoid <kronos44@mail.ru>"] |
||||
edition = "2018" |
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html |
||||
|
||||
[lib] |
||||
proc-macro = true |
||||
|
||||
[dependencies] |
||||
syn = { version = "1.0", features = ["full", "extra-traits"] } |
||||
quote = "1.0" |
||||
strum = { version = "0.20", features = ["derive"] } |
||||
proc-macro2 = "1.0" |
||||
@ -0,0 +1,179 @@ |
||||
use proc_macro::TokenStream; |
||||
use proc_macro2::{Ident, TokenStream as TokenStream2}; |
||||
use quote::quote; |
||||
use std::{collections::HashMap, str::FromStr}; |
||||
use strum::EnumString; |
||||
use syn::{ |
||||
parse_macro_input, punctuated::Punctuated, AttributeArgs, FnArg, ItemFn, Lit, NestedMeta, |
||||
ReturnType, Signature, Token, Type, |
||||
}; |
||||
|
||||
#[derive(EnumString, Debug)] |
||||
enum ReqMethod { |
||||
GET, |
||||
POST, |
||||
} |
||||
|
||||
#[derive(Debug)] |
||||
struct Endpoint { |
||||
method: ReqMethod, |
||||
} |
||||
|
||||
#[derive(Debug)] |
||||
struct FnArgs { |
||||
url_param: Option<Type>, |
||||
payload: Option<Type>, |
||||
} |
||||
|
||||
#[proc_macro_attribute] |
||||
pub fn api_route(args: TokenStream, item: TokenStream) -> TokenStream { |
||||
let args: AttributeArgs = parse_macro_input!(args); |
||||
let input: ItemFn = parse_macro_input!(item); |
||||
let Signature { |
||||
ident, |
||||
inputs, |
||||
generics, |
||||
output, |
||||
.. |
||||
} = input.sig; |
||||
let (impl_generics, _, _) = generics.split_for_impl(); |
||||
let FnArgs { url_param, payload } = parse_fn_args(inputs); |
||||
let Endpoint { method } = parse_attr_args(args); |
||||
let url_path = build_url_path(&ident, &url_param); |
||||
let return_ty = match output { |
||||
ReturnType::Type(_, ty) => quote!(#ty), |
||||
ReturnType::Default => quote!(()), |
||||
}; |
||||
let request = match method { |
||||
ReqMethod::GET => build_get(url_path), |
||||
ReqMethod::POST => build_post(url_path, &payload), |
||||
}; |
||||
let url_param = match url_param { |
||||
Some(p) => quote!(, param: #p), |
||||
None => TokenStream2::new(), |
||||
}; |
||||
let payload = match payload { |
||||
Some(p) => quote!(, payload: #p), |
||||
None => TokenStream2::new(), |
||||
}; |
||||
let q = quote! { |
||||
pub async fn #ident #impl_generics( |
||||
&self #url_param #payload |
||||
) -> UResult<#return_ty> { |
||||
let request = { |
||||
#request |
||||
}; |
||||
let response = request.send().await?; |
||||
let content_len = response.content_length(); |
||||
let is_success = match response.error_for_status_ref() { |
||||
Ok(_) => Ok(()), |
||||
Err(e) => Err(UError::from(e)) |
||||
}; |
||||
match is_success { |
||||
Ok(_) => response.json::<BaseMessage<#return_ty>>() |
||||
.await |
||||
.map(|msg| msg.into_inner()) |
||||
.or_else(|e| { |
||||
match content_len { |
||||
Some(0) => Ok(Default::default()), |
||||
_ => Err(UError::from(e)) |
||||
} |
||||
}), |
||||
Err(UError::NetError(err_src, _)) => Err( |
||||
UError::NetError( |
||||
err_src, |
||||
response.text().await.unwrap() |
||||
) |
||||
), |
||||
_ => unreachable!() |
||||
} |
||||
} |
||||
}; |
||||
//eprintln!("#!#! RESULT:\n{}", q);
|
||||
q.into() |
||||
} |
||||
|
||||
fn parse_fn_args(raw: Punctuated<FnArg, Token![,]>) -> FnArgs { |
||||
let mut arg: HashMap<String, Type> = raw |
||||
.into_iter() |
||||
.filter_map(|arg| { |
||||
if let FnArg::Typed(argt) = arg { |
||||
let mut arg_name = String::new(); |
||||
// did you think I won't overplay you? won't destroy?
|
||||
|arg_ident| -> TokenStream { |
||||
let q: TokenStream = quote!(#arg_ident).into(); |
||||
arg_name = parse_macro_input!(q as Ident).to_string(); |
||||
TokenStream::new() |
||||
}(argt.pat); |
||||
if &arg_name != "url_param" && &arg_name != "payload" { |
||||
panic!("Wrong arg name: {}", &arg_name) |
||||
} |
||||
let arg_type = *argt.ty.clone(); |
||||
Some((arg_name, arg_type)) |
||||
} else { |
||||
None |
||||
} |
||||
}) |
||||
.collect(); |
||||
FnArgs { |
||||
url_param: arg.remove("url_param"), |
||||
payload: arg.remove("payload"), |
||||
} |
||||
} |
||||
|
||||
fn build_get(url: TokenStream2) -> TokenStream2 { |
||||
quote! { |
||||
let request = self.build_get(#url); |
||||
request |
||||
} |
||||
} |
||||
|
||||
fn build_post(url: TokenStream2, payload: &Option<Type>) -> TokenStream2 { |
||||
let pld = match payload { |
||||
Some(_) => quote! { |
||||
.json(&payload.as_message()) |
||||
}, |
||||
None => TokenStream2::new(), |
||||
}; |
||||
quote! { |
||||
let request = self.build_post(#url); |
||||
request #pld |
||||
} |
||||
} |
||||
|
||||
fn build_url_path(path: &Ident, url_param: &Option<Type>) -> TokenStream2 { |
||||
let url_param = match url_param { |
||||
Some(_) => quote! { |
||||
+ &opt_to_string(param) |
||||
}, |
||||
None => TokenStream2::new(), |
||||
}; |
||||
quote! { |
||||
&format!( |
||||
"{}/{}", |
||||
stringify!(#path), |
||||
String::new() #url_param |
||||
) |
||||
} |
||||
} |
||||
|
||||
fn parse_attr_args(args: AttributeArgs) -> Endpoint { |
||||
let mut args = args.into_iter(); |
||||
let method = match args.next() { |
||||
Some(method) => match method { |
||||
NestedMeta::Lit(l) => { |
||||
if let Lit::Str(s) = l { |
||||
match ReqMethod::from_str(&s.value()) { |
||||
Ok(v) => v, |
||||
Err(_) => panic!("Unknown method"), |
||||
} |
||||
} else { |
||||
panic!("Method must be a str") |
||||
} |
||||
} |
||||
_ => panic!("Method must be on the first place"), |
||||
}, |
||||
None => panic!("Method required"), |
||||
}; |
||||
Endpoint { method } |
||||
} |
||||
@ -0,0 +1,15 @@ |
||||
/* |
||||
use std::fmt::Display; |
||||
use u_api_proc_macro::api_route; |
||||
|
||||
type UResult<T, E> = Result<T, E>; |
||||
|
||||
struct ClientHandler; |
||||
struct Paths; |
||||
|
||||
#[test] |
||||
fn test1() { |
||||
#[api_route("GET", Uuid)] |
||||
fn list<T: Display>(url_param: T) {} |
||||
} |
||||
*/ |
||||
@ -0,0 +1,146 @@ |
||||
use crate::{ |
||||
cache::JobCache, |
||||
executor::{FutRes, Waiter, DynFut}, |
||||
models::{Agent, AssignedJob, JobMeta, JobType}, |
||||
utils::{CombinedResult, OneOrMany}, |
||||
UError, |
||||
}; |
||||
use guess_host_triple::guess_host_triple; |
||||
use std::collections::HashMap; |
||||
|
||||
pub struct JobBuilder { |
||||
jobs: Waiter, |
||||
} |
||||
|
||||
impl JobBuilder { |
||||
pub fn from_request<J: OneOrMany<AssignedJob>>(job_requests: J) -> CombinedResult<Self> { |
||||
let job_requests = job_requests.into_vec(); |
||||
let mut prepared: Vec<DynFut> = vec![]; |
||||
let mut result = CombinedResult::new(); |
||||
for req in job_requests { |
||||
let job_meta = JobCache::get(&req.job_id); |
||||
if job_meta.is_none() { |
||||
result.err(UError::NoJob(req.job_id)); |
||||
continue; |
||||
} |
||||
let job_meta = job_meta.unwrap(); |
||||
let built_req = (|| { |
||||
Ok(match job_meta.exec_type { |
||||
JobType::Shell => { |
||||
let meta = JobCache::get(&req.job_id).ok_or(UError::NoJob(req.job_id))?; |
||||
let curr_platform = guess_host_triple().unwrap_or("unknown").to_string(); |
||||
if meta.platform != curr_platform { |
||||
return Err(UError::InsuitablePlatform( |
||||
meta.platform.clone(), |
||||
curr_platform, |
||||
)); |
||||
} |
||||
let job = AssignedJob::new(req.job_id, Some(&req)); |
||||
prepared.push(Box::pin(job.run())) |
||||
} |
||||
JobType::Manage => prepared.push(Box::pin(Agent::run())), |
||||
_ => todo!(), |
||||
}) |
||||
})(); |
||||
if let Err(e) = built_req { |
||||
result.err(e) |
||||
} |
||||
} |
||||
result.ok(Self { |
||||
jobs: Waiter::new(prepared), |
||||
}); |
||||
result |
||||
} |
||||
|
||||
pub fn from_meta<J: OneOrMany<JobMeta>>(job_metas: J) -> CombinedResult<Self> { |
||||
let job_requests = job_metas |
||||
.into_vec() |
||||
.into_iter() |
||||
.map(|jm| { |
||||
let j_uid = jm.id; |
||||
JobCache::insert(jm); |
||||
AssignedJob::new(j_uid, None) |
||||
}) |
||||
.collect::<Vec<AssignedJob>>(); |
||||
JobBuilder::from_request(job_requests) |
||||
} |
||||
|
||||
/// Spawn jobs and pop results later
|
||||
pub async fn spawn(mut self) -> Self { |
||||
self.jobs = self.jobs.spawn().await; |
||||
self |
||||
} |
||||
|
||||
/// Spawn jobs and wait for result
|
||||
pub async fn wait(self) -> Vec<FutRes> { |
||||
self.jobs.spawn().await.wait().await |
||||
} |
||||
|
||||
/// Spawn one job and wait for result
|
||||
pub async fn wait_one(self) -> FutRes { |
||||
self.jobs.spawn().await.wait().await.pop().unwrap() |
||||
} |
||||
} |
||||
|
||||
/// Store jobs and get results by name
|
||||
pub struct NamedJobBuilder { |
||||
builder: Option<JobBuilder>, |
||||
job_names: Vec<&'static str>, |
||||
results: HashMap<&'static str, FutRes>, |
||||
} |
||||
|
||||
impl NamedJobBuilder { |
||||
pub fn from_shell<J: OneOrMany<(&'static str, &'static str)>>( |
||||
named_jobs: J, |
||||
) -> CombinedResult<Self> { |
||||
let mut result = CombinedResult::new(); |
||||
let jobs: Vec<(&'static str, JobMeta)> = named_jobs |
||||
.into_vec() |
||||
.into_iter() |
||||
.filter_map( |
||||
|(alias, cmd)| match JobMeta::builder().with_shell(cmd).build() { |
||||
Ok(meta) => Some((alias, meta)), |
||||
Err(e) => { |
||||
result.err(e); |
||||
None |
||||
} |
||||
}, |
||||
) |
||||
.collect(); |
||||
result.ok(Self::from_meta(jobs)); |
||||
result |
||||
} |
||||
|
||||
pub fn from_meta<J: OneOrMany<(&'static str, JobMeta)>>(named_jobs: J) -> Self { |
||||
let mut job_names = vec![]; |
||||
let job_metas: Vec<JobMeta> = named_jobs |
||||
.into_vec() |
||||
.into_iter() |
||||
.map(|(alias, meta)| { |
||||
job_names.push(alias); |
||||
meta |
||||
}) |
||||
.collect(); |
||||
Self { |
||||
builder: Some(JobBuilder::from_meta(job_metas).unwrap_one()), |
||||
job_names, |
||||
results: HashMap::new(), |
||||
} |
||||
} |
||||
|
||||
pub async fn wait(mut self) -> Self { |
||||
let results = self.builder.take().unwrap().wait().await; |
||||
for (name, result) in self.job_names.iter().zip(results.into_iter()) { |
||||
self.results.insert(name, result); |
||||
} |
||||
self |
||||
} |
||||
|
||||
pub fn pop_opt(&mut self, name: &'static str) -> Option<FutRes> { |
||||
self.results.remove(name) |
||||
} |
||||
|
||||
pub fn pop(&mut self, name: &'static str) -> FutRes { |
||||
self.pop_opt(name).unwrap() |
||||
} |
||||
} |
||||
@ -1,43 +1,54 @@ |
||||
use serde::{ |
||||
Serialize, |
||||
Deserialize, |
||||
}; |
||||
use std::{ |
||||
borrow::Cow, |
||||
}; |
||||
use crate::UID; |
||||
use serde::{Deserialize, Serialize}; |
||||
use std::borrow::Cow; |
||||
use uuid::Uuid; |
||||
use crate::{UID}; |
||||
|
||||
//this is only for impl From<T> for Cow<T>
|
||||
pub struct Moo<'cow, T: Clone>(pub Cow<'cow, T>); |
||||
|
||||
pub trait ToMsg: Clone { |
||||
pub trait AsMsg: Clone + Serialize { |
||||
fn as_message<'m>(&'m self) -> BaseMessage<'m, Self> |
||||
where Moo<'m, Self>: From<&'m Self> { |
||||
where |
||||
Moo<'m, Self>: From<&'m Self>, |
||||
{ |
||||
BaseMessage::new(self) |
||||
} |
||||
} |
||||
|
||||
impl<'cow, M: AsMsg> From<M> for Moo<'cow, M> { |
||||
#[inline] |
||||
fn from(obj: M) -> Moo<'cow, M> { |
||||
Moo(Cow::Owned(obj)) |
||||
} |
||||
} |
||||
|
||||
impl<'cow, M: AsMsg> From<&'cow M> for Moo<'cow, M> { |
||||
#[inline] |
||||
fn from(obj: &'cow M) -> Moo<'cow, M> { |
||||
Moo(Cow::Borrowed(obj)) |
||||
} |
||||
} |
||||
|
||||
impl<M: AsMsg> AsMsg for Vec<M> {} |
||||
|
||||
#[derive(Serialize, Deserialize, Debug)] |
||||
pub struct BaseMessage<'cow, I> |
||||
where I: ToMsg { |
||||
pub struct BaseMessage<'cow, I: AsMsg> { |
||||
pub id: Uuid, |
||||
inner: Cow<'cow, I> |
||||
inner: Cow<'cow, I>, |
||||
} |
||||
|
||||
impl<'cow, I> BaseMessage<'cow, I> |
||||
where I: ToMsg |
||||
{ |
||||
impl<'cow, I: AsMsg> BaseMessage<'cow, I> { |
||||
pub fn new<C>(inner: C) -> Self |
||||
where C: Into<Moo<'cow, I>> { |
||||
where |
||||
C: Into<Moo<'cow, I>>, |
||||
{ |
||||
let Moo(inner) = inner.into(); |
||||
Self { |
||||
id: UID.clone(), |
||||
inner |
||||
inner, |
||||
} |
||||
} |
||||
|
||||
pub fn into_inner(self) -> I { |
||||
self.inner.into_owned() |
||||
} |
||||
} |
||||
} |
||||
|
||||
@ -0,0 +1,143 @@ |
||||
use super::JobState; |
||||
use crate::{ |
||||
cache::JobCache, |
||||
models::{schema::*, ExecResult, JobOutput}, |
||||
utils::{systime_to_string, TempFile}, |
||||
UID, |
||||
}; |
||||
use diesel::{Identifiable, Insertable, Queryable}; |
||||
use serde::{Deserialize, Serialize}; |
||||
use std::{fmt, process::Output, string::FromUtf8Error, time::SystemTime}; |
||||
use tokio::process::Command; |
||||
use uuid::Uuid; |
||||
|
||||
#[derive(
|
||||
Serialize, Deserialize, Clone, Debug, Queryable, Identifiable, Insertable, AsChangeset, |
||||
)] |
||||
#[table_name = "results"] |
||||
pub struct AssignedJob { |
||||
pub agent_id: Uuid, |
||||
pub alias: Option<String>, |
||||
pub created: SystemTime, |
||||
pub id: Uuid, |
||||
pub job_id: Uuid, |
||||
pub result: Option<Vec<u8>>, |
||||
pub state: JobState, |
||||
pub retcode: Option<i32>, |
||||
pub updated: SystemTime, |
||||
} |
||||
|
||||
impl fmt::Display for AssignedJob { |
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { |
||||
let mut out = format!("Result: {}", self.id); |
||||
out += &format!("\nAgent: {}", self.agent_id); |
||||
out += &format!("\nJob: {}", self.job_id); |
||||
if self.alias.is_some() { |
||||
out += &format!("\nAlias: {}", self.alias.as_ref().unwrap()); |
||||
} |
||||
out += &format!("\nUpdated: {}", systime_to_string(&self.updated)); |
||||
out += &format!("\nState: {}", self.state); |
||||
if self.state == JobState::Finished { |
||||
if self.retcode.is_some() { |
||||
out += &format!("\nReturn code: {}", self.retcode.unwrap()); |
||||
} |
||||
if self.result.is_some() { |
||||
out += &format!( |
||||
"\nResult: {}", |
||||
String::from_utf8_lossy(self.result.as_ref().unwrap()) |
||||
); |
||||
} |
||||
} |
||||
write!(f, "{}", out) |
||||
} |
||||
} |
||||
|
||||
impl Default for AssignedJob { |
||||
fn default() -> Self { |
||||
Self { |
||||
agent_id: Uuid::nil(), |
||||
alias: None, |
||||
created: SystemTime::now(), |
||||
id: Uuid::new_v4(), |
||||
job_id: Uuid::nil(), |
||||
result: None, |
||||
state: JobState::Queued, |
||||
retcode: None, |
||||
updated: SystemTime::now(), |
||||
} |
||||
} |
||||
} |
||||
|
||||
impl AssignedJob { |
||||
pub async fn run(mut self) -> ExecResult { |
||||
let (argv, _payload) = { |
||||
let meta = JobCache::get(&self.job_id).unwrap(); |
||||
let extracted_payload = meta |
||||
.payload |
||||
.as_ref() |
||||
.and_then(|p| TempFile::write_exec(p).ok()); |
||||
let argv = if let Some(ref p) = &extracted_payload { |
||||
meta.argv.replace("{}", &p.get_path()) |
||||
} else { |
||||
meta.argv.clone() |
||||
}; |
||||
(argv, extracted_payload) |
||||
}; |
||||
let mut split_cmd = shlex::split(&argv).unwrap().into_iter(); |
||||
let cmd = split_cmd.nth(0).unwrap(); |
||||
let args = split_cmd.collect::<Vec<String>>(); |
||||
let cmd_result = Command::new(cmd).args(args).output().await; |
||||
let (data, retcode) = match cmd_result { |
||||
Ok(Output { |
||||
status, |
||||
stdout, |
||||
stderr, |
||||
}) => ( |
||||
JobOutput::new() |
||||
.stdout(stdout) |
||||
.stderr(stderr) |
||||
.into_combined(), |
||||
status.code(), |
||||
), |
||||
Err(e) => ( |
||||
JobOutput::new() |
||||
.stderr(e.to_string().into_bytes()) |
||||
.into_combined(), |
||||
None, |
||||
), |
||||
}; |
||||
self.result = Some(data); |
||||
self.retcode = retcode; |
||||
self.updated = SystemTime::now(); |
||||
self.state = JobState::Finished; |
||||
ExecResult::Assigned(self) |
||||
} |
||||
|
||||
pub fn new(job_id: Uuid, other: Option<&Self>) -> Self { |
||||
Self { |
||||
agent_id: *UID, |
||||
job_id, |
||||
..other.unwrap_or(&Default::default()).clone() |
||||
} |
||||
} |
||||
|
||||
pub fn as_job_output(&self) -> Option<JobOutput> { |
||||
self.result |
||||
.as_ref() |
||||
.and_then(|r| JobOutput::from_combined(r)) |
||||
} |
||||
|
||||
pub fn to_raw_result(&self) -> Vec<u8> { |
||||
match self.result.as_ref() { |
||||
Some(r) => match JobOutput::from_combined(r) { |
||||
Some(o) => o.to_appropriate(), |
||||
None => r.clone(), |
||||
}, |
||||
None => b"No data".to_vec(), |
||||
} |
||||
} |
||||
|
||||
pub fn to_string_result(&self) -> Result<String, FromUtf8Error> { |
||||
String::from_utf8(self.to_raw_result()) |
||||
} |
||||
} |
||||
@ -1,116 +0,0 @@ |
||||
use super::{ExactJob, JobCache, JobMeta, JobOutput, JobState, JobType}; |
||||
use crate::{ |
||||
executor::{DynFut, Waiter}, |
||||
utils::OneOrMany, |
||||
UError, UResult, |
||||
}; |
||||
use guess_host_triple::guess_host_triple; |
||||
use std::{process::Output, time::SystemTime}; |
||||
use tokio::process::Command; |
||||
use uuid::Uuid; |
||||
|
||||
pub struct Job { |
||||
exec_type: JobType, |
||||
payload: Option<Vec<u8>>, |
||||
result: ExactJob, |
||||
} |
||||
|
||||
impl Job { |
||||
fn build(job_meta: &JobMeta, result_id: Uuid) -> UResult<Self> { |
||||
match job_meta.exec_type { |
||||
JobType::Shell => { |
||||
let curr_platform = guess_host_triple().unwrap_or("unknown").to_string(); |
||||
if job_meta.platform != curr_platform { |
||||
return Err(UError::InsuitablePlatform( |
||||
job_meta.platform.clone(), |
||||
curr_platform, |
||||
)); |
||||
} |
||||
let job_meta = job_meta.clone(); |
||||
Ok(Self { |
||||
exec_type: job_meta.exec_type, |
||||
payload: job_meta.payload, |
||||
result: ExactJob::from_meta(job_meta.id.clone(), Some(result_id)), |
||||
}) |
||||
} |
||||
_ => todo!(), |
||||
} |
||||
} |
||||
|
||||
async fn run(mut self) -> UResult<ExactJob> { |
||||
match self.exec_type { |
||||
JobType::Shell => { |
||||
let str_payload = match &self.payload { |
||||
Some(box_payload) => String::from_utf8_lossy(box_payload).into_owned(), |
||||
None => unimplemented!(), |
||||
}; |
||||
let mut cmd_parts = str_payload |
||||
.split(" ") |
||||
.map(String::from) |
||||
.collect::<Vec<String>>() |
||||
.into_iter(); |
||||
let cmd = cmd_parts.nth(0).unwrap(); |
||||
let args = cmd_parts.collect::<Vec<_>>(); |
||||
let cmd_result = Command::new(cmd).args(args).output().await; |
||||
let (data, retcode) = match cmd_result { |
||||
Ok(Output { |
||||
status, |
||||
stdout, |
||||
stderr, |
||||
}) => ( |
||||
JobOutput::new().stdout(stdout).stderr(stderr).multiline(), |
||||
status.code(), |
||||
), |
||||
Err(e) => ( |
||||
UError::JobError(e.to_string()).to_string().into_bytes(), |
||||
None, |
||||
), |
||||
}; |
||||
self.result.result = Some(data); |
||||
self.result.retcode = retcode; |
||||
self.result.updated = SystemTime::now(); |
||||
self.result.state = JobState::Finished; |
||||
} |
||||
_ => todo!(), |
||||
} |
||||
Ok(self.result) |
||||
} |
||||
} |
||||
|
||||
pub fn build_jobs_with_result<J: OneOrMany<ExactJob>>(job_requests: J) -> Waiter { |
||||
let prepared_jobs = job_requests |
||||
.into_vec() |
||||
.into_iter() |
||||
.filter_map(|jr| -> Option<DynFut> { |
||||
let job = { |
||||
let job_meta = JobCache::get(&jr.job_id); |
||||
if job_meta.is_none() { |
||||
Err(UError::NoJob(jr.job_id)) |
||||
} else { |
||||
Job::build(&*job_meta.unwrap(), jr.id) |
||||
} |
||||
}; |
||||
match job { |
||||
Ok(j) => Some(Box::pin(j.run())), |
||||
Err(e) => { |
||||
warn!("Job building error: {}", e); |
||||
None |
||||
} |
||||
} |
||||
}) |
||||
.collect::<Vec<DynFut>>(); |
||||
Waiter::new(prepared_jobs) |
||||
} |
||||
|
||||
pub fn build_jobs<J: OneOrMany<JobMeta>>(job_metas: J) -> Waiter { |
||||
let job_requests = job_metas |
||||
.into_vec() |
||||
.into_iter() |
||||
.map(|jm| { |
||||
let j_uid = jm.id; |
||||
JobCache::insert(jm); |
||||
ExactJob::from_meta(j_uid, None) |
||||
}) |
||||
.collect::<Vec<ExactJob>>(); |
||||
build_jobs_with_result(job_requests) |
||||
} |
||||
@ -1,7 +1,5 @@ |
||||
pub mod builder; |
||||
pub mod cache; |
||||
pub mod assigned; |
||||
pub mod meta; |
||||
pub mod misc; |
||||
pub mod result; |
||||
|
||||
pub use {builder::*, cache::*, meta::*, misc::*, result::*}; |
||||
pub use {assigned::*, meta::*, misc::*}; |
||||
|
||||
@ -1,155 +0,0 @@ |
||||
use super::JobState; |
||||
use crate::{models::schema::*, utils::systime_to_string, UID}; |
||||
use diesel::{Identifiable, Insertable, Queryable}; |
||||
use serde::{Deserialize, Serialize}; |
||||
use std::{fmt, time::SystemTime}; |
||||
use uuid::Uuid; |
||||
|
||||
#[derive(
|
||||
Serialize, Deserialize, Clone, Debug, Queryable, Identifiable, Insertable, AsChangeset, |
||||
)] |
||||
#[table_name = "results"] |
||||
pub struct ExactJob { |
||||
pub agent_id: Uuid, |
||||
pub created: SystemTime, |
||||
pub id: Uuid, |
||||
pub job_id: Uuid, |
||||
pub result: Option<Vec<u8>>, |
||||
pub state: JobState, |
||||
pub retcode: Option<i32>, |
||||
pub updated: SystemTime, |
||||
} |
||||
|
||||
impl Default for ExactJob { |
||||
fn default() -> Self { |
||||
Self { |
||||
agent_id: Uuid::nil(), |
||||
created: SystemTime::now(), |
||||
id: Uuid::new_v4(), |
||||
job_id: Uuid::nil(), |
||||
result: None, |
||||
state: JobState::Queued, |
||||
retcode: None, |
||||
updated: SystemTime::now(), |
||||
} |
||||
} |
||||
} |
||||
|
||||
impl fmt::Display for ExactJob { |
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { |
||||
let mut out = format!("Result {}", self.id); |
||||
out += &format!("\nAgent {}", self.agent_id); |
||||
out += &format!("\nJob: {}", self.job_id); |
||||
out += &format!("\nUpdated: {}", systime_to_string(&self.updated)); |
||||
out += &format!("\nState: {}", self.state); |
||||
if self.state == JobState::Finished { |
||||
if self.retcode.is_some() { |
||||
out += &format!("\nReturn code: {}", self.retcode.unwrap()); |
||||
} |
||||
if self.result.is_some() { |
||||
out += &format!( |
||||
"\nResult: {}", |
||||
String::from_utf8_lossy(self.result.as_ref().unwrap()) |
||||
); |
||||
} |
||||
} |
||||
write!(f, "{}", out) |
||||
} |
||||
} |
||||
|
||||
impl ExactJob { |
||||
pub fn from_meta(job_id: Uuid, result_id: Option<Uuid>) -> Self { |
||||
Self { |
||||
id: result_id.unwrap_or(Uuid::new_v4()), |
||||
agent_id: *UID, |
||||
job_id, |
||||
..Default::default() |
||||
} |
||||
} |
||||
|
||||
//pub fn as_job_output(&self) -> JobOutput {}
|
||||
} |
||||
|
||||
#[cfg(test)] |
||||
mod tests { |
||||
use super::*; |
||||
use crate::{ |
||||
models::jobs::{build_jobs, JobMeta, JobOutput}, |
||||
utils::vec_to_string, |
||||
UResult, |
||||
}; |
||||
|
||||
#[tokio::test] |
||||
async fn test_is_really_async() { |
||||
const SLEEP_SECS: u64 = 1; |
||||
let job = JobMeta::from_shell(format!("sleep {}", SLEEP_SECS)); |
||||
let sleep_jobs: Vec<JobMeta> = (0..50).map(|_| job.clone()).collect(); |
||||
let now = SystemTime::now(); |
||||
build_jobs(sleep_jobs).run_until_complete().await; |
||||
assert!(now.elapsed().unwrap().as_secs() < SLEEP_SECS + 2) |
||||
} |
||||
|
||||
#[tokio::test] |
||||
async fn test_shell_job() -> UResult<()> { |
||||
let job = JobMeta::from_shell("whoami"); |
||||
let job_result = build_jobs(job).run_one_until_complete().await; |
||||
let stdout = JobOutput::from_raw(&job_result.unwrap().result.unwrap()) |
||||
.unwrap() |
||||
.stdout; |
||||
assert_eq!(vec_to_string(&stdout).trim(), "plazmoid"); |
||||
Ok(()) |
||||
} |
||||
|
||||
#[tokio::test] |
||||
async fn test_complex_load() -> UResult<()> { |
||||
const SLEEP_SECS: u64 = 1; |
||||
let now = SystemTime::now(); |
||||
let longest_job = JobMeta::from_shell(format!("sleep {}", SLEEP_SECS)); |
||||
let longest_job = build_jobs(longest_job).spawn().await; |
||||
let ls = build_jobs(JobMeta::from_shell("ls")) |
||||
.run_one_until_complete() |
||||
.await |
||||
.unwrap(); |
||||
assert_eq!(ls.retcode.unwrap(), 0); |
||||
let result = JobOutput::from_raw(&ls.result.unwrap()).unwrap(); |
||||
let folders = String::from_utf8_lossy(&result.stdout); |
||||
let subfolders_jobs: Vec<JobMeta> = folders |
||||
.lines() |
||||
.map(|f| JobMeta::from_shell(format!("ls {}", f))) |
||||
.collect(); |
||||
let ls_subfolders = build_jobs(subfolders_jobs).run_until_complete().await; |
||||
for result in ls_subfolders { |
||||
assert_eq!(result.unwrap().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() -> UResult<()> { |
||||
let job = JobMeta::from_shell("lol_kek_puk"); |
||||
let job_result = build_jobs(job).run_one_until_complete().await.unwrap(); |
||||
let output = JobOutput::from_raw(&job_result.result.unwrap()); |
||||
assert!(output.is_none()); |
||||
assert!(job_result.retcode.is_none()); |
||||
Ok(()) |
||||
} |
||||
} |
||||
@ -1,48 +1,17 @@ |
||||
mod agent; |
||||
pub mod schema; |
||||
pub mod jobs; |
||||
mod result; |
||||
pub mod schema; |
||||
|
||||
pub use crate::{ |
||||
models::{ |
||||
agent::*, |
||||
jobs::*, |
||||
}, |
||||
messaging::*, |
||||
}; |
||||
use crate::messaging::AsMsg; |
||||
pub use crate::models::result::ExecResult; |
||||
pub use crate::models::{agent::*, jobs::*}; |
||||
use uuid::Uuid; |
||||
use std::borrow::Cow; |
||||
|
||||
// with this macro, a type can be used as message (see api)
|
||||
macro_rules! to_message { |
||||
($($type:ty),+) => { $( |
||||
|
||||
impl ToMsg for $type {} |
||||
|
||||
impl<'cow> From<$type> for Moo<'cow, $type> { |
||||
#[inline] |
||||
fn from(obj: $type) -> Moo<'cow, $type> { |
||||
Moo(Cow::Owned(obj)) |
||||
} |
||||
} |
||||
|
||||
impl<'cow> From<&'cow $type> for Moo<'cow, $type> { |
||||
#[inline] |
||||
fn from(obj: &'cow $type) -> Moo<'cow, $type> { |
||||
Moo(Cow::Borrowed(obj)) |
||||
} |
||||
} )+ |
||||
} |
||||
} |
||||
|
||||
to_message!( |
||||
Agent, |
||||
IAgent, |
||||
JobMeta, |
||||
ExactJob, |
||||
String, |
||||
Vec<Agent>, |
||||
Vec<JobMeta>, |
||||
Vec<ExactJob>, |
||||
Vec<Uuid>, |
||||
() |
||||
); |
||||
impl AsMsg for Agent {} |
||||
impl AsMsg for AssignedJob {} |
||||
impl AsMsg for ExecResult {} |
||||
impl AsMsg for JobMeta {} |
||||
impl AsMsg for String {} |
||||
impl AsMsg for Uuid {} |
||||
impl AsMsg for () {} |
||||
|
||||
@ -0,0 +1,8 @@ |
||||
use crate::models::{Agent, AssignedJob}; |
||||
use serde::{Deserialize, Serialize}; |
||||
|
||||
#[derive(Serialize, Deserialize, Clone)] |
||||
pub enum ExecResult { |
||||
Assigned(AssignedJob), |
||||
Agent(Agent), |
||||
} |
||||
@ -0,0 +1,17 @@ |
||||
/* |
||||
use std::fmt::Display; |
||||
use u_api_proc_macro::api_route; |
||||
use uuid::Uuid; |
||||
|
||||
struct Paths; |
||||
struct ClientHandler; |
||||
|
||||
#[test] |
||||
fn test_api_proc_macro() { |
||||
#[api_route("GET", Uuid)] |
||||
fn list<T: Display>(&self, msg: T) -> String {} |
||||
|
||||
#[api_route("POST", Uuid)] |
||||
fn report<T: Display>(&self, msg: T) -> String {} |
||||
} |
||||
*/ |
||||
Binary file not shown.
@ -0,0 +1,5 @@ |
||||
use std::env; |
||||
|
||||
fn main() { |
||||
println!("{}", env::args().nth(1).unwrap_or(String::new())); |
||||
} |
||||
@ -0,0 +1,161 @@ |
||||
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(()) |
||||
} |
||||
@ -0,0 +1,43 @@ |
||||
use u_lib::{models::JobOutput, utils::vec_to_string}; |
||||
|
||||
const STDOUT: &str = "<***STDOUT***>"; |
||||
const STDERR: &str = "<***STDERR***>"; |
||||
|
||||
#[test_case(
|
||||
"lol", |
||||
"kek", |
||||
&format!("{}lol{}kek", STDOUT, STDERR) |
||||
;"stdout stderr" |
||||
)] |
||||
#[test_case(
|
||||
"", |
||||
"kek", |
||||
&format!("{}kek", STDERR) |
||||
;"stderr" |
||||
)] |
||||
fn test_to_combined(stdout: &str, stderr: &str, result: &str) { |
||||
let output = JobOutput::new() |
||||
.stdout(stdout.as_bytes().to_vec()) |
||||
.stderr(stderr.as_bytes().to_vec()); |
||||
assert_eq!(&vec_to_string(&output.into_combined()), result) |
||||
} |
||||
|
||||
#[test_case(
|
||||
&format!("{}lal{}kik", STDOUT, STDERR), |
||||
"lal\nkik" |
||||
;"stdout stderr" |
||||
)] |
||||
#[test_case(
|
||||
&format!("{}qeq", STDOUT), |
||||
"qeq" |
||||
;"stdout" |
||||
)] |
||||
#[test_case(
|
||||
&format!("{}vev", STDERR), |
||||
"vev" |
||||
;"stderr" |
||||
)] |
||||
fn test_from_combined(src: &str, result: &str) { |
||||
let output = JobOutput::from_combined(src.as_bytes()).unwrap(); |
||||
assert_eq!(vec_to_string(&output.to_appropriate()).trim(), result); |
||||
} |
||||
@ -0,0 +1,10 @@ |
||||
#[macro_use] |
||||
extern crate test_case; |
||||
|
||||
#[macro_use] |
||||
extern crate u_lib; |
||||
|
||||
mod jobs { |
||||
mod execution; |
||||
mod output; |
||||
} |
||||
@ -0,0 +1,6 @@ |
||||
-- This file was automatically created by Diesel to setup helper functions |
||||
-- and other internal bookkeeping. This file is safe to edit, any future |
||||
-- changes will be added to existing projects as new migrations. |
||||
|
||||
DROP FUNCTION IF EXISTS diesel_manage_updated_at(_tbl regclass); |
||||
DROP FUNCTION IF EXISTS diesel_set_updated_at(); |
||||
@ -0,0 +1,36 @@ |
||||
-- This file was automatically created by Diesel to setup helper functions |
||||
-- and other internal bookkeeping. This file is safe to edit, any future |
||||
-- changes will be added to existing projects as new migrations. |
||||
|
||||
|
||||
|
||||
|
||||
-- Sets up a trigger for the given table to automatically set a column called |
||||
-- `updated_at` whenever the row is modified (unless `updated_at` was included |
||||
-- in the modified columns) |
||||
-- |
||||
-- # Example |
||||
-- |
||||
-- ```sql |
||||
-- CREATE TABLE users (id SERIAL PRIMARY KEY, updated_at TIMESTAMP NOT NULL DEFAULT NOW()); |
||||
-- |
||||
-- SELECT diesel_manage_updated_at('users'); |
||||
-- ``` |
||||
CREATE OR REPLACE FUNCTION diesel_manage_updated_at(_tbl regclass) RETURNS VOID AS $$ |
||||
BEGIN |
||||
EXECUTE format('CREATE TRIGGER set_updated_at BEFORE UPDATE ON %s |
||||
FOR EACH ROW EXECUTE PROCEDURE diesel_set_updated_at()', _tbl); |
||||
END; |
||||
$$ LANGUAGE plpgsql; |
||||
|
||||
CREATE OR REPLACE FUNCTION diesel_set_updated_at() RETURNS trigger AS $$ |
||||
BEGIN |
||||
IF ( |
||||
NEW IS DISTINCT FROM OLD AND |
||||
NEW.updated_at IS NOT DISTINCT FROM OLD.updated_at |
||||
) THEN |
||||
NEW.updated_at := current_timestamp; |
||||
END IF; |
||||
RETURN NEW; |
||||
END; |
||||
$$ LANGUAGE plpgsql; |
||||
Loading…
Reference in new issue