Randomly select an avatar from id.

Code cleanup/formatting

See merge request revolt/delta!2

Cargo.lock

@@ -36,6 +36,16 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "23b62fc65de8e4e7f52534fb52b0f3ed04746ae267519eef2a83941e8085068b"
 
+[[package]]
+name = "async-attributes"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "efd3d156917d94862e779f356c5acae312b08fd3121e792c857d7928c8088423"
+dependencies = [
+ "quote 1.0.8",
+ "syn 1.0.56",
+]
+
 [[package]]
 name = "async-channel"
 version = "1.5.1"
@@ -110,6 +120,7 @@ version = "1.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8f9f84f1280a2b436a2c77c2582602732b6c2f4321d5494d6e799e6c367859a8"
 dependencies = [
+ "async-attributes",
  "async-channel",
  "async-global-executor",
  "async-io",

Cargo.toml

@@ -13,13 +13,12 @@ impl_ops = "0.1.1"
 ctrlc = { version = "3.0", features = ["termination"] }
 async-tungstenite = { version = "0.10.0", features = ["async-std-runtime"] }
 rauth = { git = "https://gitlab.insrt.uk/insert/rauth" }
-async-std = { version = "1.8.0", features = ["tokio02"] }
+async-std = { version = "1.8.0", features = ["tokio02", "attributes"] }
 
 hive_pubsub = { version = "0.4.3", features = ["mongo"] }
 rocket_cors = { git = "https://github.com/lawliet89/rocket_cors", branch = "master" }
 rocket_contrib = { git = "https://github.com/SergioBenitez/Rocket", branch = "master" }
 rocket = { git = "https://github.com/SergioBenitez/Rocket", branch = "master", default-features = false }
-# ! FIXME: Switch to async-std runtime.
 mongodb = { version = "1.1.1", features = ["tokio-runtime"], default-features = false }
 
 once_cell = "1.4.1"

src/database/entities/user.rs

@@ -1,9 +1,3 @@
-use crate::database::get_collection;
-use crate::database::guards::reference::Ref;
-use mongodb::bson::{doc, from_bson, Bson};
-use rauth::auth::Session;
-use rocket::http::Status;
-use rocket::request::{self, FromRequest, Outcome, Request};
 use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -31,42 +25,10 @@ pub struct User {
     pub username: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub relations: Option<Vec<Relationship>>,
-}
-
-#[rocket::async_trait]
-impl<'a, 'r> FromRequest<'a, 'r> for User {
-    type Error = rauth::util::Error;
-
-    async fn from_request(request: &'a Request<'r>) -> request::Outcome<Self, Self::Error> {
-        let session: Session = try_outcome!(request.guard::<Session>().await);
-
-        if let Ok(result) = get_collection("users")
-            .find_one(
-                doc! {
-                    "_id": &session.user_id
-                },
-                None,
-            )
-            .await
-        {
-            if let Some(doc) = result {
-                Outcome::Success(from_bson(Bson::Document(doc)).unwrap())
-            } else {
-                Outcome::Failure((Status::Forbidden, rauth::util::Error::InvalidSession))
-            }
-        } else {
-            Outcome::Failure((
-                Status::InternalServerError,
-                rauth::util::Error::DatabaseError,
-            ))
-        }
-    }
-}
-
-impl User {
-    pub fn as_ref(&self) -> Ref {
-        Ref {
-            id: self.id.to_string(),
-        }
-    }
+
+    // ? This should never be pushed to the collection.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub relationship: Option<RelationshipStatus>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub online: Option<bool>,
 }

src/database/guards/mod.rs

@@ -1,2 +1,5 @@
 pub mod reference;
 pub mod user;
+
+pub use reference::Ref;
+// pub use user::*;

src/database/guards/reference.rs

@@ -1,6 +1,6 @@
-use crate::database::entities::*;
-use crate::database::get_collection;
+use crate::database::*;
 use crate::util::result::{Error, Result};
+
 use mongodb::bson::{doc, from_bson, Bson};
 use rocket::http::RawStr;
 use rocket::request::FromParam;
@@ -42,6 +42,14 @@ impl Ref {
     }
 }
 
+impl User {
+    pub fn as_ref(&self) -> Ref {
+        Ref {
+            id: self.id.to_string(),
+        }
+    }
+}
+
 impl<'r> FromParam<'r> for Ref {
     type Error = &'r RawStr;
 

src/database/guards/user.rs

@@ -1 +1,36 @@
+use crate::database::*;
 
+use mongodb::bson::{doc, from_bson, Bson};
+use rauth::auth::Session;
+use rocket::http::Status;
+use rocket::request::{self, FromRequest, Outcome, Request};
+
+#[rocket::async_trait]
+impl<'a, 'r> FromRequest<'a, 'r> for User {
+    type Error = rauth::util::Error;
+
+    async fn from_request(request: &'a Request<'r>) -> request::Outcome<Self, Self::Error> {
+        let session: Session = try_outcome!(request.guard::<Session>().await);
+
+        if let Ok(result) = get_collection("users")
+            .find_one(
+                doc! {
+                    "_id": &session.user_id
+                },
+                None,
+            )
+            .await
+        {
+            if let Some(doc) = result {
+                Outcome::Success(from_bson(Bson::Document(doc)).unwrap())
+            } else {
+                Outcome::Failure((Status::Forbidden, rauth::util::Error::InvalidSession))
+            }
+        } else {
+            Outcome::Failure((
+                Status::InternalServerError,
+                rauth::util::Error::DatabaseError,
+            ))
+        }
+    }
+}

src/database/mod.rs

@@ -30,3 +30,7 @@ pub mod entities;
 pub mod guards;
 pub mod migrations;
 pub mod permissions;
+
+pub use entities::*;
+pub use guards::*;
+pub use permissions::*;

src/database/permissions/mod.rs

@@ -1,49 +1,3 @@
-use crate::database::entities::User;
-use num_enum::TryFromPrimitive;
-use std::ops;
+pub mod user;
 
-#[derive(Debug, PartialEq, Eq, TryFromPrimitive, Copy, Clone)]
-#[repr(u32)]
-pub enum UserPermission {
-    Access = 1,
-    SendMessage = 2,
-    Invite = 4,
-}
-
-bitfield! {
-    pub struct UserPermissions(MSB0 [u32]);
-    u32;
-    pub get_access, _: 31;
-    pub get_send_message, _: 30;
-    pub get_invite, _: 29;
-}
-
-impl_op_ex!(+ |a: &UserPermission, b: &UserPermission| -> u32 { *a as u32 | *b as u32 });
-impl_op_ex_commutative!(+ |a: &u32, b: &UserPermission| -> u32 { *a | *b as u32 });
-
-pub async fn temp_calc_perm(_user: &User, _target: &User) -> UserPermissions<[u32; 1]> {
-    // if friends; Access + Message + Invite
-    // if mutually know each other:
-    //    and has DMs from users enabled -> Access + Message
-    //    otherwise -> Access
-    // otherwise; None
-
-    UserPermissions([UserPermission::Access + UserPermission::SendMessage + UserPermission::Invite])
-}
-
-use crate::database::entities::RelationshipStatus;
-use crate::database::guards::reference::Ref;
-
-pub fn get_relationship(a: &User, b: &Ref) -> RelationshipStatus {
-    if a.id == b.id {
-        return RelationshipStatus::Friend;
-    }
-
-    if let Some(relations) = &a.relations {
-        if let Some(relationship) = relations.iter().find(|x| x.id == b.id) {
-            return relationship.status.clone();
-        }
-    }
-
-    RelationshipStatus::None
-}
+pub use user::get_relationship;

src/database/permissions/user.rs

@@ -0,0 +1,50 @@
+use crate::database::*;
+use num_enum::TryFromPrimitive;
+use std::ops;
+
+#[derive(Debug, PartialEq, Eq, TryFromPrimitive, Copy, Clone)]
+#[repr(u32)]
+pub enum UserPermission {
+    Access = 1,
+    SendMessage = 2,
+    Invite = 4,
+}
+
+bitfield! {
+    pub struct UserPermissions(MSB0 [u32]);
+    u32;
+    pub get_access, _: 31;
+    pub get_send_message, _: 30;
+    pub get_invite, _: 29;
+}
+
+impl_op_ex!(+ |a: &UserPermission, b: &UserPermission| -> u32 { *a as u32 | *b as u32 });
+impl_op_ex_commutative!(+ |a: &u32, b: &UserPermission| -> u32 { *a | *b as u32 });
+
+pub async fn calculate(user: &User, target: &str) -> UserPermissions<[u32; 1]> {
+    // if friends; Access + Message + Invite
+    // if mutually know each other:
+    //    and has DMs from users enabled -> Access + Message
+    //    otherwise -> Access
+    // otherwise; None
+
+    if let RelationshipStatus::Friend = get_relationship(&user, &target) {
+        UserPermissions([UserPermission::Access + UserPermission::SendMessage + UserPermission::Invite])
+    } else {
+        UserPermissions([ 0 ])
+    }
+}
+
+pub fn get_relationship(a: &User, b: &str) -> RelationshipStatus {
+    if a.id == b {
+        return RelationshipStatus::Friend;
+    }
+
+    if let Some(relations) = &a.relations {
+        if let Some(relationship) = relations.iter().find(|x| x.id == b) {
+            return relationship.status.clone();
+        }
+    }
+
+    RelationshipStatus::None
+}

src/main.rs

@@ -18,17 +18,13 @@ pub mod notifications;
 pub mod routes;
 pub mod util;
 
-use async_std::task;
 use futures::join;
 use log::info;
 use rauth;
 use rocket_cors::AllowedOrigins;
 
-fn main() {
-    task::block_on(entry())
-}
-
-async fn entry() {
+#[async_std::main]
+async fn main() {
     dotenv::dotenv().ok();
     env_logger::init_from_env(env_logger::Env::default().filter_or("RUST_LOG", "info"));
 

src/notifications/events.rs

@@ -2,9 +2,8 @@ use rauth::auth::Session;
 use serde::{Deserialize, Serialize};
 use snafu::Snafu;
 
-use crate::database::entities::{RelationshipStatus, User};
-
 use super::hive::get_hive;
+use crate::database::*;
 
 #[derive(Serialize, Deserialize, Debug, Snafu)]
 #[serde(tag = "error")]
@@ -12,7 +11,7 @@ pub enum WebSocketError {
     #[snafu(display("This error has not been labelled."))]
     LabelMe,
     #[snafu(display("Internal server error."))]
-    InternalError,
+    InternalError { at: String },
     #[snafu(display("Invalid session."))]
     InvalidSession,
     #[snafu(display("User hasn't completed onboarding."))]
@@ -33,7 +32,7 @@ pub enum ClientboundNotification {
     Error(WebSocketError),
     Authenticated,
     Ready {
-        user: User
+        users: Vec<User>,
     },
 
     /*MessageCreate {
@@ -91,11 +90,17 @@ pub enum ClientboundNotification {
     GuildDelete {
         id: String,
     },*/
+
     UserRelationship {
         id: String,
         user: String,
         status: RelationshipStatus,
     },
+
+    UserPresence {
+        id: String,
+        online: bool
+    }
 }
 
 impl ClientboundNotification {

src/notifications/hive.rs

@@ -1,5 +1,5 @@
 use super::{events::ClientboundNotification, websocket};
-use crate::database::get_collection;
+use crate::database::*;
 
 use futures::FutureExt;
 use hive_pubsub::backend::mongo::MongodbPubSub;

src/notifications/mod.rs

@@ -1,4 +1,5 @@
 pub mod events;
 pub mod hive;
+pub mod payload;
 pub mod subscriptions;
 pub mod websocket;

src/notifications/payload.rs

@@ -0,0 +1,68 @@
+use crate::notifications::events::ClientboundNotification;
+use crate::{
+    database::{entities::User, get_collection},
+    util::result::{Error, Result},
+};
+use futures::StreamExt;
+use mongodb::{
+    bson::{doc, from_bson, Bson},
+    options::FindOptions,
+};
+
+use super::websocket::is_online;
+
+pub async fn generate_ready(mut user: User) -> Result<ClientboundNotification> {
+    let mut users = vec![];
+
+    if let Some(relationships) = &user.relations {
+        let user_ids: Vec<String> = relationships
+            .iter()
+            .map(|relationship| relationship.id.clone())
+            .collect();
+
+        let mut cursor = get_collection("users")
+            .find(
+                doc! {
+                    "_id": {
+                        "$in": user_ids
+                    }
+                },
+                FindOptions::builder()
+                    .projection(doc! { "_id": 1, "username": 1 })
+                    .build(),
+            )
+            .await
+            .map_err(|_| Error::DatabaseError {
+                operation: "find",
+                with: "users",
+            })?;
+
+        while let Some(result) = cursor.next().await {
+            if let Ok(doc) = result {
+                let mut user: User =
+                    from_bson(Bson::Document(doc)).map_err(|_| Error::DatabaseError {
+                        operation: "from_bson",
+                        with: "user",
+                    })?;
+
+                user.relationship = Some(
+                    relationships
+                        .iter()
+                        .find(|x| user.id == x.id)
+                        .ok_or_else(|| Error::InternalError)?
+                        .status
+                        .clone(),
+                );
+
+                user.online = Some(is_online(&user.id));
+
+                users.push(user);
+            }
+        }
+    }
+
+    user.online = Some(is_online(&user.id));
+    users.push(user);
+
+    Ok(ClientboundNotification::Ready { users })
+}

src/notifications/subscriptions.rs

@@ -1,4 +1,4 @@
-use crate::database::entities::User;
+use crate::database::*;
 
 use super::hive::get_hive;
 use hive_pubsub::PubSub;

src/notifications/websocket.rs

@@ -1,5 +1,4 @@
-use crate::database::get_collection;
-use crate::database::guards::reference::Ref;
+use crate::database::*;
 use crate::util::variables::WS_HOST;
 
 use super::subscriptions;
@@ -64,69 +63,104 @@ async fn accept(stream: TcpStream) {
         }
     };
 
-    let mut session: Option<Session> = None;
+    let session: Arc<Mutex<Option<Session>>> = Arc::new(Mutex::new(None));
+    let mutex_generator = || { session.clone() };
     let fwd = rx.map(Ok).forward(write);
-    let incoming = read.try_for_each(|msg| {
+    let incoming = read.try_for_each(async move |msg| {
+        let mutex = mutex_generator();
+        //dbg!(&mutex.lock().unwrap());
+
         if let Message::Text(text) = msg {
             if let Ok(notification) = serde_json::from_str::<ServerboundNotification>(&text) {
                 match notification {
                     ServerboundNotification::Authenticate(new_session) => {
-                        if session.is_some() {
-                            send(ClientboundNotification::Error(
-                                WebSocketError::AlreadyAuthenticated,
-                            ));
-                            return future::ok(());
+                        {
+                            if mutex.lock().unwrap().is_some() {
+                                send(ClientboundNotification::Error(
+                                    WebSocketError::AlreadyAuthenticated,
+                                ));
+                                
+                                return Ok(())
+                            }
                         }
 
-                        match task::block_on(
-                            Auth::new(get_collection("accounts")).verify_session(new_session),
-                        ) {
-                            Ok(validated_session) => {
-                                match task::block_on(
-                                    Ref {
-                                        id: validated_session.user_id.clone(),
-                                    }
-                                    .fetch_user(),
-                                ) {
-                                    Ok(user) => {
-                                        if let Ok(mut map) = USERS.write() {
-                                            map.insert(validated_session.user_id.clone(), addr);
-                                            session = Some(validated_session);
-                                            if let Ok(_) = task::block_on(
-                                                subscriptions::generate_subscriptions(&user),
-                                            ) {
-                                                send(ClientboundNotification::Authenticated);
-                                                send(ClientboundNotification::Ready { user });
-                                            } else {
-                                                send(ClientboundNotification::Error(
-                                                    WebSocketError::InternalError,
-                                                ));
-                                            }
-                                        } else {
+                        if let Ok(validated_session) = Auth::new(get_collection("accounts"))
+                            .verify_session(new_session)
+                            .await {
+                            let id = validated_session.user_id.clone();
+                            if let Ok(user) = (
+                                Ref {
+                                    id: id.clone()
+                                }
+                            )
+                            .fetch_user()
+                            .await {
+                                let was_online = is_online(&id);
+                                {
+                                    match USERS.write() {
+                                        Ok(mut map) => {
+                                            map.insert(id.clone(), addr);
+                                        }
+                                        Err(_) => {
                                             send(ClientboundNotification::Error(
-                                                WebSocketError::InternalError,
+                                                WebSocketError::InternalError { at: "Writing users map.".to_string() },
                                             ));
+
+                                            return Ok(())
+                                        }
+                                    }
+                                }
+
+                                *mutex.lock().unwrap() = Some(validated_session);
+
+                                if let Err(_) = subscriptions::generate_subscriptions(&user).await {
+                                    send(ClientboundNotification::Error(
+                                        WebSocketError::InternalError { at: "Generating subscriptions.".to_string() },
+                                    ));
+
+                                    return Ok(())
+                                }
+
+                                send(ClientboundNotification::Authenticated);
+
+                                match super::payload::generate_ready(user).await {
+                                    Ok(payload) => {
+                                        send(payload);
+
+                                        if !was_online {
+                                            ClientboundNotification::UserPresence {
+                                                id: id.clone(),
+                                                online: true
+                                            }
+                                            .publish(id)
+                                            .await
+                                            .ok();
                                         }
                                     }
                                     Err(_) => {
                                         send(ClientboundNotification::Error(
-                                            WebSocketError::OnboardingNotFinished,
+                                            WebSocketError::InternalError { at: "Generating payload.".to_string() },
                                         ));
+
+                                        return Ok(())
                                     }
                                 }
-                            }
-                            Err(_) => {
+                            } else {
                                 send(ClientboundNotification::Error(
-                                    WebSocketError::InvalidSession,
+                                    WebSocketError::OnboardingNotFinished,
                                 ));
                             }
+                        } else {
+                            send(ClientboundNotification::Error(
+                                WebSocketError::InvalidSession,
+                            ));
                         }
                     }
                 }
             }
         }
 
-        future::ok(())
+        Ok(())
     });
 
     pin_mut!(fwd, incoming);
@@ -135,7 +169,8 @@ async fn accept(stream: TcpStream) {
     info!("User {} disconnected.", &addr);
     CONNECTIONS.lock().unwrap().remove(&addr);
 
-    if let Some(session) = session {
+    let session = session.lock().unwrap();
+    if let Some(session) = session.as_ref() {
         let mut users = USERS.write().unwrap();
         users.remove(&session.user_id, &addr);
         if users.get_left(&session.user_id).is_none() {
@@ -173,3 +208,7 @@ pub fn publish(ids: Vec<String>, notification: ClientboundNotification) {
         }
     }
 }
+
+pub fn is_online(user: &String) -> bool {
+    USERS.read().unwrap().get_left(&user).is_some()
+}

src/routes/onboard/complete.rs

@@ -1,6 +1,6 @@
-use crate::database::entities::User;
-use crate::database::get_collection;
+use crate::database::*;
 use crate::util::result::{Error, Result};
+
 use mongodb::bson::doc;
 use mongodb::options::{Collation, FindOneOptions};
 use rauth::auth::Session;
@@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize};
 use validator::Validate;
 
 lazy_static! {
-    static ref RE_USERNAME: Regex = Regex::new(r"^[a-zA-Z0-9-_]+$").unwrap();
+    static ref RE_USERNAME: Regex = Regex::new(r"^[a-zA-Z0-9_]+$").unwrap();
 }
 
 #[derive(Validate, Serialize, Deserialize)]

src/routes/onboard/hello.rs

@@ -1,4 +1,5 @@
-use crate::database::entities::User;
+use crate::database::*;
+
 use rauth::auth::Session;
 use rocket_contrib::json::JsonValue;
 

src/routes/root.rs

@@ -1,4 +1,6 @@
-use crate::util::variables::{DISABLE_REGISTRATION, HCAPTCHA_SITEKEY, USE_EMAIL, USE_HCAPTCHA};
+use crate::util::variables::{
+    DISABLE_REGISTRATION, EXTERNAL_WS_URL, HCAPTCHA_SITEKEY, USE_EMAIL, USE_HCAPTCHA,
+};
 
 use mongodb::bson::doc;
 use rocket_contrib::json::JsonValue;
@@ -14,6 +16,7 @@ pub async fn root() -> JsonValue {
                 "key": HCAPTCHA_SITEKEY.to_string()
             },
             "email": *USE_EMAIL,
-        }
+        },
+        "ws": *EXTERNAL_WS_URL,
     })
 }

src/routes/users/add_friend.rs

@@ -1,41 +1,40 @@
-use crate::{
-    database::{
-        entities::{RelationshipStatus, User},
-        get_collection,
-        guards::reference::Ref,
-        permissions::get_relationship,
-    },
-    util::result::Error,
-};
-use crate::{
-    notifications::{events::ClientboundNotification, hive},
-    util::result::Result,
-};
+use crate::database::*;
+use crate::notifications::{events::ClientboundNotification, hive};
+use crate::util::result::{Error, Result};
+
 use futures::try_join;
 use mongodb::bson::doc;
-use mongodb::options::{FindOneOptions, Collation};
+use mongodb::options::{Collation, FindOneOptions};
 use rocket_contrib::json::JsonValue;
 
 #[put("/<username>/friend")]
 pub async fn req(user: User, username: String) -> Result<JsonValue> {
     let col = get_collection("users");
-    let doc = col.find_one(
-        doc! {
-            "username": username
-        },
-        FindOneOptions::builder()
-            .collation(Collation::builder().locale("en").strength(2).build())
-            .build(),
-    )
-    .await
-    .map_err(|_| Error::DatabaseError { operation: "find_one", with: "user" })?
-    .ok_or_else(|| Error::UnknownUser)?;
+    let doc = col
+        .find_one(
+            doc! {
+                "username": username
+            },
+            FindOneOptions::builder()
+                .collation(Collation::builder().locale("en").strength(2).build())
+                .build(),
+        )
+        .await
+        .map_err(|_| Error::DatabaseError {
+            operation: "find_one",
+            with: "user",
+        })?
+        .ok_or_else(|| Error::UnknownUser)?;
 
-    let target_id = doc
-        .get_str("_id")
-        .map_err(|_| Error::DatabaseError { operation: "get_str(_id)", with: "user" })?;
+    let target_id = doc.get_str("_id").map_err(|_| Error::DatabaseError {
+        operation: "get_str(_id)",
+        with: "user",
+    })?;
 
-    match get_relationship(&user, &Ref { id: target_id.to_string() }) {
+    match get_relationship(
+        &user,
+        &target_id,
+    ) {
         RelationshipStatus::User => return Err(Error::NoEffect),
         RelationshipStatus::Friend => return Err(Error::AlreadyFriends),
         RelationshipStatus::Outgoing => return Err(Error::AlreadySentRequest),

src/routes/users/block_user.rs

@@ -1,11 +1,7 @@
-use crate::{
-    database::entities::RelationshipStatus, database::entities::User, database::get_collection,
-    database::guards::reference::Ref, database::permissions::get_relationship, util::result::Error,
-};
-use crate::{
-    notifications::{events::ClientboundNotification, hive},
-    util::result::Result,
-};
+use crate::database::*;
+use crate::notifications::{events::ClientboundNotification, hive};
+use crate::util::result::{Error, Result};
+
 use futures::try_join;
 use mongodb::bson::doc;
 use rocket_contrib::json::JsonValue;
@@ -14,7 +10,7 @@ use rocket_contrib::json::JsonValue;
 pub async fn req(user: User, target: Ref) -> Result<JsonValue> {
     let col = get_collection("users");
 
-    match get_relationship(&user, &target) {
+    match get_relationship(&user, &target.id) {
         RelationshipStatus::User | RelationshipStatus::Blocked => Err(Error::NoEffect),
         RelationshipStatus::BlockedOther => {
             col.update_one(

src/routes/users/fetch_dms.rs

@@ -1,6 +1,6 @@
-use crate::database::entities::User;
-use crate::database::get_collection;
+use crate::database::*;
 use crate::util::result::{Error, Result};
+
 use futures::StreamExt;
 use mongodb::bson::doc;
 use rocket_contrib::json::JsonValue;

src/routes/users/fetch_relationship.rs

@@ -1,8 +1,9 @@
-use crate::database::{entities::User, guards::reference::Ref, permissions::get_relationship};
+use crate::database::*;
 use crate::util::result::Result;
+
 use rocket_contrib::json::JsonValue;
 
 #[get("/<target>/relationship")]
 pub async fn req(user: User, target: Ref) -> Result<JsonValue> {
-    Ok(json!({ "status": get_relationship(&user, &target) }))
+    Ok(json!({ "status": get_relationship(&user, &target.id) }))
 }

src/routes/users/fetch_relationships.rs

@@ -1,5 +1,6 @@
-use crate::database::entities::User;
+use crate::database::*;
 use crate::util::result::Result;
+
 use rocket_contrib::json::JsonValue;
 
 #[get("/relationships")]

src/routes/users/fetch_user.rs

@@ -1,6 +1,6 @@
-use crate::database::entities::User;
-use crate::database::guards::reference::Ref;
+use crate::{database::*, notifications::websocket::is_online};
 use crate::util::result::{Error, Result};
+
 use rocket_contrib::json::JsonValue;
 
 #[get("/<target>")]
@@ -9,14 +9,29 @@ pub async fn req(user: User, target: Ref) -> Result<JsonValue> {
 
     if user.id != target.id {
         // Check whether we are allowed to fetch this user.
-        let perm = crate::database::permissions::temp_calc_perm(&user, &target).await;
+        let perm = permissions::user::calculate(&user, &target.id).await;
         if !perm.get_access() {
             Err(Error::LabelMe)?
         }
 
         // Only return user relationships if the target is the caller.
         target.relations = None;
+
+        // Add relevant relationship
+        if let Some(relationships) = &user.relations {
+            target.relationship = relationships
+                .iter()
+                .find(|x| x.id == user.id)
+                .map(|x| x.status.clone())
+                .or_else(|| Some(RelationshipStatus::None));
+        } else {
+            target.relationship = Some(RelationshipStatus::None);
+        }
+    } else {
+        target.relationship = Some(RelationshipStatus::User);
     }
 
+    target.online = Some(is_online(&target.id));
+
     Ok(json!(target))
 }

src/routes/users/get_avatar.rs

@@ -1,7 +1,43 @@
 use rocket::response::NamedFile;
 use std::path::Path;
 
-#[get("/<_target>/avatar")]
-pub async fn req(_target: String) -> Option<NamedFile> {
-    NamedFile::open(Path::new("avatar.png")).await.ok()
+use crate::database::Ref;
+
+#[get("/<target>/avatar")]
+pub async fn req(target: Ref) -> Option<NamedFile> {
+    match target.id.chars().nth(25).unwrap() {
+        '0' |
+        '1' |
+        '2' |
+        '3' |
+        '4' |
+        '5' |
+        '6' |
+        '7' => NamedFile::open(Path::new("assets/user_red.png")).await.ok(),
+        '8' |
+        '9' |
+        'A' |
+        'C' |
+        'B' |
+        'D' |
+        'E' |
+        'F' => NamedFile::open(Path::new("assets/user_green.png")).await.ok(),
+        'G' |
+        'H' |
+        'J' |
+        'K' |
+        'M' |
+        'N' |
+        'P' |
+        'Q' => NamedFile::open(Path::new("assets/user_blue.png")).await.ok(),
+        'R' |
+        'S' |
+        'T' |
+        'V' |
+        'W' |
+        'X' |
+        'Y' |
+        'Z' => NamedFile::open(Path::new("assets/user_yellow.png")).await.ok(),
+        _ => unreachable!()
+    }
 }

src/routes/users/mod.rs

@@ -1,12 +1,12 @@
 use rocket::Route;
 
 mod add_friend;
-mod get_avatar;
 mod block_user;
 mod fetch_dms;
 mod fetch_relationship;
 mod fetch_relationships;
 mod fetch_user;
+mod get_avatar;
 mod open_dm;
 mod remove_friend;
 mod unblock_user;
@@ -16,11 +16,9 @@ pub fn routes() -> Vec<Route> {
         // User Information
         fetch_user::req,
         get_avatar::req,
-
         // Direct Messaging
         fetch_dms::req,
         open_dm::req,
-
         // Relationships
         fetch_relationships::req,
         fetch_relationship::req,

src/routes/users/open_dm.rs

@@ -1,7 +1,6 @@
-use crate::database::entities::{Channel, User};
-use crate::database::get_collection;
-use crate::database::guards::reference::Ref;
+use crate::database::*;
 use crate::util::result::{Error, Result};
+
 use mongodb::bson::doc;
 use rocket_contrib::json::JsonValue;
 use ulid::Ulid;

src/routes/users/remove_friend.rs

@@ -1,11 +1,7 @@
-use crate::{
-    database::entities::RelationshipStatus, database::entities::User, database::get_collection,
-    database::guards::reference::Ref, database::permissions::get_relationship, util::result::Error,
-};
-use crate::{
-    notifications::{events::ClientboundNotification, hive},
-    util::result::Result,
-};
+use crate::database::*;
+use crate::notifications::{events::ClientboundNotification, hive};
+use crate::util::result::{Error, Result};
+
 use futures::try_join;
 use hive_pubsub::PubSub;
 use mongodb::bson::doc;
@@ -15,11 +11,7 @@ use rocket_contrib::json::JsonValue;
 pub async fn req(user: User, target: Ref) -> Result<JsonValue> {
     let col = get_collection("users");
 
-    match get_relationship(&user, &target) {
-        RelationshipStatus::Blocked
-        | RelationshipStatus::BlockedOther
-        | RelationshipStatus::User
-        | RelationshipStatus::None => Err(Error::NoEffect),
+    match get_relationship(&user, &target.id) {
         RelationshipStatus::Friend
         | RelationshipStatus::Outgoing
         | RelationshipStatus::Incoming => {
@@ -80,5 +72,6 @@ pub async fn req(user: User, target: Ref) -> Result<JsonValue> {
                 }),
             }
         }
+        _ => Err(Error::NoEffect),
     }
 }

src/routes/users/unblock_user.rs

@@ -1,11 +1,7 @@
-use crate::{
-    database::entities::RelationshipStatus, database::entities::User, database::get_collection,
-    database::guards::reference::Ref, database::permissions::get_relationship, util::result::Error,
-};
-use crate::{
-    notifications::{events::ClientboundNotification, hive},
-    util::result::Result,
-};
+use crate::database::*;
+use crate::notifications::{events::ClientboundNotification, hive};
+use crate::util::result::{Error, Result};
+
 use futures::try_join;
 use hive_pubsub::PubSub;
 use mongodb::bson::doc;
@@ -15,15 +11,9 @@ use rocket_contrib::json::JsonValue;
 pub async fn req(user: User, target: Ref) -> Result<JsonValue> {
     let col = get_collection("users");
 
-    match get_relationship(&user, &target) {
-        RelationshipStatus::None
-        | RelationshipStatus::User
-        | RelationshipStatus::BlockedOther
-        | RelationshipStatus::Incoming
-        | RelationshipStatus::Outgoing
-        | RelationshipStatus::Friend => Err(Error::NoEffect),
+    match get_relationship(&user, &target.id) {
         RelationshipStatus::Blocked => {
-            match get_relationship(&target.fetch_user().await?, &user.as_ref()) {
+            match get_relationship(&target.fetch_user().await?, &user.id) {
                 RelationshipStatus::Blocked => {
                     col.update_one(
                         doc! {
@@ -115,5 +105,6 @@ pub async fn req(user: User, target: Ref) -> Result<JsonValue> {
                 _ => Err(Error::InternalError),
             }
         }
+        _ => Err(Error::NoEffect),
     }
 }

src/util/variables.rs

@@ -9,6 +9,8 @@ lazy_static! {
         env::var("REVOLT_MONGO_URI").expect("Missing REVOLT_MONGO_URI environment variable.");
     pub static ref PUBLIC_URL: String =
         env::var("REVOLT_PUBLIC_URL").expect("Missing REVOLT_PUBLIC_URL environment variable.");
+    pub static ref EXTERNAL_WS_URL: String =
+        env::var("REVOLT_EXTERNAL_WS_URL").expect("Missing REVOLT_EXTERNAL_WS_URL environment variable.");
     pub static ref HCAPTCHA_KEY: String =
         env::var("REVOLT_HCAPTCHA_KEY").unwrap_or_else(|_| "0x0000000000000000000000000000000000000000".to_string());
     pub static ref HCAPTCHA_SITEKEY: String =
@@ -40,12 +42,8 @@ lazy_static! {
 pub fn preflight_checks() {
     if *USE_EMAIL == false {
         #[cfg(not(debug_assertions))]
-        {
-            if !env::var("REVOLT_UNSAFE_NO_EMAIL").map_or(false, |v| v == *"1") {
-                panic!(
-                    "Not letting you run this in production, set REVOLT_UNSAFE_NO_EMAIL=1 to run."
-                );
-            }
+        if !env::var("REVOLT_UNSAFE_NO_EMAIL").map_or(false, |v| v == *"1") {
+            panic!("Running in production without email is not recommended, set REVOLT_UNSAFE_NO_EMAIL=1 to override.");
         }
 
         #[cfg(debug_assertions)]
@@ -54,10 +52,8 @@ pub fn preflight_checks() {
 
     if *USE_HCAPTCHA == false {
         #[cfg(not(debug_assertions))]
-        {
-            if !env::var("REVOLT_UNSAFE_NO_CAPTCHA").map_or(false, |v| v == *"1") {
-                panic!("Not letting you run this in production, set REVOLT_UNSAFE_NO_CAPTCHA=1 to run.");
-            }
+        if !env::var("REVOLT_UNSAFE_NO_CAPTCHA").map_or(false, |v| v == *"1") {
+            panic!("Running in production without CAPTCHA is not recommended, set REVOLT_UNSAFE_NO_CAPTCHA=1 to override.");
         }
 
         #[cfg(debug_assertions)]