blatherskite

a toy discord-like chat app backend written for a swe class
Log | Files | Refs | README

commit 01af1b510f0a5deef8f3c9a1ce3662d2d55b7a9f
parent 77941234e41781b75f21693b4b012ad0c1693d6c
Author: quantumish <freifeld.david@gmail.com>
Date:   Wed, 28 Sep 2022 14:26:43 -0700

Improve authentication system

Diffstat:
Mscuttlebutt/Cargo.toml | 1+
Mscuttlebutt/src/main.rs | 74++++++++++++++++++++++++++++++++++++++++++++++++++++++--------------------
Mscuttlebutt/src/responses.rs | 16++++++++++++++++
3 files changed, 71 insertions(+), 20 deletions(-)

diff --git a/scuttlebutt/Cargo.toml b/scuttlebutt/Cargo.toml @@ -17,6 +17,7 @@ more-asserts = "0.3.0" poem = { version = "1.3.42", features = ["test"] } poem-openapi = { version = "2.0.12", features = ["swagger-ui"] } pretty_assertions = "1.3.0" +rand = "0.8.5" rustflake = "0.1.1" serde = "1.0.144" serde_json = "1.0.85" diff --git a/scuttlebutt/src/main.rs b/scuttlebutt/src/main.rs @@ -1,6 +1,6 @@ use cassandra_cpp::*; use chrono::{DateTime, Duration, Local}; -use hmac::{Hmac, Mac}; +use hmac::Hmac; use jwt::{SignWithKey, VerifyWithKey}; use poem::{listener::TcpListener, web::Data, EndpointExt, Request, Result, Route, Server}; use poem_openapi::{ @@ -9,6 +9,7 @@ use poem_openapi::{ payload::{Base64, Json, PlainText}, *, }; +use rand::{distributions::Alphanumeric, Rng}; use serde::{Deserialize, Serialize}; use sha2::Sha256; use std::sync::{Arc, Mutex}; @@ -41,9 +42,11 @@ async fn api_checker(req: &Request, api_key: ApiKey) -> Option<User> { .unwrap(), ) .unwrap(); - let key = todo!(); // Query DB here... - // let server_key = req.data::<ServerKey>().unwrap(); - // VerifyWithKey::<User>::verify_with_key(api_key.key.as_str(), server_key).ok() + if claims.exp < Local::now() { + return None; + } + let server_key = req.data::<ServerKey>().unwrap(); + VerifyWithKey::<User>::verify_with_key(api_key.key.as_str(), server_key).ok() } struct Api { @@ -70,10 +73,11 @@ impl Api { cluster.set_load_balance_round_robin(); let session = cluster.connect().unwrap(); - let space_stmt = &stmt!(&format!("CREATE KEYSPACE IF NOT EXISTS {keyspc} WITH replication = {{'class':'SimpleStrategy', 'replication_factor': 1}}")); - let table_stmt = &stmt!(&format!("CREATE TABLE IF NOT EXISTS {keyspc}.users (id bigint PRIMARY KEY, name text, email text, hash text);")); - session.execute(space_stmt).wait().unwrap(); - session.execute(table_stmt).wait().unwrap(); + session.execute(&stmt!(&format!("CREATE KEYSPACE IF NOT EXISTS {keyspc} WITH replication = {{'class':'SimpleStrategy', 'replication_factor': 1}}"))).wait().unwrap(); + session.execute(&stmt!(&format!("CREATE TABLE IF NOT EXISTS {keyspc}.users (id bigint PRIMARY KEY, name text, email text, hash text);"))).wait().unwrap(); + session.execute(&stmt!(&format!("CREATE TABLE IF NOT EXISTS {keyspc}.channels (id bigint PRIMARY KEY, name text, group bigint, members list<bigint>, channels list<bigint>);"))).wait().unwrap(); + session.execute(&stmt!(&format!("CREATE TABLE IF NOT EXISTS {keyspc}.user_groups (id bigint PRIMARY KEY, groups list<bigint>);"))).wait().unwrap(); + session.execute(&stmt!(&format!("CREATE TABLE IF NOT EXISTS {keyspc}.messages (group bigint, channel bigint, author bigint, time timestamp, content text, PRIMARY KEY ((group, channel)));"))).wait().unwrap(); Api { sess: session, @@ -82,16 +86,39 @@ impl Api { } #[oai(path = "/login", method = "post")] - async fn login(&self, id: Query<i64>, hash: Base64<Vec<u8>>) -> Result<PlainText<String>> { - let key = - Hmac::<Sha256>::new_from_slice(&hash.0).map_err(poem::error::InternalServerError)?; - let token = Claims { - id: id.0, - exp: Local::now() + Duration::days(1), + async fn login( + &self, + key: Data<&ServerKey>, + id: Query<i64>, + hash: Base64<Vec<u8>>, + ) -> LoginResponse { + use LoginResponse::*; + let hash_stmt = &stmt!(&format!( + "SELECT id, name, email, hash FROM {}.users WHERE id={};", + self.kspc, id.0 + )); + let res = self.sess.execute(hash_stmt).wait().unwrap(); + if res.row_count() == 1 { + let row = res.first_row().unwrap(); + let db_hash: String = row.get(3).unwrap(); + if base64::decode(db_hash).unwrap() != hash.0 { + Unauthorized + } else { + let row = res.first_row().unwrap(); + let token = Claims { + id: id.0, + exp: Local::now() + Duration::days(1), + } + .sign_with_key(key.0); + Success(PlainText(token.unwrap().to_string())) + } + } else if res.row_count() == 0 { + NotFound + } else { + InternalError(PlainText( + "Found multiple (or negative?) number of rows.".to_string(), + )) } - .sign_with_key(&key) - .map_err(poem::error::InternalServerError)?; - Ok(PlainText(token)) } #[oai(path = "/user", method = "get")] @@ -331,6 +358,7 @@ impl Api { #[tokio::main] async fn main() -> Result<(), std::io::Error> { + use hmac::Mac; if std::env::var_os("RUST_LOG").is_none() { std::env::set_var("RUST_LOG", "poem=debug"); } @@ -346,12 +374,18 @@ async fn main() -> Result<(), std::io::Error> { let ui = api_service.swagger_ui(); let spec = api_service.spec(); + let key: String = rand::thread_rng() + .sample_iter(&Alphanumeric) + .take(7) + .map(char::from) + .collect(); Server::new(TcpListener::bind("127.0.0.1:3000")) .run( Route::new() .nest("/api", api_service) .nest("/", ui) - .at("/spec", poem::endpoint::make_sync(move |_| spec.clone())), + .at("/spec", poem::endpoint::make_sync(move |_| spec.clone())) + .data(ServerKey::new_from_slice(&key.as_bytes()).unwrap()), ) .await } @@ -368,8 +402,8 @@ mod tests { use sha2::Digest; fn setup() -> TestClient<Route> { - let app = - OpenApiService::new(Api, "Scuttlebutt", "1.0").server("http://localhost:3000/api"); + let app = OpenApiService::new(Api::new("test"), "Scuttlebutt", "1.0") + .server("http://localhost:3000/api"); TestClient::new(Route::new().nest("/api", app)) } diff --git a/scuttlebutt/src/responses.rs b/scuttlebutt/src/responses.rs @@ -37,6 +37,22 @@ pub struct Message { } #[derive(ApiResponse)] +pub enum LoginResponse { + /// Returns the authentication token requested. + #[oai(status = 200)] + Success(PlainText<String>), + /// User ID not found + #[oai(status = 404)] + NotFound, + /// Incorrect hash provided + #[oai(status = 401)] + Unauthorized, + /// Internal server error when attempting to access database/sign key + #[oai(status = 500)] + InternalError(PlainText<String>), +} + +#[derive(ApiResponse)] pub enum UserResponse { /// Returns the user requested. #[oai(status = 200)]