tegakituesday.com/src/main.rs

#[macro_use]
extern crate rocket;

use core::panic;
use rocket::fs::{relative, FileServer};
use rocket::http::{Cookie, CookieJar};
use rocket::request::{FromRequest};
use rocket::response::{Redirect};
use rocket::{request, Request};
use rocket_dyn_templates::{context, Template};
use sass_rocket_fairing::SassFairing;
use std::collections::HashMap;
use std::convert::Infallible;
use std::env;
use std::fs;

mod challenge;
use challenge::Challenge;

mod kyujitai;

#[get("/<challenge>")]
fn get_challenge(challenge: u32, cookies: &CookieJar<'_>) -> Template {
    let value = cookies
        .get_private(TOKEN_COOKIE)
        .map(|cookie| cookie.value().to_owned());
    let logged_in = value.is_some();
    Template::render(
        "index",
        context! {
            challenge,
            logged_in,
            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!(
        "https://discord.com/api/oauth2/authorize?client_id={client_id}&redirect_uri={redirect_uri}&response_type=code&scope=identify%20guilds.join%20guilds",
        client_id = env::var("CLIENT_ID").unwrap(),
        redirect_uri = format!("{}login", env::var("DOMAIN").unwrap()),
    ))
    // TODO: After returning from Discord go to previous page (with Referer?)
}

#[get("/login?<code>")]
fn login_success(code: String, cookies: &CookieJar<'_>) -> Redirect {
    cookies.add_private(Cookie::new(TOKEN_COOKIE, code));
    Redirect::to("/")
}

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";

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

#[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, login_success, logout])
        .mount("/css", FileServer::from(relative!("styles/css")))
        .attach(Template::fairing())
        .attach(SassFairing::default())
}