use oiseau::{cache::Cache, query_row, query_rows}; use tetratto_shared::unix_epoch_timestamp; use crate::model::{ Error, Result, auth::{User, InviteCode}, permissions::FinePermission, }; use crate::{auto_method, DataManager}; use oiseau::{PostgresRow, execute, get, params}; impl DataManager { /// Get a [`InviteCode`] from an SQL row. pub(crate) fn get_invite_code_from_row(x: &PostgresRow) -> InviteCode { InviteCode { id: get!(x->0(i64)) as usize, created: get!(x->1(i64)) as usize, owner: get!(x->2(i64)) as usize, code: get!(x->3(String)), is_used: get!(x->4(i32)) as i8 == 1, } } auto_method!(get_invite_code_by_id()@get_invite_code_from_row -> "SELECT * FROM invite_codes WHERE id = $1" --name="invite code" --returns=InviteCode --cache-key-tmpl="atto.invite_code:{}"); auto_method!(get_invite_code_by_code(&str)@get_invite_code_from_row -> "SELECT * FROM invite_codes WHERE code = $1" --name="invite code" --returns=InviteCode); /// Get invite_codes by `owner`. pub async fn get_invite_codes_by_owner( &self, owner: usize, batch: usize, page: usize, ) -> Result> { let conn = match self.0.connect().await { Ok(c) => c, Err(e) => return Err(Error::DatabaseConnection(e.to_string())), }; let res = query_rows!( &conn, "SELECT * FROM invite_codes WHERE owner = $1 ORDER BY created DESC LIMIT $2 OFFSET $3", &[&(owner as i64), &(batch as i64), &((page * batch) as i64)], |x| { Self::get_invite_code_from_row(x) } ); if res.is_err() { return Err(Error::GeneralNotFound("invite_code".to_string())); } Ok(res.unwrap()) } /// Get invite_codes by `owner`. pub async fn get_invite_codes_by_owner_count(&self, owner: usize) -> Result { let conn = match self.0.connect().await { Ok(c) => c, Err(e) => return Err(Error::DatabaseConnection(e.to_string())), }; let res = query_row!( &conn, "SELECT COUNT(*)::int FROM invite_codes WHERE owner = $1", &[&(owner as i64)], |x| Ok(x.get::(0)) ); if res.is_err() { return Err(Error::GeneralNotFound("invite_code".to_string())); } Ok(res.unwrap()) } /// Fill a vector of invite codes with the user that used them. pub async fn fill_invite_codes( &self, codes: Vec, ) -> Result, InviteCode)>> { let mut out = Vec::new(); for code in codes { if code.is_used { out.push(( match self.get_user_by_invite_code(code.id as i64).await { Ok(u) => Some(u), Err(_) => None, }, code, )) } else { out.push((None, code)) } } Ok(out) } const MAXIMUM_FREE_INVITE_CODES: usize = 4; const MAXIMUM_SUPPORTER_INVITE_CODES: usize = 48; const MINIMUM_ACCOUNT_AGE_FOR_INVITE_CODES: usize = 2_629_800_000; // 1mo /// Create a new invite_code in the database. /// /// # Arguments /// * `data` - a mock [`InviteCode`] object to insert pub async fn create_invite_code(&self, data: InviteCode, user: &User) -> Result { // check account creation date (if we aren't a supporter OR this is a purchased account) if !user.permissions.check(FinePermission::SUPPORTER) | user.was_purchased { if unix_epoch_timestamp() - user.created < Self::MINIMUM_ACCOUNT_AGE_FOR_INVITE_CODES { return Err(Error::MiscError( "Your account is too young to do this".to_string(), )); } } // ... if !user.permissions.check(FinePermission::SUPPORTER) { // our account is old enough, but we need to make sure we don't already have // 2 invite codes if (self.get_invite_codes_by_owner_count(user.id).await? as usize) >= Self::MAXIMUM_FREE_INVITE_CODES { return Err(Error::MiscError( "You already have the maximum number of invite codes you can create" .to_string(), )); } } else if !user.permissions.check(FinePermission::MANAGE_USERS) { // check count since we're also not a moderator with MANAGE_USERS if (self.get_invite_codes_by_owner_count(user.id).await? as usize) >= Self::MAXIMUM_SUPPORTER_INVITE_CODES { return Err(Error::MiscError( "You already have the maximum number of invite codes you can create" .to_string(), )); } } let conn = match self.0.connect().await { Ok(c) => c, Err(e) => return Err(Error::DatabaseConnection(e.to_string())), }; let res = execute!( &conn, "INSERT INTO invite_codes VALUES ($1, $2, $3, $4, $5)", params![ &(data.id as i64), &(data.created as i64), &(data.owner as i64), &data.code, &{ if data.is_used { 1 } else { 0 } } ] ); if let Err(e) = res { return Err(Error::DatabaseError(e.to_string())); } Ok(data) } pub async fn delete_invite_code(&self, id: usize, user: &User) -> Result<()> { if !user.permissions.check(FinePermission::MANAGE_USERS) { return Err(Error::NotAllowed); } let conn = match self.0.connect().await { Ok(c) => c, Err(e) => return Err(Error::DatabaseConnection(e.to_string())), }; let res = execute!( &conn, "DELETE FROM invite_codes WHERE id = $1", &[&(id as i64)] ); if let Err(e) = res { return Err(Error::DatabaseError(e.to_string())); } self.0.1.remove(format!("atto.invite_code:{}", id)).await; Ok(()) } pub async fn update_invite_code_is_used(&self, id: usize, new_is_used: bool) -> Result<()> { let conn = match self.0.connect().await { Ok(c) => c, Err(e) => return Err(Error::DatabaseConnection(e.to_string())), }; let res = execute!( &conn, "UPDATE invite_codes SET is_used = $1 WHERE id = $2", params![&{ if new_is_used { 1 } else { 0 } }, &(id as i64)] ); if let Err(e) = res { return Err(Error::DatabaseError(e.to_string())); } self.0.1.remove(format!("atto.invite_code:{}", id)).await; Ok(()) } }