#[macro_use] extern crate tracing; mod db; mod error; mod handlers; use crate::handlers::{Endpoints, PayloadFlags}; use db::PgRepo; use error::{Error as ServerError, RejResponse}; use std::{convert::Infallible, sync::Arc}; use u_lib::{ config, db::async_pool, messaging::{AsMsg, Reportable}, models::*, types::Id, }; use warp::{ body, log::{custom, Info}, reply::{json, Json, Response}, Filter, Rejection, Reply, }; const DEFAULT_RESPONSE: &str = "null"; fn into_message(msg: M) -> Json { json(&msg) } pub fn init_endpoints( auth_token: &str, db: PgRepo, ) -> impl Filter + Clone { fn make_optional( f: impl Filter + Clone, ) -> impl Filter,), Error = Infallible> + Clone { f.map(Some) .or_else(|_| async { Ok::<(Option,), Infallible>((None,)) }) } let path = |p: &'static str| warp::post().and(warp::path(p)); let create_qs_cfg = || serde_qs::Config::new(1, true); let with_db = { let adb = Arc::new(db); warp::any().map(move || adb.clone()) }; let get_agents = path("get_agents") .and(with_db.clone()) .and(make_optional(warp::path::param::())) .and_then(Endpoints::get_agents) .map(into_message); let upload_jobs = path("upload_jobs") .and(with_db.clone()) .and(body::json::>()) .and_then(Endpoints::upload_jobs) .map(into_message); let get_job = path("get_job") .and(with_db.clone()) .and(warp::path::param::()) .and(make_optional(serde_qs::warp::query::( create_qs_cfg(), ))) .and_then(Endpoints::get_job) .map(into_message); let get_jobs = path("get_jobs") .and(with_db.clone()) .and_then(Endpoints::get_jobs) .map(into_message); let get_assigned_jobs = path("get_assigned_jobs") .and(with_db.clone()) .and(make_optional(warp::path::param::())) .and_then(Endpoints::get_assigned_jobs) .map(into_message); let get_personal_jobs = path("get_personal_jobs") .and(with_db.clone()) .and(warp::path::param::()) .and_then(Endpoints::get_personal_jobs) .map(into_message); let del = path("del") .and(with_db.clone()) .and(warp::path::param::()) .and_then(Endpoints::del) .map(ok); let assign_jobs = path("assign_jobs") .and(with_db.clone()) .and(body::json::>()) .and_then(Endpoints::assign_jobs) .map(into_message); let report = path("report") .and(with_db.clone()) .and(body::json::>()) .and(warp::header("User-Agent")) .and_then(Endpoints::report) .map(ok); let update_agent = path("update_agent") .and(with_db.clone()) .and(body::json::()) .and_then(Endpoints::update_agent) .map(ok); let update_job = path("update_job") .and(with_db.clone()) .and(body::json::()) .and_then(Endpoints::update_job) .map(ok); let update_assigned_job = path("update_result") .and(with_db.clone()) .and(body::json::()) .and_then(Endpoints::update_assigned_job) .map(ok); let get_payloads = path("get_payloads") .and(with_db.clone()) .and_then(Endpoints::get_payloads) .map(into_message); let get_payload = path("get_payload") .and(with_db.clone()) .and(warp::path::param::()) .and(make_optional(serde_qs::warp::query::( create_qs_cfg(), ))) .and_then(Endpoints::get_payload) .map(into_message); let upload_payload = path("upload_payload") .and(with_db.clone()) .and(body::json::()) .and_then(Endpoints::upload_payload) .map(ok); let update_payload = path("update_payload") .and(with_db.clone()) .and(body::json::()) .and_then(Endpoints::update_payload) .map(ok); let ping = path("ping").map(|| DEFAULT_RESPONSE); let auth_token = format!("Bearer {auth_token}",).into_boxed_str(); let auth_header = warp::header::exact("authorization", Box::leak(auth_token)); let auth_zone = (get_agents .or(get_job.clone()) .or(get_jobs.clone()) .or(get_payloads) .or(get_payload) .or(upload_jobs) .or(upload_payload) .or(del) .or(assign_jobs) .or(get_assigned_jobs) .or(update_agent) .or(update_job) .or(update_assigned_job) .or(update_payload) .or(ping)) .and(auth_header); let agent_zone = get_job.or(get_jobs).or(get_personal_jobs).or(report); auth_zone.or(agent_zone) } pub async fn preload_jobs(repo: &PgRepo) -> Result<(), ServerError> { repo.interact(|mut db| { let job_alias = "agent_hello"; let if_job_exists = db.get_job_by_alias(job_alias)?; if if_job_exists.is_none() { let agent_hello = RawJob::default() .with_type(JobType::Init) .with_alias(job_alias) .try_into_job() .unwrap(); db.insert_jobs(&[agent_hello.meta])?; } Ok(()) }) .await } pub async fn serve() -> Result<(), ServerError> { let env = config::DBEnv::load()?; let pool = async_pool(&env); let db = PgRepo::new(pool); preload_jobs(&db).await?; let env = config::AccessEnv::load()?; let routes = init_endpoints(&env.admin_auth_token, db) .recover(handle_rejection) .with(custom(logger)); let server_cert = include_bytes!("../../../certs/server.crt"); let server_key = include_bytes!("../../../certs/server.key"); let ca = include_bytes!("../../../certs/ca.crt"); warp::serve(routes) .tls() .cert(server_cert) .key(server_key) .client_auth_required(ca) .run(([0, 0, 0, 0], config::MASTER_PORT)) .await; Ok(()) } async fn handle_rejection(rej: Rejection) -> Result { let resp = if let Some(err) = rej.find::() { error!("{:?}", err); RejResponse::bad_request(err.to_string()) } else if rej.is_not_found() { RejResponse::not_found("not found placeholder") } else { error!("{:?}", rej); RejResponse::internal() }; Ok(resp.into_response()) } fn logger(info: Info<'_>) { info!(target: "warp", "{raddr} {agent_id} \"{path}\" {status}", raddr = info.remote_addr().unwrap_or(([0, 0, 0, 0], 0).into()), path = info.path(), agent_id = info.user_agent() .map(|id: &str| id.splitn(3, '-') .take(2) .collect::() ) .unwrap_or_else(|| "NO_AGENT_UID".to_string()), status = info.status() ); } fn ok(_: T) -> impl Reply { DEFAULT_RESPONSE } /* #[cfg(test)] mod tests { use super::*; use crate::handlers::Endpoints; use handlers::build_ok; use u_lib::messaging::{AsMsg, BaseMessage, Reportable}; use uuid::Uuid; use warp::test; #[rstest] #[case(Some(Uuid::new_v4()))] #[should_panic] #[case(None)] #[tokio::test] async fn test_get_agent_jobs_unauthorized(#[case] uid: Option) { let mock = Endpoints::faux(); when!(mock.get_agent_jobs).then_return(Ok(build_ok(""))); //mock.expect().with(eq(uid)).returning(|_| Ok(build_ok(""))); test::request() .path(&format!( "/get_agent_jobs/{}", uid.map(|u| u.simple().to_string()).unwrap_or(String::new()) )) .method("GET") .filter(&init_filters("")) .await .unwrap(); } #[tokio::test] async fn test_report_unauth_successful() { let mock = Endpoints::report(); mock.expect() .withf(|msg: &BaseMessage<'_, Vec>| msg.inner_ref()[0] == Reportable::Dummy) .returning(|_| Ok(build_ok(""))); test::request() .path("/report/") .method("POST") .json(&vec![Reportable::Dummy].as_message()) .filter(&init_filters("")) .await .unwrap(); mock.checkpoint(); } } */