fixed integration tests, improve server code

- mount cargo registry in tests container
- perform cargo update on host machine
- add context to db errors
- use db.conn field only in UDB
- handle rejections
- prettier logs
- pass agent_uid in user-agent header
pull/1/head
plazmoid 2 years ago
parent 6fe0e71959
commit 7b59031bfe
  1. 1
      Cargo.toml
  2. 9
      Makefile.toml
  3. 2
      bin/u_agent/src/lib.rs
  4. 2
      bin/u_panel/Cargo.toml
  5. 4
      bin/u_panel/src/argparse.rs
  6. 1
      bin/u_server/Cargo.toml
  7. 165
      bin/u_server/src/db.rs
  8. 59
      bin/u_server/src/error.rs
  9. 22
      bin/u_server/src/errors.rs
  10. 38
      bin/u_server/src/handlers.rs
  11. 72
      bin/u_server/src/u_server.rs
  12. 3
      integration/docker-compose.yml
  13. 2
      integration/integration_tests.py
  14. 2
      integration/tests/fixtures/agent.rs
  15. 4
      integration/tests/helpers/panel.rs
  16. 2
      lib/u_lib/Cargo.toml
  17. 22
      lib/u_lib/src/api.rs
  18. 2
      lib/u_lib/src/models/agent.rs
  19. 4
      lib/u_lib/src/models/jobs/meta.rs
  20. 4
      lib/u_lib/src/utils/platform.rs
  21. 90
      migrations/2020-10-24-111622_create_all/up.sql
  22. 4
      spec.txt

@ -9,6 +9,7 @@ members = [
] ]
[workspace.dependencies] [workspace.dependencies]
anyhow = "1.0.58"
reqwest = { version = "0.11", features = ["json"] } reqwest = { version = "0.11", features = ["json"] }
serde = { version = "1.0", features = ["derive"] } serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0" serde_json = "1.0"

@ -48,8 +48,12 @@ dependencies = ["build_static_libs", "build_frontend"]
command = "${CARGO}" command = "${CARGO}"
args = ["build", "--target", "${TARGET}", "${@}"] args = ["build", "--target", "${TARGET}", "${@}"]
[tasks.cargo_update]
command = "${CARGO}"
args = ["update"]
[tasks.release_tasks] [tasks.release_tasks]
condition = { env = { "PROFILE_OVERRIDE" = "release"} } condition = { env = { PROFILE_OVERRIDE = "release"} }
script = ''' script = '''
BINS=$(ls ./target/${TARGET}/${PROFILE_OVERRIDE}/u_* -1 | grep -v ".d") BINS=$(ls ./target/${TARGET}/${PROFILE_OVERRIDE}/u_* -1 | grep -v ".d")
echo "Stripping..." echo "Stripping..."
@ -59,7 +63,7 @@ upx -9 $BINS
''' '''
[tasks.build] [tasks.build]
dependencies = ["cargo_build", "release_tasks"] dependencies = ["cargo_update", "cargo_build", "release_tasks"]
clear = true clear = true
[tasks.run] [tasks.run]
@ -70,6 +74,7 @@ command = "${CARGO}"
args = ["test", "--target", "${TARGET}", "--lib", "--", "${@}"] args = ["test", "--target", "${TARGET}", "--lib", "--", "${@}"]
[tasks.integration] [tasks.integration]
dependencies = ["cargo_update"]
script = ''' script = '''
[[ ! -d "./target/${TARGET}/${PROFILE_OVERRIDE}" ]] && echo 'No target folder. Build project first' && exit 1 [[ ! -d "./target/${TARGET}/${PROFILE_OVERRIDE}" ]] && echo 'No target folder. Build project first' && exit 1
cd ./integration cd ./integration

@ -115,6 +115,6 @@ pub async fn run_forever() -> ! {
// ErrChan::send(UError::Runtime(e.to_string()), "deeeemon").await // ErrChan::send(UError::Runtime(e.to_string()), "deeeemon").await
// } // }
} }
info!("Startup"); info!("Starting agent {}", get_self_uid());
agent_loop(client).await agent_loop(client).await
} }

@ -9,7 +9,7 @@ edition = "2021"
[dependencies] [dependencies]
actix-cors = "0.6.1" actix-cors = "0.6.1"
actix-web = "4.1" actix-web = "4.1"
anyhow = "1.0.44" anyhow = { workspace = true }
futures-util = "0.3.21" futures-util = "0.3.21"
mime_guess = "2.0.4" mime_guess = "2.0.4"
once_cell = "1.8.0" once_cell = "1.8.0"

@ -78,11 +78,11 @@ fn parse_uuid(src: &str) -> Result<Uuid, String> {
pub async fn process_cmd(client: ClientHandler, args: Args) -> UResult<String> { pub async fn process_cmd(client: ClientHandler, args: Args) -> UResult<String> {
fn to_json<Msg: AsMsg>(data: AnyResult<Msg>) -> String { fn to_json<Msg: AsMsg>(data: AnyResult<Msg>) -> String {
let data = match data { let result = match data {
Ok(r) => PanelResult::Ok(r), Ok(r) => PanelResult::Ok(r),
Err(e) => PanelResult::Err(e.downcast().expect("unknown error type")), Err(e) => PanelResult::Err(e.downcast().expect("unknown error type")),
}; };
serde_json::to_string(&data).unwrap() serde_json::to_string(&result).unwrap()
} }
Ok(match args.cmd { Ok(match args.cmd {

@ -5,6 +5,7 @@ name = "u_server"
version = "0.1.0" version = "0.1.0"
[dependencies] [dependencies]
anyhow = { workspace = true }
diesel = { version = "1.4.5", features = ["postgres", "uuid"] } diesel = { version = "1.4.5", features = ["postgres", "uuid"] }
hyper = "0.14" hyper = "0.14"
once_cell = "1.7.2" once_cell = "1.7.2"

@ -1,4 +1,4 @@
use crate::errors::{Error, SResult}; use crate::error::Error as ServerError;
use diesel::{pg::PgConnection, prelude::*, result::Error as DslError}; use diesel::{pg::PgConnection, prelude::*, result::Error as DslError};
use once_cell::sync::OnceCell; use once_cell::sync::OnceCell;
use serde::Deserialize; use serde::Deserialize;
@ -9,9 +9,10 @@ use u_lib::{
}; };
use uuid::Uuid; use uuid::Uuid;
type Result<T> = std::result::Result<T, ServerError>;
pub struct UDB { pub struct UDB {
pub conn: PgConnection, conn: PgConnection,
__p: (),
} }
static DB: OnceCell<Mutex<UDB>> = OnceCell::new(); static DB: OnceCell<Mutex<UDB>> = OnceCell::new();
@ -34,7 +35,6 @@ impl UDB {
); );
let instance = UDB { let instance = UDB {
conn: PgConnection::establish(&db_url).unwrap(), conn: PgConnection::establish(&db_url).unwrap(),
__p: (),
}; };
Mutex::new(instance) Mutex::new(instance)
}) })
@ -42,69 +42,89 @@ impl UDB {
.unwrap() .unwrap()
} }
pub fn insert_jobs(&self, job_metas: &[JobMeta]) -> SResult<()> { pub fn insert_jobs(&self, job_metas: &[JobMeta]) -> Result<()> {
use schema::jobs; use schema::jobs;
diesel::insert_into(jobs::table) diesel::insert_into(jobs::table)
.values(job_metas) .values(job_metas)
.execute(&self.conn)?; .execute(&self.conn)
.map_err(with_err_ctx("Can't insert jobs"))?;
Ok(()) Ok(())
} }
pub fn get_jobs(&self, uid: Option<Uuid>) -> SResult<Vec<JobMeta>> { pub fn get_jobs(&self, ouid: Option<Uuid>) -> Result<Vec<JobMeta>> {
use schema::jobs; use schema::jobs;
let result = if uid.is_some() {
jobs::table match ouid {
.filter(jobs::id.eq(uid.unwrap())) Some(uid) => jobs::table
.get_results::<JobMeta>(&self.conn)? .filter(jobs::id.eq(uid))
} else { .get_results::<JobMeta>(&self.conn),
jobs::table.load::<JobMeta>(&self.conn)? None => jobs::table.load::<JobMeta>(&self.conn),
}; }
Ok(result) .map_err(with_err_ctx("Can't get exact jobs"))
} }
pub fn find_job_by_alias(&self, alias: &str) -> SResult<JobMeta> { pub fn find_job_by_alias(&self, alias: &str) -> Result<Option<JobMeta>> {
use schema::jobs; use schema::jobs;
let result = jobs::table let result = jobs::table
.filter(jobs::alias.eq(alias)) .filter(jobs::alias.eq(alias))
.first::<JobMeta>(&self.conn)?; .first::<JobMeta>(&self.conn)
.optional()
.map_err(with_err_ctx(format!("Can't find job by alias {alias}")))?;
Ok(result) Ok(result)
} }
pub fn insert_agent(&self, agent: &Agent) -> SResult<()> { pub fn insert_agent(&self, agent: &Agent) -> Result<()> {
use schema::agents; use schema::agents;
diesel::insert_into(agents::table) diesel::insert_into(agents::table)
.values(agent) .values(agent)
.on_conflict(agents::id) .on_conflict(agents::id)
.do_update() .do_update()
.set(agent) .set(agent)
.execute(&self.conn)?; .execute(&self.conn)
.map_err(with_err_ctx(format!("Can't insert agent {agent:x?}")))?;
Ok(())
}
pub fn insert_result(&self, result: &AssignedJob) -> Result<()> {
use schema::results;
diesel::insert_into(results::table)
.values(result)
.execute(&self.conn)
.map_err(with_err_ctx(format!("Can't insert result {result:x?}")))?;
Ok(()) Ok(())
} }
pub fn get_agents(&self, uid: Option<Uuid>) -> SResult<Vec<Agent>> { pub fn get_agents(&self, ouid: Option<Uuid>) -> Result<Vec<Agent>> {
use schema::agents; use schema::agents;
let result = if uid.is_some() {
agents::table match ouid {
.filter(agents::id.eq(uid.unwrap())) Some(uid) => agents::table
.load::<Agent>(&self.conn)? .filter(agents::id.eq(uid))
} else { .load::<Agent>(&self.conn),
agents::table.load::<Agent>(&self.conn)? None => agents::table.load::<Agent>(&self.conn),
}; }
Ok(result) .map_err(with_err_ctx(format!("Can't get agent(s) {ouid:?}")))
} }
pub fn update_job_status(&self, uid: Uuid, status: JobState) -> SResult<()> { pub fn update_job_status(&self, uid: Uuid, status: JobState) -> Result<()> {
use schema::results; use schema::results;
diesel::update(results::table) diesel::update(results::table)
.filter(results::id.eq(uid)) .filter(results::id.eq(uid))
.set(results::state.eq(status)) .set(results::state.eq(status))
.execute(&self.conn)?; .execute(&self.conn)
.map_err(with_err_ctx(format!("Can't update status of job {uid}")))?;
Ok(()) Ok(())
} }
//TODO: filters possibly could work in a wrong way, check //TODO: filters possibly could work in a wrong way, check
pub fn get_exact_jobs(&self, uid: Option<Uuid>, personal: bool) -> SResult<Vec<AssignedJob>> { pub fn get_exact_jobs(&self, uid: Option<Uuid>, personal: bool) -> Result<Vec<AssignedJob>> {
use schema::results; use schema::results;
let mut q = results::table.into_boxed(); let mut q = results::table.into_boxed();
/*if uid.is_some() { /*if uid.is_some() {
q = q.filter(results::agent_id.eq(uid.unwrap())) q = q.filter(results::agent_id.eq(uid.unwrap()))
@ -121,32 +141,19 @@ impl UDB {
.or_filter(results::job_id.eq(uid.unwrap())) .or_filter(results::job_id.eq(uid.unwrap()))
.or_filter(results::id.eq(uid.unwrap())) .or_filter(results::id.eq(uid.unwrap()))
} }
let result = q.load::<AssignedJob>(&self.conn)?; let result = q
.load::<AssignedJob>(&self.conn)
.map_err(with_err_ctx("Can't get exact jobs"))?;
Ok(result) Ok(result)
} }
pub fn set_jobs_for_agent(&self, agent_uid: &Uuid, job_uids: &[Uuid]) -> SResult<Vec<Uuid>> { pub fn set_jobs_for_agent(&self, agent_uid: &Uuid, job_uids: &[Uuid]) -> Result<Vec<Uuid>> {
use schema::{agents::dsl::agents, jobs::dsl::jobs, results}; use schema::results;
if let Err(DslError::NotFound) = agents.find(agent_uid).first::<Agent>(&self.conn) {
return Err(Error::NotFound(agent_uid.to_string()));
}
let not_found_jobs = job_uids
.iter()
.filter_map(|job_uid| {
if let Err(DslError::NotFound) = jobs.find(job_uid).first::<JobMeta>(&self.conn) {
Some(job_uid.to_string())
} else {
None
}
})
.collect::<Vec<String>>();
if !not_found_jobs.is_empty() {
return Err(Error::NotFound(not_found_jobs.join(", ")));
}
let job_requests = job_uids let job_requests = job_uids
.iter() .iter()
.map(|job_uid| { .map(|job_uid| {
info!("set_jobs_for_agent: set {} for {}", job_uid, agent_uid); debug!("set_jobs_for_agent: set {} for {}", job_uid, agent_uid);
AssignedJob { AssignedJob {
job_id: *job_uid, job_id: *job_uid,
agent_id: *agent_uid, agent_id: *agent_uid,
@ -154,46 +161,84 @@ impl UDB {
} }
}) })
.collect::<Vec<AssignedJob>>(); .collect::<Vec<AssignedJob>>();
diesel::insert_into(results::table) diesel::insert_into(results::table)
.values(&job_requests) .values(&job_requests)
.execute(&self.conn)?; .execute(&self.conn)
let assigned_uids = job_requests.iter().map(|aj| aj.id).collect(); .map_err(with_err_ctx(format!(
Ok(assigned_uids) "Can't setup jobs {job_uids:?} for agent {agent_uid:?}"
)))?;
Ok(job_requests.iter().map(|aj| aj.id).collect())
} }
pub fn del_jobs(&self, uids: &[Uuid]) -> SResult<usize> { pub fn del_jobs(&self, uids: &[Uuid]) -> Result<usize> {
use schema::jobs; use schema::jobs;
let mut affected = 0; let mut affected = 0;
for &uid in uids { for &uid in uids {
let deleted = diesel::delete(jobs::table) let deleted = diesel::delete(jobs::table)
.filter(jobs::id.eq(uid)) .filter(jobs::id.eq(uid))
.execute(&self.conn)?; .execute(&self.conn)
.map_err(with_err_ctx("Can't delete jobs"))?;
affected += deleted; affected += deleted;
} }
Ok(affected) Ok(affected)
} }
pub fn del_results(&self, uids: &[Uuid]) -> SResult<usize> { pub fn del_results(&self, uids: &[Uuid]) -> Result<usize> {
use schema::results; use schema::results;
let mut affected = 0; let mut affected = 0;
for &uid in uids { for &uid in uids {
let deleted = diesel::delete(results::table) let deleted = diesel::delete(results::table)
.filter(results::id.eq(uid)) .filter(results::id.eq(uid))
.execute(&self.conn)?; .execute(&self.conn)
.map_err(with_err_ctx("Can't delete results"))?;
affected += deleted; affected += deleted;
} }
Ok(affected) Ok(affected)
} }
pub fn del_agents(&self, uids: &[Uuid]) -> SResult<usize> { pub fn del_agents(&self, uids: &[Uuid]) -> Result<usize> {
use schema::agents; use schema::agents;
let mut affected = 0; let mut affected = 0;
for &uid in uids { for &uid in uids {
let deleted = diesel::delete(agents::table) let deleted = diesel::delete(agents::table)
.filter(agents::id.eq(uid)) .filter(agents::id.eq(uid))
.execute(&self.conn)?; .execute(&self.conn)
.map_err(with_err_ctx("Can't delete agents"))?;
affected += deleted; affected += deleted;
} }
Ok(affected) Ok(affected)
} }
pub fn update_agent(&self, agent: &Agent) -> Result<()> {
agent
.save_changes::<Agent>(&self.conn)
.map_err(with_err_ctx(format!("Can't update agent {agent:x?}")))?;
Ok(())
}
pub fn update_job(&self, job: &JobMeta) -> Result<()> {
job.save_changes::<JobMeta>(&self.conn)
.map_err(with_err_ctx(format!("Can't update job {job:x?}")))?;
Ok(())
}
pub fn update_result(&self, result: &AssignedJob) -> Result<()> {
debug!(
"updating result: id = {}, job_id = {}, agent_id = {}",
result.id, result.job_id, result.agent_id
);
result
.save_changes::<AssignedJob>(&self.conn)
.map_err(with_err_ctx(format!("Can't update result {result:x?}")))?;
Ok(())
}
}
fn with_err_ctx(msg: impl AsRef<str>) -> impl Fn(DslError) -> ServerError {
move |err| ServerError::DBErrorCtx(format!("{}, reason: {err}", msg.as_ref()))
} }

@ -0,0 +1,59 @@
use diesel::result::Error as DslError;
use thiserror::Error;
use warp::{
http::StatusCode,
reject::Reject,
reply::{with_status, Response},
Reply,
};
#[derive(Error, Debug)]
pub enum Error {
#[error("Error processing {0}")]
ProcessingError(String),
#[error(transparent)]
DBError(#[from] DslError),
#[error("DB error: {0}")]
DBErrorCtx(String),
#[error("General error: {0}")]
Other(String),
}
impl Reject for Error {}
pub struct RejResponse {
message: String,
status: StatusCode,
}
impl RejResponse {
pub fn not_found(msg: impl Into<String>) -> Self {
Self {
message: msg.into(),
status: StatusCode::NOT_FOUND,
}
}
pub fn bad_request(msg: impl Into<String>) -> Self {
Self {
message: msg.into(),
status: StatusCode::BAD_REQUEST,
}
}
pub fn internal() -> Self {
Self {
message: "INTERNAL_SERVER_ERROR".to_string(),
status: StatusCode::INTERNAL_SERVER_ERROR,
}
}
}
impl Reply for RejResponse {
fn into_response(self) -> Response {
with_status(self.message, self.status).into_response()
}
}

@ -1,22 +0,0 @@
use diesel::result::Error as DslError;
use thiserror::Error;
use warp::reject::Reject;
pub type SResult<T> = Result<T, Error>;
#[derive(Error, Debug)]
pub enum Error {
#[error("{0} is not found")]
NotFound(String),
#[error("Error processing {0}")]
ProcessingError(String),
#[error(transparent)]
DBError(#[from] DslError),
#[error("General error: {0}")]
Other(String),
}
impl Reject for Error {}

@ -1,8 +1,7 @@
use std::time::SystemTime; use std::time::SystemTime;
use crate::db::UDB; use crate::db::UDB;
use crate::errors::Error; use crate::error::Error;
use diesel::SaveChangesDsl;
use u_lib::{ use u_lib::{
messaging::{AsMsg, BaseMessage, Reportable}, messaging::{AsMsg, BaseMessage, Reportable},
models::*, models::*,
@ -39,7 +38,9 @@ impl Endpoints {
if agents.is_empty() { if agents.is_empty() {
let db = UDB::lock_db(); let db = UDB::lock_db();
db.insert_agent(&Agent::with_id(uid.unwrap()))?; db.insert_agent(&Agent::with_id(uid.unwrap()))?;
let job = db.find_job_by_alias("agent_hello")?; let job = db
.find_job_by_alias("agent_hello")?
.expect("agent_hello job not found");
db.set_jobs_for_agent(&uid.unwrap(), &[job.id])?; db.set_jobs_for_agent(&uid.unwrap(), &[job.id])?;
} }
let result = UDB::lock_db().get_exact_jobs(uid, true)?; let result = UDB::lock_db().get_exact_jobs(uid, true)?;
@ -76,8 +77,16 @@ impl Endpoints {
msg.into_inner() msg.into_inner()
.into_iter() .into_iter()
.map(|ident| { .map(|ident| {
Uuid::parse_str(&ident) Uuid::parse_str(&ident).or_else(|_| {
.or_else(|_| UDB::lock_db().find_job_by_alias(&ident).map(|j| j.id)) let job_from_db = UDB::lock_db().find_job_by_alias(&ident);
match job_from_db {
Ok(job) => match job {
Some(j) => Ok(j.id),
None => Err(Error::ProcessingError(format!("unknown ident {ident}"))),
},
Err(e) => Err(e),
}
})
}) })
.collect::<Result<Vec<Uuid>, Error>>() .collect::<Result<Vec<Uuid>, Error>>()
.and_then(|j| UDB::lock_db().set_jobs_for_agent(&agent_uid, &j)) .and_then(|j| UDB::lock_db().set_jobs_for_agent(&agent_uid, &j))
@ -117,10 +126,7 @@ impl Endpoints {
JobType::Terminate => todo!(), JobType::Terminate => todo!(),
JobType::Update => todo!(), JobType::Update => todo!(),
} }
let db = UDB::lock_db(); UDB::lock_db().update_result(&result)?;
result
.save_changes::<AssignedJob>(&db.conn)
.map_err(Error::from)?;
} }
Reportable::Error(e) => { Reportable::Error(e) => {
warn!("{} reported an error: {}", id, e); warn!("{} reported an error: {}", id, e);
@ -132,27 +138,19 @@ impl Endpoints {
} }
pub async fn update_agent(agent: BaseMessage<'static, Agent>) -> EndpResult<()> { pub async fn update_agent(agent: BaseMessage<'static, Agent>) -> EndpResult<()> {
agent UDB::lock_db().update_agent(&agent.into_inner())?;
.into_inner()
.save_changes::<Agent>(&UDB::lock_db().conn)
.map_err(Error::from)?;
Ok(()) Ok(())
} }
pub async fn update_job(job: BaseMessage<'static, JobMeta>) -> EndpResult<()> { pub async fn update_job(job: BaseMessage<'static, JobMeta>) -> EndpResult<()> {
job.into_inner() UDB::lock_db().update_job(&job.into_inner())?;
.save_changes::<JobMeta>(&UDB::lock_db().conn)
.map_err(Error::from)?;
Ok(()) Ok(())
} }
pub async fn update_assigned_job( pub async fn update_assigned_job(
assigned: BaseMessage<'static, AssignedJob>, assigned: BaseMessage<'static, AssignedJob>,
) -> EndpResult<()> { ) -> EndpResult<()> {
assigned UDB::lock_db().update_result(&assigned.into_inner())?;
.into_inner()
.save_changes::<AssignedJob>(&UDB::lock_db().conn)
.map_err(Error::from)?;
Ok(()) Ok(())
} }

@ -12,12 +12,12 @@ extern crate diesel;
// in this block // in this block
mod db; mod db;
mod errors; mod error;
mod handlers; mod handlers;
use errors::{Error, SResult}; use error::{Error as ServerError, RejResponse};
use serde::{de::DeserializeOwned, Deserialize}; use serde::{de::DeserializeOwned, Deserialize};
use std::path::PathBuf; use std::{convert::Infallible, path::PathBuf};
use u_lib::{ use u_lib::{
config::MASTER_PORT, config::MASTER_PORT,
logging::init_logger, logging::init_logger,
@ -28,7 +28,8 @@ use u_lib::{
use uuid::Uuid; use uuid::Uuid;
use warp::{ use warp::{
body, body,
reply::{json, reply, Json}, log::{custom, Info},
reply::{json, reply, Json, Response},
Filter, Rejection, Reply, Filter, Rejection, Reply,
}; };
@ -55,7 +56,7 @@ pub fn init_endpoints(
auth_token: &str, auth_token: &str,
) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone { ) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone {
let path = |p: &'static str| warp::post().and(warp::path(p)); let path = |p: &'static str| warp::post().and(warp::path(p));
let infallible_none = |_| async { Ok::<(Option<Uuid>,), std::convert::Infallible>((None,)) }; let infallible_none = |_| async { Ok::<_, Infallible>((None::<Uuid>,)) };
let get_agents = path("get_agents") let get_agents = path("get_agents")
.and( .and(
@ -131,7 +132,7 @@ pub fn init_endpoints(
.map(ok); .map(ok);
let auth_token = format!("Bearer {auth_token}",).into_boxed_str(); let auth_token = format!("Bearer {auth_token}",).into_boxed_str();
let auth_header = warp::header::exact("Authorization", Box::leak(auth_token)); let auth_header = warp::header::exact("authorization", Box::leak(auth_token));
let auth_zone = (get_agents let auth_zone = (get_agents
.or(get_jobs) .or(get_jobs)
@ -150,32 +151,31 @@ pub fn init_endpoints(
auth_zone.or(agent_zone) auth_zone.or(agent_zone)
} }
pub fn prefill_jobs() -> SResult<()> { pub fn prefill_jobs() -> Result<(), ServerError> {
let job_alias = "agent_hello"; let job_alias = "agent_hello";
let if_job_exists = UDB::lock_db().find_job_by_alias(job_alias); let if_job_exists = UDB::lock_db().find_job_by_alias(job_alias)?;
match if_job_exists { if if_job_exists.is_none() {
Ok(_) => Ok(()), let agent_hello = JobMeta::builder()
Err(Error::DBError(diesel::result::Error::NotFound)) => { .with_type(JobType::Init)
let agent_hello = JobMeta::builder() .with_alias(job_alias)
.with_type(JobType::Init) .build()
.with_alias(job_alias) .unwrap();
.build() UDB::lock_db().insert_jobs(&[agent_hello])?
.unwrap();
UDB::lock_db().insert_jobs(&[agent_hello])
}
Err(e) => Err(e),
} }
Ok(())
} }
pub async fn serve() -> SResult<()> { pub async fn serve() -> Result<(), ServerError> {
init_logger(Some("u_server")); init_logger(Some("u_server"));
prefill_jobs()?; prefill_jobs()?;
let env = load_env::<ServEnv>().map_err(|e| Error::Other(e.to_string()))?;
let routes = init_endpoints(&env.admin_auth_token);
let certs_dir = PathBuf::from("certs"); let certs_dir = PathBuf::from("certs");
let env = load_env::<ServEnv>().map_err(|e| ServerError::Other(e.to_string()))?;
let routes = init_endpoints(&env.admin_auth_token)
.recover(handle_rejection)
.with(custom(logger));
warp::serve(routes.with(warp::log("warp"))) warp::serve(routes)
.tls() .tls()
.cert_path(certs_dir.join("server.crt")) .cert_path(certs_dir.join("server.crt"))
.key_path(certs_dir.join("server.key")) .key_path(certs_dir.join("server.key"))
@ -185,6 +185,32 @@ pub async fn serve() -> SResult<()> {
Ok(()) Ok(())
} }
async fn handle_rejection(rej: Rejection) -> Result<Response, Infallible> {
let resp = if let Some(err) = rej.find::<ServerError>() {
error!("{:x?}", err);
RejResponse::bad_request(err.to_string())
} else if rej.is_not_found() {
RejResponse::not_found("not found placeholder")
} else {
RejResponse::internal()
};
Ok(resp.into_response())
}
fn logger(info: Info<'_>) {
info!(target: "warp",
"{raddr} {agent_uid} \"{path}\"",
raddr = info.remote_addr().unwrap_or(([0, 0, 0, 0], 0).into()),
path = info.path(),
agent_uid = info.user_agent()
.map(|uid: &str| uid.splitn(3, '-')
.take(2)
.collect::<String>()
)
.unwrap_or_else(|| "NO_AGENT".to_string())
);
}
fn ok<T>(_: T) -> impl Reply { fn ok<T>(_: T) -> impl Reply {
reply() reply()
} }

@ -29,7 +29,7 @@ services:
- ../.env - ../.env
- ../.env.private - ../.env.private
environment: environment:
RUST_LOG: info RUST_LOG: warp=info,u_server_lib=debug
healthcheck: healthcheck:
test: ss -tlpn | grep 63714 test: ss -tlpn | grep 63714
interval: 5s interval: 5s
@ -80,6 +80,7 @@ services:
networks: networks:
- u_net - u_net
volumes: volumes:
- ${HOME}/.cargo/registry/:/usr/local/cargo/registry/
- ../__Cargo_integration.toml:/tests/Cargo.toml - ../__Cargo_integration.toml:/tests/Cargo.toml
- ./:/tests/integration/ - ./:/tests/integration/
- ../certs:/tests/certs - ../certs:/tests/certs

@ -59,7 +59,7 @@ def run_tests():
if not only_setup_cluster: if not only_setup_cluster:
CLUSTER.run('cargo test --test integration') CLUSTER.run('cargo test --test integration')
except Exception as e: except Exception as e:
# CLUSTER.print_containers_logs() CLUSTER.print_containers_logs()
fail(e) fail(e)
finally: finally:
_cleanup() _cleanup()

@ -17,6 +17,7 @@ impl RegisteredAgent {
pub async fn register_agent() -> RegisteredAgent { pub async fn register_agent() -> RegisteredAgent {
let cli = ClientHandler::new(&ENV.u_server, None); let cli = ClientHandler::new(&ENV.u_server, None);
let agent_uid = Uuid::new_v4(); let agent_uid = Uuid::new_v4();
println!("registering agent {agent_uid}");
let resp = cli let resp = cli
.get_personal_jobs(agent_uid) .get_personal_jobs(agent_uid)
.await .await
@ -27,6 +28,7 @@ pub async fn register_agent() -> RegisteredAgent {
let job = cli.get_jobs(Some(job_id)).await.unwrap().pop().unwrap(); let job = cli.get_jobs(Some(job_id)).await.unwrap().pop().unwrap();
assert_eq!(job.alias, Some("agent_hello".to_string())); assert_eq!(job.alias, Some("agent_hello".to_string()));
let mut agent_data = AssignedJob::from(&job); let mut agent_data = AssignedJob::from(&job);
agent_data.agent_id = agent_uid;
agent_data.set_result(&Agent { agent_data.set_result(&Agent {
id: agent_uid, id: agent_uid,
..Default::default() ..Default::default()

@ -43,8 +43,8 @@ impl Panel {
.as_ref(), .as_ref(),
); );
match &result { match &result {
PanelResult::Ok(r) => eprintln!("<<<+ {r:?}"), PanelResult::Ok(r) => eprintln!("<<<+ {r:02x?}"),
PanelResult::Err(e) => eprintln!("<<<! {e:?}"), PanelResult::Err(e) => eprintln!("<<<! {e:02x?}"),
} }
result result
} }

@ -7,7 +7,7 @@ edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies] [dependencies]
anyhow = "1.0.58" anyhow = { workspace = true }
chrono = "0.4.19" chrono = "0.4.19"
diesel = { version = "1.4.5", features = ["postgres", "uuid"], optional = true } diesel = { version = "1.4.5", features = ["postgres", "uuid"], optional = true }
diesel-derive-enum = { version = "1", features = ["postgres"], optional = true } diesel-derive-enum = { version = "1", features = ["postgres"], optional = true }

@ -2,7 +2,7 @@ use std::collections::HashMap;
use std::fmt::Debug; use std::fmt::Debug;
use crate::{ use crate::{
config::MASTER_PORT, config::{get_self_uid, MASTER_PORT},
messaging::{self, AsMsg, BaseMessage, Empty}, messaging::{self, AsMsg, BaseMessage, Empty},
models::{self}, models::{self},
utils::{opt_to_string, OneOrVec}, utils::{opt_to_string, OneOrVec},
@ -26,20 +26,22 @@ pub struct ClientHandler {
impl ClientHandler { impl ClientHandler {
pub fn new(server: &str, password: Option<String>) -> Self { pub fn new(server: &str, password: Option<String>) -> Self {
let identity = Identity::from_pkcs12_der(AGENT_IDENTITY, "").unwrap(); let identity = Identity::from_pkcs12_der(AGENT_IDENTITY, "").unwrap();
let mut client = Client::builder().identity(identity); let mut default_headers = HashMap::from([(
"user-agent".to_string(),
get_self_uid().hyphenated().to_string(),
)]);
if let Some(pwd) = password { if let Some(pwd) = password {
client = client.default_headers( default_headers.insert("authorization".to_string(), format!("Bearer {pwd}"));
HeaderMap::try_from(&HashMap::from([(
"Authorization".to_string(),
format!("Bearer {pwd}"),
)]))
.unwrap(),
)
} }
let client = client
let client = Client::builder()
.identity(identity)
.default_headers(HeaderMap::try_from(&default_headers).unwrap())
.add_root_certificate(Certificate::from_pem(ROOT_CA_CERT).unwrap()) .add_root_certificate(Certificate::from_pem(ROOT_CA_CERT).unwrap())
.build() .build()
.unwrap(); .unwrap();
Self { Self {
client, client,
base_url: Url::parse(&format!("https://{}:{}", server, MASTER_PORT)).unwrap(), base_url: Url::parse(&format!("https://{}:{}", server, MASTER_PORT)).unwrap(),

@ -96,7 +96,7 @@ impl Agent {
hostname: decoder(builder.pop("hostname")), hostname: decoder(builder.pop("hostname")),
is_root: &decoder(builder.pop("is_root")) == "0", is_root: &decoder(builder.pop("is_root")) == "0",
username: decoder(builder.pop("username")), username: decoder(builder.pop("username")),
platform: Platform::current().into_string(), platform: Platform::current_as_string(),
..Default::default() ..Default::default()
} }
} }

@ -31,7 +31,7 @@ pub struct JobMeta {
pub exec_type: JobType, pub exec_type: JobType,
///target triple ///target triple
#[serde(default)] #[serde(default = "Platform::current_as_string")]
pub platform: String, pub platform: String,
#[serde(default)] #[serde(default)]
@ -67,7 +67,7 @@ impl Default for JobMeta {
alias: None, alias: None,
argv: String::new(), argv: String::new(),
exec_type: JobType::Shell, exec_type: JobType::Shell,
platform: Platform::current().into_string(), platform: Platform::current_as_string(),
payload: None, payload: None,
schedule: None, schedule: None,
payload_path: None, payload_path: None,

@ -15,6 +15,10 @@ impl Platform {
Self(guess_host_triple().unwrap_or("unknown").to_string()) Self(guess_host_triple().unwrap_or("unknown").to_string())
} }
pub fn current_as_string() -> String {
Self::current().into_string()
}
pub fn matches(&self, pf: impl AsRef<str>) -> bool { pub fn matches(&self, pf: impl AsRef<str>) -> bool {
match PlatformReq::from_str(pf.as_ref()) { match PlatformReq::from_str(pf.as_ref()) {
Ok(p) => p.matches(&_Platform::find(&self.0).unwrap()), Ok(p) => p.matches(&_Platform::find(&self.0).unwrap()),

@ -4,60 +4,58 @@ CREATE TYPE JobState AS ENUM ('queued', 'running', 'finished');
CREATE TYPE AgentState AS ENUM ('new', 'active', 'banned'); CREATE TYPE AgentState AS ENUM ('new', 'active', 'banned');
CREATE TABLE IF NOT EXISTS agents ( CREATE TABLE IF NOT EXISTS agents (
alias TEXT alias TEXT,
, hostname TEXT NOT NULL hostname TEXT NOT NULL,
, id UUID NOT NULL DEFAULT uuid_generate_v4() id UUID NOT NULL DEFAULT uuid_generate_v4(),
, ip_gray TEXT ip_gray TEXT,
, ip_white TEXT ip_white TEXT,
, is_root BOOLEAN NOT NULL DEFAULT false is_root BOOLEAN NOT NULL DEFAULT false,
, is_root_allowed BOOLEAN NOT NULL DEFAULT false is_root_allowed BOOLEAN NOT NULL DEFAULT false,
, last_active TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP last_active TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
-- target triplet platform TEXT NOT NULL,
, platform TEXT NOT NULL regtime TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
, regtime TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP state AgentState NOT NULL DEFAULT 'new',
, state AgentState NOT NULL DEFAULT 'new' token TEXT,
-- is needed to processing requests username TEXT NOT NULL,
, token TEXT
, username TEXT NOT NULL PRIMARY KEY(id)
, PRIMARY KEY(id)
); );
CREATE TABLE IF NOT EXISTS jobs ( CREATE TABLE IF NOT EXISTS jobs (
alias TEXT alias TEXT,
, argv TEXT NOT NULL argv TEXT NOT NULL,
, id UUID NOT NULL DEFAULT uuid_generate_v4() id UUID NOT NULL DEFAULT uuid_generate_v4(),
, exec_type JobType NOT NULL DEFAULT 'shell' exec_type JobType NOT NULL DEFAULT 'shell',
, platform TEXT NOT NULL platform TEXT NOT NULL,
, payload BYTEA payload BYTEA,
, payload_path TEXT payload_path TEXT,
, schedule TEXT schedule TEXT,
, PRIMARY KEY(id) PRIMARY KEY(id)
); );
CREATE TABLE IF NOT EXISTS results ( CREATE TABLE IF NOT EXISTS results (
agent_id UUID NOT NULL agent_id UUID NOT NULL,
, alias TEXT alias TEXT,
, created TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP created TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
, id UUID NOT NULL DEFAULT uuid_generate_v4() id UUID NOT NULL DEFAULT uuid_generate_v4(),
, job_id UUID NOT NULL job_id UUID NOT NULL,
, result BYTEA result BYTEA,
, state JobState NOT NULL DEFAULT 'queued' state JobState NOT NULL DEFAULT 'queued',
, exec_type JobType NOT NULL DEFAULT 'shell' exec_type JobType NOT NULL DEFAULT 'shell',
, retcode INTEGER retcode INTEGER,
, updated TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP updated TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
, FOREIGN KEY(agent_id) REFERENCES agents(id) ON DELETE CASCADE FOREIGN KEY(agent_id) REFERENCES agents(id) ON DELETE CASCADE,
, FOREIGN KEY(job_id) REFERENCES jobs(id) ON DELETE CASCADE FOREIGN KEY(job_id) REFERENCES jobs(id) ON DELETE CASCADE,
, PRIMARY KEY(id) PRIMARY KEY(id)
); );
CREATE TABLE IF NOT EXISTS certificates ( CREATE TABLE IF NOT EXISTS certificates (
agent_id UUID NOT NULL agent_id UUID NOT NULL,
, id UUID NOT NULL DEFAULT uuid_generate_v4() id UUID NOT NULL DEFAULT uuid_generate_v4(),
, is_revoked BOOLEAN NOT NULL DEFAULT FALSE is_revoked BOOLEAN NOT NULL DEFAULT FALSE,
, PRIMARY KEY(id) PRIMARY KEY(id),
, FOREIGN KEY(agent_id) REFERENCES agents(id) FOREIGN KEY(agent_id) REFERENCES agents(id)
); );

@ -1,7 +1,9 @@
todos:
Upload/download files Upload/download files
More tests More tests
Agent update (use more JobType's) Agent update (use more JobType's)
Erase log macros in release mode Erase log macros in release mode
Bump wine version to test agent on windows Bump wine version to test agent on windows
Store downloaded payload on disk instead of memory Store downloaded payload on disk instead of ram
Improve web interface Improve web interface
Migrator binary
Loading…
Cancel
Save