Skip to main content

max / makenotwork

4.7 KB · 164 lines History Blame Raw
1 //! Invite code generation and redemption queries.
2
3 use rand::RngExt;
4 use sqlx::PgPool;
5
6 use super::models::DbInviteCode;
7 use super::{InviteCodeId, UserId};
8 use crate::error::Result;
9
10 /// Charset for invite codes: uppercase alphanumeric minus ambiguous chars (I/O/0/1).
11 const CODE_CHARSET: &[u8] = b"ABCDEFGHJKLMNPQRSTUVWXYZ23456789";
12 const CODE_LENGTH: usize = 12;
13
14 /// Generate a random 12-character invite code from the unambiguous charset.
15 #[tracing::instrument(skip_all)]
16 pub fn generate_invite_code() -> String {
17 let mut rng = rand::rng();
18 (0..CODE_LENGTH)
19 .map(|_| CODE_CHARSET[rng.random_range(0..CODE_CHARSET.len())] as char)
20 .collect()
21 }
22
23 /// Format a raw 12-char code as `XXXX-XXXX-XXXX` for display.
24 #[tracing::instrument(skip_all)]
25 pub fn format_invite_code(code: &str) -> String {
26 let chars: Vec<char> = code.chars().collect();
27 if chars.len() != 12 {
28 return code.to_string();
29 }
30 format!("{}-{}-{}", &code[..4], &code[4..8], &code[8..12])
31 }
32
33 /// Insert a new invite code for a creator.
34 #[tracing::instrument(skip_all)]
35 pub async fn create_invite_code(
36 pool: &PgPool,
37 creator_id: UserId,
38 code: &str,
39 ) -> Result<DbInviteCode> {
40 // Store the hash, never the plaintext: a DB read must not yield a usable
41 // invite. The caller keeps the raw `code` to show/share with the creator.
42 let invite = sqlx::query_as::<_, DbInviteCode>(
43 r"
44 INSERT INTO invite_codes (creator_id, code)
45 VALUES ($1, $2)
46 RETURNING *
47 ",
48 )
49 .bind(creator_id)
50 .bind(crate::crypto::invite_code_hash(code))
51 .fetch_one(pool)
52 .await?;
53
54 Ok(invite)
55 }
56
57 /// Count unredeemed (active) invite codes for a creator.
58 #[tracing::instrument(skip_all)]
59 pub async fn count_active_invites(pool: &PgPool, creator_id: UserId) -> Result<i64> {
60 let count: (i64,) = sqlx::query_as(
61 "SELECT COUNT(*) FROM invite_codes WHERE creator_id = $1 AND redeemed_by_id IS NULL",
62 )
63 .bind(creator_id)
64 .fetch_one(pool)
65 .await?;
66
67 Ok(count.0)
68 }
69
70 /// List all invite codes created by a specific creator, newest first.
71 #[tracing::instrument(skip_all)]
72 pub async fn get_invites_by_creator(
73 pool: &PgPool,
74 creator_id: UserId,
75 ) -> Result<Vec<DbInviteCode>> {
76 let invites = sqlx::query_as::<_, DbInviteCode>(
77 "SELECT * FROM invite_codes WHERE creator_id = $1 ORDER BY created_at DESC",
78 )
79 .bind(creator_id)
80 .fetch_all(pool)
81 .await?;
82
83 Ok(invites)
84 }
85
86 /// Look up an unredeemed invite code by its raw code string. The raw code is
87 /// hashed before lookup (codes are stored hashed); the caller normalizes it
88 /// (uppercase, dashes stripped) first.
89 #[tracing::instrument(skip_all)]
90 pub async fn get_valid_invite_code(pool: &PgPool, code: &str) -> Result<Option<DbInviteCode>> {
91 let invite = sqlx::query_as::<_, DbInviteCode>(
92 "SELECT * FROM invite_codes WHERE code = $1 AND redeemed_by_id IS NULL",
93 )
94 .bind(crate::crypto::invite_code_hash(code))
95 .fetch_optional(pool)
96 .await?;
97
98 Ok(invite)
99 }
100
101 /// Atomically claim an invite code for a user. Returns `true` if this call won
102 /// the claim, `false` if the code was already redeemed (lost race / reused).
103 ///
104 /// The `AND redeemed_by_id IS NULL` guard makes the read-then-write in the
105 /// caller race-free: two concurrent redeemers both pass `get_valid_invite_code`
106 /// but only one satisfies the guarded UPDATE, so a single-use invite can never
107 /// onboard two users.
108 #[tracing::instrument(skip_all)]
109 pub async fn redeem_invite_code(
110 pool: &PgPool,
111 code_id: InviteCodeId,
112 user_id: UserId,
113 ) -> Result<bool> {
114 let claimed = sqlx::query_scalar::<_, InviteCodeId>(
115 "UPDATE invite_codes SET redeemed_by_id = $2, redeemed_at = NOW() \
116 WHERE id = $1 AND redeemed_by_id IS NULL RETURNING id",
117 )
118 .bind(code_id)
119 .bind(user_id)
120 .fetch_optional(pool)
121 .await?;
122
123 Ok(claimed.is_some())
124 }
125
126 #[cfg(test)]
127 mod tests {
128 use super::*;
129
130 #[test]
131 fn generate_code_length() {
132 let code = generate_invite_code();
133 assert_eq!(code.len(), 12);
134 }
135
136 #[test]
137 fn generate_code_uses_valid_charset() {
138 let code = generate_invite_code();
139 for c in code.chars() {
140 assert!(
141 CODE_CHARSET.contains(&(c as u8)),
142 "Invalid char '{c}' in generated code"
143 );
144 }
145 }
146
147 #[test]
148 fn generate_codes_are_unique() {
149 let a = generate_invite_code();
150 let b = generate_invite_code();
151 assert_ne!(a, b);
152 }
153
154 #[test]
155 fn format_code_xxxx_xxxx_xxxx() {
156 assert_eq!(format_invite_code("ABCD1234EFGH"), "ABCD-1234-EFGH");
157 }
158
159 #[test]
160 fn format_code_passthrough_on_wrong_length() {
161 assert_eq!(format_invite_code("SHORT"), "SHORT");
162 }
163 }
164