From 33937354b72e4605511fe3d3419e7a1fbff1ae3e Mon Sep 17 00:00:00 2001 From: Hardcore Sushi Date: Mon, 26 Apr 2021 18:23:26 +0200 Subject: [PATCH] Websocket authentication --- src/frontend/index.html | 18 ++++ src/main.rs | 233 ++++++++++++++++++---------------------- src/ui_interface.rs | 14 ++- 3 files changed, 133 insertions(+), 132 deletions(-) diff --git a/src/frontend/index.html b/src/frontend/index.html index 925f505..5e8919e 100644 --- a/src/frontend/index.html +++ b/src/frontend/index.html @@ -515,10 +515,28 @@ }); } } + + //source: https://www.w3schools.com/js/js_cookies.asp + function getCookie(cname) { + var name = cname + "="; + var decodedCookie = decodeURIComponent(document.cookie); + var ca = decodedCookie.split(';'); + for(var i = 0; i >) -> u16 { loop { let (stream, _addr) = server.accept().await.unwrap(); if *worker_done.read().unwrap() { - let stream = stream.into_std().unwrap(); - stream.set_nonblocking(false).unwrap(); - match tungstenite::accept(stream.try_clone().unwrap()) { - Ok(websocket) => { - let ui_connection = UiConnection::new(websocket); - let global_vars = global_vars.clone(); - global_vars.read().unwrap().session_manager.set_ui_connection(ui_connection); - *worker_done.write().unwrap() = false; - websocket_worker(stream, global_vars, worker_done.clone()).await; + let ui_auth_token = { + global_vars.clone().read().unwrap().ui_auth_token.clone() + }; + if let Some(ui_auth_token) = ui_auth_token { + let stream = stream.into_std().unwrap(); + stream.set_nonblocking(false).unwrap(); + match tungstenite::accept(stream.try_clone().unwrap()) { + Ok(mut websocket) => { + if let Ok(message) = websocket.read_message() { //waiting for auth token + match message.into_text() { + Ok(token) => { + if token == ui_auth_token { + let ui_connection = UiConnection::new(websocket); + let global_vars = global_vars.clone(); + global_vars.read().unwrap().session_manager.set_ui_connection(ui_connection.clone()); + *worker_done.write().unwrap() = false; + websocket_worker(ui_connection, global_vars, worker_done.clone()).await; + } + } + Err(e) => print_error!(e) + } + } + } + Err(e) => print_error!(e) } - Err(e) => print_error!(e) } } } @@ -78,8 +92,7 @@ fn load_msgs(session_manager: Arc, ui_connection: &mut UiConnect } } -async fn websocket_worker(websocket_strem: TcpStream, global_vars: Arc>, worker_done: Arc>) { - let mut ui_connection = UiConnection::from_raw_socket(websocket_strem.try_clone().unwrap()); +async fn websocket_worker(mut ui_connection: UiConnection, global_vars: Arc>, worker_done: Arc>) { let session_manager = global_vars.read().unwrap().session_manager.clone(); ui_connection.set_name(&session_manager.get_my_name()); session_manager.list_contacts().into_iter().for_each(|contact|{ @@ -99,7 +112,7 @@ async fn websocket_worker(websocket_strem: TcpStream, global_vars: Arc { @@ -107,8 +120,8 @@ async fn websocket_worker(websocket_strem: TcpStream, global_vars: Arc = msg.split(" ").collect(); match args[0] { @@ -229,6 +242,16 @@ async fn websocket_worker(websocket_strem: TcpStream, global_vars: Arc bool { + if let Some(cookie) = req.cookie(constants::HTTP_COOKIE_NAME) { + let global_vars = req.app_data::>>>().unwrap(); + if let Some(token) = &global_vars.read().unwrap().ui_auth_token { + return token == cookie.value(); + } + } + false +} + #[derive(Deserialize, Serialize, Debug)] struct FileInfo { uuid: String, @@ -236,65 +259,54 @@ struct FileInfo { } fn handle_load_file(req: HttpRequest, file_info: web::Query) -> HttpResponse { - match req.cookie(constants::HTTP_COOKIE_NAME) { - Some(cookie) => { - let global_vars = req.app_data::>>>().unwrap(); - let global_vars = global_vars.read().unwrap(); - if global_vars.http_session_manager.is_registered(cookie.value()) { - match Uuid::from_str(&file_info.uuid) { - Ok(uuid) => { - match global_vars.session_manager.load_file(uuid) { - Some(buffer) => { - return HttpResponse::Ok().header("Content-Disposition", format!("attachment; filename=\"{}\"", escape_double_quote(html_escape::decode_html_entities(&file_info.file_name).to_string()))).content_type("application/octet-stream").body(buffer); - } - None => {} - } + if is_authenticated(&req) { + match Uuid::from_str(&file_info.uuid) { + Ok(uuid) => { + let global_vars = req.app_data::>>>().unwrap(); + match global_vars.read().unwrap().session_manager.load_file(uuid) { + Some(buffer) => { + return HttpResponse::Ok().header("Content-Disposition", format!("attachment; filename=\"{}\"", escape_double_quote(html_escape::decode_html_entities(&file_info.file_name).to_string()))).content_type("application/octet-stream").body(buffer); } - Err(e) => print_error!(e) + None => {} } } + Err(e) => print_error!(e) } - None => {} } HttpResponse::NotFound().finish() } async fn handle_send_file(req: HttpRequest, mut payload: Multipart) -> HttpResponse { - let cookie = req.cookie(constants::HTTP_COOKIE_NAME); - if cookie.is_some() { - let global_vars = req.app_data::>>>().unwrap(); - let global_vars_read = global_vars.read().unwrap(); - if global_vars_read.http_session_manager.is_registered(cookie.unwrap().value()) { - drop(global_vars_read); //releasing mutex while uploading the file - let mut session_id: Option = None; - while let Ok(Some(mut field)) = payload.try_next().await { - let content_disposition = field.content_disposition().unwrap(); - if let Some(name) = content_disposition.get_name() { - if name == "session_id" { - if let Some(Ok(raw_id)) = field.next().await { - session_id = Some(std::str::from_utf8(&raw_id).unwrap().parse().unwrap()); - } - } else if session_id.is_some() { - let filename = content_disposition.get_filename().unwrap(); - let mut buffer = Vec::new(); - while let Some(chunk) = field.next().await { - buffer.extend(chunk.unwrap()); - } - let session_id = session_id.unwrap(); - let global_vars_read = global_vars.read().unwrap(); - match global_vars_read.session_manager.send_to(&session_id, protocol::file(filename, &buffer)).await { - Ok(_) => { - match global_vars_read.session_manager.store_file(&session_id, &buffer) { - Ok(file_uuid) => { - let msg = [&[protocol::Headers::FILE][..], file_uuid.as_bytes(), filename.as_bytes()].concat(); - global_vars_read.session_manager.store_msg(&session_id, true, msg); - return HttpResponse::Ok().body(file_uuid.to_string()); - } - Err(e) => print_error!(e) + if is_authenticated(&req) { + let mut session_id: Option = None; + while let Ok(Some(mut field)) = payload.try_next().await { + let content_disposition = field.content_disposition().unwrap(); + if let Some(name) = content_disposition.get_name() { + if name == "session_id" { + if let Some(Ok(raw_id)) = field.next().await { + session_id = Some(std::str::from_utf8(&raw_id).unwrap().parse().unwrap()); + } + } else if session_id.is_some() { + let filename = content_disposition.get_filename().unwrap(); + let mut buffer = Vec::new(); + while let Some(chunk) = field.next().await { + buffer.extend(chunk.unwrap()); + } + let session_id = session_id.unwrap(); + let global_vars = req.app_data::>>>().unwrap(); + let global_vars_read = global_vars.read().unwrap(); + match global_vars_read.session_manager.send_to(&session_id, protocol::file(filename, &buffer)).await { + Ok(_) => { + match global_vars_read.session_manager.store_file(&session_id, &buffer) { + Ok(file_uuid) => { + let msg = [&[protocol::Headers::FILE][..], file_uuid.as_bytes(), filename.as_bytes()].concat(); + global_vars_read.session_manager.store_msg(&session_id, true, msg); + return HttpResponse::Ok().body(file_uuid.to_string()); } + Err(e) => print_error!(e) } - Err(e) => print_error!(e) } + Err(e) => print_error!(e) } } } @@ -304,27 +316,25 @@ async fn handle_send_file(req: HttpRequest, mut payload: Multipart) -> HttpRespo } async fn handle_logout(req: HttpRequest) -> HttpResponse { - match req.cookie(constants::HTTP_COOKIE_NAME) { - Some(cookie) => { - let global_vars = req.app_data::>>>().unwrap(); - let mut global_vars_write = global_vars.write().unwrap(); - if global_vars_write.session_manager.is_identity_loaded() { - global_vars_write.http_session_manager.remove(cookie.value()); - global_vars_write.session_manager.stop().await; - } - if Identity::is_protected().unwrap_or(true) { - HttpResponse::Found().header(header::LOCATION, "/").finish() - } else { - HttpResponse::Ok().body(include_str!("frontend/logout.html")) - } + if is_authenticated(&req) { + let global_vars = req.app_data::>>>().unwrap(); + let mut global_vars_write = global_vars.write().unwrap(); + if global_vars_write.session_manager.is_identity_loaded() { + global_vars_write.ui_auth_token = None; + global_vars_write.session_manager.stop().await; } - None => HttpResponse::Unauthorized().finish() + if Identity::is_protected().unwrap_or(true) { + HttpResponse::Found().header(header::LOCATION, "/").finish() + } else { + HttpResponse::Ok().body(include_str!("frontend/logout.html")) + } + } else { + HttpResponse::Unauthorized().finish() } } fn login(identity: Identity, global_vars: &Arc>) -> HttpResponse { let mut global_vars_write = global_vars.write().unwrap(); - let cookie_value = global_vars_write.http_session_manager.register(); let session_manager = global_vars_write.session_manager.clone(); if !session_manager.is_identity_loaded() { global_vars_write.session_manager.set_identity(Some(identity)); @@ -334,11 +344,15 @@ fn login(identity: Identity, global_vars: &Arc>) -> HttpRespo } }); } - let cookie = CookieBuilder::new(constants::HTTP_COOKIE_NAME, cookie_value) - .http_only(true) - .max_age(time::Duration::hours(4) - ).finish(); - HttpResponse::Found().header(header::LOCATION, "/").set_header(header::SET_COOKIE, cookie.to_string()).finish() + let mut raw_cookie = [0; 32]; + OsRng.fill_bytes(&mut raw_cookie); + let cookie_value = base64::encode(raw_cookie); + global_vars_write.ui_auth_token = Some(cookie_value.clone()); + let cookie = CookieBuilder::new(constants::HTTP_COOKIE_NAME, cookie_value).max_age(time::Duration::hours(4)).finish(); + HttpResponse::Found() + .header(header::LOCATION, "/") + .set_header(header::SET_COOKIE, cookie.to_string()) + .finish() } fn on_identity_loaded(identity: Identity, global_vars: &Arc>) -> HttpResponse { @@ -428,21 +442,14 @@ fn index_not_logged_in(global_vars: &Arc>) -> HttpResponse { async fn handle_index(req: HttpRequest) -> HttpResponse { let global_vars = req.app_data::>>>().unwrap(); - match req.cookie(constants::HTTP_COOKIE_NAME) { - Some(cookie) => { - let global_vars_read = global_vars.read().unwrap(); - if global_vars_read.http_session_manager.is_registered(cookie.value()) { - HttpResponse::Ok().body( - include_str!("frontend/index.html") - .replace("WEBSOCKET_PORT", &global_vars_read.websocket_port.to_string()) - .replace("IS_IDENTITY_PROTECTED", &Identity::is_protected().unwrap().to_string()) - ) - } else { - drop(global_vars_read); - index_not_logged_in(global_vars) - } - } - None => index_not_logged_in(global_vars) + if is_authenticated(&req) { + HttpResponse::Ok().body( + include_str!("frontend/index.html") + .replace("WEBSOCKET_PORT", &global_vars.read().unwrap().websocket_port.to_string()) + .replace("IS_IDENTITY_PROTECTED", &Identity::is_protected().unwrap().to_string()) + ) + } else { + index_not_logged_in(global_vars) } } @@ -564,38 +571,10 @@ struct CreateParams { password_confirm: String, } -struct HttpSessionsManager { - http_sessions: Vec, -} - -impl HttpSessionsManager { - fn get_index(&self, cookie: &str) -> Option { - self.http_sessions.iter().position(|c| c == cookie) - } - pub fn new() -> HttpSessionsManager { - HttpSessionsManager { - http_sessions: Vec::new() - } - } - pub fn register(&mut self) -> String { - let mut raw_cookie = [0; 32]; - OsRng.fill_bytes(&mut raw_cookie); - let cookie = base64::encode(raw_cookie); - self.http_sessions.push(cookie.clone()); - cookie - } - pub fn remove(&mut self, cookie: &str) -> Option { - Some(self.http_sessions.remove(self.get_index(cookie)?)) - } - pub fn is_registered(&self, cookie: &str) -> bool { - self.get_index(cookie).is_some() - } -} - struct GlobalVars { session_manager: Arc, websocket_port: u16, - http_session_manager: HttpSessionsManager, + ui_auth_token: Option, tokio_handle: Handle, } @@ -612,7 +591,7 @@ async fn main() { let global_vars = Arc::new(RwLock::new(GlobalVars { session_manager: Arc::new(SessionManager::new()), websocket_port: 0, - http_session_manager: HttpSessionsManager::new(), + ui_auth_token: None, tokio_handle: Handle::current(), })); let websocket_port = start_websocket_server(global_vars.clone()).await; diff --git a/src/ui_interface.rs b/src/ui_interface.rs index 5964921..bf58603 100644 --- a/src/ui_interface.rs +++ b/src/ui_interface.rs @@ -91,11 +91,6 @@ pub struct UiConnection{ } impl UiConnection { - pub fn from_raw_socket(stream: TcpStream) -> UiConnection { - let websocket = WebSocket::from_raw_socket(stream, Role::Server, None); - UiConnection::new(websocket) - } - pub fn new(websocket: WebSocket) -> UiConnection { UiConnection { websocket: websocket, @@ -161,4 +156,13 @@ impl UiConnection { pub fn logout(&mut self) { self.write_message(Message::from("logout")); } +} + +impl Clone for UiConnection { + fn clone(&self) -> Self { + UiConnection { + websocket: WebSocket::from_raw_socket(self.websocket.get_ref().try_clone().unwrap(), Role::Server, None), + is_valid: self.is_valid + } + } } \ No newline at end of file