Get notifications working properly.

src/notifications/events.rs

@@ -1,15 +1,25 @@
 use rauth::auth::Session;
 use serde::{Deserialize, Serialize};
 use snafu::Snafu;
+use hive_pubsub::PubSub;
+
+use crate::database::entities::RelationshipStatus;
+
+use super::hive::get_hive;
 
 #[derive(Serialize, Deserialize, Debug, Snafu)]
 #[serde(tag = "type")]
 pub enum WebSocketError {
     #[snafu(display("This error has not been labelled."))]
     LabelMe,
-
     #[snafu(display("Internal server error."))]
     InternalError,
+    #[snafu(display("Invalid session."))]
+    InvalidSession,
+    #[snafu(display("User hasn't completed onboarding."))]
+    OnboardingNotFinished,
+    #[snafu(display("Already authenticated with server."))]
+    AlreadyAuthenticated,
 }
 
 #[derive(Deserialize, Debug)]
@@ -22,6 +32,7 @@ pub enum ServerboundNotification {
 #[serde(tag = "type")]
 pub enum ClientboundNotification {
     Error(WebSocketError),
+    Authenticated,
 
     /*MessageCreate {
         id: String,
@@ -78,9 +89,16 @@ pub enum ClientboundNotification {
     GuildDelete {
         id: String,
     },*/
+    
     UserRelationship {
         id: String,
         user: String,
-        status: i32,
+        status: RelationshipStatus,
     },
 }
+
+impl ClientboundNotification {
+    pub async fn publish(self, topic: String) -> Result<(), String> {
+        hive_pubsub::backend::mongo::publish(get_hive(), &topic, self).await
+    }
+}

src/notifications/hive.rs

@@ -1,4 +1,4 @@
-use super::events::ClientboundNotification;
+use super::{events::ClientboundNotification, websocket};
 use crate::database::get_collection;
 
 use futures::FutureExt;
@@ -8,14 +8,15 @@ use log::{debug, error};
 use once_cell::sync::OnceCell;
 use serde_json::to_string;
 
-static HIVE: OnceCell<MongodbPubSub<String, String, ClientboundNotification>> = OnceCell::new();
+type Hive = MongodbPubSub<String, String, ClientboundNotification>;
+static HIVE: OnceCell<Hive> = OnceCell::new();
 
 pub async fn init_hive() {
     let hive = MongodbPubSub::new(
-        |_ids, notification| {
+        |ids, notification| {
             if let Ok(data) = to_string(&notification) {
                 debug!("Pushing out notification. {}", data);
-                // ! FIXME: push to websocket
+                websocket::publish(ids, notification);
             } else {
                 error!("Failed to serialise notification.");
             }
@@ -39,12 +40,7 @@ pub async fn listen() {
     dbg!("a");
 }
 
-pub fn publish(topic: &String, data: ClientboundNotification) -> Result<(), String> {
-    let hive = HIVE.get().unwrap();
-    hive.publish(topic, data)
-}
-
-pub fn subscribe(user: String, topics: Vec<String>) -> Result<(), String> {
+pub fn subscribe_multiple(user: String, topics: Vec<String>) -> Result<(), String> {
     let hive = HIVE.get().unwrap();
     for topic in topics {
         hive.subscribe(user.clone(), topic)?;
@@ -53,16 +49,15 @@ pub fn subscribe(user: String, topics: Vec<String>) -> Result<(), String> {
     Ok(())
 }
 
-pub fn drop_user(user: &String) -> Result<(), String> {
+pub fn subscribe_if_exists(user: String, topic: String) -> Result<(), String> {
     let hive = HIVE.get().unwrap();
-    hive.drop_client(user)?;
+    if hive.hive.map.lock().unwrap().get_left(&user).is_some() {
+        hive.subscribe(user, topic)?;
+    }
 
     Ok(())
 }
 
-pub fn drop_topic(topic: &String) -> Result<(), String> {
-    let hive = HIVE.get().unwrap();
-    hive.drop_topic(topic)?;
-
-    Ok(())
+pub fn get_hive() -> &'static Hive {
+    HIVE.get().unwrap()
 }

src/notifications/mod.rs

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

src/notifications/subscriptions.rs

@@ -0,0 +1,17 @@
+use crate::database::entities::User;
+
+use super::hive::get_hive;
+use hive_pubsub::PubSub;
+
+pub async fn generate_subscriptions(user: &User) -> Result<(), String> {
+    let hive = get_hive();
+    hive.subscribe(user.id.clone(), user.id.clone())?;
+
+    if let Some(relations) = &user.relations {
+        for relation in relations {
+            hive.subscribe(user.id.clone(), relation.id.clone())?;
+        }
+    }
+
+    Ok(())
+}

src/notifications/websocket.rs

@@ -1,27 +1,31 @@
+use crate::database::get_collection;
 use crate::{database::entities::User, util::variables::WS_HOST};
+use crate::database::guards::reference::Ref;
+
+use super::subscriptions;
 
 use async_std::net::{TcpListener, TcpStream};
 use async_std::task;
 use async_tungstenite::tungstenite::Message;
 use futures::channel::mpsc::{unbounded, UnboundedSender};
+use futures::stream::TryStreamExt;
 use futures::{pin_mut, prelude::*};
-use log::info;
+use log::{debug, info};
 use many_to_many::ManyToMany;
-use rauth::auth::Session;
+use rauth::auth::{Auth, Session};
 use std::collections::HashMap;
 use std::net::SocketAddr;
-use std::str::from_utf8;
 use std::sync::{Arc, Mutex, RwLock};
-use ulid::Ulid;
+use hive_pubsub::PubSub;
 
-use super::events::ServerboundNotification;
+use super::{events::{ClientboundNotification, ServerboundNotification, WebSocketError}, hive::get_hive};
 
 type Tx = UnboundedSender<Message>;
 type PeerMap = Arc<Mutex<HashMap<SocketAddr, Tx>>>;
 
 lazy_static! {
     static ref CONNECTIONS: PeerMap = Arc::new(Mutex::new(HashMap::new()));
-    static ref USERS: Arc<RwLock<ManyToMany<String, String>>> =
+    static ref USERS: Arc<RwLock<ManyToMany<String, SocketAddr>>> =
         Arc::new(RwLock::new(ManyToMany::new()));
 }
 
@@ -45,32 +49,111 @@ async fn accept(stream: TcpStream) {
 
     info!("User established WebSocket connection from {}.", &addr);
 
-    let id = Ulid::new().to_string();
     let (write, read) = ws_stream.split();
-
     let (tx, rx) = unbounded();
-    CONNECTIONS.lock().unwrap().insert(addr, tx);
+    CONNECTIONS.lock().unwrap().insert(addr, tx.clone());
 
-    let session: Option<Session> = None;
-    let user: Option<User> = None;
+    let send = |notification: ClientboundNotification| {
+        if let Ok(response) = serde_json::to_string(
+            &notification,
+        ) {
+            if let Err(_) = tx.unbounded_send(Message::Text(response)) {
+                debug!("Failed unbounded_send to websocket stream.");
+            }
+        }
+    };
 
+    let mut session: Option<Session> = None;
     let fwd = rx.map(Ok).forward(write);
-    let reading = read.for_each(|message| async {
-        let data = message.unwrap().into_data();
-        // if you mess with the data, you get the bazooki
-        let string = from_utf8(&data).unwrap();
+    let incoming = read.try_for_each(|msg| {
+        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 let Ok(notification) = serde_json::from_str::<ServerboundNotification>(string) {
-            match notification {
-                ServerboundNotification::Authenticate(a) => {
-                    dbg!(a);
+                        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);
+                                            } else {
+                                                send(ClientboundNotification::Error(WebSocketError::InternalError));
+                                            }
+                                        } else {
+                                            send(ClientboundNotification::Error(WebSocketError::InternalError));
+                                        }
+                                    },
+                                    Err(_) => {
+                                        send(ClientboundNotification::Error(WebSocketError::OnboardingNotFinished));
+                                    }
+                                }
+                            }
+                            Err(_) => {
+                                send(ClientboundNotification::Error(WebSocketError::InvalidSession));
+                            }
+                        }
+                    }
                 }
             }
         }
+
+        future::ok(())
     });
 
-    pin_mut!(fwd, reading);
-    future::select(fwd, reading).await;
+    pin_mut!(fwd, incoming);
+    future::select(fwd, incoming).await;
 
-    println!("User {} disconnected.", &addr);
+    info!("User {} disconnected.", &addr);
+    CONNECTIONS.lock().unwrap().remove(&addr);
+
+    if let Some(session) = session {
+        let mut users = USERS.write().unwrap();
+        users.remove(&session.user_id, &addr);
+        if users.get_left(&session.user_id).is_none() {
+            get_hive().drop_client(&session.user_id).unwrap();
+        }
+    }
+}
+
+pub fn publish(ids: Vec<String>, notification: ClientboundNotification) {
+    let mut targets = vec![];
+    {
+        let users = USERS.read().unwrap();
+        for id in ids {
+            // Block certain notifications from reaching users that aren't meant to see them.
+            if let ClientboundNotification::UserRelationship { id: user_id, .. } = &notification {
+                if &id != user_id {
+                    continue;
+                }
+            }
+
+            if let Some(mut arr) = users.get_left(&id) {
+                targets.append(&mut arr);
+            }
+        }
+    }
+
+    let msg = Message::Text(serde_json::to_string(&notification).unwrap());
+
+    let connections = CONNECTIONS.lock().unwrap();
+    for target in targets {
+        if let Some(conn) = connections.get(&target) {
+            if let Err(_) = conn.unbounded_send(msg.clone()) {
+                debug!("Failed unbounded_send.");
+            }
+        }
+    }
 }

src/routes/users/add_friend.rs

@@ -1,4 +1,4 @@
-use crate::util::result::Result;
+use crate::{notifications::{events::ClientboundNotification, hive}, util::result::Result};
 use crate::{
     database::{
         entities::{RelationshipStatus, User},
@@ -49,7 +49,22 @@ pub async fn req(user: User, target: Ref) -> Result<JsonValue> {
                     None
                 )
             ) {
-                Ok(_) => Ok(json!({ "status": "Friend" })),
+                Ok(_) => {
+                    try_join!(
+                        ClientboundNotification::UserRelationship {
+                            id: user.id.clone(),
+                            user: target.id.clone(),
+                            status: RelationshipStatus::Friend
+                        }.publish(user.id.clone()),
+                        ClientboundNotification::UserRelationship {
+                            id: target.id.clone(),
+                            user: user.id.clone(),
+                            status: RelationshipStatus::Friend
+                        }.publish(target.id.clone())
+                    ).ok();
+
+                    Ok(json!({ "status": "Friend" }))
+                },
                 Err(_) => Err(Error::DatabaseError {
                     operation: "update_one",
                     with: "user",
@@ -87,7 +102,25 @@ pub async fn req(user: User, target: Ref) -> Result<JsonValue> {
                     None
                 )
             ) {
-                Ok(_) => Ok(json!({ "status": "Outgoing" })),
+                Ok(_) => {
+                    try_join!(
+                        ClientboundNotification::UserRelationship {
+                            id: user.id.clone(),
+                            user: target.id.clone(),
+                            status: RelationshipStatus::Outgoing
+                        }.publish(user.id.clone()),
+                        ClientboundNotification::UserRelationship {
+                            id: target.id.clone(),
+                            user: user.id.clone(),
+                            status: RelationshipStatus::Incoming
+                        }.publish(target.id.clone())
+                    ).ok();
+
+                    hive::subscribe_if_exists(user.id.clone(), target.id.clone()).ok();
+                    hive::subscribe_if_exists(target.id.clone(), user.id.clone()).ok();
+
+                    Ok(json!({ "status": "Outgoing" }))
+                },
                 Err(_) => Err(Error::DatabaseError {
                     operation: "update_one",
                     with: "user",