diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..dcacc4e --- /dev/null +++ b/src/config.rs @@ -0,0 +1,26 @@ +use std::env; + +pub fn max_cost() -> u64 { + env::var("MAX_COST") + .unwrap_or("1000000000".to_owned()) + .parse::() + .unwrap_or(1000000000) +} + +pub fn max_memory() -> u32 { + env::var("MAX_MEMORY") + .unwrap_or("4096".to_owned()) + .parse::() + .unwrap_or(4096) +} + +pub fn server_port() -> u16 { + env::var("PORT") + .unwrap_or("33000".to_owned()) + .parse::() + .unwrap_or(33000) +} + +pub fn app_secret() -> String { + env::var("APP_SECRET").unwrap_or("APP_SECRET".to_owned()) +} diff --git a/src/main.rs b/src/main.rs index 7fde71f..73fe089 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,5 @@ mod cli; +mod config; mod cost; mod memory; mod read; diff --git a/src/server/core.rs b/src/server/core.rs index 9270c77..fdf12af 100644 --- a/src/server/core.rs +++ b/src/server/core.rs @@ -5,6 +5,7 @@ use rocket::Build; use rocket::Config; use rocket::Rocket; +use crate::config::*; use crate::server::execute; use crate::server::jwt; @@ -37,10 +38,7 @@ pub fn rocket() -> Rocket { rocket::build() .configure(Config { address: Ipv4Addr::new(0, 0, 0, 0).into(), - port: std::env::var("PORT") - .unwrap_or_else(|_| "33000".to_string()) - .parse() - .unwrap(), + port: server_port(), ..Config::default() }) .mount("/", routes![index, info, jwt::validate, execute::execute]) diff --git a/src/server/execute.rs b/src/server/execute.rs index ec6d6ad..c75ce31 100644 --- a/src/server/execute.rs +++ b/src/server/execute.rs @@ -1,3 +1,4 @@ +use crate::config::*; use crate::run; use crate::server::jwt; use base64::{engine::general_purpose, Engine as _}; @@ -50,7 +51,7 @@ pub async fn execute( } }; - if submission.cost > 1000000000 { + if submission.cost > max_cost() { return Json(ExecutionResult { success: false, cost: None, @@ -61,7 +62,7 @@ pub async fn execute( }); } - if submission.memory > 1024 { + if submission.memory > max_memory() { return Json(ExecutionResult { success: false, cost: None, diff --git a/src/server/jwt.rs b/src/server/jwt.rs index 1ccb62f..768350b 100644 --- a/src/server/jwt.rs +++ b/src/server/jwt.rs @@ -2,7 +2,8 @@ use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation}; use rocket::request::{self, FromRequest, Request}; use rocket::serde::json::Json; use rocket::serde::{Deserialize, Serialize}; -use std::env; + +use crate::config::*; #[derive(Debug, Serialize, Deserialize)] #[serde(crate = "rocket::serde")] @@ -11,10 +12,7 @@ struct Claims { } pub fn is_valid_token(token: &str) -> bool { - let secret = match env::var("JWT_SECRET") { - Ok(secret) => secret, - Err(_) => "SECRET_TOKEN".to_string(), - }; + let secret = app_secret(); let validation = Validation::new(Algorithm::HS256);