From b87f396f40d4f6b50b2f0d66f1c275ff190e36a0 Mon Sep 17 00:00:00 2001
From: Paul Makles <paulmakles@gmail.com>
Date: Sun, 3 Jan 2021 17:43:20 +0000
Subject: [PATCH] Make ready payload send all users we have a relationship
with.
---
src/database/entities/user.rs | 56 ++++++++++++++++++++++++++++++++--
src/notifications/events.rs | 2 +-
src/notifications/websocket.rs | 14 ++++++++-
3 files changed, 68 insertions(+), 4 deletions(-)
diff --git a/src/database/entities/user.rs b/src/database/entities/user.rs
index 1d97b7f..53bd531 100644
--- a/src/database/entities/user.rs
+++ b/src/database/entities/user.rs
@@ -1,10 +1,11 @@
-use crate::database::get_collection;
+use crate::{database::get_collection, notifications::events::ClientboundNotification, util::result::{Error, Result}};
use crate::database::guards::reference::Ref;
-use mongodb::bson::{doc, from_bson, Bson};
+use mongodb::{bson::{doc, from_bson, Bson}, options::FindOptions};
use rauth::auth::Session;
use rocket::http::Status;
use rocket::request::{self, FromRequest, Outcome, Request};
use serde::{Deserialize, Serialize};
+use futures::StreamExt;
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum RelationshipStatus {
@@ -31,6 +32,10 @@ pub struct User {
pub username: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub relations: Option<Vec<Relationship>>,
+
+ // ? This should never be pushed to the collection.
+ #[serde(skip_serializing_if = "Option::is_none")]
+ pub relationship: Option<RelationshipStatus>
}
#[rocket::async_trait]
@@ -69,4 +74,51 @@ impl User {
id: self.id.to_string(),
}
}
+
+ pub async fn generate_ready_payload(self) -> Result<ClientboundNotification> {
+ let mut users = vec![];
+
+ if let Some(relationships) = &self.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()
+ );
+
+ users.push(user);
+ }
+ }
+ }
+
+ users.push(self);
+
+ Ok(ClientboundNotification::Ready { users })
+ }
}
diff --git a/src/notifications/events.rs b/src/notifications/events.rs
index 0cda340..eb289fe 100644
--- a/src/notifications/events.rs
+++ b/src/notifications/events.rs
@@ -33,7 +33,7 @@ pub enum ClientboundNotification {
Error(WebSocketError),
Authenticated,
Ready {
- user: User
+ users: Vec<User>
},
/*MessageCreate {
diff --git a/src/notifications/websocket.rs b/src/notifications/websocket.rs
index dd47235..6bc8936 100644
--- a/src/notifications/websocket.rs
+++ b/src/notifications/websocket.rs
@@ -96,7 +96,19 @@ async fn accept(stream: TcpStream) {
subscriptions::generate_subscriptions(&user),
) {
send(ClientboundNotification::Authenticated);
- send(ClientboundNotification::Ready { user });
+
+ match task::block_on(
+ user.generate_ready_payload()
+ ) {
+ Ok(payload) => {
+ send(payload);
+ }
+ Err(_) => {
+ send(ClientboundNotification::Error(
+ WebSocketError::InternalError,
+ ));
+ }
+ }
} else {
send(ClientboundNotification::Error(
WebSocketError::InternalError,
--
GitLab