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