From fc22ed413f1e7e45b2dead5623064c7153002262 Mon Sep 17 00:00:00 2001
From: alemi <me@alemi.dev>
Date: Mon, 6 Mar 2023 18:52:23 +0100
Subject: [PATCH] feat: improved access token security (JWT sorta)

---
 Cargo.toml         |  4 ++++
 src/main.rs        | 13 ++++++++----
 src/persistence.rs | 52 ++++++++++++++++++++++++++++++++++++++++++++++
 src/routes/auth.rs | 24 ++++++++++++++++-----
 4 files changed, 84 insertions(+), 9 deletions(-)
 create mode 100644 src/persistence.rs

diff --git a/Cargo.toml b/Cargo.toml
index ef0c485..de5ab66 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -14,6 +14,10 @@ reqwest = { version = "0.11", features = ["json"] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 tracing-subscriber = "0.3"
+rand = { version = "0.8", features = ["getrandom"]}
 uuid = "1"
 chrono = "0.4"
 tracing = "0.1"
+jwt = "0.16"
+sha2 = "0.10"
+hmac = "0.12"
diff --git a/src/main.rs b/src/main.rs
index bd04cea..ed9a102 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,7 @@
 mod proto;
 mod entities;
 mod routes;
+mod persistence;
 
 use std::{collections::HashMap, sync::Arc};
 
@@ -14,10 +15,10 @@ use uuid::Uuid;
 use tracing_subscriber::prelude::*;
 use tracing::{info, metadata::LevelFilter};
 
-use crate::routes::{
+use crate::{routes::{
 	auth::{authenticate, validate, refresh},
 	session::{join, has_joined_wrapper, profile}, register::register_unmigrated,
-};
+}, persistence::load_secret};
 
 /// Reimplementation of legacy auth server for minecraft
 #[derive(Parser, Debug, Clone)]
@@ -56,6 +57,7 @@ pub struct AppState {
 	store: Arc<Mutex<HashMap<Uuid, JoinAttempt>>>,
 	db: DatabaseConnection,
 	cfg: ConfigArgs,
+	secret: String,
 }
 
 #[tokio::main]
@@ -69,7 +71,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 	let cfg = ConfigArgs::parse();
 
 	let db = Database::connect(cfg.database.clone()).await?;
-	let store = Arc::new(Mutex::new(HashMap::new()));
+
+	let secret = load_secret(&db).await?;
+
+	let store = Arc::new(Mutex::new(HashMap::new())); // TODO do this as an Actor
 
 	let addr = cfg.bind_addr.parse()?;
 
@@ -85,7 +90,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 		// CUSTOM
 		.route("/register/unmigrated", post(register_unmigrated))
 		.fallback(fallback_route)
-		.with_state(AppState { store, db, cfg });
+		.with_state(AppState { store, db, cfg, secret });
 
 	info!(target: "MAIN", "serving Yggdrasil on {}", &addr);
 
diff --git a/src/persistence.rs b/src/persistence.rs
new file mode 100644
index 0000000..603e8a2
--- /dev/null
+++ b/src/persistence.rs
@@ -0,0 +1,52 @@
+use chrono::Utc;
+use hmac::{Hmac, Mac};
+use jwt::SignWithKey;
+use rand::{rngs::OsRng, Rng, distributions::Alphanumeric};
+use sea_orm::{EntityTrait, DatabaseConnection, ActiveValue::NotSet, Set, DbErr};
+use sha2::Sha384;
+use tracing::info;
+use std::collections::BTreeMap;
+
+use crate::entities;
+
+/// Since we check against our db for each validate request, there's no real use for a JWT here. It
+/// is still used to allow easier eventual future migration, and to make tokens look like true old
+/// Mojang tokens
+pub fn new_auth_token(secret: &[u8], fields: Vec<(&str, &str)>) -> Result<String, jwt::Error> {
+	let key: Hmac<Sha384> = Hmac::new_from_slice(secret)?;
+	let mut claims : BTreeMap<&str, &str> = BTreeMap::new();
+	for (key, value) in fields {
+		claims.insert(key, value);
+	}
+	let token = claims.sign_with_key(&key)?;
+	Ok(token)
+}
+
+pub async fn load_secret(db: &DatabaseConnection) -> Result<String, DbErr>  {
+	let secret;
+
+	if let Some(state) = entities::persistence::Entity::find().one(db).await? {
+		secret = state.secret;
+	} else {
+		info!(target: "SECRET", "generating new secret");
+		secret = OsRng::default()
+			.sample_iter(&Alphanumeric)
+			.take(64)
+			.map(char::from)
+			.collect();
+
+		entities::persistence::Entity::delete_many().exec(db).await?;
+
+		entities::persistence::Entity::insert(
+			entities::persistence::ActiveModel {
+				id: NotSet,
+				secret: Set(secret.clone()),
+				last_edit: Set(Utc::now()),
+			}
+		).exec(db).await?;
+	}
+
+	Ok(secret)
+}
+
+
diff --git a/src/routes/auth.rs b/src/routes/auth.rs
index 38a397b..7922c39 100644
--- a/src/routes/auth.rs
+++ b/src/routes/auth.rs
@@ -1,10 +1,10 @@
 use axum::{extract::State, Json, http::StatusCode};
-use chrono::Utc;
+use chrono::{Utc, Duration};
 use sea_orm::{EntityTrait, QueryFilter, ColumnTrait, ActiveValue::NotSet, Set};
-use uuid::Uuid;
 use tracing::{info, warn};
+use uuid::Uuid;
 
-use crate::{entities, AppState, proto};
+use crate::{entities, AppState, proto, persistence::new_auth_token};
 
 
 pub async fn validate(State(state): State<AppState>, Json(payload): Json<proto::ValidateRequest>) -> Result<StatusCode, StatusCode> {
@@ -33,9 +33,13 @@ pub async fn refresh(State(state): State<AppState>, Json(payload): Json<proto::R
 		let user = entities::user::Entity::find_by_id(t.user_id).one(&state.db).await
 			.map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, proto::Error::simple("db error").json()))?
 			.ok_or((StatusCode::NOT_FOUND, proto::Error::simple("no user owns this token").json()))?;
+
 		entities::token::Entity::delete_by_id(t.id).exec(&state.db).await
 			.map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, proto::Error::simple("db error").json()))?;
-		let new_access_token = Uuid::new_v4(); // TODO same as with authenticate
+
+		let new_access_token = new_auth_token(state.secret.as_bytes(), vec![("uuid", &user.uuid.to_string())])
+			.map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, proto::Error::simple("invalid secret").json()))?;
+
 		entities::token::Entity::insert(
 			entities::token::ActiveModel{
 				id: NotSet,
@@ -45,12 +49,14 @@ pub async fn refresh(State(state): State<AppState>, Json(payload): Json<proto::R
 			}
 		).exec(&state.db).await
 			.map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, proto::Error::simple("db error").json()))?;
+
 		let response = proto::RefreshResponse {
 			accessToken: new_access_token.to_string(),
 			clientToken: payload.clientToken,
 			selectedProfile: proto::Profile { id: user.uuid, name: user.name },
 			user: None,
 		};
+
 		info!(target: "AUTH", "[REFRESH] answering with {:?}", response);
 		Ok(Json(response))
 	} else {
@@ -70,12 +76,17 @@ pub async fn authenticate(State(state): State<AppState>, Json(payload): Json<pro
 			let s = entities::property::Entity::find().filter(
 				entities::property::Column::UserId.eq(u.id)
 			).one(&state.db).await.map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, proto::Error::simple("db error").json()))?;
+
 			let skin = match s {
 				Some(s) => proto::Property::from(s),
 				None    => proto::Property::default_skin(),
 			};
+
 			// make new token
-			let access_token = Uuid::new_v4().to_string(); // TODO maybe use a JWT?
+
+			let access_token = new_auth_token(state.secret.as_bytes(), vec![("uuid", &u.uuid.to_string())])
+				.map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, proto::Error::simple("invalid secret").json()))?;
+
 			entities::token::Entity::insert(entities::token::ActiveModel {
 				id: NotSet,
 				user_id: Set(u.id),
@@ -83,11 +94,13 @@ pub async fn authenticate(State(state): State<AppState>, Json(payload): Json<pro
 				created_at: Set(Utc::now()),
 			}).exec(&state.db).await
 				.map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, proto::Error::simple("db error").json()))?;
+
 			let client_token = payload.clientToken.unwrap_or(Uuid::new_v4().to_string());
 			let profile = proto::Profile {
 				name: u.name.clone(),
 				id: u.uuid,
 			};
+
 			let response =proto::AuthenticateResponse {
 				accessToken: access_token,
 				user: proto::User { id: u.uuid, username: u.name, properties: Some(vec![ skin ]) },
@@ -95,6 +108,7 @@ pub async fn authenticate(State(state): State<AppState>, Json(payload): Json<pro
 				availableProfiles: vec![profile.clone()],
 				selectedProfile: profile,
 			};
+
 			info!(target: "AUTH", "[AUTHENTICATE] answering with {:?}", response);
 			Ok(Json(response))
 		} else {