Implement basic login system

src/main.rs

@@ -2,7 +2,14 @@
 extern crate rocket;
 
 use core::panic;
+use std::collections::HashMap;
+use std::convert::Infallible;
+use std::env;
 use rocket::fs::{FileServer, relative};
+use rocket::{Request, request};
+use rocket::http::{Cookie, CookieJar};
+use rocket::request::{FromRequest, FlashMessage};
+use rocket::response::{Redirect, Flash};
 use rocket_dyn_templates::{context, Template};
 use std::fs;
 use sass_rocket_fairing::SassFairing;
@@ -13,11 +20,16 @@ use challenge::Challenge;
 mod kyujitai;
 
 #[get("/<challenge>")]
-fn get_challenge(challenge: u32) -> Template {
+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 = {
@@ -51,11 +63,70 @@ fn get_challenge(challenge: u32) -> Template {
     )
 }
 
+#[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])
+        .mount("/", routes![get_challenge, login, login_success, logout])
         .mount("/css", FileServer::from(relative!("styles/css")))
         .attach(Template::fairing())
         .attach(SassFairing::default())