#[macro_use] extern crate rocket; use core::panic; use chrono::{Duration, Utc}; use derive_more::From; use reqwest::StatusCode; use rocket::form::Form; use rocket::fs::{relative, FileServer}; use rocket::http::{Cookie, CookieJar}; use rocket::request::{FromRequest}; use rocket::response::{Redirect, content::RawHtml}; use rocket::{request, Request}; use rocket_dyn_templates::{context, Template}; use sass_rocket_fairing::SassFairing; use serde::ser::SerializeStruct; use serde::{Deserialize, Serialize, Serializer}; use std::collections::HashMap; use std::convert::Infallible; use std::env; use std::fs; mod challenge; use challenge::Challenge; 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! { challenge, user: User::get(cookies).await.unwrap(), content: { use comrak::{parse_document, Arena, ComrakOptions}; let options = { let mut options = ComrakOptions::default(); options.extension.front_matter_delimiter = Some("---".to_owned()); options }; let arena = Arena::new(); let root = parse_document( &arena, &fs::read_to_string(format!("content/challenges/{challenge}.md")).expect("Couldn't find challenge file"), &options, ); if let Some(node) = root.children().next() { if let comrak::nodes::NodeValue::FrontMatter(frontmatter) = &node.data.borrow().value { let frontmatter = { // Trim starting and ending fences let lines: Vec<&str> = frontmatter.trim().lines().collect(); lines[1..lines.len() - 1].join("\n") }; let challenge: Challenge = serde_yaml::from_str(&frontmatter).unwrap(); challenge } else { panic!("No frontmatter!") } } else { panic!("Empty document!") } } }, ) } #[get("/login")] fn login() -> Redirect { Redirect::to(format!( // Switch from response_type=code to response_type=token from URL generator "https://discord.com/api/oauth2/authorize?client_id={client_id}&redirect_uri={redirect_uri}&response_type=token&scope=identify%20guilds.join%20guilds", client_id = env::var("CLIENT_ID").unwrap(), redirect_uri = format!("{}success", env::var("DOMAIN").unwrap()), )) } #[derive(FromForm)] struct Login<'r> { token_type: &'r str, access_token: &'r str, expires_in: u64, scope: &'r str, } #[post("/login", data = "")] async fn post_login(login: Form>, cookies: &CookieJar<'_>) -> Redirect { if (login.token_type != "Bearer" || login.scope != "guilds.join+identify+guilds") && 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("/") } #[get("/success")] fn success() -> RawHtml<&'static str> { RawHtml("
") } struct Referer(Option); #[rocket::async_trait] impl<'r> FromRequest<'r> for Referer { type Error = Infallible; async fn from_request(req: &'r Request<'_>) -> request::Outcome { let referer = req.headers().get_one("Referer"); request::Outcome::Success(Referer(referer.map(|referer| referer.to_owned()))) } } const TOKEN_COOKIE: &str = "token"; const TOKEN_EXPIRE_COOKIE: &str = "token_expire"; #[get("/logout")] fn logout(cookies: &CookieJar<'_>, referer: Referer) -> Redirect { let token = match cookies.get_private(TOKEN_COOKIE) { Some(cookie) => cookie.value().to_owned(), None => return Redirect::to("/"), }; rocket::tokio::spawn(async { let client = reqwest::Client::new(); let params = { let mut params = HashMap::new(); params.insert("client_id", env::var("CLIENT_ID").unwrap()); params.insert("client_secret", env::var("CLIENT_SECRET").unwrap()); params.insert("token", token); params }; match client .post("https://discord.com/api/oauth2/token/revoke") .header("Content-Type", "application/x-www-form-urlencoded") .form(¶ms) .send() .await { Ok(_) => println!("Successfully revoked token"), Err(error) => println!("Failed to revoke token: {:?}", error), }; }); User::purge(cookies); let redirect_url = referer.0.unwrap_or("/".to_owned()); Redirect::to(redirect_url) } #[derive(Default, Deserialize)] struct User { #[serde(deserialize_with = "deserialize_id")] id: u64, #[serde(rename = "username")] name: String, #[serde(deserialize_with = "deserialize_discriminator")] discriminator: u16, avatar: String, } fn deserialize_id<'de, D>(deserializer: D) -> Result where D: serde::Deserializer<'de>, { let id_str: &str = serde::Deserialize::deserialize(deserializer)?; id_str.parse().map_err(serde::de::Error::custom) } fn deserialize_discriminator<'de, D>(deserializer: D) -> Result where D: serde::Deserializer<'de>, { let id_str: &str = serde::Deserialize::deserialize(deserializer)?; id_str.parse().map_err(serde::de::Error::custom) } impl Serialize for User { fn serialize(&self, serializer: S) -> Result where S: Serializer, { let mut state = serializer.serialize_struct("User", 5)?; 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("username", &self.username())?; state.end() } } #[derive(From, Debug)] enum GetUserError { ReqwestError(reqwest::Error), DeserializeError(serde_json::Error), #[allow(unused)] 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 { return self.name.clone(); } format!("{}#{:0>4}", self.name, self.discriminator) } async fn init(token: &str, cookies: &CookieJar<'_>) -> Result { let (status, text) = { let response = reqwest::Client::new() .get("https://discord.com/api/users/@me") .header("Authorization", format!("Bearer {token}")) .send() .await?; (response.status(), response.text().await) }; if !status.is_success() { return Err(GetUserError::DiscordError { status, message: text.ok(), }) } 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)) } } #[launch] fn rocket() -> _ { let config = rocket::Config::figment().merge(("port", 1313)); dotenv::dotenv().expect("Failed to load .env file"); rocket::custom(config) .mount("/", routes![get_challenge, login, post_login, success, logout]) .mount("/css", FileServer::from(relative!("styles/css"))) .attach(Template::fairing()) .attach(SassFairing::default()) } #[cfg(test)] mod tests { use crate::User; fn test_user(name: &str, discriminator: u16) -> User { User { name: name.to_owned(), discriminator, ..Default::default() } } #[test] fn test_legacy_username() { let user = test_user("test", 123); assert_eq!(user.username(), "test#0123"); } #[test] fn test_new_username() { let user = test_user("test", 0); assert_eq!(user.username(), "test"); } }