Просмотр исходного кода

darkirc: Implement seen message tracking in USER scope

parazyd 2 лет назад
Родитель
Сommit
bc60781225
3 измененных файлов с 30 добавлено и 10 удалено
  1. 13 3
      bin/darkirc/src/irc/client.rs
  2. 6 4
      bin/darkirc/src/irc/command.rs
  3. 11 3
      bin/darkirc/src/main.rs

+ 13 - 3
bin/darkirc/src/irc/client.rs

@@ -20,7 +20,7 @@ use std::{
     collections::{HashMap, HashSet},
     sync::{
         atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
-        Arc,
+        Arc, OnceLock,
     },
 };
 
@@ -55,7 +55,7 @@ pub enum ReplyType {
     Cap(String),
 }
 
-/// Stateful IRC client, used for each client connection
+/// Stateful IRC client handler, used for each client connection
 pub struct Client {
     /// Pointer to parent `IrcServer`
     pub server: Arc<IrcServer>,
@@ -81,6 +81,9 @@ pub struct Client {
     pub realname: RwLock<String>,
     /// Client caps
     pub caps: RwLock<HashMap<String, bool>>,
+    /// Set of seen messages for the user
+    /// TODO: It grows indefinitely, needs to be pruned.
+    pub seen: OnceLock<sled::Tree>,
 }
 
 impl Client {
@@ -105,6 +108,7 @@ impl Client {
             nickname: RwLock::new(String::from("*")),
             realname: RwLock::new(String::from("*")),
             caps: RwLock::new(caps),
+            seen: OnceLock::new(),
         })
     }
 
@@ -155,6 +159,9 @@ impl Client {
                             if let Err(e) = self.server.darkirc.event_graph.dag_insert(event.clone()).await {
                                 error!("[IRC CLIENT] Failed inserting new event to DAG: {}", e);
                             } else {
+                                // We sent this, so it should be considered seen.
+                                self.seen.get().unwrap().insert(event_id.as_bytes(), &[]).unwrap();
+
                                 // Otherwise, broadcast it
                                 self.server.darkirc.p2p.broadcast(&EventPut(event)).await;
                             }
@@ -214,6 +221,9 @@ impl Client {
                             error!("[IRC CLIENT] Failed writing PRIVMSG to client: {}", e);
                             continue
                         }
+
+                        // Mark the message as seen for this USER
+                        self.seen.get().unwrap().insert(event_id.as_bytes(), &[]).unwrap();
                     }
                 }
             }
@@ -267,7 +277,7 @@ impl Client {
         // Commands can begin with :garbage, but we will reject clients
         // doing that for now to keep the protocol simple and focused.
         let cmd = tokens.next().ok_or(Error::ParseFailed("Invalid command line"))?;
-        let args = line.replacen(cmd, "", 1).trim().to_string();
+        let args = line.replacen(cmd, "", 1);
         let cmd = cmd.to_uppercase();
 
         debug!("[{}] --> {}{}", self.addr, cmd, args);

+ 6 - 4
bin/darkirc/src/irc/command.rs

@@ -50,10 +50,7 @@
 //! Some of the above commands could actually be implemented and could
 //! work in respect to the P2P network.
 
-use std::{
-    collections::HashSet,
-    sync::{atomic::Ordering::SeqCst, Arc},
-};
+use std::{collections::HashSet, sync::atomic::Ordering::SeqCst};
 
 use darkfi::Result;
 use log::{error, info};
@@ -789,6 +786,11 @@ impl Client {
         *self.username.write().await = username.to_string();
         *self.realname.write().await = realname.to_string();
 
+        // The username is now set, we can open the sled tree for seen messages
+        self.seen
+            .set(self.server.darkirc.sled.open_tree(format!("darkirc_user_{}", username)).unwrap())
+            .unwrap();
+
         // If the nickname is set, we can complete the registration
         if nick != "*" {
             self.registered.store(true, SeqCst);

+ 11 - 3
bin/darkirc/src/main.rs

@@ -105,6 +105,8 @@ struct Args {
 pub struct DarkIrc {
     /// P2P network pointer
     p2p: P2pPtr,
+    /// Sled DB (also used in event_graph)
+    sled: sled::Db,
     /// Event Graph instance
     event_graph: EventGraphPtr,
     /// JSON-RPC connection tracker
@@ -114,8 +116,13 @@ pub struct DarkIrc {
 }
 
 impl DarkIrc {
-    fn new(p2p: P2pPtr, event_graph: EventGraphPtr, dnet_sub: JsonSubscriber) -> Self {
-        Self { p2p, event_graph, rpc_connections: Mutex::new(HashSet::new()), dnet_sub }
+    fn new(
+        p2p: P2pPtr,
+        sled: sled::Db,
+        event_graph: EventGraphPtr,
+        dnet_sub: JsonSubscriber,
+    ) -> Self {
+        Self { p2p, sled, event_graph, rpc_connections: Mutex::new(HashSet::new()), dnet_sub }
     }
 }
 
@@ -208,7 +215,8 @@ async fn realmain(args: Args, ex: Arc<Executor<'static>>) -> Result<()> {
     );
 
     info!("Starting JSON-RPC server");
-    let darkirc = Arc::new(DarkIrc::new(p2p.clone(), event_graph.clone(), dnet_sub));
+    let darkirc =
+        Arc::new(DarkIrc::new(p2p.clone(), sled_db.clone(), event_graph.clone(), dnet_sub));
     let darkirc_ = Arc::clone(&darkirc);
     let rpc_task = StoppableTask::new();
     rpc_task.clone().start(