diff --git a/Cargo.lock b/Cargo.lock index 50032d7..3f33aa5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1903,6 +1903,7 @@ dependencies = [ "actix-web", "clap", "env_logger", + "futures", "log", "serde_json", "zoidberg_lib", diff --git a/zoidberg_client/src/main.rs b/zoidberg_client/src/main.rs index 8026a07..f437b4d 100644 --- a/zoidberg_client/src/main.rs +++ b/zoidberg_client/src/main.rs @@ -12,7 +12,7 @@ use zoidberg_lib::types::{FetchResponse, Job, RegisterResponse, Status, Update}; const VERSION: &str = env!("CARGO_PKG_VERSION"); fn build_client(secret: &str) -> Client { - let cookie = format!("secret={}", secret); + let cookie = format!("{}", secret); let mut headers = header::HeaderMap::new(); headers.insert( @@ -31,11 +31,12 @@ fn build_client(secret: &str) -> Client { #[derive(Debug)] struct Worker { id: i32, + secret: String, } impl Worker { - async fn new(server: &str) -> Result> { - let res = build_client("some_secret") + async fn new(server: &str, secret: &str) -> Result> { + let res = build_client(secret) .get(format!("{}/register", server)) .send() .await?; @@ -43,7 +44,10 @@ impl Worker { let body = res.text().await?; let r: RegisterResponse = serde_json::from_str(&body)?; log::info!("registered worker with id: {}", &r.id); - Ok(Worker { id: r.id }) + Ok(Worker { + id: r.id, + secret: secret.to_string(), + }) } async fn update(self: &Self, jobs: &[Job]) -> Result<(), Box> { @@ -56,7 +60,7 @@ impl Worker { }) .collect(); - let body = build_client("some_secret") + let body = build_client(&self.secret) .post("http://localhost:8080/update") .json(&updates) .send() @@ -69,7 +73,7 @@ impl Worker { } async fn fetch(self: &Self) -> Result> { - let res = build_client("some_secret") + let res = build_client(&self.secret) .get("http://localhost:8080/fetch") .send() .await?; @@ -125,8 +129,12 @@ async fn main() -> Result<(), Box> { ) .get_matches(); let server = matches.value_of("server").unwrap(); + let secret = std::env::var("ZOIDBERG_SECRET") + .expect("Please set the $ZOIDBERG_SECRET environment variable"); - let client = Worker::new(server).await.expect("Could not create client"); + let client = Worker::new(server, &secret) + .await + .expect("Could not create client"); let pause = time::Duration::from_secs(1); let long_pause = time::Duration::from_secs(20); diff --git a/zoidberg_server/Cargo.toml b/zoidberg_server/Cargo.toml index 723dab4..c0e6220 100644 --- a/zoidberg_server/Cargo.toml +++ b/zoidberg_server/Cargo.toml @@ -15,3 +15,4 @@ serde_json = "1.0" clap = "3.2" env_logger = "0.9" log = "0.4" +futures = "0.3.24" diff --git a/zoidberg_server/src/main.rs b/zoidberg_server/src/main.rs index 5ad6e0c..c5b5caa 100644 --- a/zoidberg_server/src/main.rs +++ b/zoidberg_server/src/main.rs @@ -1,6 +1,11 @@ -use actix_web::{get, post, web, App, HttpResponse, HttpServer, Responder, Result}; +use actix_web::error::ErrorBadRequest; +use actix_web::{ + dev, get, middleware::Logger, post, web, App, Error, FromRequest, HttpRequest, HttpResponse, + HttpServer, Responder, Result, +}; use clap; use env_logger::Env; +use futures::future::{err, ok, Ready}; use log; use std::sync::Mutex; use zoidberg_lib::types::{FetchResponse, Job, RegisterResponse, StatusRequest, Update, Worker}; @@ -17,6 +22,41 @@ struct State { jobs: Mutex>, } +impl State { + fn new() -> Self { + Self { + counter_workers: Mutex::new(0), + counter_jobs: Mutex::new(0), + workers: Mutex::new(Vec::new()), + new_jobs: Mutex::new(Vec::new()), + jobs: Mutex::new(Vec::new()), + } + } +} + +struct Authorization {} + +impl FromRequest for Authorization { + type Error = Error; + type Future = Ready>; + + fn from_request(req: &HttpRequest, _payload: &mut dev::Payload) -> Self::Future { + if let Some(head) = req.headers().get("cookie") { + if let Ok(cookie) = head.to_str() { + if let Some(secret) = req.app_data::() { + println!("{} == {}", secret, cookie); + if secret == cookie { + return ok(Authorization {}); + } else { + return err(ErrorBadRequest("no auth")); + } + } + } + } + err(ErrorBadRequest("no auth")) + } +} + #[get("/")] async fn index(data: web::Data) -> impl Responder { let workers = data.workers.lock().unwrap(); @@ -26,7 +66,7 @@ async fn index(data: web::Data) -> impl Responder { } #[get("/register")] -async fn register(data: web::Data) -> Result { +async fn register(data: web::Data, _: Authorization) -> Result { let mut counter_workers = data.counter_workers.lock().unwrap(); *counter_workers += 1; @@ -42,7 +82,7 @@ async fn register(data: web::Data) -> Result { } #[get("/fetch")] -async fn fetch(data: web::Data) -> Result { +async fn fetch(data: web::Data, _: Authorization) -> Result { let mut new_jobs = data.new_jobs.lock().unwrap(); if let Some(j) = new_jobs.pop() { return Ok(web::Json(FetchResponse::Jobs(vec![j]))); @@ -54,6 +94,7 @@ async fn fetch(data: web::Data) -> Result { async fn status( s: web::Json>, data: web::Data, + _: Authorization, ) -> Result { let jobs = data.jobs.lock().unwrap(); let status_updates: Vec = jobs @@ -66,7 +107,11 @@ async fn status( } #[post("/update")] -async fn update(updates: web::Json>, data: web::Data) -> Result { +async fn update( + updates: web::Json>, + data: web::Data, + _: Authorization, +) -> Result { let mut jobs = data.jobs.lock().unwrap(); let mut n = 0; for update in updates.iter() { @@ -87,7 +132,11 @@ async fn update(updates: web::Json>, data: web::Data) -> Resu } #[post("/submit")] -async fn submit(data: web::Data, js: web::Json>) -> Result { +async fn submit( + data: web::Data, + js: web::Json>, + _: Authorization, +) -> Result { let mut new_jobs = data.new_jobs.lock().unwrap(); let mut jobs = data.jobs.lock().unwrap(); let mut counter_jobs = data.counter_jobs.lock().unwrap(); @@ -113,21 +162,20 @@ async fn submit(data: web::Data, js: web::Json>) -> Result std::io::Result<()> { env_logger::Builder::from_env(Env::default().default_filter_or("zoidberg_server=info")).init(); + let secret = std::env::var("ZOIDBERG_SECRET") + .expect("Please set the $ZOIDBERG_SECRET environment variable"); + let _matches = clap::App::new("Zoidberg server") .version(VERSION) .author("Johannes Heuel") .get_matches(); - let state = web::Data::new(State { - counter_workers: Mutex::new(0), - counter_jobs: Mutex::new(0), - workers: Mutex::new(Vec::new()), - new_jobs: Mutex::new(Vec::new()), - jobs: Mutex::new(Vec::new()), - }); + let state = web::Data::new(State::new()); HttpServer::new(move || { App::new() + .wrap(Logger::default()) + .app_data(secret.clone()) .app_data(state.clone()) .service(index) .service(register) @@ -151,13 +199,7 @@ mod tests { async fn test_index() { let app = test::init_service( App::new() - .app_data(web::Data::new(State { - counter_workers: Mutex::new(0), - counter_jobs: Mutex::new(0), - workers: Mutex::new(Vec::new()), - new_jobs: Mutex::new(Vec::new()), - jobs: Mutex::new(Vec::new()), - })) + .app_data(web::Data::new(State::new())) .service(index), ) .await; @@ -170,17 +212,15 @@ mod tests { async fn test_register() { let app = test::init_service( App::new() - .app_data(web::Data::new(State { - counter_workers: Mutex::new(0), - counter_jobs: Mutex::new(0), - workers: Mutex::new(Vec::new()), - new_jobs: Mutex::new(Vec::new()), - jobs: Mutex::new(Vec::new()), - })) + .app_data(String::from("secret")) + .app_data(web::Data::new(State::new())) .service(register), ) .await; - let req = test::TestRequest::get().uri("/register").to_request(); + let req = test::TestRequest::get() + .append_header(("cookie", "secret")) + .uri("/register") + .to_request(); let resp: RegisterResponse = test::call_and_read_body_json(&app, req).await; assert_eq!(resp.id, 1); } @@ -190,6 +230,7 @@ mod tests { let cmd = String::from("hi"); let app = test::init_service( App::new() + .app_data(String::from("secret")) .app_data(web::Data::new(State { counter_workers: Mutex::new(0), counter_jobs: Mutex::new(0), @@ -204,7 +245,10 @@ mod tests { .service(fetch), ) .await; - let req = test::TestRequest::get().uri("/fetch").to_request(); + let req = test::TestRequest::get() + .append_header(("cookie", "secret")) + .uri("/fetch") + .to_request(); let resp: FetchResponse = test::call_and_read_body_json(&app, req).await; match resp { FetchResponse::Nop => { @@ -225,6 +269,7 @@ mod tests { let cmd = String::from("hi"); let app = test::init_service( App::new() + .app_data(String::from("secret")) .app_data(web::Data::new(State { counter_workers: Mutex::new(0), counter_jobs: Mutex::new(0), @@ -240,6 +285,7 @@ mod tests { ) .await; let req = test::TestRequest::post() + .append_header(("cookie", "secret")) .set_json(vec![StatusRequest { id: 1 }]) .uri("/status") .to_request(); @@ -251,17 +297,13 @@ mod tests { async fn test_update() { let app = test::init_service( App::new() - .app_data(web::Data::new(State { - counter_workers: Mutex::new(0), - counter_jobs: Mutex::new(0), - workers: Mutex::new(Vec::new()), - new_jobs: Mutex::new(Vec::new()), - jobs: Mutex::new(Vec::new()), - })) + .app_data(String::from("secret")) + .app_data(web::Data::new(State::new())) .service(update), ) .await; let req = test::TestRequest::post() + .append_header(("cookie", "secret")) .set_json(vec![Update { worker: 0, job: 0, @@ -277,17 +319,13 @@ mod tests { async fn test_submit() { let app = test::init_service( App::new() - .app_data(web::Data::new(State { - counter_workers: Mutex::new(0), - counter_jobs: Mutex::new(0), - workers: Mutex::new(Vec::new()), - new_jobs: Mutex::new(Vec::new()), - jobs: Mutex::new(Vec::new()), - })) + .app_data(String::from("secret")) + .app_data(web::Data::new(State::new())) .service(submit), ) .await; let req = test::TestRequest::post() + .append_header(("cookie", "secret")) .set_json(vec![Job { id: 0, cmd: String::from("hi"),