tetratto_core/database/
auth.rs

1use super::common::NAME_REGEX;
2use oiseau::cache::Cache;
3use crate::model::{
4    Error, Result,
5    auth::{Token, User, UserSettings},
6    permissions::{FinePermission, SecondaryPermission},
7    oauth::AuthGrant,
8    moderation::AuditLogEntry,
9    auth::{
10        Achievement, AchievementName, AchievementRarity, Notification, UserConnections,
11        ACHIEVEMENTS,
12    },
13};
14use tetratto_shared::{
15    hash::{hash_salted, salt},
16    unix_epoch_timestamp,
17};
18use crate::{auto_method, DataManager};
19use oiseau::{PostgresRow, execute, get, query_row, params};
20
21macro_rules! update_role_fn {
22    ($name:ident, $role_ty:ty, $col:literal) => {
23        pub async fn $name(
24            &self,
25            id: usize,
26            role: $role_ty,
27            user: &User,
28            force: bool,
29        ) -> Result<()> {
30            let other_user = self.get_user_by_id(id).await?;
31
32            if !force {
33                // check permission
34                if !user.permissions.check(FinePermission::MANAGE_USERS) {
35                    return Err(Error::NotAllowed);
36                }
37
38                if other_user.permissions.check_manager() && !user.permissions.check_admin() {
39                    return Err(Error::MiscError(
40                        "Cannot manage the role of other managers".to_string(),
41                    ));
42                }
43
44                if other_user.permissions == user.permissions {
45                    return Err(Error::MiscError(
46                        "Cannot manage users of equal level to you".to_string(),
47                    ));
48                }
49            }
50
51            // ...
52            let conn = match self.0.connect().await {
53                Ok(c) => c,
54                Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
55            };
56
57            let res = execute!(
58                &conn,
59                &format!("UPDATE users SET {} = $1 WHERE id = $2", $col),
60                params![&(role.bits() as i32), &(id as i64)]
61            );
62
63            if let Err(e) = res {
64                return Err(Error::DatabaseError(e.to_string()));
65            }
66
67            self.cache_clear_user(&other_user).await;
68
69            // create audit log entry
70            self.create_audit_log_entry(AuditLogEntry::new(
71                user.id,
72                format!(
73                    "invoked `{}` with x value `{}` and y value `{}`",
74                    $col,
75                    other_user.id,
76                    role.bits()
77                ),
78            ))
79            .await?;
80
81            // ...
82            Ok(())
83        }
84    };
85}
86
87impl DataManager {
88    /// Get a [`User`] from an SQL row.
89    pub(crate) fn get_user_from_row(x: &PostgresRow) -> User {
90        User {
91            id: get!(x->0(i64)) as usize,
92            created: get!(x->1(i64)) as usize,
93            username: get!(x->2(String)),
94            password: get!(x->3(String)),
95            salt: get!(x->4(String)),
96            settings: serde_json::from_str(&get!(x->5(String)).to_string()).unwrap(),
97            tokens: serde_json::from_str(&get!(x->6(String)).to_string()).unwrap(),
98            permissions: FinePermission::from_bits(get!(x->7(i32)) as u32).unwrap(),
99            is_verified: get!(x->8(i32)) as i8 == 1,
100            notification_count: {
101                let x = get!(x->9(i32)) as usize;
102                // we're a little too close to the maximum count, clearly something's gone wrong
103                if x > usize::MAX - 1000 { 0 } else { x }
104            },
105            follower_count: get!(x->10(i32)) as usize,
106            following_count: get!(x->11(i32)) as usize,
107            last_seen: get!(x->12(i64)) as usize,
108            totp: get!(x->13(String)),
109            recovery_codes: serde_json::from_str(&get!(x->14(String)).to_string()).unwrap(),
110            post_count: get!(x->15(i32)) as usize,
111            request_count: {
112                let x = get!(x->16(i32)) as usize;
113                if x > usize::MAX - 1000 { 0 } else { x }
114            },
115            connections: serde_json::from_str(&get!(x->17(String)).to_string()).unwrap(),
116            stripe_id: get!(x->18(String)),
117            grants: serde_json::from_str(&get!(x->19(String)).to_string()).unwrap(),
118            associated: serde_json::from_str(&get!(x->20(String)).to_string()).unwrap(),
119            invite_code: get!(x->21(i64)) as usize,
120            secondary_permissions: SecondaryPermission::from_bits(get!(x->22(i32)) as u32).unwrap(),
121            achievements: serde_json::from_str(&get!(x->23(String)).to_string()).unwrap(),
122            awaiting_purchase: get!(x->24(i32)) as i8 == 1,
123            was_purchased: get!(x->25(i32)) as i8 == 1,
124            ban_reason: get!(x->26(String)),
125            is_deactivated: get!(x->27(i32)) as i8 == 1,
126            ban_expire: get!(x->28(i64)) as usize,
127            checkouts: serde_json::from_str(&get!(x->29(String)).to_string()).unwrap(),
128            last_policy_consent: get!(x->30(i64)) as usize,
129            close_friends_stack: get!(x->31(i64)) as usize,
130            missed_messages_count: get!(x->32(i32)) as usize,
131            views: get!(x->33(i32)) as usize,
132            shrimpcamp_link: get!(x->34(i64)) as usize,
133        }
134    }
135
136    auto_method!(get_user_by_id(usize as i64)@get_user_from_row -> "SELECT * FROM users WHERE id = $1" --name="user" --returns=User --cache-key-tmpl="atto.user:{}");
137    auto_method!(get_user_by_username(&str)@get_user_from_row -> "SELECT * FROM users WHERE username = $1" --name="user" --returns=User --cache-key-tmpl="atto.user:{}");
138    auto_method!(get_user_by_username_no_cache(&str)@get_user_from_row -> "SELECT * FROM users WHERE username = $1" --name="user" --returns=User);
139    auto_method!(get_user_by_browser_session(&str)@get_user_from_row -> "SELECT * FROM users WHERE browser_session = $1" --name="user" --returns=User);
140
141    /// Get a user given just their ID. Returns the void user if the user doesn't exist.
142    ///
143    /// # Arguments
144    /// * `id` - the ID of the user
145    pub async fn get_user_by_id_with_void(&self, id: usize) -> Result<User> {
146        let conn = match self.0.connect().await {
147            Ok(c) => c,
148            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
149        };
150
151        let res = query_row!(
152            &conn,
153            "SELECT * FROM users WHERE id = $1",
154            &[&(id as i64)],
155            |x| Ok(Self::get_user_from_row(x))
156        );
157
158        if res.is_err() {
159            return Ok(User::deleted());
160            // return Err(Error::UserNotFound);
161        }
162
163        Ok(res.unwrap())
164    }
165
166    /// Get a user given just their auth token.
167    ///
168    /// # Arguments
169    /// * `token` - the token of the user
170    pub async fn get_user_by_token(&self, token: &str) -> Result<User> {
171        let conn = match self.0.connect().await {
172            Ok(c) => c,
173            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
174        };
175
176        let res = query_row!(
177            &conn,
178            "SELECT * FROM users WHERE tokens LIKE $1",
179            &[&format!("%\"{token}\"%")],
180            |x| Ok(Self::get_user_from_row(x))
181        );
182
183        if res.is_err() {
184            return Err(Error::UserNotFound);
185        }
186
187        Ok(res.unwrap())
188    }
189
190    /// Get a user given just their grant token.
191    ///
192    /// Also returns the auth grant this token is associated with from the user.
193    ///
194    /// # Arguments
195    /// * `token` - the token of the user
196    pub async fn get_user_by_grant_token(
197        &self,
198        token: &str,
199        check_expiration: bool,
200    ) -> Result<(AuthGrant, User)> {
201        let conn = match self.0.connect().await {
202            Ok(c) => c,
203            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
204        };
205
206        let res = query_row!(
207            &conn,
208            "SELECT * FROM users WHERE grants LIKE $1",
209            &[&format!("%\"token\":\"{token}\"%")],
210            |x| Ok(Self::get_user_from_row(x))
211        );
212
213        if res.is_err() {
214            return Err(Error::UserNotFound);
215        }
216
217        let user = res.unwrap();
218        let grant = user
219            .grants
220            .iter()
221            .find(|x| x.token == token)
222            .unwrap()
223            .clone();
224
225        // check token expiry
226        if check_expiration {
227            let now = unix_epoch_timestamp();
228            let delta = now - grant.last_updated;
229
230            if delta > 604_800_000 {
231                return Err(Error::MiscError("Token expired".to_string()));
232            }
233        }
234
235        // ...
236        Ok((grant, user))
237    }
238
239    /// Create a new user in the database.
240    ///
241    /// # Arguments
242    /// * `data` - a mock [`User`] object to insert
243    pub async fn create_user(&self, mut data: User) -> Result<()> {
244        if !self.0.0.security.registration_enabled {
245            return Err(Error::RegistrationDisabled);
246        }
247
248        data.username = data.username.to_lowercase();
249
250        // check values
251        if data.username.len() < 2 {
252            return Err(Error::DataTooShort("username".to_string()));
253        } else if data.username.len() > 32 {
254            return Err(Error::DataTooLong("username".to_string()));
255        }
256
257        if data.password.len() < 6 {
258            return Err(Error::DataTooShort("password".to_string()));
259        }
260
261        if self.0.0.banned_usernames.contains(&data.username) {
262            return Err(Error::MiscError("This username cannot be used".to_string()));
263        }
264
265        let regex = regex::RegexBuilder::new(NAME_REGEX)
266            .multi_line(true)
267            .build()
268            .unwrap();
269
270        if regex.captures(&data.username).is_some() {
271            return Err(Error::MiscError(
272                "This username contains invalid characters".to_string(),
273            ));
274        }
275
276        // make sure username isn't taken
277        if self.get_user_by_username(&data.username).await.is_ok() {
278            return Err(Error::UsernameInUse);
279        }
280
281        // ...
282        let conn = match self.0.connect().await {
283            Ok(c) => c,
284            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
285        };
286
287        let res = execute!(
288            &conn,
289            "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, $23, $24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35)",
290            params![
291                &(data.id as i64),
292                &(data.created as i64),
293                &data.username.to_lowercase(),
294                &data.password,
295                &data.salt,
296                &serde_json::to_string(&data.settings).unwrap(),
297                &serde_json::to_string(&data.tokens).unwrap(),
298                &(FinePermission::DEFAULT.bits() as i32),
299                &if data.is_verified { 1_i32 } else { 0_i32 },
300                &0_i32,
301                &0_i32,
302                &0_i32,
303                &(data.last_seen as i64),
304                &String::new(),
305                "[]",
306                &0_i32,
307                &0_i32,
308                &serde_json::to_string(&data.connections).unwrap(),
309                &"",
310                &serde_json::to_string(&data.grants).unwrap(),
311                &serde_json::to_string(&data.associated).unwrap(),
312                &(data.invite_code as i64),
313                &(SecondaryPermission::DEFAULT.bits() as i32),
314                &serde_json::to_string(&data.achievements).unwrap(),
315                &if data.awaiting_purchase { 1_i32 } else { 0_i32 },
316                &if data.was_purchased { 1_i32 } else { 0_i32 },
317                &data.ban_reason,
318                &if data.is_deactivated { 1_i32 } else { 0_i32 },
319                &(data.ban_expire as i64),
320                &serde_json::to_string(&data.checkouts).unwrap(),
321                &(data.last_policy_consent as i64),
322                &(data.close_friends_stack as i64),
323                &(data.missed_messages_count as i32),
324                &(data.views as i32),
325                &(data.shrimpcamp_link as i64),
326            ]
327        );
328
329        if let Err(e) = res {
330            return Err(Error::DatabaseError(e.to_string()));
331        }
332
333        Ok(())
334    }
335
336    /// Delete an existing user in the database.
337    ///
338    /// # Arguments
339    /// * `id` - the ID of the user
340    /// * `password` - the current password of the user
341    /// * `force` - if we should delete even if the given password is incorrect
342    pub async fn delete_user(&self, id: usize, password: &str, force: bool) -> Result<User> {
343        let user = self.get_user_by_id(id).await?;
344
345        if (hash_salted(password.to_string(), user.salt.clone()) != user.password) && !force {
346            return Err(Error::IncorrectPassword);
347        }
348
349        let conn = match self.0.connect().await {
350            Ok(c) => c,
351            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
352        };
353
354        let res = execute!(&conn, "DELETE FROM users WHERE id = $1", &[&(id as i64)]);
355
356        if let Err(e) = res {
357            return Err(Error::DatabaseError(e.to_string()));
358        }
359
360        self.cache_clear_user(&user).await;
361
362        // delete communities
363        for community in self.get_communities_by_owner(user.id).await? {
364            self.delete_community(community.id, &user).await?;
365        }
366
367        // delete memberships
368        // member counts will remain the same... but that should probably be changed
369        let res = execute!(
370            &conn,
371            "DELETE FROM memberships WHERE owner = $1",
372            &[&(id as i64)]
373        );
374
375        if let Err(e) = res {
376            return Err(Error::DatabaseError(e.to_string()));
377        }
378
379        // delete notifications
380        let res = execute!(
381            &conn,
382            "DELETE FROM notifications WHERE owner = $1",
383            &[&(id as i64)]
384        );
385
386        if let Err(e) = res {
387            return Err(Error::DatabaseError(e.to_string()));
388        }
389
390        // delete requests
391        let res = execute!(
392            &conn,
393            "DELETE FROM requests WHERE owner = $1",
394            &[&(id as i64)]
395        );
396
397        if let Err(e) = res {
398            return Err(Error::DatabaseError(e.to_string()));
399        }
400
401        // delete warnings
402        let res = execute!(
403            &conn,
404            "DELETE FROM user_warnings WHERE receiver = $1",
405            &[&(id as i64)]
406        );
407
408        if let Err(e) = res {
409            return Err(Error::DatabaseError(e.to_string()));
410        }
411
412        // delete blocks
413        let res = execute!(
414            &conn,
415            "DELETE FROM userblocks WHERE initiator = $1 OR receiver = $1",
416            &[&(id as i64)]
417        );
418
419        if let Err(e) = res {
420            return Err(Error::DatabaseError(e.to_string()));
421        }
422
423        let res = execute!(
424            &conn,
425            "DELETE FROM ipblocks WHERE initiator = $1",
426            &[&(id as i64)]
427        );
428
429        if let Err(e) = res {
430            return Err(Error::DatabaseError(e.to_string()));
431        }
432
433        // delete reactions
434        // reactions counts will remain the same :)
435        let res = execute!(
436            &conn,
437            "DELETE FROM reactions WHERE owner = $1",
438            &[&(id as i64)]
439        );
440
441        if let Err(e) = res {
442            return Err(Error::DatabaseError(e.to_string()));
443        }
444
445        // delete stacks
446        let res = execute!(
447            &conn,
448            "DELETE FROM stacks WHERE owner = $1",
449            &[&(id as i64)]
450        );
451
452        if let Err(e) = res {
453            return Err(Error::DatabaseError(e.to_string()));
454        }
455
456        // delete drafts
457        let res = execute!(
458            &conn,
459            "DELETE FROM drafts WHERE owner = $1",
460            &[&(id as i64)]
461        );
462
463        if let Err(e) = res {
464            return Err(Error::DatabaseError(e.to_string()));
465        }
466
467        // delete posts
468        let res = execute!(&conn, "DELETE FROM posts WHERE owner = $1", &[&(id as i64)]);
469
470        if let Err(e) = res {
471            return Err(Error::DatabaseError(e.to_string()));
472        }
473
474        // delete polls
475        let res = execute!(&conn, "DELETE FROM polls WHERE owner = $1", &[&(id as i64)]);
476
477        if let Err(e) = res {
478            return Err(Error::DatabaseError(e.to_string()));
479        }
480
481        // delete poll votes
482        let res = execute!(
483            &conn,
484            "DELETE FROM pollvotes WHERE owner = $1",
485            &[&(id as i64)]
486        );
487
488        if let Err(e) = res {
489            return Err(Error::DatabaseError(e.to_string()));
490        }
491
492        // delete stackblocks
493        let res = execute!(
494            &conn,
495            "DELETE FROM stackblocks WHERE initiator = $1",
496            &[&(id as i64)]
497        );
498
499        if let Err(e) = res {
500            return Err(Error::DatabaseError(e.to_string()));
501        }
502
503        // delete invite codes
504        let res = execute!(
505            &conn,
506            "DELETE FROM invite_codes WHERE owner = $1",
507            &[&(id as i64)]
508        );
509
510        if let Err(e) = res {
511            return Err(Error::DatabaseError(e.to_string()));
512        }
513
514        // delete guest_logs
515        let res = execute!(
516            &conn,
517            "DELETE FROM guest_logs WHERE owner = $1",
518            &[&(id as i64)]
519        );
520
521        if let Err(e) = res {
522            return Err(Error::DatabaseError(e.to_string()));
523        }
524
525        // delete transfers
526        let res = execute!(
527            &conn,
528            "DELETE FROM transfers WHERE sender = $1 OR receiver = $1",
529            &[&(id as i64)]
530        );
531
532        if let Err(e) = res {
533            return Err(Error::DatabaseError(e.to_string()));
534        }
535
536        // delete products
537        let res = execute!(
538            &conn,
539            "DELETE FROM products WHERE owner = $1",
540            &[&(id as i64)]
541        );
542
543        if let Err(e) = res {
544            return Err(Error::DatabaseError(e.to_string()));
545        }
546
547        // delete letters
548        let res = execute!(
549            &conn,
550            "DELETE FROM letters WHERE owner = $1",
551            &[&(id as i64)]
552        );
553
554        if let Err(e) = res {
555            return Err(Error::DatabaseError(e.to_string()));
556        }
557
558        // delete ads
559        let res = execute!(&conn, "DELETE FROM ads WHERE owner = $1", &[&(id as i64)]);
560
561        if let Err(e) = res {
562            return Err(Error::DatabaseError(e.to_string()));
563        }
564
565        // delete user follows... individually since it requires updating user counts
566        for follow in self.get_userfollows_by_receiver_all(id).await? {
567            self.delete_userfollow(follow.id, &user, true).await?;
568        }
569
570        for follow in self.get_userfollows_by_initiator_all(id).await? {
571            self.delete_userfollow(follow.id, &user, true).await?;
572        }
573
574        // delete apps
575        for app in self.get_apps_by_owner(id).await? {
576            self.delete_app(app.id, &user).await?;
577        }
578
579        // delete uploads
580        for upload in match self.2.get_uploads_by_owner_all(user.id).await {
581            Ok(x) => x,
582            Err(e) => return Err(Error::MiscError(e.to_string())),
583        } {
584            if let Err(e) = self.2.delete_upload(upload.id).await {
585                return Err(Error::MiscError(e.to_string()));
586            }
587        }
588
589        // delete polls
590        for poll in self.get_polls_by_owner_all(user.id).await? {
591            self.delete_poll(poll.id, &user).await?;
592        }
593
594        // free up invite code
595        if self.0.0.security.enable_invite_codes
596            && user.invite_code != 0
597            && self.get_invite_code_by_id(user.invite_code).await.is_ok()
598        {
599            // we're checking if the code is ok because the owner might've deleted their account,
600            // deleting all of their invite codes as well
601            self.update_invite_code_is_used(user.invite_code, false)
602                .await?;
603        }
604
605        // ...
606        Ok(user)
607    }
608
609    pub async fn update_user_verified_status(&self, id: usize, x: bool, user: User) -> Result<()> {
610        if !user.permissions.check(FinePermission::MANAGE_VERIFIED) {
611            return Err(Error::NotAllowed);
612        }
613
614        let other_user = self.get_user_by_id(id).await?;
615
616        let conn = match self.0.connect().await {
617            Ok(c) => c,
618            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
619        };
620
621        let res = execute!(
622            &conn,
623            "UPDATE users SET verified = $1 WHERE id = $2",
624            params![&{ if x { 1 } else { 0 } }, &(id as i64)]
625        );
626
627        if let Err(e) = res {
628            return Err(Error::DatabaseError(e.to_string()));
629        }
630
631        self.cache_clear_user(&other_user).await;
632
633        // create audit log entry
634        self.create_audit_log_entry(AuditLogEntry::new(
635            user.id,
636            format!(
637                "invoked `update_user_verified_status` with x value `{}` and y value `{}`",
638                other_user.id, x
639            ),
640        ))
641        .await?;
642
643        // ...
644        Ok(())
645    }
646
647    pub async fn update_user_is_deactivated(&self, id: usize, x: bool, user: User) -> Result<()> {
648        if id != user.id && !user.permissions.check(FinePermission::MANAGE_USERS) {
649            return Err(Error::NotAllowed);
650        }
651
652        let other_user = self.get_user_by_id(id).await?;
653
654        let conn = match self.0.connect().await {
655            Ok(c) => c,
656            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
657        };
658
659        let res = execute!(
660            &conn,
661            "UPDATE users SET is_deactivated = $1 WHERE id = $2",
662            params![&{ if x { 1 } else { 0 } }, &(id as i64)]
663        );
664
665        if let Err(e) = res {
666            return Err(Error::DatabaseError(e.to_string()));
667        }
668
669        self.cache_clear_user(&other_user).await;
670
671        // create audit log entry (if we aren't the user that is being updated)
672        if user.id != other_user.id {
673            self.create_audit_log_entry(AuditLogEntry::new(
674                user.id,
675                format!(
676                    "invoked `update_user_is_deactivated` with x value `{}` and y value `{}`",
677                    other_user.id, x
678                ),
679            ))
680            .await?;
681        }
682
683        // ...
684        Ok(())
685    }
686
687    pub async fn update_user_password(
688        &self,
689        id: usize,
690        from: String,
691        to: String,
692        user: User,
693        force: bool,
694    ) -> Result<()> {
695        // verify password
696        if !user.check_password(from.clone()) && !force {
697            return Err(Error::MiscError("Password does not match".to_string()));
698        }
699
700        // ...
701        let conn = match self.0.connect().await {
702            Ok(c) => c,
703            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
704        };
705
706        let new_salt = salt();
707        let new_password = hash_salted(to, new_salt.clone());
708        let res = execute!(
709            &conn,
710            "UPDATE users SET password = $1, salt = $2 WHERE id = $3",
711            params![&new_password.as_str(), &new_salt.as_str(), &(id as i64)]
712        );
713
714        if let Err(e) = res {
715            return Err(Error::DatabaseError(e.to_string()));
716        }
717
718        self.cache_clear_user(&user).await;
719        Ok(())
720    }
721
722    pub async fn update_user_username(&self, id: usize, to: String, user: User) -> Result<()> {
723        // check value
724        if to.len() < 2 {
725            return Err(Error::DataTooShort("username".to_string()));
726        } else if to.len() > 32 {
727            return Err(Error::DataTooLong("username".to_string()));
728        }
729
730        if self.0.0.banned_usernames.contains(&to) {
731            return Err(Error::MiscError("This username cannot be used".to_string()));
732        }
733
734        let regex = regex::RegexBuilder::new(r"[^\w_\-\.!]+")
735            .multi_line(true)
736            .build()
737            .unwrap();
738
739        if regex.captures(&to).is_some() {
740            return Err(Error::MiscError(
741                "This username contains invalid characters".to_string(),
742            ));
743        }
744
745        // ...
746        let conn = match self.0.connect().await {
747            Ok(c) => c,
748            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
749        };
750
751        let res = execute!(
752            &conn,
753            "UPDATE users SET username = $1 WHERE id = $2",
754            params![&to.to_lowercase(), &(id as i64)]
755        );
756
757        if let Err(e) = res {
758            return Err(Error::DatabaseError(e.to_string()));
759        }
760
761        self.cache_clear_user(&user).await;
762        Ok(())
763    }
764
765    pub async fn update_user_awaiting_purchased_status(
766        &self,
767        id: usize,
768        x: bool,
769        user: User,
770        require_permission: bool,
771    ) -> Result<()> {
772        if (user.id != id) | require_permission
773            && !user.permissions.check(FinePermission::MANAGE_USERS)
774        {
775            return Err(Error::NotAllowed);
776        }
777
778        let other_user = self.get_user_by_id(id).await?;
779
780        let conn = match self.0.connect().await {
781            Ok(c) => c,
782            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
783        };
784
785        let res = execute!(
786            &conn,
787            "UPDATE users SET awaiting_purchase = $1 WHERE id = $2",
788            params![&{ if x { 1 } else { 0 } }, &(id as i64)]
789        );
790
791        if let Err(e) = res {
792            return Err(Error::DatabaseError(e.to_string()));
793        }
794
795        self.cache_clear_user(&other_user).await;
796
797        // create audit log entry
798        if user.id != other_user.id {
799            self.create_audit_log_entry(AuditLogEntry::new(
800                user.id,
801                format!(
802                    "invoked `update_user_purchased_status` with x value `{}` and y value `{}`",
803                    other_user.id, x
804                ),
805            ))
806            .await?;
807        }
808
809        // ...
810        Ok(())
811    }
812
813    pub async fn seen_user(&self, user: &User) -> Result<()> {
814        let conn = match self.0.connect().await {
815            Ok(c) => c,
816            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
817        };
818
819        let res = execute!(
820            &conn,
821            "UPDATE users SET last_seen = $1 WHERE id = $2",
822            params![&(unix_epoch_timestamp() as i64), &(user.id as i64)]
823        );
824
825        if let Err(e) = res {
826            return Err(Error::DatabaseError(e.to_string()));
827        }
828
829        self.cache_clear_user(user).await;
830
831        Ok(())
832    }
833
834    /// Add an achievement to a user.
835    ///
836    /// Still returns `Ok` if the user already has the achievement.
837    #[async_recursion::async_recursion]
838    pub async fn add_achievement(
839        &self,
840        user: &mut User,
841        achievement: Achievement,
842        check_for_final: bool,
843    ) -> Result<()> {
844        if user.settings.disable_achievements {
845            return Ok(());
846        }
847
848        if user.achievements.iter().any(|x| x.name == achievement.name) {
849            return Ok(());
850        }
851
852        // send notif
853        self.create_notification(Notification::new(
854            "You've earned a new achievement!".to_string(),
855            format!(
856                "You've earned the \"{}\" [achievement](/achievements)!",
857                achievement.name.title()
858            ),
859            user.id,
860        ))
861        .await?;
862
863        // add achievement
864        user.achievements.push(achievement);
865        self.update_user_achievements(user.id, user.achievements.to_owned())
866            .await?;
867
868        // check for final
869        if check_for_final && user.achievements.len() + 1 == ACHIEVEMENTS {
870            self.add_achievement(user, AchievementName::GetAllOtherAchievements.into(), false)
871                .await?;
872        }
873
874        // ...
875        Ok(())
876    }
877
878    /// Fill achievements with their title and description.
879    ///
880    /// # Returns
881    /// `(name, description, rarity, achievement)`
882    pub fn fill_achievements(
883        &self,
884        mut list: Vec<Achievement>,
885    ) -> Vec<(String, String, AchievementRarity, Achievement)> {
886        let mut out = Vec::new();
887
888        // sort by unlocked desc
889        list.sort_by(|a, b| a.unlocked.cmp(&b.unlocked));
890        list.reverse();
891
892        // ...
893        for x in list {
894            out.push((
895                x.name.title().to_string(),
896                x.name.description().to_string(),
897                x.name.rarity(),
898                x,
899            ))
900        }
901
902        out
903    }
904
905    /// Validate a given TOTP code for the given profile.
906    pub fn check_totp(&self, ua: &User, code: &str) -> bool {
907        let totp = ua.totp(Some(
908            self.0
909                .0
910                .host
911                .replace("http://", "")
912                .replace("https://", "")
913                .replace(":", "_"),
914        ));
915
916        if let Some(totp) = totp {
917            return !code.is_empty()
918                && (totp.check_current(code).unwrap()
919                    | ua.recovery_codes.contains(&code.to_string()));
920        }
921
922        true
923    }
924
925    /// Generate 8 random recovery codes for TOTP.
926    pub fn generate_totp_recovery_codes() -> Vec<String> {
927        let mut out: Vec<String> = Vec::new();
928
929        for _ in 0..9 {
930            out.push(salt())
931        }
932
933        out
934    }
935
936    /// Update the profile's TOTP secret.
937    ///
938    /// # Arguments
939    /// * `id` - the ID of the user
940    /// * `secret` - the TOTP secret
941    /// * `recovery` - the TOTP recovery codes
942    pub async fn update_user_totp(
943        &self,
944        id: usize,
945        secret: &str,
946        recovery: &Vec<String>,
947    ) -> Result<()> {
948        let user = self.get_user_by_id(id).await?;
949
950        // update
951        let conn = match self.0.connect().await {
952            Ok(c) => c,
953            Err(e) => return Err(Error::DatabaseConnection(e.to_string())),
954        };
955
956        let res = execute!(
957            &conn,
958            "UPDATE users SET totp = $1, recovery_codes = $2 WHERE id = $3",
959            params![
960                &secret,
961                &serde_json::to_string(recovery).unwrap(),
962                &(id as i64)
963            ]
964        );
965
966        if let Err(e) = res {
967            return Err(Error::DatabaseError(e.to_string()));
968        }
969
970        self.cache_clear_user(&user).await;
971        Ok(())
972    }
973
974    /// Enable TOTP for a profile.
975    ///
976    /// # Arguments
977    /// * `id` - the ID of the user to enable TOTP for
978    /// * `user` - the user doing this
979    ///
980    /// # Returns
981    /// `Result<(secret, qr base64)>`
982    pub async fn enable_totp(
983        &self,
984        id: usize,
985        user: User,
986    ) -> Result<(String, String, Vec<String>)> {
987        let other_user = self.get_user_by_id(id).await?;
988
989        if other_user.id != user.id {
990            if other_user.permissions.check(FinePermission::MANAGE_USERS) {
991                // create audit log entry
992                self.create_audit_log_entry(AuditLogEntry::new(
993                    user.id,
994                    format!("invoked `enable_totp` with x value `{}`", other_user.id,),
995                ))
996                .await?;
997            } else {
998                return Err(Error::NotAllowed);
999            }
1000        }
1001
1002        let secret = totp_rs::Secret::default().to_string();
1003        let recovery = Self::generate_totp_recovery_codes();
1004        self.update_user_totp(id, &secret, &recovery).await?;
1005
1006        // fetch profile again (with totp information)
1007        let other_user = self.get_user_by_id(id).await?;
1008
1009        // get totp
1010        let totp = other_user.totp(Some(
1011            self.0
1012                .0
1013                .host
1014                .replace("http://", "")
1015                .replace("https://", "")
1016                .replace(":", "_"),
1017        ));
1018
1019        if totp.is_none() {
1020            return Err(Error::MiscError("Failed to get TOTP code".to_string()));
1021        }
1022
1023        let totp = totp.unwrap();
1024
1025        // generate qr
1026        let qr = match totp.get_qr_base64() {
1027            Ok(q) => q,
1028            Err(e) => return Err(Error::MiscError(e.to_string())),
1029        };
1030
1031        // return
1032        Ok((totp.get_secret_base32(), qr, recovery))
1033    }
1034
1035    pub async fn cache_clear_user(&self, user: &User) {
1036        self.0.1.remove(format!("atto.user:{}", user.id)).await;
1037        self.0
1038            .1
1039            .remove(format!("atto.user:{}", user.username))
1040            .await;
1041    }
1042
1043    update_role_fn!(update_user_role, FinePermission, "permissions");
1044    update_role_fn!(
1045        update_user_secondary_role,
1046        SecondaryPermission,
1047        "secondary_permissions"
1048    );
1049
1050    auto_method!(update_user_tokens(Vec<Token>)@get_user_by_id -> "UPDATE users SET tokens = $1 WHERE id = $2" --serde --cache-key-tmpl=cache_clear_user);
1051    auto_method!(update_user_grants(Vec<AuthGrant>)@get_user_by_id -> "UPDATE users SET grants = $1 WHERE id = $2" --serde --cache-key-tmpl=cache_clear_user);
1052    auto_method!(update_user_settings(UserSettings)@get_user_by_id -> "UPDATE users SET settings = $1 WHERE id = $2" --serde --cache-key-tmpl=cache_clear_user);
1053    auto_method!(update_user_connections(UserConnections)@get_user_by_id -> "UPDATE users SET connections = $1 WHERE id = $2" --serde --cache-key-tmpl=cache_clear_user);
1054    auto_method!(update_user_associated(Vec<usize>)@get_user_by_id -> "UPDATE users SET associated = $1 WHERE id = $2" --serde --cache-key-tmpl=cache_clear_user);
1055    auto_method!(update_user_achievements(Vec<Achievement>)@get_user_by_id -> "UPDATE users SET achievements = $1 WHERE id = $2" --serde --cache-key-tmpl=cache_clear_user);
1056    auto_method!(update_user_invite_code(i64)@get_user_by_id -> "UPDATE users SET invite_code = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
1057    auto_method!(update_user_browser_session(&str)@get_user_by_id -> "UPDATE users SET browser_session = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
1058    auto_method!(update_user_ban_reason(&str)@get_user_by_id -> "UPDATE users SET ban_reason = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
1059    auto_method!(update_user_channel_mutes(Vec<usize>)@get_user_by_id -> "UPDATE users SET channel_mutes = $1 WHERE id = $2" --serde --cache-key-tmpl=cache_clear_user);
1060    auto_method!(update_user_ban_expire(i64)@get_user_by_id -> "UPDATE users SET ban_expire = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
1061    auto_method!(update_user_checkouts(Vec<String>)@get_user_by_id -> "UPDATE users SET checkouts = $1 WHERE id = $2" --serde --cache-key-tmpl=cache_clear_user);
1062    auto_method!(update_user_last_policy_consent(i64)@get_user_by_id -> "UPDATE users SET last_policy_consent = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
1063    auto_method!(update_user_close_friends_stack(i64)@get_user_by_id -> "UPDATE users SET close_friends_stack = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
1064    auto_method!(update_user_shrimpcamp_link(i64)@get_user_by_id -> "UPDATE users SET shrimpcamp_link = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
1065
1066    auto_method!(get_user_by_stripe_id(&str)@get_user_from_row -> "SELECT * FROM users WHERE stripe_id = $1" --name="user" --returns=User);
1067    auto_method!(update_user_stripe_id(&str)@get_user_by_id -> "UPDATE users SET stripe_id = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
1068
1069    auto_method!(update_user_notification_count(i32)@get_user_by_id -> "UPDATE users SET notification_count = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
1070    auto_method!(incr_user_notifications()@get_user_by_id -> "UPDATE users SET notification_count = notification_count + 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --incr);
1071    auto_method!(decr_user_notifications()@get_user_by_id -> "UPDATE users SET notification_count = notification_count - 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --decr=notification_count);
1072
1073    auto_method!(incr_user_follower_count()@get_user_by_id -> "UPDATE users SET follower_count = follower_count + 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --incr);
1074    auto_method!(decr_user_follower_count()@get_user_by_id -> "UPDATE users SET follower_count = follower_count - 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --decr=follower_count);
1075
1076    auto_method!(incr_user_following_count()@get_user_by_id -> "UPDATE users SET following_count = following_count + 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --incr);
1077    auto_method!(decr_user_following_count()@get_user_by_id -> "UPDATE users SET following_count = following_count - 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --decr=following_count);
1078
1079    auto_method!(incr_user_post_count()@get_user_by_id -> "UPDATE users SET post_count = post_count + 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --incr);
1080    auto_method!(decr_user_post_count()@get_user_by_id -> "UPDATE users SET post_count = post_count - 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --decr=post_count);
1081
1082    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);
1083    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);
1084    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);
1085
1086    auto_method!(get_user_by_invite_code(i64)@get_user_from_row -> "SELECT * FROM users WHERE invite_code = $1" --name="user" --returns=User);
1087
1088    auto_method!(update_user_missed_messages_count(i32)@get_user_by_id -> "UPDATE users SET missed_messages_count = $1 WHERE id = $2" --cache-key-tmpl=cache_clear_user);
1089    auto_method!(incr_user_missed_messages()@get_user_by_id -> "UPDATE users SET missed_messages_count = missed_messages_count + 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --incr);
1090    auto_method!(decr_user_missed_messages()@get_user_by_id -> "UPDATE users SET missed_messages_count = missed_messages_count - 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --decr=notification_count);
1091
1092    auto_method!(incr_profile_views()@get_user_by_id -> "UPDATE users SET views = views + 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --incr);
1093    auto_method!(decr_profile_views()@get_user_by_id -> "UPDATE users SET views = views - 1 WHERE id = $1" --cache-key-tmpl=cache_clear_user --decr=views);
1094}
tetratto_core/database/auth.rs

tetratto_core/database/
auth.rs
