t/tetratto
1
0
Fork 0
Code Issues Pull requests Projects Releases 4 Packages Wiki Activity Actions

add: invite codes

Browse source
This commit is contained in:
trisua 2025-06-22 13:03:02 -04:00
parent d1a074eaeb
commit 626c6711ef
19 changed files with 410 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

9
crates/core/src/config.rs
View file

@ -13,6 +13,9 @@ pub struct SecurityConfig {
/// The name of the header which will contain the real IP of the connecting user.
#[serde(default = "default_real_ip_header")]
pub real_ip_header: String,
/// If users require an invite code to register. Invite codes can be generated by supporters.
#[serde(default = "default_enable_invite_codes")]
pub enable_invite_codes: bool,
}
fn default_security_registration_enabled() -> bool {
@ -23,11 +26,16 @@ fn default_real_ip_header() -> String {
"CF-Connecting-IP".to_string()
}
fn default_enable_invite_codes() -> bool {
false
}
impl Default for SecurityConfig {
fn default() -> Self {
Self {
registration_enabled: default_security_registration_enabled(),
real_ip_header: default_real_ip_header(),
enable_invite_codes: default_enable_invite_codes(),
}
}
}
@ -341,6 +349,7 @@ fn default_banned_usernames() -> Vec<String> {
"stacks".to_string(),
"stack".to_string(),
"search".to_string(),
"journals".to_string(),
]
}

6
crates/core/src/database/auth.rs
View file

@ -45,6 +45,7 @@ impl DataManager {
stripe_id: get!(x->18(String)),
grants: serde_json::from_str(&get!(x->19(String)).to_string()).unwrap(),
associated: serde_json::from_str(&get!(x->20(String)).to_string()).unwrap(),
invite_code: get!(x->21(i64)) as usize,
}
}
@ -200,7 +201,7 @@ impl DataManager {
let res = execute!(
&conn,
"INSERT INTO users VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21)",
"INSERT INTO users VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22)",
params![
&(data.id as i64),
&(data.created as i64),
@ -223,6 +224,7 @@ impl DataManager {
&"",
&serde_json::to_string(&data.grants).unwrap(),
&serde_json::to_string(&data.associated).unwrap(),
&(data.invite_code as i64)
]
);
@ -842,4 +844,6 @@ impl DataManager {
auto_method!(update_user_request_count(i32)@get_user_by_id -> "UPDATE users SET request_count = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
auto_method!(incr_user_request_count()@get_user_by_id -> "UPDATE users SET request_count = request_count + 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --incr);
auto_method!(decr_user_request_count()@get_user_by_id -> "UPDATE users SET request_count = request_count - 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --decr=request_count);
auto_method!(get_user_by_invite_code(i64)@get_user_from_row -> "SELECT * FROM users WHERE invite_code = $1" --name="user" --returns=User);
}

4
crates/core/src/database/common.rs
View file

@ -39,6 +39,7 @@ impl DataManager {
execute!(&conn, common::CREATE_TABLE_JOURNALS).unwrap();
execute!(&conn, common::CREATE_TABLE_NOTES).unwrap();
execute!(&conn, common::CREATE_TABLE_MESSAGE_REACTIONS).unwrap();
execute!(&conn, common::CREATE_TABLE_INVITE_CODES).unwrap();
self.0
.1
@ -115,7 +116,8 @@ macro_rules! auto_method {
Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
};
let res = query_row!(&conn, $query, &[&selector], |x| { Ok(Self::$select_fn(x)) });
let res =
oiseau::query_row!(&conn, $query, &[&selector], |x| { Ok(Self::$select_fn(x)) });
if res.is_err() {
return Err(Error::GeneralNotFound($name_.to_string()));

1
crates/core/src/database/drivers/common.rs
View file

@ -26,3 +26,4 @@ pub const CREATE_TABLE_STACKBLOCKS: &str = include_str!("./sql/create_stackblock
pub const CREATE_TABLE_JOURNALS: &str = include_str!("./sql/create_journals.sql");
pub const CREATE_TABLE_NOTES: &str = include_str!("./sql/create_notes.sql");
pub const CREATE_TABLE_MESSAGE_REACTIONS: &str = include_str!("./sql/create_message_reactions.sql");
pub const CREATE_TABLE_INVITE_CODES: &str = include_str!("./sql/create_invite_codes.sql");

7
crates/core/src/database/drivers/sql/create_invite_codes.sql Normal file
View file

@ -0,0 +1,7 @@
CREATE TABLE IF NOT EXISTS invite_codes (
id BIGINT NOT NULL PRIMARY KEY,
created BIGINT NOT NULL,
owner BIGINT NOT NULL,
code TEXT NOT NULL,
is_used INT NOT NULL
)

159
crates/core/src/database/invite_codes.rs Normal file
View file

@ -0,0 +1,159 @@
use oiseau::{cache::Cache, query_rows};
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) -> Result<Vec<InviteCode>> {
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",
&[&(owner 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())
}
/// Fill a vector of invite codes with the user that used them.
pub async fn fill_invite_codes(
&self,
codes: Vec<InviteCode>,
) -> Result<Vec<(Option<User>, 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_SUPPORTER_INVITE_CODES: usize = 15;
/// 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<InviteCode> {
if !user.permissions.check(FinePermission::SUPPORTER) {
return Err(Error::RequiresSupporter);
} else {
// check count
if self.get_invite_codes_by_owner(user.id).await?.len()
>= 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_INVITES) {
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(())
}
}

1
crates/core/src/database/mod.rs
View file

@ -8,6 +8,7 @@ pub mod connections;
mod drafts;
mod drivers;
mod emojis;
mod invite_codes;
mod ipbans;
mod ipblocks;
mod journals;

2
crates/core/src/model/addr.rs
View file

@ -1,7 +1,7 @@
use std::net::SocketAddr;
/// How many bytes should be taken as the prefix (from the begining of the address).
pub(crate) const IPV6_PREFIX_BYTES: usize = 16;
pub(crate) const IPV6_PREFIX_BYTES: usize = 11;
/// The protocol of a [`RemoteAddr`].
#[derive(Clone, Debug, PartialEq, Eq)]

26
crates/core/src/model/auth.rs
View file

@ -49,6 +49,9 @@ pub struct User {
/// A list of the IDs of all accounts the user has signed into through the UI.
#[serde(default)]
pub associated: Vec<usize>,
/// The ID of the [`InviteCode`] this user provided during registration.
#[serde(default)]
pub invite_code: usize,
}
pub type UserConnections =
@ -283,6 +286,7 @@ impl User {
stripe_id: String::new(),
grants: Vec::new(),
associated: Vec::new(),
invite_code: 0,
}
}
@ -591,3 +595,25 @@ impl UserWarning {
}
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct InviteCode {
pub id: usize,
pub created: usize,
pub owner: usize,
pub code: String,
pub is_used: bool,
}
impl InviteCode {
/// Create a new [`InviteCode`].
pub fn new(owner: usize) -> Self {
Self {
id: Snowflake::new().to_string().parse::<usize>().unwrap(),
created: unix_epoch_timestamp(),
owner,
code: salt(),
is_used: false,
}
}
}

1
crates/core/src/model/permissions.rs
View file

@ -39,6 +39,7 @@ bitflags! {
const MANAGE_APPS = 1 << 28;
const MANAGE_JOURNALS = 1 << 29;
const MANAGE_NOTES = 1 << 30;
const MANAGE_INVITES = 1 << 31;
const _ = !0;
}