Get notifications working properly.

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.");
+            }
+        }
+    }
 }