diff --git a/src/data/communication.rs b/src/data/communication.rs index e21dcae..4dd6d43 100644 --- a/src/data/communication.rs +++ b/src/data/communication.rs @@ -81,6 +81,8 @@ pub enum DataTypes { online_status, omikron_id, omikron_connections, + reset_token, + new_token, } impl DataTypes { @@ -161,6 +163,8 @@ impl DataTypes { "onlinestatus" => DataTypes::online_status, "omikronid" => DataTypes::omikron_id, "omikronconnections" => DataTypes::omikron_connections, + "resettoken" => DataTypes::reset_token, + "newtoken" => DataTypes::new_token, _ => DataTypes::error_type, // fallback if unknown } } @@ -170,6 +174,7 @@ impl DataTypes { #[allow(non_camel_case_types, dead_code)] pub enum CommunicationType { error, + error_invalid_data, error_invalid_user_id, error_invalid_omikron_id, error_not_found, @@ -244,6 +249,12 @@ pub enum CommunicationType { change_user_data, change_iota_data, + get_register, + complete_register_user, + complete_register_iota, + delete_user, + delete_iota, + start_register, complete_register, } @@ -262,6 +273,7 @@ impl CommunicationType { "function" => CommunicationType::function, "update" => CommunicationType::update, "createuser" => CommunicationType::create_user, + "errorinvaliddata" => CommunicationType::error_invalid_data, "errorinvaliduserid" => CommunicationType::error_invalid_user_id, "errorinvalidomikronid" => CommunicationType::error_invalid_omikron_id, "errornotfound" => CommunicationType::error_not_found, @@ -326,6 +338,12 @@ impl CommunicationType { "changeuserdata" => CommunicationType::change_user_data, "changeiotadata" => CommunicationType::change_iota_data, + "getregister" => CommunicationType::get_register, + "completeregisteruser" => CommunicationType::complete_register_user, + "completeregisteriota" => CommunicationType::complete_register_iota, + "deleteuser" => CommunicationType::delete_user, + "deleteiota" => CommunicationType::delete_iota, + "startregister" => CommunicationType::start_register, "completeregister" => CommunicationType::complete_register, diff --git a/src/server/omikron_connection.rs b/src/server/omikron_connection.rs index f1bd08f..bebab4a 100644 --- a/src/server/omikron_connection.rs +++ b/src/server/omikron_connection.rs @@ -562,8 +562,270 @@ impl OmikronConnection { return; } - if cv.is_type(CommunicationType::change_user_data) {} - if cv.is_type(CommunicationType::change_iota_data) {} + + if cv.is_type(CommunicationType::get_register) { + let register_id = sql::get_register_id().await; + let response = CommunicationValue::new(CommunicationType::get_register) + .with_id(cv.get_id()) + .add_data( + DataTypes::user_id, + JsonValue::Number(Number::from(register_id)), + ); + self.send_message(&response).await; + return; + } + if cv.is_type(CommunicationType::complete_register_iota) { + if let (Some(iota_id), Some(public_key)) = ( + cv.get_data(DataTypes::iota_id).and_then(|v| v.as_i64()), + cv.get_data(DataTypes::public_key).and_then(|v| v.as_str()), + ) { + match sql::register_complete_iota(iota_id, public_key.to_string()).await { + Ok(_) => { + let response = CommunicationValue::new(CommunicationType::success) + .with_id(cv.get_id()); + self.send_message(&response).await; + } + Err(e) => { + self.send_message( + &CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data_str(DataTypes::error_type, e.to_string()), + ) + .await; + } + } + } else { + self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_data) + .await; + } + return; + } + if cv.is_type(CommunicationType::complete_register_user) { + if let ( + Some(user_id), + Some(username), + Some(public_key), + Some(iota_id), + Some(reset_token), + ) = ( + cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()), + cv.get_data(DataTypes::username).and_then(|v| v.as_str()), + cv.get_data(DataTypes::public_key).and_then(|v| v.as_str()), + cv.get_data(DataTypes::iota_id).and_then(|v| v.as_i64()), + cv.get_data(DataTypes::reset_token).and_then(|v| v.as_str()), // Assuming reset_token is sent as token + ) { + // The documentation does not specify a private_key_hash, using an empty string. + match sql::register_complete_user( + user_id, + username.to_string(), + public_key.to_string(), + "".to_string(), // private_key_hash + iota_id, + reset_token.to_string(), + ) + .await + { + Ok(_) => { + let response = CommunicationValue::new(CommunicationType::success) + .with_id(cv.get_id()); + self.send_message(&response).await; + } + Err(e) => { + self.send_message( + &CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data_str(DataTypes::error_type, e.to_string()), + ) + .await; + } + } + } else { + self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_data) + .await; + } + return; + } + if cv.is_type(CommunicationType::change_user_data) { + if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) { + let mut success = true; + let mut error_message = String::new(); + + if let Some(username) = cv.get_data(DataTypes::username).and_then(|v| v.as_str()) { + if let Err(e) = sql::change_username(user_id, username.to_string()).await { + success = false; + error_message = e.to_string(); + } + } + if let Some(display) = cv.get_data(DataTypes::display).and_then(|v| v.as_str()) { + if let Err(e) = sql::change_display_name(user_id, display.to_string()).await { + success = false; + error_message = e.to_string(); + } + } + if let Some(avatar) = cv.get_data(DataTypes::avatar).and_then(|v| v.as_str()) { + if let Err(e) = sql::change_avatar(user_id, avatar.to_string()).await { + success = false; + error_message = e.to_string(); + } + } + if let Some(about) = cv.get_data(DataTypes::about).and_then(|v| v.as_str()) { + if let Err(e) = sql::change_about(user_id, about.to_string()).await { + success = false; + error_message = e.to_string(); + } + } + if let Some(status) = cv.get_data(DataTypes::status).and_then(|v| v.as_str()) { + if let Err(e) = sql::change_status(user_id, status.to_string()).await { + success = false; + error_message = e.to_string(); + } + } + if let Some(public_key) = + cv.get_data(DataTypes::public_key).and_then(|v| v.as_str()) + { + if let Some(private_key_hash) = cv + .get_data(DataTypes::private_key_hash) + .and_then(|v| v.as_str()) + { + if let Err(e) = sql::change_keys( + user_id, + public_key.to_string(), + private_key_hash.to_string(), + ) + .await + { + success = false; + error_message = e.to_string(); + } + } else { + success = false; + error_message = + "private_key_hash is required when changing public_key".to_string(); + } + } + + if success { + let response = + CommunicationValue::new(CommunicationType::success).with_id(cv.get_id()); + self.send_message(&response).await; + } else { + self.send_message( + &CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data_str(DataTypes::error_type, error_message), + ) + .await; + } + } else { + self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_data) + .await; + } + return; + } + if cv.is_type(CommunicationType::change_iota_data) { + if let (Some(user_id), Some(iota_id), Some(reset_token), Some(new_token)) = ( + cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()), + cv.get_data(DataTypes::iota_id).and_then(|v| v.as_i64()), + cv.get_data(DataTypes::reset_token).and_then(|v| v.as_str()), + cv.get_data(DataTypes::new_token).and_then(|v| v.as_str()), + ) { + match sql::get_by_user_id(user_id).await { + Ok(user) => { + let current_token = user.11; // token is the 12th element (index 11) + if current_token == reset_token { + let mut success = true; + let mut error_message = String::new(); + if let Err(e) = sql::change_iota_id(user_id, iota_id).await { + success = false; + error_message = e.to_string(); + } + if success { + if let Err(e) = + sql::change_token(user_id, new_token.to_string()).await + { + success = false; + error_message = e.to_string(); + } + } + + if success { + let response = CommunicationValue::new(CommunicationType::success) + .with_id(cv.get_id()); + self.send_message(&response).await; + } else { + self.send_message( + &CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data_str(DataTypes::error_type, error_message), + ) + .await; + } + } else { + self.send_error_response( + &cv.get_id(), + // Using this for invalid token + CommunicationType::error_invalid_challenge, + ) + .await; + } + } + Err(_) => { + self.send_error_response(&cv.get_id(), CommunicationType::error_not_found) + .await; + } + } + } else { + self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_data) + .await; + } + return; + } + if cv.is_type(CommunicationType::delete_user) { + if let Some(user_id) = cv.get_data(DataTypes::user_id).and_then(|v| v.as_i64()) { + match sql::delete_user(user_id).await { + Ok(_) => { + let response = CommunicationValue::new(CommunicationType::success) + .with_id(cv.get_id()); + self.send_message(&response).await; + } + Err(e) => { + self.send_message( + &CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data_str(DataTypes::error_type, e.to_string()), + ) + .await; + } + } + } else { + self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_data) + .await; + } + return; + } + if cv.is_type(CommunicationType::delete_iota) { + if let Some(iota_id) = cv.get_data(DataTypes::iota_id).and_then(|v| v.as_i64()) { + match sql::delete_iota(iota_id).await { + Ok(_) => { + let response = CommunicationValue::new(CommunicationType::success) + .with_id(cv.get_id()); + self.send_message(&response).await; + } + Err(e) => { + self.send_message( + &CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data_str(DataTypes::error_type, e.to_string()), + ) + .await; + } + } + } else { + self.send_error_response(&cv.get_id(), CommunicationType::error_invalid_data) + .await; + } + return; + } } async fn send_error_response(&self, message_id: &Uuid, error_type: CommunicationType) { diff --git a/src/server/server.rs b/src/server/server.rs index da1763d..6c5102b 100644 --- a/src/server/server.rs +++ b/src/server/server.rs @@ -58,7 +58,7 @@ impl Service> for HttpService { let headers = parts.headers.clone(); let fut = async move { - let is_websocket_upgrade = path == "/ws" + let is_websocket_upgrade = path.starts_with("/ws") && method == Method::GET && headers .get("connection") diff --git a/src/sql/sql.rs b/src/sql/sql.rs index 251ecb8..bd3d283 100644 --- a/src/sql/sql.rs +++ b/src/sql/sql.rs @@ -427,6 +427,18 @@ pub async fn change_keys( Ok(()) } +pub async fn change_token(id: i64, new_token: String) -> Result<(), sqlx::Error> { + let db_lock = SQL_DB.read().await; + let pool = db_lock.as_ref().expect("Database pool is not initialized"); + + sqlx::query("UPDATE users SET token = ? WHERE id = ?") + .bind(new_token) + .bind(id) + .execute(pool) + .await?; + + Ok(()) +} pub async fn register_complete_user( id: i64, username: String, @@ -470,9 +482,6 @@ pub async fn print_users() -> Result<(), Box> { let about: Vec = row.get("about"); let sub_level: i32 = row.get("sub_level"); let sub_end: i64 = row.get("sub_end"); - let public_key: String = row.get("public_key"); - let private_key_hash: String = row.get("private_key_hash"); - let token: Vec = row.get("token"); log!( "User: {:?}", @@ -484,10 +493,7 @@ pub async fn print_users() -> Result<(), Box> { String::from_utf8_lossy(&status), String::from_utf8_lossy(&about), sub_level, - sub_end, - public_key, - private_key_hash, - String::from_utf8_lossy(&token) + sub_end ) ); }