unki/bin/u_server/src/main.rs

mod db;
mod handlers;

use warp::{body, Filter, Rejection};

#[macro_use]
extern crate log;
extern crate env_logger;

use db::lock_db;
use serde::de::DeserializeOwned;
use u_lib::{
    config::MASTER_PORT,
    messaging::{AsMsg, BaseMessage},
    models::*,
};
use uuid::Uuid;

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 prefill_jobs() {
    let agent_hello = JobMeta::builder()
        .with_type(misc::JobType::Manage)
        .with_alias("agent_hello")
        .build()
        .unwrap();
    lock_db().insert_jobs(&[agent_hello]).ok();
}

fn init() {
    env_logger::init();
    prefill_jobs();
}

#[tokio::main]
async fn main() {
    init();
    let infallible_none = |_| async { Ok::<(Option<Uuid>,), std::convert::Infallible>((None,)) };

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

    let upload_jobs = warp::post()
        .and(warp::path("upload_jobs"))
        .and(get_content::<Vec<JobMeta>>())
        .and_then(handlers::upload_jobs);

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

    let get_agent_jobs = warp::get()
        .and(warp::path("get_agent_jobs"))
        .and(
            warp::path::param::<Uuid>()
                .map(Some)
                .or_else(infallible_none),
        )
        .and_then(|uid| handlers::get_agent_jobs(uid, false));

    let get_personal_jobs = warp::get()
        .and(warp::path("get_agent_jobs"))
        .and(warp::path::param::<Uuid>().map(Some))
        .and_then(|uid| handlers::get_agent_jobs(uid, true));

    let del = warp::get()
        .and(warp::path("del"))
        .and(warp::path::param::<Uuid>())
        .and_then(handlers::del);

    let set_jobs = warp::post()
        .and(warp::path("set_jobs"))
        .and(warp::path::param::<Uuid>())
        .and(get_content::<Vec<Uuid>>())
        .and_then(handlers::set_jobs);

    let report = warp::post()
        .and(warp::path("report"))
        .and(get_content::<Vec<ExecResult>>().and_then(handlers::report));

    let auth_token = warp::header::exact("authorization", "Bearer 123qwe");

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

    let auth_zone = auth_token.and(
        get_agents
            .or(get_jobs)
            .or(upload_jobs)
            .or(del)
            .or(set_jobs)
            .or(get_agent_jobs),
    );

    let routes = agent_zone.or(auth_zone);
    warp::serve(routes.with(warp::log("warp")))
        .run(([0, 0, 0, 0], MASTER_PORT))
        .await;
}
/*
#[cfg(test)]
mod tests {
    use super::*;
    /*
        #[tokio::test]
        async fn test_gather() {
        }
    */
}
*/