Refactoring

rust
Elnu 2 years ago
parent 244e66cfa3
commit 94bd2da27b

@ -0,0 +1,10 @@
pub mod token {
pub const TOKEN_COOKIE: &str = "token";
pub const TOKEN_EXPIRE_COOKIE: &str = "token_expire";
}
pub mod user {
pub const USER_ID_COOKIE: &str = "user_id";
pub const USER_NAME_COOKIE: &str = "user_name";
pub const USER_DISCRIMINATOR_COOKIE: &str = "user_discriminator";
pub const USER_AVATAR_COOKIE: &str = "user_avatar";
}

@ -1,327 +1,22 @@
#[macro_use]
extern crate rocket;
use chrono::{Duration, Utc};
use poise::serenity_prelude::{GatewayIntents, Client, Http};
use core::panic;
use std::ops::Deref;
use derive_more::From;
use reqwest::StatusCode;
use rocket::form::Form;
use poise::serenity_prelude::Http;
use rocket::fs::{relative, FileServer};
use rocket::http::{Cookie, CookieJar};
use rocket::request::FromRequest;
use rocket::response::{content::RawHtml, Redirect};
use rocket::{request, Request, State};
use rocket_dyn_templates::{context, Template};
use rocket_dyn_templates::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 models;
mod kyujitai;
mod utils;
#[get("/<challenge>")]
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 = "<login>")]
async fn post_login(login: Form<Login<'_>>, 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(
"<form action=\"/login\" method=\"post\"></form>
<script>
try {
const params = new URLSearchParams(location.hash.slice(1));
const form = document.querySelector(\"form\");
[\"token_type\", \"access_token\", \"expires_in\", \"scope\"].forEach(field => {
const input = document.createElement(\"input\");
input.type = \"hidden\";
input.name = field;
input.value = params.get(field);
form.appendChild(input);
});
form.submit();
} catch {
location.href = \"/\";
}
</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),
};
});
User::purge(cookies);
let redirect_url = referer.0.unwrap_or("/".to_owned());
Redirect::to(redirect_url)
}
#[get("/testing")]
async fn testing(cookies: &CookieJar<'_>, http: &State<Http>) -> String {
// Get logged in user's join date in 字ちゃん server
format!("{:?}", http.get_guild(814700630958276649).await
.expect("Failed to get testing guild")
.member(http.deref(), User::get(cookies).await
.expect("Failed to get logged in user data")
.expect("No logged in user")
.id
).await
.expect("Failed to fetch user in server")
.joined_at)
}
#[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,
}
mod cookies;
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>,
},
}
fn parse_cookie_value<T: std::str::FromStr>(cookies: &CookieJar<'_>, name: &str) -> Option<T> {
cookies.get_private(name)?.value().parse().ok()
}
mod routes;
use routes::*;
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<Self, GetUserError> {
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<Self> {
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<Option<Self>, 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::<i64>())
.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))
}
}
mod prelude;
#[launch]
async fn rocket() -> _ {
@ -343,28 +38,3 @@ async fn rocket() -> _ {
.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");
}
}

@ -4,7 +4,7 @@ use std::str::FromStr;
use serde::{Deserialize, Deserializer, Serialize};
use serde_yaml::Value;
use crate::kyujitai::Kyujitai;
use crate::prelude::*;
#[derive(Serialize, Deserialize)]
pub struct Challenge {
@ -18,6 +18,40 @@ pub struct Challenge {
pub date: chrono::NaiveDate,
}
impl Challenge {
pub fn get(number: u32) -> Self {
use comrak::{parse_document, Arena, ComrakOptions};
use std::fs;
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/{number}.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!")
}
}
}
#[derive(Serialize, Deserialize)]
pub struct Song {
pub japanese: Option<String>,

@ -0,0 +1,5 @@
mod user;
pub use user::User;
mod challenge;
pub use challenge::Challenge;

@ -0,0 +1,116 @@
mod serial;
#[cfg(test)]
mod tests;
use chrono::Utc;
use derive_more::From;
use reqwest::StatusCode;
use rocket::http::{CookieJar, Cookie};
use serial::*;
use serde::Deserialize;
use crate::cookies::{token::*, user::*};
#[derive(Default, Deserialize)]
pub struct User {
#[serde(deserialize_with = "deserialize_id")]
pub id: u64,
#[serde(rename = "username")]
pub name: String,
#[serde(deserialize_with = "deserialize_discriminator")]
pub discriminator: u16,
pub avatar: String,
}
impl User {
pub fn username(&self) -> String {
if self.discriminator == 0 {
return self.name.clone();
}
format!("{}#{:0>4}", self.name, self.discriminator)
}
pub async fn init(token: &str, cookies: &CookieJar<'_>) -> Result<Self, GetUserError> {
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_COOKIE, user.id.to_string()));
cookies.add_private(Cookie::new(USER_NAME_COOKIE, user.name.clone()));
cookies.add_private(Cookie::new(
USER_DISCRIMINATOR_COOKIE,
user.discriminator.to_string(),
));
cookies.add_private(Cookie::new(USER_AVATAR_COOKIE, user.avatar.clone()));
Ok(user)
}
pub fn purge(cookies: &CookieJar<'_>) {
cookies.remove_private(Cookie::named(TOKEN_COOKIE));
cookies.remove_private(Cookie::named(USER_ID_COOKIE));
cookies.remove_private(Cookie::named(USER_NAME_COOKIE));
cookies.remove_private(Cookie::named(USER_DISCRIMINATOR_COOKIE));
cookies.remove_private(Cookie::named(USER_AVATAR_COOKIE));
cookies.remove(Cookie::named(TOKEN_EXPIRE_COOKIE));
}
fn from_cookies(cookies: &CookieJar<'_>) -> Option<Self> {
Some(Self {
id: parse_cookie_value(cookies, USER_ID_COOKIE)?,
name: parse_cookie_value(cookies, USER_NAME_COOKIE)?,
discriminator: parse_cookie_value(cookies, USER_DISCRIMINATOR_COOKIE)?,
avatar: parse_cookie_value(cookies, USER_AVATAR_COOKIE)?,
})
}
pub async fn get(cookies: &CookieJar<'_>) -> Result<Option<Self>, 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::<i64>())
.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_COOKIE));
cookies.remove_private(Cookie::named(USER_NAME_COOKIE));
cookies.remove_private(Cookie::named(USER_DISCRIMINATOR_COOKIE));
cookies.remove_private(Cookie::named(USER_AVATAR_COOKIE));
cookies.remove(Cookie::named(TOKEN_EXPIRE_COOKIE));
return Ok(None);
}
Ok(Some(user))
}
}
#[derive(From, Debug)]
pub enum GetUserError {
ReqwestError(reqwest::Error),
DeserializeError(serde_json::Error),
#[allow(unused)]
DiscordError {
status: StatusCode,
message: Option<String>,
},
}
fn parse_cookie_value<T: std::str::FromStr>(cookies: &CookieJar<'_>, name: &str) -> Option<T> {
cookies.get_private(name)?.value().parse().ok()
}

@ -0,0 +1,34 @@
use serde::{Serialize, Serializer, ser::SerializeStruct};
use super::User;
pub 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)
}
pub 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()
}
}

@ -0,0 +1,21 @@
use super::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");
}

@ -0,0 +1 @@
pub use crate::utils::Kyujitai;

@ -0,0 +1,22 @@
use rocket::http::CookieJar;
use rocket_dyn_templates::{Template, context};
use crate::models::{User, Challenge};
#[get("/<challenge>")]
pub 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: Challenge::get(challenge),
},
)
}

@ -0,0 +1,13 @@
use std::env;
use rocket::response::Redirect;
#[get("/login")]
pub 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()),
))
}

@ -0,0 +1,36 @@
use std::{collections::HashMap, env};
use rocket::{http::CookieJar, response::Redirect};
use crate::{utils::Referer, cookies::token::TOKEN_COOKIE, models::User};
#[get("/logout")]
pub 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),
};
});
User::purge(cookies);
let redirect_url = referer.0.unwrap_or("/".to_owned());
Redirect::to(redirect_url)
}

@ -0,0 +1,23 @@
#[path = "get_challenge.rs"]
mod _get_challenge;
pub use _get_challenge::get_challenge;
#[path = "login.rs"]
mod _login;
pub use _login::login;
#[path = "logout.rs"]
mod _logout;
pub use _logout::logout;
#[path = "post_login.rs"]
mod _post_login;
pub use _post_login::post_login;
#[path = "success.rs"]
mod _success;
pub use _success::success;
#[path = "testing.rs"]
mod _testing;
pub use _testing::testing;

@ -0,0 +1,27 @@
use chrono::{Utc, Duration};
use rocket::{form::Form, http::{CookieJar, Cookie}, response::Redirect};
use crate::{cookies::token::TOKEN_EXPIRE_COOKIE, models::User};
#[derive(FromForm)]
pub struct Login<'r> {
token_type: &'r str,
access_token: &'r str,
expires_in: u64,
scope: &'r str,
}
#[post("/login", data = "<login>")]
pub async fn post_login(login: Form<Login<'_>>, 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("/")
}

@ -0,0 +1,24 @@
use rocket::response::content::RawHtml;
#[get("/success")]
pub fn success() -> RawHtml<&'static str> {
RawHtml(
"<form action=\"/login\" method=\"post\"></form>
<script>
try {
const params = new URLSearchParams(location.hash.slice(1));
const form = document.querySelector(\"form\");
[\"token_type\", \"access_token\", \"expires_in\", \"scope\"].forEach(field => {
const input = document.createElement(\"input\");
input.type = \"hidden\";
input.name = field;
input.value = params.get(field);
form.appendChild(input);
});
form.submit();
} catch {
location.href = \"/\";
}
</script>",
)
}

@ -0,0 +1,20 @@
use std::ops::Deref;
use poise::serenity_prelude::Http;
use rocket::{http::CookieJar, State};
use crate::models::User;
#[get("/testing")]
pub async fn testing(cookies: &CookieJar<'_>, http: &State<Http>) -> String {
// Get logged in user's join date in 字ちゃん server
format!("{:?}", http.get_guild(814700630958276649).await
.expect("Failed to get testing guild")
.member(http.deref(), User::get(cookies).await
.expect("Failed to get logged in user data")
.expect("No logged in user")
.id
).await
.expect("Failed to fetch user in server")
.joined_at)
}

@ -0,0 +1,5 @@
mod kyujitai;
pub use kyujitai::Kyujitai;
mod referer;
pub use referer::Referer;

@ -0,0 +1,15 @@
use std::convert::Infallible;
use rocket::{request::{FromRequest, self}, Request};
pub struct Referer(pub 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())))
}
}
Loading…
Cancel
Save