Database implementation, migration, user rework, animated server icons

src/main.rs

@@ -2,13 +2,14 @@
 extern crate rocket;
 
 use poise::serenity_prelude::Http;
-use rocket::fs::{relative, FileServer};
+use rocket::{fs::{relative, FileServer}, Rocket, Ignite};
 use rocket_dyn_templates::{tera, Template};
 use sass_rocket_fairing::SassFairing;
 use std::{collections::HashMap, env};
+use clap::{Parser, Subcommand};
 
 mod models;
-use models::Settings;
+use models::{Settings, Database};
 
 mod utils;
 
@@ -24,12 +25,52 @@ use crate::i18n::langs_filter;
 
 mod prelude;
 
-#[launch]
-async fn rocket() -> _ {
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+#[command(propagate_version = true)]
+struct Args {
+    #[command(subcommand)]
+    cmd: Option<Command>,
+}
+
+#[derive(Subcommand)]
+enum Command {
+    LoadLegacy,
+}
+
+#[tokio::main]
+async fn main() {
     dotenv::dotenv().expect("Failed to load .env file");
 
+    let args = Args::parse();
+
+    match &args.cmd {
+        Some(cmd) => match cmd {
+            Command::LoadLegacy => {
+                if Database::file_exists() {
+                    println!("Cannot load legacy submissions to an existing database");
+                    return;
+                }
+                let http = http();
+                Database::new(false)
+                    .expect("Failed to load database")
+                    .load_legacy(&http).await
+                    .expect("Failed to load legacy submissions");
+            },
+        },
+        None => {
+            rocket().await.expect("Failed to launch rocket");
+        },
+    }
+}
+
+fn http() -> Http {
     let token = env::var("DISCORD_TOKEN").expect("Expected a token in the environment");
-    let http = Http::new(&token);
+    Http::new(&token)
+}
+
+async fn rocket() -> Result<Rocket<Ignite>, rocket::Error> {
+    let http = http();
 
     let config = rocket::Config::figment().merge(("port", 1313)).merge((
         "secret_key",
@@ -63,4 +104,7 @@ async fn rocket() -> _ {
             )
         }))
         .attach(SassFairing::default())
+        .ignite().await
+        .unwrap()
+        .launch().await
 }

src/models/database.rs

@@ -0,0 +1,153 @@
+use std::{fs::File, io::Read, collections::HashMap, path::Path};
+
+use poise::serenity_prelude::Http;
+use rusqlite::{Connection, params};
+
+use crate::{utils::get_challenge_number, models::User};
+
+use super::{LegacySubmission, Submission};
+
+pub struct Database {
+    conn: Connection,
+}
+
+const DATABASE_FILENAME: &str = "database.db";
+
+impl Database {
+    pub fn file_exists() -> bool {
+        Path::new(DATABASE_FILENAME).exists()
+    }
+
+    pub fn new(
+        testing: bool,
+    ) -> rusqlite::Result<Self> {
+        let conn = if testing {
+            Connection::open_in_memory()
+        } else {
+            Connection::open(DATABASE_FILENAME)
+        }?;
+        conn.execute(
+            "CREATE TABLE IF NOT EXISTS Submission (
+                id        INTEGER PRIMARY KEY,
+                author_id INTEGER NOT NULL,
+                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
+                image     TEXT NOT NULL,
+                challenge INTEGER NOT NULL,
+                FOREIGN KEY (author_id) REFERENCES User(id)
+            )",
+            params![],
+        )?;
+        conn.execute(
+            "CREATE TABLE IF NOT EXISTS User (
+                id            INTEGER PRIMARY KEY,
+                name          TEXT NOT NULL,
+                discriminator INTEGER NOT NULL,
+                avatar        TEXT
+            )",
+            params![],
+        )?;
+        Ok(Self { conn })
+    }
+
+    pub fn has_submitted(&self, user_id: u64) -> rusqlite::Result<bool> {
+        Ok(self.conn
+            .prepare("SELECT 1 FROM User WHERE id = ?1 LIMIT 1")?
+            .query_row(params![user_id], |_row| Ok(()))
+            .is_ok())
+    }
+
+    pub async fn load_legacy(&self, http: &Http) -> rusqlite::Result<()> {
+        let latest_challenge = get_challenge_number();
+        // HashMap of archived users that are no longer sharing a server with 字ちゃん
+        // Their historical usernames and discriminators will be used
+        let mut archived_users = HashMap::new();
+        for n in 1..=latest_challenge {
+            println!("Loading legacy challenge {n}/{latest_challenge}...");
+            let mut file = File::open(format!("data/challenges/{n}.json")).unwrap();
+            let mut contents = String::new();
+            file.read_to_string(&mut contents).unwrap();
+            for (legacy, submissions) in serde_json::from_str::<Vec<LegacySubmission>>(&contents)
+                .unwrap()
+                .iter()
+                .map(|legacy| (legacy, legacy.parse().unwrap())) {
+                    let mut already_updated = false;
+                    for submission in submissions {
+                        self.conn.execute(
+                            "INSERT INTO Submission(author_id, timestamp, image, challenge) VALUES (?1, ?2, ?3, ?4)",
+                            params![
+                                &submission.author_id,
+                                &submission.timestamp,
+                                &submission.image,
+                                n,
+                            ]
+                        )?;
+                        let id = submission.author_id;
+                        if !self.has_submitted(id)? {
+                            println!("Fetching user {id}...");
+                            let previously_archived = archived_users.contains_key(&id);
+                            // Parse archived user out of legacy and insert into HashMap
+                            let mut archive = || {
+                                if already_updated {
+                                    return;
+                                }
+                                if previously_archived {
+                                    println!("Updating archived data for user {id}");
+                                } else {
+                                    println!("Adding archived data for user {id}");
+                                }
+                                let (name, discriminator) = {
+                                    let mut iter = legacy.username.split("#");
+                                    let name = iter.next().unwrap().to_owned();
+                                    let discriminator = iter
+                                        .next()
+                                        .map(|str| str.parse().unwrap())
+                                        .unwrap_or(0);
+                                    (name, discriminator)
+                                };
+                                archived_users.insert(id, User {
+                                    id,
+                                    name,
+                                    discriminator,
+                                    avatar: None,
+                                });
+                                already_updated = true;
+                            };
+                            if previously_archived {
+                                // If it already contains the archived user,
+                                // overwrite write their data since they may have updated
+                                // their username/discriminator since their previous submission
+                                archive();
+                            } else {
+                                match User::fetch(&http, submission.author_id).await {
+                                    Ok(user) => {
+                                        self.conn.execute(
+                                            "INSERT INTO User(id, name, discriminator, avatar) VALUES (?1, ?2, ?3, ?4)",
+                                            params![user.id, user.name, user.discriminator, user.avatar]
+                                        )?;
+                                    },
+                                    Err(error) => {
+                                        println!("Failed to fetch user {}, may update archived data: {error}", submission.author_id);
+                                        archive();
+                                    },
+                                };
+                            }
+                        }
+                    }
+                }
+        }
+        Ok(())
+    }
+
+    #[allow(dead_code)]
+    pub fn refresh_users(&self) -> rusqlite::Result<()> {
+        // Periodically refresh all changable user data (name, discriminator, avatar)
+        // Ideally this should run periodically.
+        todo!()
+    }
+
+    #[allow(dead_code, unused_variables)]
+    pub fn insert_submission(&self, submission: &Submission) -> rusqlite::Result<()> {
+        // For new submissions only
+        todo!()
+    }
+}

src/models/mod.rs

@@ -1,8 +1,14 @@
 mod user;
-pub use user::User;
+pub use user::{Username, User, SessionUser};
 
 mod challenge;
 pub use challenge::Challenge;
 
 mod settings;
 pub use settings::Settings;
+
+mod submission;
+pub use submission::*;
+
+mod database;
+pub use database::Database;

src/models/settings.rs

@@ -50,7 +50,12 @@ pub struct Guild {
 impl Guild {
     pub async fn load(&mut self, http: &Http) -> poise::serenity_prelude::Result<()> {
         let server = http.get_guild(self.id).await?;
-        self.icon = Some(server.icon);
+        self.icon = Some(server.icon_url().map(|icon| if icon.contains("/a_") {
+            // serenity only gives non-animated URL
+            icon.replace("webp", "gif")
+        } else {
+            icon
+        }));
         self.name = Some(server.name);
         Ok(())
     }

src/models/submission/legacy_submission.rs

@@ -0,0 +1,52 @@
+use chrono::{Utc, TimeZone};
+use serde::Deserialize;
+use derive_more::From;
+
+use super::Submission;
+
+#[derive(Deserialize)]
+pub struct LegacySubmission {
+    pub id: String,
+    pub images: Vec<String>,
+    pub username: String,
+}
+
+#[derive(From, Debug)]
+pub enum LegacySubmissionParseError {
+    BadAuthorId(std::num::ParseIntError),
+}
+
+impl LegacySubmission {
+    pub fn parse(&self) -> Result<Vec<Submission>, LegacySubmissionParseError> {
+        let author_id = self.id.parse()?;
+        Ok(self.images
+            .iter()
+            .map(|image| {
+                Submission {
+                    author_id,
+                    timestamp: {
+                        // The last number of the filename is either a discriminator or a datestamp
+                        // We can split apart the filename and check if it's >9999. In that case,
+                        // it's a datestamp.
+                        (|| image
+                            // Get filename without extension
+                            .split(".")
+                            .nth(0)?
+                            // Get last number
+                            .split("-")
+                            .last()?
+                            .parse()
+                            .ok()
+                            // Check if discriminator or timestamp, then convert
+                            .map(|number| if number > 9999 {
+                                Utc.timestamp_millis_opt(number).single()
+                            } else {
+                                None
+                            })?)()
+                    },
+                    image: image.clone(),
+                }
+            })
+            .collect())
+    }
+}

src/models/submission/mod.rs

@@ -0,0 +1,5 @@
+mod submission;
+pub use submission::Submission;
+
+mod legacy_submission;
+pub use legacy_submission::LegacySubmission;

src/models/submission/submission.rs

@@ -0,0 +1,15 @@
+use chrono::{Utc, DateTime};
+
+// Challenge submission
+// In the legacy site version, one submission held 1 or more images.
+// Now, 1 submission = 1 image, and the leaderboard count will be labeled as "participations"
+pub struct Submission {
+    pub author_id: u64,
+    // Some fields might be empty for legacy submissions
+    // Starting during challenge #87, submissions have a datestamp appended to their filename. 
+    // TODO: Determine whether this datestamp is local time or UTC
+    pub timestamp: Option<DateTime<Utc>>,
+    // Image path relative to submission folder.
+    // Thumbnail image path will be determined from this.
+    pub image: String,
+}

src/models/user/mod.rs

@@ -5,14 +5,19 @@ mod tests;
 
 use chrono::Utc;
 use derive_more::From;
+use poise::serenity_prelude::{self, UserId, Http};
 use reqwest::StatusCode;
 use rocket::http::{Cookie, CookieJar};
 use serial::*;
 
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 
 use crate::cookies::{token::*, user::*};
 
+pub trait Username {
+    fn username(&self) -> String;
+}
+
 #[derive(Default, Deserialize)]
 pub struct User {
     #[serde(deserialize_with = "deserialize_id")]
@@ -21,17 +26,56 @@ pub struct User {
     pub name: String,
     #[serde(deserialize_with = "deserialize_discriminator")]
     pub discriminator: u16,
-    pub avatar: String,
+    pub avatar: Option<String>,
 }
 
-impl User {
-    pub fn username(&self) -> String {
+impl Username for User {
+    fn username(&self) -> String {
         if self.discriminator == 0 {
             return self.name.clone();
         }
         format!("{}#{:0>4}", self.name, self.discriminator)
     }
+}
 
+impl User {
+    pub async fn fetch(http: &Http, id: u64) -> serenity_prelude::Result<Self> {
+        let user = UserId(id).to_user(http).await?;
+        Ok(Self {
+            id,
+            avatar: user.avatar,
+            name: user.name,
+            discriminator: user.discriminator,
+        })
+    }
+
+    fn avatar(&self) -> String {
+        match &self.avatar {
+            // https://cdn.discordapp.com/avatars/67795786229878784/e58524afe21b5058fc6f3cdc19aea8e1.webp?size=1024
+            Some(avatar) => format!(
+                "https://cdn.discordapp.com/avatars/{}/{}.{}?size=1024",
+                self.id,
+                avatar,
+                if avatar.starts_with("a_") { "gif "} else { "webp" }
+            ),
+            // Archived user or user with no avatar, calculate default avatar
+            // https://www.reddit.com/r/discordapp/comments/au6v4e/comment/eh61dm6/
+            // https://docs.rs/serenity/0.11.5/serenity/model/user/struct.User.html#method.default_avatar_url
+            None => format!("https://cdn.discordapp.com/embed/avatars/{}.png", self.discriminator % 5),
+        }
+    }
+}
+
+#[derive(Default, Serialize)]
+pub struct SessionUser(pub User);
+
+impl Username for SessionUser {
+    fn username(&self) -> String {
+        self.0.username()
+    }
+}
+
+impl SessionUser {
     pub async fn init(token: &str, cookies: &CookieJar<'_>) -> Result<Self, GetUserError> {
         let (status, text) = {
             let response = reqwest::Client::new()
@@ -47,7 +91,7 @@ impl User {
                 message: text.ok(),
             });
         }
-        let user: Self = serde_json::from_str(&text?)?;
+        let user: User = serde_json::from_str(&text?)?;
         cookies.add_private(Cookie::new(TOKEN_COOKIE, token.to_owned()));
         cookies.add_private(Cookie::new(USER_ID_COOKIE, user.id.to_string()));
         cookies.add_private(Cookie::new(USER_NAME_COOKIE, user.name.clone()));
@@ -55,8 +99,8 @@ impl User {
             USER_DISCRIMINATOR_COOKIE,
             user.discriminator.to_string(),
         ));
-        cookies.add_private(Cookie::new(USER_AVATAR_COOKIE, user.avatar.clone()));
-        Ok(user)
+        cookies.add_private(Cookie::new(USER_AVATAR_COOKIE, user.avatar()));
+        Ok(Self(user))
     }
 
     pub fn purge(cookies: &CookieJar<'_>) {
@@ -69,12 +113,12 @@ impl User {
     }
 
     fn from_cookies(cookies: &CookieJar<'_>) -> Option<Self> {
-        Some(Self {
+        Some(Self(User {
             id: parse_cookie_value(cookies, USER_ID_COOKIE)?,
             name: parse_cookie_value(cookies, USER_NAME_COOKIE)?,
             discriminator: parse_cookie_value(cookies, USER_DISCRIMINATOR_COOKIE)?,
-            avatar: parse_cookie_value(cookies, USER_AVATAR_COOKIE)?,
-        })
+            avatar: Some(parse_cookie_value(cookies, USER_AVATAR_COOKIE)?),
+        }))
     }
 
     pub async fn get(cookies: &CookieJar<'_>) -> Result<Option<Self>, GetUserError> {

src/models/user/serial.rs

@@ -1,6 +1,6 @@
 use serde::{ser::SerializeStruct, Serialize, Serializer};
 
-use super::User;
+use super::{User, Username};
 
 pub fn deserialize_id<'de, D>(deserializer: D) -> Result<u64, D::Error>
 where
@@ -27,7 +27,7 @@ impl Serialize for User {
         state.serialize_field("id", &self.id)?;
         state.serialize_field("name", &self.name)?;
         state.serialize_field("discriminator", &self.discriminator)?;
-        state.serialize_field("avatar", &self.avatar)?;
+        state.serialize_field("avatar", &self.avatar())?;
         state.serialize_field("username", &self.username())?;
         state.end()
     }

src/models/user/tests.rs

@@ -1,4 +1,4 @@
-use super::User;
+use super::{User, Username};
 
 fn test_user(name: &str, discriminator: u16) -> User {
     User {

src/routes/get_challenge.rs

@@ -6,7 +6,7 @@ use rocket_dyn_templates::{context, Template};
 use crate::{
     cookies::LANG_COOKIE,
     i18n::DEFAULT as DEFAULT_LANG,
-    models::{Challenge, Settings, User},
+    models::{Challenge, Settings, SessionUser},
     utils::AcceptLanguage,
 };
 
@@ -27,7 +27,7 @@ pub async fn get_challenge(
                 .map(|cookie| vec![cookie.value().to_owned()])
                 .or(accept_language.0)
                 .unwrap_or_else(|| vec![DEFAULT_LANG.to_owned()]),
-            user: User::get(cookies).await.unwrap(),
+            user: SessionUser::get(cookies).await.unwrap(),
             content: Challenge::get(challenge),
         },
     )

src/routes/logout.rs

@@ -2,7 +2,7 @@ use std::{collections::HashMap, env};
 
 use rocket::{http::CookieJar, response::Redirect};
 
-use crate::{cookies::token::TOKEN_COOKIE, models::User, utils::Referer};
+use crate::{cookies::token::TOKEN_COOKIE, models::SessionUser, utils::Referer};
 
 #[get("/logout")]
 pub fn logout(cookies: &CookieJar<'_>, referer: Referer) -> Redirect {
@@ -29,7 +29,7 @@ pub fn logout(cookies: &CookieJar<'_>, referer: Referer) -> Redirect {
             println!("Failed to revoke token: {:?}", error);
         }
     });
-    User::purge(cookies);
+    SessionUser::purge(cookies);
     let redirect_url = referer.0.unwrap_or("/".to_owned());
     Redirect::to(redirect_url)
 }

src/routes/post_login.rs

@@ -5,7 +5,7 @@ use rocket::{
     response::Redirect,
 };
 
-use crate::{cookies::{token::TOKEN_EXPIRE_COOKIE, WELCOMED_COOKIE}, models::User};
+use crate::{cookies::{token::TOKEN_EXPIRE_COOKIE, WELCOMED_COOKIE}, models::SessionUser};
 
 #[derive(FromForm)]
 pub struct Login<'r> {
@@ -18,7 +18,7 @@ pub struct Login<'r> {
 #[post("/login", data = "<login>")]
 pub async fn post_login(login: Form<Login<'_>>, cookies: &CookieJar<'_>) -> Redirect {
     if (login.token_type != "Bearer" || login.scope.split("+").any(|scope| scope == "identify"))
-        && User::init(login.access_token, cookies).await.is_ok()
+        && SessionUser::init(login.access_token, cookies).await.is_ok()
     {
         cookies.add(Cookie::new(
             TOKEN_EXPIRE_COOKIE,

src/routes/testing.rs

@@ -3,7 +3,7 @@ use std::ops::Deref;
 use poise::serenity_prelude::Http;
 use rocket::{http::CookieJar, State};
 
-use crate::models::User;
+use crate::models::SessionUser;
 
 #[get("/testing")]
 pub async fn testing(cookies: &CookieJar<'_>, http: &State<Http>) -> String {
@@ -15,11 +15,11 @@ pub async fn testing(cookies: &CookieJar<'_>, http: &State<Http>) -> String {
             .expect("Failed to get testing guild")
             .member(
                 http.deref(),
-                User::get(cookies)
+                SessionUser::get(cookies)
                     .await
                     .expect("Failed to get logged in user data")
                     .expect("No logged in user")
-                    .id
+                    .0.id
             )
             .await
             .expect("Failed to fetch user in server")

src/utils/challenge.rs

@@ -0,0 +1,21 @@
+use std::fs;
+
+pub fn get_challenge_number() -> i32 {
+    let paths = fs::read_dir("content/challenges").unwrap();
+    let mut max = 0;
+    for path in paths {
+        let number = path
+            .unwrap()
+            .path()
+            .file_stem()
+            .unwrap()
+            .to_str()
+            .unwrap()
+            .parse::<i32>()
+            .unwrap();
+        if number > max {
+            max = number;
+        }
+    }
+    max
+}

src/utils/mod.rs

@@ -3,3 +3,6 @@ pub use kyujitai::Kyujitai;
 
 mod headers;
 pub use headers::*;
+
+mod challenge;
+pub use challenge::*;