Websocket authentication

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 <ca.length; i++) {
+                    var c = ca[i];
+                    while (c.charAt(0) == ' ') {
+                        c = c.substring(1);
+                    }
+                    if (c.indexOf(name) == 0) {
+                        return c.substring(name.length, c.length);
+                    }
+                }
+                return "";
+            }
             
             socket = new WebSocket("ws://"+location.hostname+":WEBSOCKET_PORT/ws");
             socket.onopen = function() {
                 console.log("Connected");
+                socket.send(getCookie("aira_auth")); //authenticating websocket connection
                 window.onfocus = function() {
                     if (current_chat_index != -1) {
                         socket.send("set_seen "+current_chat_index);

src/main.rs

@@ -7,7 +7,7 @@ mod ui_interface;
 mod constants;
 mod discovery;
 
-use std::{env, fs, io, net::{SocketAddr, TcpStream}, path::Path, str::FromStr, sync::{Arc, RwLock}};
+use std::{env, fs, io, net::SocketAddr, path::Path, str::FromStr, sync::{Arc, RwLock}};
 use tokio::{net::TcpListener, runtime::Handle};
 use actix_web::{App, HttpMessage, HttpRequest, HttpResponse, HttpServer, http::{header, CookieBuilder}, web, web::Data};
 use actix_multipart::Multipart;
@@ -32,17 +32,31 @@ async fn start_websocket_server(global_vars: Arc<RwLock<GlobalVars>>) -> u16 {
         loop {
             let (stream, _addr) = server.accept().await.unwrap();
             if *worker_done.read().unwrap() {
-                let stream = stream.into_std().unwrap();
+                let ui_auth_token = {
-                stream.set_nonblocking(false).unwrap();
+                    global_vars.clone().read().unwrap().ui_auth_token.clone()
-                match tungstenite::accept(stream.try_clone().unwrap()) {
+                };
-                    Ok(websocket) => {
+                if let Some(ui_auth_token) = ui_auth_token {
-                        let ui_connection = UiConnection::new(websocket);
+                    let stream = stream.into_std().unwrap();
-                        let global_vars = global_vars.clone();
+                    stream.set_nonblocking(false).unwrap();
-                        global_vars.read().unwrap().session_manager.set_ui_connection(ui_connection);
+                    match tungstenite::accept(stream.try_clone().unwrap()) {
-                        *worker_done.write().unwrap() = false;
+                        Ok(mut websocket) => {
-                        websocket_worker(stream, global_vars, worker_done.clone()).await;
+                            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<SessionManager>, ui_connection: &mut UiConnect
     }
 }
 
-async fn websocket_worker(websocket_strem: TcpStream, global_vars: Arc<RwLock<GlobalVars>>, worker_done: Arc<RwLock<bool>>) {
+async fn websocket_worker(mut ui_connection: UiConnection, global_vars: Arc<RwLock<GlobalVars>>, worker_done: Arc<RwLock<bool>>) {
-    let mut ui_connection = UiConnection::from_raw_socket(websocket_strem.try_clone().unwrap());
     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<RwLock<Gl
     });
     discover_peers(session_manager.clone());
     let handle = Handle::current();
-    std::thread::spawn(move || {
+    std::thread::spawn(move || { //new thread needed to block on read_message() without blocking tokio tasks
         loop {
             match ui_connection.websocket.read_message() {
                 Ok(msg) => {
@@ -107,8 +120,8 @@ async fn websocket_worker(websocket_strem: TcpStream, global_vars: Arc<RwLock<Gl
                         ui_connection.write_message(Message::Pong(Vec::new())); //not sure if I'm doing this right
                     } else if msg.is_text() {
                         let msg = msg.into_text().unwrap();
+                        let mut ui_connection = ui_connection.clone();
                         let session_manager = session_manager.clone();
-                        let mut ui_connection = UiConnection::from_raw_socket(websocket_strem.try_clone().unwrap());
                         handle.spawn(async move {
                             let args: Vec<&str> = msg.split(" ").collect();
                             match args[0] {
@@ -229,6 +242,16 @@ async fn websocket_worker(websocket_strem: TcpStream, global_vars: Arc<RwLock<Gl
     });
 }
 
+fn is_authenticated(req: &HttpRequest) -> bool {
+    if let Some(cookie) = req.cookie(constants::HTTP_COOKIE_NAME) {
+        let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().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<FileInfo>) -> HttpResponse {
-    match req.cookie(constants::HTTP_COOKIE_NAME) {
+    if is_authenticated(&req) {
-        Some(cookie) => {
+        match Uuid::from_str(&file_info.uuid) {
-            let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
+            Ok(uuid) => {
-            let global_vars = global_vars.read().unwrap();
+                let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
-            if global_vars.http_session_manager.is_registered(cookie.value()) {
+                match global_vars.read().unwrap().session_manager.load_file(uuid) {
-                match Uuid::from_str(&file_info.uuid) {
+                    Some(buffer) => {
-                    Ok(uuid) => {
+                        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);
-                        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 => {}
-                        }
                     }
-                    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 is_authenticated(&req) {
-    if cookie.is_some() {
+        let mut session_id: Option<usize> = None;
-        let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
+        while let Ok(Some(mut field)) = payload.try_next().await {
-        let global_vars_read = global_vars.read().unwrap();
+            let content_disposition = field.content_disposition().unwrap();
-        if global_vars_read.http_session_manager.is_registered(cookie.unwrap().value()) {
+            if let Some(name) = content_disposition.get_name() {
-            drop(global_vars_read); //releasing mutex while uploading the file
+                if name == "session_id" {
-            let mut session_id: Option<usize> = None;
+                    if let Some(Ok(raw_id)) = field.next().await {
-            while let Ok(Some(mut field)) = payload.try_next().await {
+                        session_id = Some(std::str::from_utf8(&raw_id).unwrap().parse().unwrap());
-                let content_disposition = field.content_disposition().unwrap();
+                    }
-                if let Some(name) = content_disposition.get_name() {
+                } else if session_id.is_some() {
-                    if name == "session_id" {
+                    let filename = content_disposition.get_filename().unwrap();
-                        if let Some(Ok(raw_id)) = field.next().await {
+                    let mut buffer = Vec::new();
-                            session_id = Some(std::str::from_utf8(&raw_id).unwrap().parse().unwrap());
+                    while let Some(chunk) = field.next().await {
-                        }
+                        buffer.extend(chunk.unwrap());
-                    } else if session_id.is_some() {
+                    }
-                        let filename = content_disposition.get_filename().unwrap();
+                    let session_id = session_id.unwrap();
-                        let mut buffer = Vec::new();
+                    let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
-                        while let Some(chunk) = field.next().await {
+                    let global_vars_read = global_vars.read().unwrap();
-                            buffer.extend(chunk.unwrap());
+                    match global_vars_read.session_manager.send_to(&session_id,  protocol::file(filename, &buffer)).await {
-                        }
+                        Ok(_) => {
-                        let session_id = session_id.unwrap();
+                            match global_vars_read.session_manager.store_file(&session_id, &buffer) {
-                        let global_vars_read = global_vars.read().unwrap();
+                                Ok(file_uuid) => {
-                        match global_vars_read.session_manager.send_to(&session_id,  protocol::file(filename, &buffer)).await {
+                                    let msg = [&[protocol::Headers::FILE][..], file_uuid.as_bytes(), filename.as_bytes()].concat();
-                            Ok(_) => {
+                                    global_vars_read.session_manager.store_msg(&session_id, true, msg);
-                                match global_vars_read.session_manager.store_file(&session_id, &buffer) {
+                                    return HttpResponse::Ok().body(file_uuid.to_string());
-                                    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)
                         }
+                        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) {
+    if is_authenticated(&req) {
-        Some(cookie) => {
+        let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
-            let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
+        let mut global_vars_write = global_vars.write().unwrap();
-            let mut global_vars_write = global_vars.write().unwrap();
+        if global_vars_write.session_manager.is_identity_loaded() {
-            if global_vars_write.session_manager.is_identity_loaded() {
+            global_vars_write.ui_auth_token = None;
-                global_vars_write.http_session_manager.remove(cookie.value());
+            global_vars_write.session_manager.stop().await;
-                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"))
-            }
         }
-        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<RwLock<GlobalVars>>) -> 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<RwLock<GlobalVars>>) -> HttpRespo
             }
         });
     }
-    let cookie = CookieBuilder::new(constants::HTTP_COOKIE_NAME, cookie_value)
+    let mut raw_cookie = [0; 32];
-        .http_only(true)
+    OsRng.fill_bytes(&mut raw_cookie);
-        .max_age(time::Duration::hours(4)
+    let cookie_value = base64::encode(raw_cookie);
-    ).finish();
+    global_vars_write.ui_auth_token = Some(cookie_value.clone());
-    HttpResponse::Found().header(header::LOCATION, "/").set_header(header::SET_COOKIE, cookie.to_string()).finish()
+    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<RwLock<GlobalVars>>) -> HttpResponse {
@@ -428,21 +442,14 @@ fn index_not_logged_in(global_vars: &Arc<RwLock<GlobalVars>>) -> HttpResponse {
 
 async fn handle_index(req: HttpRequest) -> HttpResponse {
     let global_vars = req.app_data::<Data<Arc<RwLock<GlobalVars>>>>().unwrap();
-    match req.cookie(constants::HTTP_COOKIE_NAME) {
+    if is_authenticated(&req) {
-        Some(cookie) => {
+        HttpResponse::Ok().body(
-            let global_vars_read = global_vars.read().unwrap();
+            include_str!("frontend/index.html")
-            if global_vars_read.http_session_manager.is_registered(cookie.value()) {
+                .replace("WEBSOCKET_PORT", &global_vars.read().unwrap().websocket_port.to_string())
-                HttpResponse::Ok().body(
+                .replace("IS_IDENTITY_PROTECTED", &Identity::is_protected().unwrap().to_string())
-                    include_str!("frontend/index.html")
+        )
-                        .replace("WEBSOCKET_PORT", &global_vars_read.websocket_port.to_string())
+    } else {
-                        .replace("IS_IDENTITY_PROTECTED", &Identity::is_protected().unwrap().to_string())
+        index_not_logged_in(global_vars)
-                )
-            } else {
-                drop(global_vars_read);
-                index_not_logged_in(global_vars)
-            }
-        }
-        None => index_not_logged_in(global_vars)
     }
 }
 
@@ -564,38 +571,10 @@ struct CreateParams {
     password_confirm: String,
 }
 
-struct HttpSessionsManager {
-    http_sessions: Vec<String>,
-}
-
-impl HttpSessionsManager {
-    fn get_index(&self, cookie: &str) -> Option<usize> {
-        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<String> {
-        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<SessionManager>,
     websocket_port: u16,
-    http_session_manager: HttpSessionsManager,
+    ui_auth_token: Option<String>,
     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;

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<TcpStream>) -> 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
+        }
+    }
 }