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("/")] 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 = "")] -fn post_login(login: Form>, cookies: &CookieJar<'_>) -> Redirect { +async fn post_login(login: Form>, 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 }, } +fn parse_cookie_value(cookies: &CookieJar<'_>, name: &str) -> Option { + 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, 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::()) - .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 { 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 { + 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, 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::()) + .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)) } }