unki/bin/u_server/src/u_server.rs

#[macro_use]
extern crate tracing;

#[cfg(test)]
#[macro_use]
extern crate rstest;

// due to linking errors
extern crate openssl;
// don't touch anything
extern crate diesel;
// in this block

mod db;
mod errors;
mod handlers;

use errors::{Error, SResult};
use serde::{de::DeserializeOwned, Deserialize};
use std::path::PathBuf;
use u_lib::{
    config::MASTER_PORT,
    logging::init_logger,
    messaging::{AsMsg, BaseMessage, Reportable},
    models::*,
    utils::load_env,
};
use uuid::Uuid;
use warp::{
    body,
    reply::{json, reply, Json},
    Filter, Rejection, Reply,
};

use crate::db::UDB;
use crate::handlers::Endpoints;

#[derive(Deserialize)]
struct ServEnv {
    admin_auth_token: String,
}

fn get_content<M>() -> impl Filter<Extract = (BaseMessage<'static, M>,), Error = Rejection> + Clone
where
    M: AsMsg + Sync + Send + DeserializeOwned + 'static,
{
    body::content_length_limit(1024 * 64).and(body::json::<BaseMessage<M>>())
}

fn into_message<M: AsMsg>(msg: M) -> Json {
    json(&msg.as_message())
}

pub fn init_endpoints(
    auth_token: &str,
) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone {
    let path = |p: &'static str| warp::post().and(warp::path(p));
    let infallible_none = |_| async { Ok::<(Option<Uuid>,), std::convert::Infallible>((None,)) };

    let get_agents = path("get_agents")
        .and(
            warp::path::param::<Uuid>()
                .map(Some)
                .or_else(infallible_none),
        )
        .and_then(Endpoints::get_agents)
        .map(into_message);

    let upload_jobs = path("upload_jobs")
        .and(get_content::<Vec<JobMeta>>())
        .and_then(Endpoints::upload_jobs)
        .map(ok);

    let get_jobs = path("get_jobs")
        .and(
            warp::path::param::<Uuid>()
                .map(Some)
                .or_else(infallible_none),
        )
        .and_then(Endpoints::get_jobs)
        .map(into_message);

    let get_agent_jobs = path("get_agent_jobs")
        .and(
            warp::path::param::<Uuid>()
                .map(Some)
                .or_else(infallible_none),
        )
        .and_then(Endpoints::get_agent_jobs)
        .map(into_message);

    let get_personal_jobs = path("get_personal_jobs")
        .and(warp::path::param::<Uuid>().map(Some))
        .and_then(Endpoints::get_personal_jobs)
        .map(into_message);

    let del = path("del")
        .and(warp::path::param::<Uuid>())
        .and_then(Endpoints::del)
        .map(ok);

    let set_jobs = path("set_jobs")
        .and(warp::path::param::<Uuid>())
        .and(get_content::<Vec<String>>())
        .and_then(Endpoints::set_jobs)
        .map(into_message);

    let report = path("report")
        .and(get_content::<Vec<Reportable>>())
        .and_then(Endpoints::report)
        .map(ok);

    let update_agent = path("update_item")
        .and(get_content::<Agent>())
        .and_then(Endpoints::update_agent)
        .map(ok);

    let update_job = path("update_item")
        .and(get_content::<JobMeta>())
        .and_then(Endpoints::update_job)
        .map(ok);

    let update_assigned_job = path("update_item")
        .and(get_content::<AssignedJob>())
        .and_then(Endpoints::update_assigned_job)
        .map(ok);

    let download = path("download")
        .and(warp::path::param::<String>())
        .and_then(Endpoints::download)
        .map(ok);

    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_jobs)
        .or(upload_jobs)
        .or(del)
        .or(set_jobs)
        .or(get_agent_jobs)
        .or(update_agent)
        .or(update_job)
        .or(update_assigned_job)
        .or(download))
    .and(auth_header);

    let agent_zone = get_jobs.or(get_personal_jobs).or(report).or(download);

    auth_zone.or(agent_zone)
}

pub fn prefill_jobs() -> SResult<()> {
    let agent_hello = RawJobMeta::builder()
        .with_type(misc::JobType::Manage)
        .with_alias("agent_hello")
        .build()
        .unwrap();
    UDB::lock_db().insert_jobs(&[agent_hello])
}

pub async fn serve() -> SResult<()> {
    init_logger(Some("u_server"));
    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");

    warp::serve(routes.with(warp::log("warp")))
        .tls()
        .cert_path(certs_dir.join("server.crt"))
        .key_path(certs_dir.join("server.key"))
        .client_auth_required_path(certs_dir.join("ca.crt"))
        .run(([0, 0, 0, 0], MASTER_PORT))
        .await;
    Ok(())
}

fn ok<T>(_: T) -> impl Reply {
    reply()
}

/*
#[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<Uuid>) {
        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<Reportable>>| 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();
    }
}
*/