use axum::{extract::State, response::IntoResponse, routing::get, Router}; use axum_extra::extract::cookie::{Cookie, CookieJar}; use clap::Parser; use morethantext::MoreThanText; use tokio::{spawn, sync::mpsc::channel}; const LOCALHOST: &str = "127.0.0.1"; const SESSION_KEY: &str = "sessionid"; #[derive(Parser, Debug)] #[command(version, about, long_about = None)] struct Args { /// Post used #[arg(short, long, default_value_t = 3000)] port: u16, /// IP used #[arg(short, long, default_value_t = LOCALHOST.to_string())] address: String, /// cluster host #[arg(short, long, num_args(0..))] node: Vec, } #[tokio::main] async fn main() { let args = Args::parse(); let addr = format!("{}:{}", args.address, args.port); let state = MoreThanText::new(); let app = create_app(state).await; let listener = tokio::net::TcpListener::bind(&addr).await.unwrap(); axum::serve(listener, app.into_make_service()) .await .unwrap(); } async fn create_app(state: MoreThanText) -> Router { Router::new().route("/", get(mtt_conn)).with_state(state) } async fn mtt_conn(jar: CookieJar, state: State) -> impl IntoResponse { let sid = match jar.get(SESSION_KEY) { Some(cookie) => Some(cookie.value().to_string()), None => None, }; let sess_info = sid.clone(); let (tx, mut rx) = channel(5); spawn(async move { tx.send(state.clone().request(sess_info)).await.unwrap(); }); let reply = rx.recv().await.unwrap(); let cookie = Cookie::build((SESSION_KEY, reply.get_session().to_string())); let cookies = jar.add(cookie); (cookies, reply.get_content()) } #[cfg(test)] mod servers { use super::*; use axum::{ body::Body, http::{Request, StatusCode}, }; use tower::ServiceExt; #[tokio::test] async fn get_home_page() { let app = create_app(MoreThanText::new()).await; let response = app .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap()) .await .unwrap(); assert_eq!(response.status(), StatusCode::OK); let sessid = format!("{:?}", response.headers().get("set-cookie").unwrap()); assert!(sessid.contains(SESSION_KEY), "did not set session id"); } #[tokio::test] async fn session_ids_are_unique() { let app = create_app(MoreThanText::new()).await; let mut holder: Vec = Vec::new(); for _ in 0..5 { let response = app .clone() .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap()) .await .unwrap(); let sessid = format!("{:?}", response.headers().get("set-cookie").unwrap()); assert!( !holder.contains(&sessid), "found duplicate entry: {:?}", holder ); holder.push(sessid); } } #[tokio::test] async fn receive_file_not_found() { let app = create_app(MoreThanText::new()).await; let response = app .oneshot( Request::builder() .uri("/isomething") .body(Body::empty()) .unwrap(), ) .await .unwrap(); assert_eq!(response.status(), StatusCode::NOT_FOUND); } }