From 5ef44a787e8e3cd9495497673d7c9144bba17f31 Mon Sep 17 00:00:00 2001
From: ElnuDev <elnu@elnu.com>
Date: Fri, 16 Jun 2023 17:19:37 -0700
Subject: [PATCH] Implement user data cookie storage

---
 src/main.rs | 75 ++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 57 insertions(+), 18 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index b944251..32e9c37 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -27,6 +27,7 @@ mod kyujitai;
 
 #[get("/<challenge>")]
 async fn get_challenge(challenge: u32, cookies: &CookieJar<'_>) -> Template { 
+    println!("{:?}", cookies.get_private("user_name").map(|cookie| cookie.value().to_owned()));
     Template::render(
         "index",
         context! {
@@ -85,10 +86,11 @@ struct Login<'r> {
 }
 
 #[post("/login", data = "<login>")]
-fn post_login(login: Form<Login<'_>>, cookies: &CookieJar<'_>) -> Redirect {
+async fn post_login(login: Form<Login<'_>>, cookies: &CookieJar<'_>) -> Redirect {
     if login.token_type != "Bearer" || login.scope != "guilds.join+identify+guilds" {
-        cookies.add_private(Cookie::new(TOKEN_COOKIE, login.access_token.to_owned()));
-        cookies.add(Cookie::new(TOKEN_EXPIRE_COOKIE, (Utc::now() + Duration::seconds(login.expires_in as i64)).timestamp().to_string()));
+        if User::init(login.access_token, cookies).await.is_ok() {
+            cookies.add(Cookie::new(TOKEN_EXPIRE_COOKIE, (Utc::now() + Duration::seconds(login.expires_in as i64)).timestamp().to_string()));
+        }
     }
     Redirect::to("/")
 }
@@ -151,7 +153,7 @@ fn logout(cookies: &CookieJar<'_>, referer: Referer) -> Redirect {
             Err(error) => println!("Failed to revoke token: {:?}", error),
         };
     });
-    cookies.remove_private(Cookie::named(TOKEN_COOKIE));
+    User::purge(cookies);
     let redirect_url = referer.0.unwrap_or("/".to_owned());
     Redirect::to(redirect_url)
 }
@@ -199,6 +201,7 @@ impl Serialize for User {
 }
 
 #[derive(From, Debug)]
+
 enum GetUserError {
     ReqwestError(reqwest::Error),
     DeserializeError(serde_json::Error),
@@ -206,6 +209,10 @@ enum GetUserError {
     DiscordError { status: StatusCode, message: Option<String> },
 }
 
+fn parse_cookie_value<T: std::str::FromStr>(cookies: &CookieJar<'_>, name: &str) -> Option<T> {
+    cookies.get_private(name)?.value().parse().ok()
+}
+
 impl User {
     fn username(&self) -> String {
         if self.discriminator == 0 {
@@ -214,19 +221,7 @@ impl User {
         format!("{}#{:0>4}", self.name, self.discriminator)
     }
 
-    async fn get(cookies: &CookieJar<'_>) -> Result<Option<Self>, GetUserError> {
-        let token = match cookies.get_private(TOKEN_COOKIE) {
-            Some(cookie) => cookie.value().to_owned(),
-            None => return Ok(None),
-        };
-        if cookies.get(TOKEN_EXPIRE_COOKIE)
-            .map(|expire| expire.value().parse::<i64>())
-            .and_then(Result::ok)
-            .map_or(true, |timestamp| Utc::now().timestamp() >= timestamp) {
-                cookies.remove_private(Cookie::named(TOKEN_COOKIE));
-                cookies.remove(Cookie::named(TOKEN_EXPIRE_COOKIE));
-                return Ok(None);
-            }
+    async fn init(token: &str, cookies: &CookieJar<'_>) -> Result<Self, GetUserError> {
         let (status, text) = {
             let response = reqwest::Client::new()
                 .get("https://discord.com/api/users/@me")
@@ -241,7 +236,51 @@ impl User {
                 message: text.ok(),
             })
         }
-        Ok(Some(serde_json::from_str(&text?)?))
+        let user: Self = serde_json::from_str(&text?)?;
+        cookies.add_private(Cookie::new(TOKEN_COOKIE, token.to_owned()));
+        cookies.add_private(Cookie::new("user_id", user.id.to_string()));
+        cookies.add_private(Cookie::new("user_name", user.name.clone()));
+        cookies.add_private(Cookie::new("user_discriminator", user.discriminator.to_string()));
+        cookies.add_private(Cookie::new("user_avatar", user.avatar.clone()));
+        Ok(user)
+    }
+
+    fn purge(cookies: &CookieJar<'_>) {
+        cookies.remove_private(Cookie::named(TOKEN_COOKIE));
+        cookies.remove_private(Cookie::named("user_id"));
+        cookies.remove_private(Cookie::named("user_name"));
+        cookies.remove_private(Cookie::named("user_discriminator"));
+        cookies.remove_private(Cookie::named("user_avatar"));
+        cookies.remove(Cookie::named(TOKEN_EXPIRE_COOKIE));
+    }
+
+    fn from_cookies(cookies: &CookieJar<'_>) -> Option<Self> {
+        Some(Self {
+            id: parse_cookie_value(cookies, "user_id")?,
+            name: parse_cookie_value(cookies, "user_name")?,
+            discriminator: parse_cookie_value(cookies, "user_discriminator")?,
+            avatar: parse_cookie_value(cookies, "user_avatar")?,
+        })
+    }
+
+    async fn get(cookies: &CookieJar<'_>) -> Result<Option<Self>, GetUserError> {
+        let user = match Self::from_cookies(cookies) {
+            Some(user) => user,
+            None => return Ok(None),
+        };
+        if cookies.get(TOKEN_EXPIRE_COOKIE)
+            .map(|expire| expire.value().parse::<i64>())
+            .and_then(Result::ok)
+            .map_or(true, |timestamp| Utc::now().timestamp() >= timestamp) {
+                cookies.remove_private(Cookie::named(TOKEN_COOKIE));
+                cookies.remove_private(Cookie::named("user_id"));
+                cookies.remove_private(Cookie::named("user_name"));
+                cookies.remove_private(Cookie::named("user_discriminator"));
+                cookies.remove_private(Cookie::named("user_avatar"));
+                cookies.remove(Cookie::named(TOKEN_EXPIRE_COOKIE));
+                return Ok(None);
+            }
+        Ok(Some(user))
     }
 }