use crate::{ db::pool, error::{OmegaError, Result}, models::{IotaId, User, UserId}, sql::connection_status::UserStatus, }; use mtp::crypto::PublicKeyBundle; use sqlx::{FromRow, MySql, QueryBuilder, Row}; use std::collections::HashMap; pub const MAX_PROTOCOL_ID: i64 = (1_i64 << 48) - 1; const ID_ALLOCATION_ATTEMPTS: usize = 16; pub async fn get_register_id() -> Result { use std::time::{SystemTime, UNIX_EPOCH}; let timestamp = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_secs()) .unwrap_or(0); let ts = timestamp as i64; if ts >= 1 && ts <= MAX_PROTOCOL_ID { return Ok(UserId::from(ts)); } // Fall back to random if the timestamp is outside the 48-bit range. let id = (rand::random::() & ((1_u64 << 48) - 1)) as i64; Ok(UserId::from(id.max(1))) } pub fn valid_protocol_id(id: i64) -> bool { (1..=MAX_PROTOCOL_ID).contains(&id) } /// Allocate an ID that is durable, short-lived, and bound to the connected /// Iota. `token` is presented again when completing the registration. pub async fn allocate_registration(iota_id: IotaId, request_id: u32) -> Result<(UserId, String)> { if !valid_protocol_id(iota_id.0) { return Err(OmegaError::Validation( "Iota ID is outside the 48-bit protocol range".into(), )); } for _ in 0..ID_ALLOCATION_ATTEMPTS { let id = get_register_id().await?; let token = uuid::Uuid::new_v4().to_string(); let result = sqlx::query( "INSERT INTO registration_leases (token, user_id, iota_id, request_id, expires_at) \ VALUES (?, ?, ?, ?, DATE_ADD(UTC_TIMESTAMP(), INTERVAL 10 MINUTE))", ) .bind(&token) .bind(id.0) .bind(iota_id.0) .bind(request_id) .execute(&pool().await?) .await; match result { Ok(_) => return Ok((id, token)), Err(error) if is_duplicate_key(&error) => { let existing = sqlx::query( "SELECT user_id, token, expires_at >= UTC_TIMESTAMP() AS current \ FROM registration_leases WHERE iota_id = ? AND request_id = ?", ) .bind(iota_id.0) .bind(request_id) .fetch_optional(&pool().await?) .await?; if let Some(existing) = existing { let current: i8 = existing.get("current"); if current != 0 { return Ok(( UserId::from(existing.get::("user_id")), existing.get::("token"), )); } } continue; } Err(error) => return Err(error.into()), } } Err(OmegaError::Validation( "could not allocate a unique registration ID".into(), )) } pub(crate) fn is_duplicate_key(error: &sqlx::Error) -> bool { error.as_database_error().is_some_and(|database| { // MySQL's generic database-error API exposes SQLSTATE (23000) as // `code()`. The driver-specific duplicate-key number is retained in // the diagnostic message. database.code().as_deref() == Some("23000") && database.message().contains("1062") }) } const USER_BY_USERNAME_QUERY: &str = "SELECT id, iota_id, username, display, status, presence_preference, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE username = ?"; const USER_BY_ID_QUERY: &str = "SELECT id, iota_id, username, display, status, presence_preference, about, avatar, sub_level, sub_end, public_key, token FROM users WHERE id = ?"; const USER_COLUMNS: &str = "SELECT id, iota_id, username, display, status, presence_preference, about, avatar, sub_level, sub_end, public_key, token FROM users"; #[derive(FromRow)] struct UserRow { id: i64, iota_id: i64, username: Vec, display: Option>, status: Option>, presence_preference: Vec, about: Option>, avatar: Option>, sub_level: i32, sub_end: i64, public_key: Vec, token: Vec, } impl TryFrom for User { type Error = sqlx::Error; fn try_from(row: UserRow) -> std::result::Result { let public_key = PublicKeyBundle::from_bytes(&row.public_key) .map_err(|error| sqlx::Error::Decode(Box::new(error)))?; let decode = |value| String::from_utf8(value).map_err(|error| sqlx::Error::Decode(Box::new(error))); Ok(User { id: row.id.into(), iota_id: row.iota_id.into(), username: decode(row.username)?, display: row.display.map(decode).transpose()?, status: row.status.map(decode).transpose()?, presence_preference: decode(row.presence_preference)?, about: row.about.map(decode).transpose()?, avatar: row.avatar, sub_level: row.sub_level, sub_end: row.sub_end, public_key, token: decode(row.token)?, }) } } pub async fn get_by_username(username: &str) -> Result { let row = sqlx::query_as::<_, UserRow>(USER_BY_USERNAME_QUERY) .bind(username) .fetch_optional(&pool().await?) .await? .ok_or(OmegaError::NotFound)?; row.try_into().map_err(OmegaError::from) } pub async fn get_by_user_id(id: UserId) -> Result { let row = sqlx::query_as::<_, UserRow>(USER_BY_ID_QUERY) .bind(id.0) .fetch_optional(&pool().await?) .await? .ok_or(OmegaError::NotFound)?; row.try_into().map_err(OmegaError::from) } pub async fn get_users_by_iota_id(id: IotaId) -> Result> { get_users_by_iota_ids(&[id.0]).await } fn normalized_ids(ids: &[i64]) -> Vec { let mut ids = ids .iter() .copied() .filter(|id| valid_protocol_id(*id)) .collect::>(); ids.sort_unstable(); ids.dedup(); ids } fn append_in_clause(query: &mut QueryBuilder<'_, MySql>, ids: &[i64]) { query.push("("); for (index, id) in ids.iter().enumerate() { if index > 0 { query.push(", "); } query.push_bind(*id); } query.push(")"); } async fn fetch_users(mut query: QueryBuilder<'_, MySql>) -> Result> { query .build_query_as::() .fetch_all(&pool().await?) .await? .into_iter() .map(|row| row.try_into().map_err(OmegaError::from)) .collect() } pub async fn get_users_by_ids(ids: &[i64]) -> Result> { let ids = normalized_ids(ids); if ids.is_empty() { return Ok(Vec::new()); } let mut query = QueryBuilder::::new(USER_COLUMNS); query.push(" WHERE id IN "); append_in_clause(&mut query, &ids); fetch_users(query).await } pub async fn get_users_by_iota_ids(ids: &[i64]) -> Result> { let ids = normalized_ids(ids); if ids.is_empty() { return Ok(Vec::new()); } let mut query = QueryBuilder::::new(USER_COLUMNS); query.push(" WHERE iota_id IN "); append_in_clause(&mut query, &ids); fetch_users(query).await } pub async fn get_users_by_ids_and_iota_ids( user_ids: &[i64], iota_ids: &[i64], ) -> Result> { let user_ids = normalized_ids(user_ids); let iota_ids = normalized_ids(iota_ids); if user_ids.is_empty() && iota_ids.is_empty() { return Ok(Vec::new()); } let mut query = QueryBuilder::::new(USER_COLUMNS); query.push(" WHERE "); if !user_ids.is_empty() { query.push("id IN "); append_in_clause(&mut query, &user_ids); } if !iota_ids.is_empty() { if !user_ids.is_empty() { query.push(" OR "); } query.push("iota_id IN "); append_in_clause(&mut query, &iota_ids); } fetch_users(query).await } #[derive(FromRow)] struct PresencePreferenceRow { id: i64, presence_preference: Vec, } pub async fn get_presence_preferences(ids: &[i64]) -> Result> { let ids = normalized_ids(ids); if ids.is_empty() { return Ok(HashMap::new()); } let mut query = QueryBuilder::::new("SELECT id, presence_preference FROM users WHERE id IN "); append_in_clause(&mut query, &ids); let rows = query .build_query_as::() .fetch_all(&pool().await?) .await?; let mut preferences = HashMap::with_capacity(rows.len()); for row in rows { let status = String::from_utf8(row.presence_preference) .ok() .and_then(|value| UserStatus::from_client_preference(&value)); let status = match status { Some(status) => status, None => { crate::log_in!( crate::util::logger::PrintType::General, "Invalid persisted presence preference for user {}, using user_online", row.id ); UserStatus::user_online } }; preferences.insert(row.id, status); } Ok(preferences) } async fn update( id: UserId, query: &'static str, value: impl Send + sqlx::Encode<'static, sqlx::MySql> + sqlx::Type + 'static, ) -> Result<()> { sqlx::query(query) .bind(value) .bind(id.0) .execute(&pool().await?) .await?; Ok(()) } pub async fn change_username(id: UserId, value: String) -> Result<()> { update( id, "UPDATE users SET username = ? WHERE id = ?", value.into_bytes(), ) .await } pub async fn change_display_name(id: UserId, value: String) -> Result<()> { update( id, "UPDATE users SET display = ? WHERE id = ?", value.into_bytes(), ) .await } pub async fn change_avatar(id: UserId, value: String) -> Result<()> { update( id, "UPDATE users SET avatar = ? WHERE id = ?", value.into_bytes(), ) .await } pub async fn change_about(id: UserId, value: String) -> Result<()> { update( id, "UPDATE users SET about = ? WHERE id = ?", value.into_bytes(), ) .await } pub async fn change_status(id: UserId, value: String) -> Result<()> { update( id, "UPDATE users SET status = ? WHERE id = ?", value.into_bytes(), ) .await } pub async fn change_presence_preference(id: UserId, value: String) -> Result<()> { update( id, "UPDATE users SET presence_preference = ? WHERE id = ?", value.into_bytes(), ) .await } pub async fn change_iota_id(id: UserId, value: IotaId) -> Result<()> { sqlx::query("UPDATE users SET iota_id = ? WHERE id = ?") .bind(value.0) .bind(id.0) .execute(&pool().await?) .await?; Ok(()) } pub async fn change_token(id: UserId, value: String) -> Result<()> { update( id, "UPDATE users SET token = ? WHERE id = ?", value.into_bytes(), ) .await } pub async fn delete_user(id: UserId) -> Result<()> { sqlx::query("DELETE FROM users WHERE id = ?") .bind(id.0) .execute(&pool().await?) .await?; Ok(()) } pub async fn change_keys(id: UserId, public_key: PublicKeyBundle) -> Result<()> { sqlx::query("UPDATE users SET public_key = ? WHERE id = ?") .bind(public_key.as_bytes()) .bind(id.0) .execute(&pool().await?) .await?; Ok(()) } pub async fn register_complete_user( id: UserId, username: String, public_key: PublicKeyBundle, iota_id: IotaId, token: String, registration_token: String, ) -> Result<()> { if !valid_protocol_id(id.0) || !valid_protocol_id(iota_id.0) { return Err(OmegaError::Validation( "user or Iota ID is outside the 48-bit protocol range".into(), )); } if !valid_username(&username) { return Err(OmegaError::Validation("invalid username".into())); } let mut transaction = pool().await?.begin().await?; let lease = sqlx::query( "SELECT iota_id, completed_at IS NOT NULL AS completed, \ expires_at >= UTC_TIMESTAMP() AS current FROM registration_leases \ WHERE token = ? AND user_id = ? FOR UPDATE", ) .bind(®istration_token) .bind(id.0) .fetch_optional(&mut *transaction) .await? .ok_or_else(|| OmegaError::Validation("unknown registration lease".into()))?; let lease_iota_id: i64 = lease.get("iota_id"); let completed: i8 = lease.get("completed"); let current: i8 = lease.get("current"); if lease_iota_id != iota_id.0 || (completed == 0 && current == 0) { return Err(OmegaError::Validation( "expired or mismatched registration lease".into(), )); } let insert_result = sqlx::query( "INSERT INTO users (id, username, public_key, iota_id, token) VALUES (?, ?, ?, ?, ?)", ) .bind(id.0) .bind(username.as_bytes()) .bind(public_key.as_bytes()) .bind(iota_id.0) .bind(token.as_bytes()) .execute(&mut *transaction) .await; let result: Result<()> = match insert_result { Ok(_) => Ok(()), Err(insert_error) => { let existing = sqlx::query_as::<_, UserRow>(USER_BY_ID_QUERY) .bind(id.0) .fetch_optional(&mut *transaction) .await?; match existing .map(User::try_from) .transpose() .map_err(OmegaError::from)? { Some(existing) if existing.iota_id == iota_id && existing.username == username && existing.public_key.as_bytes() == public_key.as_bytes() && existing.token == token => { Ok(()) } _ => Err(insert_error.into()), } } }; result?; sqlx::query("UPDATE registration_leases SET completed_at = UTC_TIMESTAMP() WHERE token = ?") .bind(®istration_token) .execute(&mut *transaction) .await?; transaction.commit().await?; Ok(()) } fn valid_username(username: &str) -> bool { !username.is_empty() && username.chars().count() <= 15 && !username.chars().any(char::is_control) && !username.contains(['/', '\\']) } #[cfg(test)] mod tests { use super::{MAX_PROTOCOL_ID, get_register_id, valid_protocol_id}; #[tokio::test] async fn generated_registration_ids_fit_the_mtp_wire_range() { for _ in 0..128 { assert!(valid_protocol_id(get_register_id().await.unwrap().0)); } assert!(!valid_protocol_id(0)); assert!(valid_protocol_id(MAX_PROTOCOL_ID)); assert!(!valid_protocol_id(MAX_PROTOCOL_ID + 1)); } }