use bson::{bson, doc, from_bson, Document}; use mongodb::options::FindOneOptions; use rocket::http::{RawStr, Status}; use rocket::request::{self, FromParam, FromRequest, Request}; use rocket::Outcome; use serde::{Deserialize, Serialize}; use crate::database; use database::user::{User, UserRelationship}; #[derive(Serialize, Deserialize, Debug, Clone)] pub struct UserRef { pub id: String, pub username: String, pub email_verified: bool, } impl UserRef { pub fn fetch_data(&self, projection: Document) -> Option<Document> { database::get_collection("users") .find_one( doc! { "_id": &self.id }, FindOneOptions::builder().projection(projection).build(), ) .expect("Failed to fetch user from database.") } pub fn fetch_relationships(&self) -> Option<Vec<UserRelationship>> { let user = database::get_collection("users") .find_one( doc! { "_id": &self.id }, FindOneOptions::builder() .projection(doc! { "relations": 1 }) .build(), ) .expect("Failed to fetch user relationships from database.") .expect("Missing user document."); if let Ok(arr) = user.get_array("relations") { let mut relationships = vec![]; for item in arr { relationships.push(from_bson(item.clone()).unwrap()); } Some(relationships) } else { None } } } #[derive(Debug)] pub enum AuthError { BadCount, Missing, Invalid, } impl<'a, 'r> FromRequest<'a, 'r> for UserRef { type Error = AuthError; fn from_request(request: &'a Request<'r>) -> request::Outcome<Self, Self::Error> { let keys: Vec<_> = request.headers().get("x-auth-token").collect(); match keys.len() { 0 => Outcome::Failure((Status::Forbidden, AuthError::Missing)), 1 => { let key = keys[0]; let result = database::get_collection("users") .find_one( doc! { "access_token": key }, FindOneOptions::builder() .projection(doc! { "_id": 1, "username": 1, "email_verification.verified": 1, }) .build(), ) .unwrap(); if let Some(user) = result { Outcome::Success(UserRef { id: user.get_str("_id").unwrap().to_string(), username: user.get_str("username").unwrap().to_string(), email_verified: user .get_document("email_verification") .unwrap() .get_bool("verified") .unwrap(), }) } else { Outcome::Failure((Status::Forbidden, AuthError::Invalid)) } } _ => Outcome::Failure((Status::BadRequest, AuthError::BadCount)), } } } impl<'a, 'r> FromRequest<'a, 'r> for User { type Error = AuthError; fn from_request(request: &'a Request<'r>) -> request::Outcome<Self, Self::Error> { let keys: Vec<_> = request.headers().get("x-auth-token").collect(); match keys.len() { 0 => Outcome::Failure((Status::Forbidden, AuthError::Missing)), 1 => { let key = keys[0]; let col = database::get_db().collection("users"); let result = col.find_one(doc! { "access_token": key }, None).unwrap(); if let Some(user) = result { Outcome::Success( from_bson(bson::Bson::Document(user)).expect("Failed to unwrap user."), ) } else { Outcome::Failure((Status::Forbidden, AuthError::Invalid)) } } _ => Outcome::Failure((Status::BadRequest, AuthError::BadCount)), } } } impl<'r> FromParam<'r> for UserRef { type Error = &'r RawStr; fn from_param(param: &'r RawStr) -> Result<Self, Self::Error> { let col = database::get_db().collection("users"); let result = database::get_collection("users") .find_one( doc! { "_id": param.to_string() }, FindOneOptions::builder() .projection(doc! { "_id": 1, "username": 1, "email_verification.verified": 1, }) .build(), ) .unwrap(); if let Some(user) = result { Ok(UserRef { id: user.get_str("_id").unwrap().to_string(), username: user.get_str("username").unwrap().to_string(), email_verified: user .get_document("email_verification") .unwrap() .get_bool("verified") .unwrap(), }) } else { Err(param) } } } impl<'r> FromParam<'r> for User { type Error = &'r RawStr; fn from_param(param: &'r RawStr) -> Result<Self, Self::Error> { let col = database::get_db().collection("users"); let result = col .find_one(doc! { "_id": param.to_string() }, None) .unwrap(); if let Some(user) = result { Ok(from_bson(bson::Bson::Document(user)).expect("Failed to unwrap user.")) } else { Err(param) } } }