json web tokens (auth)

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"

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"

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

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

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<EncodingKey> = 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<DecodingKey> = 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<AppState>) -> Result<String, StatusCode> {
+    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<Claims> {
+    let key = &*AUTH_PUBLIC_KEY;
+
+    decode::<Claims>(token, key, &Validation::new(Algorithm::RS512))
+        .map(|token_data| token_data.claims)
+        .map_err(|e| {
+            println!("{e:?}");
+            e
+        })
+        .ok()
+}

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

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<WebsocketHandlerParams>,
     State(state): State<AppState>,
 ) -> 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) {
+async fn handler(mut socket: WebSocket, state: AppState, params: WebsocketHandlerParams) {
-    let mut next_client_id = state.next_client_id.lock().await;
+    let Some(claims) = verify_token(&params.token) else {
+        let _ = socket.close();
+        return;
+    };
 
-    *next_client_id += 1;
+    let _ = socket.send((claims.user_id).to_string().into()).await;
-    let _ = socket.send((*next_client_id).to_string().into()).await;
-
-    drop(next_client_id);
 
     let (sender, receiver) = socket.split();