diff --git a/Cargo.lock b/Cargo.lock index 6118184..40fa885 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -338,6 +338,7 @@ dependencies = [ "axum_typed_multipart", "chrono", "futures-util", + "jsonwebtoken", "serde", "serde_json", "tokio", @@ -425,6 +426,15 @@ version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a2330da5de22e8a3cb63252ce2abb30116bf5265e89c0e01bc17015ce30a476" +[[package]] +name = "deranged" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c9e6a11ca8224451684bc0d7d5a7adbf8f2fd6887261a1cfc3c0432f9d4068e" +dependencies = [ + "powerfmt", +] + [[package]] name = "digest" version = "0.10.7" @@ -474,7 +484,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cea14ef9355e3beab063703aa9dab15afd25f0667c341310c1e5274bb1d0da18" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -574,8 +584,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" dependencies = [ "cfg-if", + "js-sys", "libc", "wasi 0.11.0+wasi-snapshot-preview1", + "wasm-bindgen", ] [[package]] @@ -872,6 +884,21 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonwebtoken" +version = "9.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a87cc7a48537badeae96744432de36f4be2b4a34a05a5ef32e9dd8a1c169dde" +dependencies = [ + "base64", + "js-sys", + "pem", + "ring", + "serde", + "serde_json", + "simple_asn1", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -998,6 +1025,31 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-conv" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -1051,6 +1103,16 @@ dependencies = [ "windows-targets", ] +[[package]] +name = "pem" +version = "3.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38af38e8470ac9dee3ce1bae1af9c1671fffc44ddfd8bd1d0a3445bf349a8ef3" +dependencies = [ + "base64", + "serde", +] + [[package]] name = "percent-encoding" version = "2.3.1" @@ -1098,6 +1160,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -1345,6 +1413,20 @@ dependencies = [ "bytecheck", ] +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.16", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + [[package]] name = "rkyv" version = "0.7.45" @@ -1406,7 +1488,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -1519,6 +1601,18 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" +[[package]] +name = "simple_asn1" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "297f631f50729c8c99b84667867963997ec0b50f32b2a7dbcab828ef0541e8bb" +dependencies = [ + "num-bigint", + "num-traits", + "thiserror 2.0.12", + "time", +] + [[package]] name = "slab" version = "0.4.9" @@ -1639,7 +1733,7 @@ dependencies = [ "getrandom 0.3.3", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -1682,6 +1776,37 @@ dependencies = [ "syn 2.0.101", ] +[[package]] +name = "time" +version = "0.3.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7619e19bc266e0f9c5e6686659d394bc57973859340060a69221e57dbc0c40" +dependencies = [ + "deranged", + "itoa", + "num-conv", + "powerfmt", + "serde", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9e9a38711f559d9e3ce1cdb06dd7c5b8ea546bc90052da6d06bb76da74bb07c" + +[[package]] +name = "time-macros" +version = "0.2.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3526739392ec93fd8b359c8e98514cb3e8e021beb4e5f597b00a0221f8ed8a49" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tinystr" version = "0.8.1" @@ -1904,6 +2029,12 @@ dependencies = [ "tinyvec", ] +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + [[package]] name = "url" version = "2.5.4" diff --git a/Cargo.toml b/Cargo.toml index a782790..a891261 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,3 +14,4 @@ validify = "2.0.0" axum_typed_multipart = "0.16.2" serde_json = "1.0.140" chrono = { version = "0.4.41", features = ["serde"] } +jsonwebtoken = "9.3.1" diff --git a/private_key.pem b/private_key.pem new file mode 100644 index 0000000..c163fd0 --- /dev/null +++ b/private_key.pem @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQDa/BTgQpggjW0s +grn2ZFBetA+LhOo8tKrMZRQrwnfcSMJEBEDDv6uRaBXalfibJthbq7RuxmRorfrf +5KhQMbgZwAJxOv+u4syRKUMvr9/7W2ZkLiPiEBEhwmvsxiF9Xve8gDx1AT/BNghG +yQbXQmbjzDmOR78btQkySNp988tXN0zDlDAvL4MRp/dlue7ndC+zvi6cTH+pUZTj +o6EQNFqQJIAzx53sP9Gdlt+MWc4YVdx2xZVQT+DqbPIUfTjTH3oj5Dft0VIWwgbS +d6aCiExml4SVy48bJJ4UrvC/HfRuztm5731pTbhqFSa43rAumzUvFEJQW4X0GVuv +4Ab10SQPAgMBAAECggEAHcK56OskAZ3hgyBruttR5jenuoO10co77CCl58yQvZSl +0i3sdK332IBszT1/K+wmO08lQRN1Ue9THlp/sBERvsN5kX891bwSUSgVGXCtirDr +7sbWhEKVZsvQHxL45NujOjBeBwjEbAhhLOg6exwVykpXCB8iudaWCV9yp5YabHUM +e/eWvtOyOQaH23BACLK+5PbpJoBsOGGVu319HuDo0bl+kUsDqSE7lM57DvmJ8ruU +mU4Yefk3CVKgUDnw03EyftG5X25vVB7uju7lJxgaDWVL1SNxmzb9b9MdBlFI/N6x +mVeBSaIWE2sIFeOe2djv7pLQE7aVeYaUpv3eox75AQKBgQDw6Y7im2N65IH9FLao +fmkDDP6R/Ls2TpPEXZ+XT/r3WUnEZNxGfkoN9epTu12PePHjpWjUQQNrT9qUUWv/ +hGQNamovadZipCwQ59BgeD5aEjBg5gXqPdg05HggZlcVZM/Iys+ttGpvlgNlUbMc +1UBpRnNxQwL8WL2VFPpxUMv0owKBgQDosveU+x+L8xrqUT1ZnxOKZguL1d1lJ3OF +fmlsk4Tkq+dy1k7+PzfaA/lZm/eK6qbnWvE+h0dfJmd5RdaisgITWp691BTgqHEX +2xc2sSan+kj7YIpHZN6Os+CZqBP/b951arlCVlyzhASzlIPjv2ClrRJK4thtPpN5 +Qvy5LiIdpQKBgBNfVKQYRrmdIm7NeZq2/47FPrtSjWNikjNaK/ko709wmKP4UFZC +1YWl3r/48x8UZK8fn/xhdfrtIvCGU57qcHk7s2o4ooqwFPMPEQyQ3sKzb7x7AsLB +Ul9+d8JzOFF6F9NXq7eWKHDsqT0+gp6w70sH5USYVOch46zXUQRHsP7tAoGAW4r8 +ZzmboXzT29Z14in4j/uxEyPhYwfwJFYktzgZkbQm60Bv9wl7uWLS13PKIHZLjrxy +J30kkMQU9NqKo5qPwnll9ZvXUsIxK4dfTH7IjFCkIXBLNebyvwT4mLpMk7K4u0xu +8S/0bBLFnrgQYqpJJSQ/DALNzbxOW+gEb/08IpkCgYBdq9IDct8dTqjw0WBAcltO +iCN5Wr+fvTv7FgxEbqw8BZvGRSUI2/ItZuet9EzphEfnp78VKigRQvg9Q3QiLerX +q/k8WIp3rLWV0tkJwdiIWHyRD+Z5vUPGvafV/J3wxOjNKrL3uSuUlwPC/KXAQylQ +pXeXbob1DZxzf5XVWXEXew== +-----END PRIVATE KEY----- diff --git a/public_key.pem b/public_key.pem new file mode 100644 index 0000000..2886737 --- /dev/null +++ b/public_key.pem @@ -0,0 +1,9 @@ +-----BEGIN PUBLIC KEY----- +MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA2vwU4EKYII1tLIK59mRQ +XrQPi4TqPLSqzGUUK8J33EjCRARAw7+rkWgV2pX4mybYW6u0bsZkaK363+SoUDG4 +GcACcTr/ruLMkSlDL6/f+1tmZC4j4hARIcJr7MYhfV73vIA8dQE/wTYIRskG10Jm +48w5jke/G7UJMkjaffPLVzdMw5QwLy+DEaf3Zbnu53Qvs74unEx/qVGU46OhEDRa +kCSAM8ed7D/RnZbfjFnOGFXcdsWVUE/g6mzyFH040x96I+Q37dFSFsIG0nemgohM +ZpeElcuPGySeFK7wvx30bs7Zue99aU24ahUmuN6wLps1LxRCUFuF9Blbr+AG9dEk +DwIDAQAB +-----END PUBLIC KEY----- diff --git a/src/auth.rs b/src/auth.rs new file mode 100644 index 0000000..ede045b --- /dev/null +++ b/src/auth.rs @@ -0,0 +1,65 @@ +use std::sync::LazyLock; + +use axum::{extract::State, http::StatusCode}; +use chrono::Utc; +use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, Validation, decode, encode}; +use serde::{Deserialize, Serialize}; + +use crate::state::AppState; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Claims { + pub user_id: u32, + pub iat: usize, + pub exp: usize, +} + +pub static AUTH_SECRET_KEY: LazyLock = LazyLock::new(|| { + let bytes = std::fs::read("./private_key.pem").unwrap(); + let encoding_key = EncodingKey::from_rsa_pem(&bytes).unwrap(); + + encoding_key +}); + +pub static AUTH_PUBLIC_KEY: LazyLock = LazyLock::new(|| { + let bytes = std::fs::read("./public_key.pem").unwrap(); + let decoding_key = DecodingKey::from_rsa_pem(&bytes).unwrap(); + + decoding_key +}); + +impl Claims { + pub fn new(user_id: u32) -> Self { + let now = (Utc::now().timestamp_millis() / 1000) as usize; + Self { + user_id, + iat: now, + // Should be about 1 year + exp: now + 31540000, + } + } +} + +pub async fn get_auth_token(State(state): State) -> Result { + let mut next_client_id = state.next_client_id.lock().await; + let claims = Claims::new(*next_client_id); + + *next_client_id += 1; + + encode(&Header::new(Algorithm::RS512), &claims, &AUTH_SECRET_KEY).map_err(|e| { + dbg!(&e); + StatusCode::INTERNAL_SERVER_ERROR + }) +} + +pub fn verify_token(token: &str) -> Option { + let key = &*AUTH_PUBLIC_KEY; + + decode::(token, key, &Validation::new(Algorithm::RS512)) + .map(|token_data| token_data.claims) + .map_err(|e| { + println!("{e:?}"); + e + }) + .ok() +} diff --git a/src/main.rs b/src/main.rs index 94ff645..f68c922 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +mod auth; mod message; mod send_message_handler; mod state; @@ -16,6 +17,7 @@ async fn main() -> anyhow::Result<()> { let state = AppState::new(sender); let router: Router = Router::new() + .route("/auth", get(auth::get_auth_token)) .route("/message", post(send_message_handler::handler)) .route("/ws", get(websockets::websocket_handler)) .with_state(state); diff --git a/src/websockets.rs b/src/websockets.rs index 2867bca..841b846 100644 --- a/src/websockets.rs +++ b/src/websockets.rs @@ -1,6 +1,6 @@ use axum::{ extract::{ - State, WebSocketUpgrade, + Query, State, WebSocketUpgrade, ws::{Message, WebSocket}, }, response::Response, @@ -10,24 +10,31 @@ use futures_util::{ SinkExt as _, StreamExt, stream::{SplitSink, SplitStream}, }; +use serde::Deserialize; use tokio::sync::broadcast::Receiver; -use crate::{message::ChatMessage, state::AppState}; +use crate::{auth::verify_token, message::ChatMessage, state::AppState}; + +#[derive(Debug, Deserialize)] +pub struct WebsocketHandlerParams { + token: String, +} pub async fn websocket_handler( upgrade: WebSocketUpgrade, + Query(params): Query, State(state): State, ) -> Response { - upgrade.on_upgrade(move |websocket| handler(websocket, state)) + upgrade.on_upgrade(move |websocket| handler(websocket, state, params)) } -async fn handler(mut socket: WebSocket, state: AppState) { - let mut next_client_id = state.next_client_id.lock().await; +async fn handler(mut socket: WebSocket, state: AppState, params: WebsocketHandlerParams) { + let Some(claims) = verify_token(¶ms.token) else { + let _ = socket.close(); + return; + }; - *next_client_id += 1; - let _ = socket.send((*next_client_id).to_string().into()).await; - - drop(next_client_id); + let _ = socket.send((claims.user_id).to_string().into()).await; let (sender, receiver) = socket.split();