Add unsubscribe endpoint

Add action to output envelope
2022-08-27 22:13:19 +02:00 · 2022-08-27 18:19:29 +02:00
9 changed files with 226 additions and 109 deletions
--- a/src/error.rs
+++ b/src/error.rs
@@ -41,6 +41,7 @@ impl InputEnvelope {
        OutputEnvelope {
            request_nr: self.request_nr,
            error: Some(Error::new(kind, class, msg)),
+            action: None,
            data: Value::Null,
        }
    }
@@ -49,6 +50,7 @@ impl InputEnvelope {
        OutputEnvelope {
            request_nr: self.request_nr,
            error: Some(error),
+            action: None,
            data: Value::Null,
        }
    }
@@ -85,8 +87,9 @@ impl From<Error> for OutputEnvelope {
    fn from(error: Error) -> Self {
        OutputEnvelope {
            error: Some(error),
-            data: Value::Null,
            request_nr: None,
+            action: None,
+            data: Value::Null,
        }
    }
 }
--- a/src/http.rs
+++ b/src/http.rs
@@ -3,7 +3,7 @@ use crate::usimp;
 use crate::usimp::*;
 use crate::websocket;
 use hyper::{body, header, Body, Method, Request, Response, StatusCode};
-use serde_json::{Map, Value};
+use serde_json::Value;
 use std::str::FromStr;
 use uuid::Uuid;

@@ -146,10 +146,7 @@ pub async fn handler(mut req: Request<Body>) -> Result<Response<Body>, hyper::Er
            };

            if let Some(output) = output {
-                let mut data = Value::Object(Map::new());
-
-                match output.error {
-                    Some(error) => {
+                if let Some(ref error) = output.error {
                    res = match error.class {
                        ErrorClass::ClientProtocolError => res.status(StatusCode::BAD_REQUEST),
                        ErrorClass::ServerError => {
@@ -157,21 +154,11 @@ pub async fn handler(mut req: Request<Body>) -> Result<Response<Body>, hyper::Er
                        }
                        _ => res.status(StatusCode::OK),
                    };
-                        data["status"] = Value::from("error");
-                        data["error"] = Value::from(error);
-                    }
-                    None => {
-                        data["status"] = Value::from("success");
-                    }
                }

-                data["request_nr"] = match output.request_nr {
-                    Some(nr) => Value::from(nr),
-                    None => Value::Null,
-                };
-                data["data"] = output.data;
+                let data: Value = output.into();
                return Ok(res
-                    .body(Body::from(serde_json::to_string(&data).unwrap() + "\r\n"))
+                    .body(Body::from(data.to_string() + "\r\n"))
                    .unwrap());
            } else {
                res = res.status(StatusCode::NO_CONTENT);
--- a/src/usimp/handler/mod.rs
+++ b/src/usimp/handler/mod.rs
@@ -2,6 +2,7 @@ mod authenticate;
 mod new_event;
 mod ping;
 mod subscribe;
+mod unsubscribe;

 use crate::usimp::*;
 use tokio::sync::mpsc;
@@ -28,10 +29,11 @@ pub async fn endpoint(

    println!("Endpoint: {}", input.endpoint);
    Ok(match input.endpoint.as_str() {
-        "ping" => input.respond(ping::handle(&input, session).await?),
-        "authenticate" => input.respond(authenticate::handle(&input, session).await?),
-        "subscribe" => input.respond(subscribe::handle(&input, session, tx).await?),
-        "new_event" => input.respond(new_event::handle(&input, session).await?),
+        "ping" => input.respond(ping::handle(&input, session).await?, None),
+        "authenticate" => input.respond(authenticate::handle(&input, session).await?, None),
+        "subscribe" => input.respond(subscribe::handle(&input, session, tx).await?, Some(OutputAction::Subscribe)),
+        "unsubscribe" => input.respond(unsubscribe::handle(&input, session).await?, Some(OutputAction::Unsubscribe)),
+        "new_event" => input.respond(new_event::handle(&input, session).await?, None),
        _ => input.new_error(
            ErrorKind::UsimpError,
            ErrorClass::ClientProtocolError,
--- a/src/usimp/handler/new_event.rs
+++ b/src/usimp/handler/new_event.rs
@@ -22,14 +22,14 @@ pub async fn handle(input: &InputEnvelope, session: Option<Session>) -> Result<V
 }

 async fn new_event(input: Input, session: Option<Session>) -> Result<Output, Error> {
-    let _account = get_account(&session)?;
+    let _account = get_account_opt(&session)?;
    let mut uuids = vec![];
    // TODO check permissions
    for mut event in input.events {
        let uuid = Uuid::new_v4();
        event.id = Some(uuid);
        uuids.push(uuid);
-        subscription::push(&input.room_id, event).await?;
+        subscription::push_room(&input.room_id, event).await?;
    }
    Ok(Output {events: uuids})
 }
--- a/src/usimp/handler/subscribe.rs
+++ b/src/usimp/handler/subscribe.rs
@@ -25,8 +25,7 @@ pub async fn handle(
            session,
            input.request_nr,
            tx,
-        )
-        .await?,
+        ).await?,
    )?)
 }

@@ -36,8 +35,16 @@ async fn subscribe(
    req_nr: Option<u64>,
    tx: Option<mpsc::Sender<WebSocketEnvelope>>,
 ) -> Result<Output, Error> {
-    let account = get_account(&session)?;
-    let mut rx = subscription::subscribe_account(account).await;
+    let session = match session {
+        Some(s) => s,
+        None => return Err(Error::new(
+            ErrorKind::SubscriptionError,
+            ErrorClass::ClientError,
+            None,
+        )),
+    };
+
+    let mut rx = subscription::subscribe_account(&session, req_nr).await?;
    match tx {
        Some(tx) => {
            tokio::spawn(async move {
@@ -46,6 +53,7 @@ async fn subscribe(
                        .send(OutputEnvelope {
                            error: None,
                            request_nr: req_nr,
+                            action: Some(OutputAction::Push),
                            data: serde_json::json![{"events": [event]}],
                        }.into())
                        .await;
--- a/src/usimp/handler/unsubscribe.rs
+++ b/src/usimp/handler/unsubscribe.rs
@@ -0,0 +1,45 @@
+use crate::usimp::subscription;
+use crate::usimp::*;
+
+use serde::{Deserialize, Serialize};
+use serde_json::{from_value, to_value, Value};
+use tokio::sync::mpsc;
+use crate::websocket::WebSocketEnvelope;
+
+#[derive(Serialize, Deserialize, Clone)]
+struct Input {}
+
+#[derive(Serialize, Deserialize, Clone)]
+struct Output {}
+
+pub async fn handle(
+    input: &InputEnvelope,
+    session: Option<Session>,
+) -> Result<Value, Error> {
+    Ok(to_value(
+        unsubscribe(
+            from_value(input.data.clone())?,
+            session,
+            input.request_nr,
+        ).await?,
+    )?)
+}
+
+async fn unsubscribe(
+    _input: Input,
+    session: Option<Session>,
+    req_nr: Option<u64>,
+) -> Result<Output, Error> {
+    let session = match session {
+        Some(s) => s,
+        None => return Err(Error::new(
+            ErrorKind::SubscriptionError,
+            ErrorClass::ClientError,
+            None,
+        )),
+    };
+
+    subscription::unsubscribe_account(&session, req_nr).await?;
+
+    Ok(Output {})
+}
--- a/src/usimp/mod.rs
+++ b/src/usimp/mod.rs
@@ -9,9 +9,25 @@ use base64_url;
 use crypto::digest::Digest;
 use crypto::sha2::Sha256;
 use serde::{Deserialize, Serialize};
-use serde_json::Value;
+use serde_json::{Map, Value};
 use uuid::Uuid;

+pub enum OutputAction {
+    Subscribe,
+    Unsubscribe,
+    Push,
+}
+
+impl From<OutputAction> for Value {
+    fn from(action: OutputAction) -> Self {
+        Value::from(match action {
+            OutputAction::Subscribe => "subscribe",
+            OutputAction::Unsubscribe => "unsubscribe",
+            OutputAction::Push => "push",
+        })
+    }
+}
+
 #[derive(Serialize, Deserialize)]
 pub struct InputEnvelope {
    pub endpoint: String,
@@ -25,9 +41,40 @@ pub struct InputEnvelope {
 pub struct OutputEnvelope {
    pub error: Option<Error>,
    pub request_nr: Option<u64>,
+    pub action: Option<OutputAction>,
    pub data: Value,
 }

+impl From<OutputEnvelope> for Value {
+    fn from(msg: OutputEnvelope) -> Self {
+        let mut envelope = Value::Object(Map::new());
+
+        envelope["request_nr"] = match msg.request_nr {
+            Some(nr) => Value::from(nr),
+            None => Value::Null,
+        };
+
+        match msg.error {
+            Some(error) => {
+                envelope["status"] = Value::from("error");
+                envelope["error"] = Value::from(error);
+            }
+            None => {
+                envelope["status"] = Value::from("success");
+            }
+        }
+
+        envelope["action"] = match msg.action {
+            Some(a) => Value::from(a),
+            None => Value::Null,
+        };
+
+        envelope["data"] = msg.data;
+
+        envelope
+    }
+}
+
 #[derive(Clone, Serialize, Deserialize)]
 pub struct Event {
    data: Value,
@@ -46,10 +93,9 @@ pub struct Session {
    account: Option<Account>,
 }

-pub fn get_account(session: &Option<Session>) -> Result<&Account, Error> {
+pub fn get_account_opt(session: &Option<Session>) -> Result<&Account, Error> {
    match session {
-        Some(session) => match &session.account {
-            Some(account) => Ok(&account),
+        Some(session) => get_account(session),
        None => {
            return Err(Error::new(
                ErrorKind::UsimpError,
@@ -57,7 +103,12 @@ pub fn get_account(session: &Option<Session>) -> Result<&Account, Error> {
                None,
            ))
        }
-        },
+    }
+}
+
+pub fn get_account(session: &Session) -> Result<&Account, Error> {
+    match &session.account {
+        Some(account) => Ok(&account),
        None => {
            return Err(Error::new(
                ErrorKind::UsimpError,
@@ -69,10 +120,11 @@ pub fn get_account(session: &Option<Session>) -> Result<&Account, Error> {
 }

 impl InputEnvelope {
-    pub fn respond(&self, data: Value) -> OutputEnvelope {
+    pub fn respond(&self, data: Value, action: Option<OutputAction>) -> OutputEnvelope {
        OutputEnvelope {
            error: None,
            request_nr: self.request_nr,
+            action,
            data,
        }
    }
--- a/src/usimp/subscription.rs
+++ b/src/usimp/subscription.rs
@@ -1,36 +1,80 @@
 use crate::database;
 use crate::usimp::*;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
+use std::ops::Deref;
 use std::sync::Arc;
 use tokio::sync::{mpsc, Mutex};
+use tokio::sync::mpsc::Sender;

-static mut ROOMS: Option<Arc<Mutex<HashMap<Uuid, Vec<mpsc::Sender<Event>>>>>> = None;
-static mut ACCOUNTS: Option<Arc<Mutex<HashMap<Uuid, Vec<mpsc::Sender<Event>>>>>> = None;
+#[derive(Clone, Eq, Hash, PartialEq)]
+struct Subscription {
+    session: Uuid,
+    req_nr: Option<u64>,
+}
+
+static mut SUBSCRIPTIONS: Option<Arc<Mutex<HashMap<Subscription, mpsc::Sender<Event>>>>> = None;
+static mut ACCOUNTS: Option<Arc<Mutex<HashMap<Uuid, HashSet<Subscription>>>>> = None;

 pub fn init() {
    unsafe {
-        ROOMS = Some(Arc::new(Mutex::new(HashMap::new())));
+        SUBSCRIPTIONS = Some(Arc::new(Mutex::new(HashMap::new())));
        ACCOUNTS = Some(Arc::new(Mutex::new(HashMap::new())));
    }
 }

-pub async fn subscribe_account(account: &Account) -> mpsc::Receiver<Event> {
+pub async fn subscribe_account(session: &Session, req_nr: Option<u64>) -> Result<mpsc::Receiver<Event>, Error> {
+    let account = get_account(session)?;
+    let sub = Subscription {session: session.id, req_nr};
+
    let (tx, rx) = mpsc::channel::<Event>(64);
    unsafe {
-        let mut acc = ACCOUNTS.as_ref().unwrap().lock().await;
-        match acc.get_mut(&account.id) {
-            Some(vec) => {
-                vec.push(tx);
+        let mut subs = SUBSCRIPTIONS.as_ref().unwrap().lock().await;
+        let mut accs = ACCOUNTS.as_ref().unwrap().lock().await;
+        match accs.get_mut(&account.id) {
+            Some(set) => {
+                set.insert(sub.clone());
            }
            None => {
-                acc.insert(account.id, vec![tx]);
+                let mut set = HashSet::new();
+                set.insert(sub.clone());
+                accs.insert(account.id, set);
            }
        }
+        match subs.get_mut(&sub) {
+            Some(_) => return Err(Error::new(
+                ErrorKind::SubscriptionError,
+                ErrorClass::ClientError,
+                None
+            )),
+            None => {
+                subs.insert(sub.clone(), tx);
+            }
        }
-    rx
    }

-pub async fn push(room_id: &Uuid, event: Event) -> Result<(), Error> {
+    Ok(rx)
+}
+
+pub async fn unsubscribe_account(session: &Session, req_nr: Option<u64>) -> Result<(), Error> {
+    let account = get_account(session)?;
+    let sub = Subscription {session: session.id, req_nr};
+
+    unsafe {
+        let mut subs = SUBSCRIPTIONS.as_ref().unwrap().lock().await;
+        let mut accs = ACCOUNTS.as_ref().unwrap().lock().await;
+        match accs.get_mut(&account.id) {
+            Some(set) => {
+                set.remove(&sub);
+            }
+            None => {},
+        }
+        subs.remove(&sub);
+    }
+
+    Ok(())
+}
+
+pub async fn push_room(room_id: &Uuid, event: Event) -> Result<(), Error> {
    let backend = database::client().await?;
    let accounts = match backend {
        database::Client::Postgres(client) => {
@@ -50,26 +94,26 @@ pub async fn push(room_id: &Uuid, event: Event) -> Result<(), Error> {
        }
    };

-    let mut rooms = unsafe {
-        let mut rooms = ROOMS.as_ref().unwrap().lock().await;
-        if let Some(rooms) = rooms.get_mut(room_id) {
-            rooms.clone()
-        } else {
-            Vec::new()
-        }
-    };
+    let mut room: Vec<mpsc::Sender<Event>> = Vec::new();

    for account in accounts {
        unsafe {
-            let mut accounts = ACCOUNTS.as_ref().unwrap().lock().await;
-            if let Some(acc) = accounts.get_mut(&account) {
-                let mut acc = acc.clone();
-                rooms.append(&mut acc);
+            let subs = SUBSCRIPTIONS.as_ref().unwrap().lock().await;
+            let accs = ACCOUNTS.as_ref().unwrap().lock().await;
+            if let Some(acc_subs) = accs.get(&account) {
+                for sub in acc_subs {
+                    match subs.get(sub) {
+                        None => {}
+                        Some(tx) => {
+                            room.push(tx.clone());
+                        }
+                    }
+                }
            }
        }
    }

-    for tx in rooms {
+    for tx in room {
        let _res = tx.send(event.clone()).await;
    }

--- a/src/websocket.rs
+++ b/src/websocket.rs
@@ -8,7 +8,7 @@ use hyper::{header, Body, Request, StatusCode};
 use hyper_tungstenite::hyper::upgrade::Upgraded;
 use hyper_tungstenite::tungstenite::{handshake, Message};
 use hyper_tungstenite::{tungstenite::protocol::Role, WebSocketStream};
-use serde_json::{Map, Value};
+use serde_json::Value;
 use tokio::sync::mpsc;

 pub enum WebSocketEnvelope {
@@ -50,21 +50,7 @@ async fn sender(
                break;
            }
            WebSocketEnvelope::Text(msg) => {
-                let mut envelope = Value::Object(Map::new());
-                envelope["data"] = msg.data;
-                envelope["request_nr"] = match msg.request_nr {
-                    Some(nr) => Value::from(nr),
-                    None => Value::Null,
-                };
-                match msg.error {
-                    Some(error) => {
-                        envelope["status"] = Value::from("error");
-                        envelope["error"] = Value::from(error);
-                    }
-                    None => {
-                        envelope["status"] = Value::from("success");
-                    }
-                }
+                let envelope: Value = msg.into();
                if let Err(error) = sink.send(Message::Text(envelope.to_string())).await {
                    eprintln!("{:?}", error);
                    return;
@@ -90,16 +76,11 @@ async fn receiver(
                    _res = tx.send(WebSocketEnvelope::Close).await;
                    break;
                } else if msg.is_binary() {
-                    _res = tx.send(WebSocketEnvelope::Text(OutputEnvelope {
-                        error: Some(Error {
-                            kind: ErrorKind::WebSocketError,
-                            class: ErrorClass::ClientProtocolError,
-                            msg: Some("Binary frames are not allowed".to_string()),
-                            desc: None,
-                        }),
-                        request_nr: None,
-                        data: Value::Null,
-                    })).await;
+                    _res = tx.send(WebSocketEnvelope::Text(OutputEnvelope::from(Error::new(
+                        ErrorKind::WebSocketError,
+                        ErrorClass::ClientProtocolError,
+                        Some("Binary frames are not allowed".to_string())
+                    )))).await;
                } else if msg.is_text() {
                    let input: InputEnvelope = serde_json::from_slice(&msg.into_data()[..]).unwrap();
                    let output = match usimp::endpoint(&input, Some(tx.clone())).await {
@@ -108,16 +89,11 @@ async fn receiver(
                    };
                    _res = tx.send(WebSocketEnvelope::Text(output)).await;
                } else {
-                    _res = tx.send(WebSocketEnvelope::Text(OutputEnvelope {
-                        error: Some(Error {
-                            kind: ErrorKind::WebSocketError,
-                            class: ErrorClass::ClientProtocolError,
-                            msg: Some("Unknown frame opcode".to_string()),
-                            desc: None,
-                        }),
-                        request_nr: None,
-                        data: Value::Null,
-                    })).await;
+                    _res = tx.send(WebSocketEnvelope::Text(OutputEnvelope::from(Error::new(
+                        ErrorKind::WebSocketError,
+                        ErrorClass::ClientProtocolError,
+                        Some("Unknown frame opcode".to_string())
+                    )))).await;
                }
            }
            Err(error) => println!("{:?}", error),
Author	SHA1	Message	Date
Lorenz Stechauner	d9672bdac0	Add unsubscribe endpoint	2022-08-27 22:13:19 +02:00
Lorenz Stechauner	97ae71d553	Add action to output envelope	2022-08-27 18:19:29 +02:00