tegakituesday.com/src/main.rs

#[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("/<challenge>")]
async fn get_challenge(challenge: u32, cookies: &CookieJar<'_>) -> Template { 
    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> {
    access_token: &'r str,
    expires_in: u64,
}

#[post("/login", data = "<login>")]
fn post_login(login: Form<Login<'_>>, cookies: &CookieJar<'_>) -> Redirect {
    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()));
    Redirect::to("/")
}

#[get("/success")]
fn success() -> RawHtml<&'static str> {
    RawHtml("<form action=\"/login\" method=\"post\">
    <input type=\"hidden\" name=\"access_token\">
    <input type=\"hidden\" name=\"expires_in\">
</form>
<script>
    const params = new URLSearchParams(location.hash.slice(1));
    document.querySelector(\"[name=access_token]\").value = params.get(\"access_token\");
    document.querySelector(\"[name=expires_in]\").value = params.get(\"expires_in\");
    document.querySelector(\"form\").submit();
</script>")
}

struct Referer(Option<String>);

#[rocket::async_trait]
impl<'r> FromRequest<'r> for Referer {
    type Error = Infallible;

    async fn from_request(req: &'r Request<'_>) -> request::Outcome<Self, Self::Error> {
        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(&params)
            .send()
            .await
        {
            Ok(_) => println!("Successfully revoked token"),
            Err(error) => println!("Failed to revoke token: {:?}", error),
        };
    });
    cookies.remove_private(Cookie::named(TOKEN_COOKIE));
    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<u64, D::Error>
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<u16, D::Error>
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<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    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<String> },
}

impl User {
    fn username(&self) -> String {
        if self.discriminator == 0 {
            return self.name.clone();
        }
        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);
            }
        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(),
            })
        }
        Ok(Some(serde_json::from_str(&text?)?))
    }
}

#[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 {
        let mut user = User::default();
        user.name = name.to_owned();
        user.discriminator = discriminator;
        user
    }

    #[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");
    }
}