wip

crates/relay_auth/src/lib.rs

@@ -13,7 +13,7 @@ use gpui::{
 use nostr_sdk::prelude::*;
 use settings::{AppSettings, AuthMode};
 use smallvec::{smallvec, SmallVec};
-use state::{tracker, NostrRegistry};
+use state::NostrRegistry;
 use theme::ActiveTheme;
 use ui::button::{Button, ButtonVariants};
 use ui::notification::Notification;
@@ -28,7 +28,7 @@ pub fn init(window: &mut Window, cx: &mut App) {
 
 /// Authentication request
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
-pub struct AuthRequest {
+struct AuthRequest {
     url: RelayUrl,
     challenge: String,
 }
@@ -56,6 +56,12 @@ impl AuthRequest {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
+enum Signal {
+    Auth(Arc<AuthRequest>),
+    Pending((EventId, RelayUrl)),
+}
+
 struct GlobalRelayAuth(Entity<RelayAuth>);
 
 impl Global for GlobalRelayAuth {}
@@ -63,6 +69,9 @@ impl Global for GlobalRelayAuth {}
 // Relay authentication
 #[derive(Debug)]
 pub struct RelayAuth {
+    /// Pending events waiting for resend after authentication
+    pending_events: HashSet<(EventId, RelayUrl)>,
+
     /// Tasks for asynchronous operations
     tasks: SmallVec<[Task<()>; 2]>,
 
@@ -84,89 +93,113 @@ impl RelayAuth {
     /// Create a new relay auth instance
     fn new(window: &mut Window, cx: &mut Context<Self>) -> Self {
         let nostr = NostrRegistry::global(cx);
-        // Channel for communication between nostr and gpui
-        let (tx, rx) = flume::bounded::<Arc<AuthRequest>>(100);
-
         let mut subscriptions = smallvec![];
-        let mut tasks = smallvec![];
 
         subscriptions.push(
-            // Observe the current state
-            cx.observe(&nostr, move |this, state, cx| {
+            // Observe the nostr state
+            cx.observe_in(&nostr, window, move |this, state, window, cx| {
                 if state.read(cx).connected() {
-                    this.handle_notifications(tx.clone(), cx)
-                }
-            }),
-        );
-
-        tasks.push(
-            // Update GPUI states
-            cx.spawn_in(window, async move |this, cx| {
-                while let Ok(req) = rx.recv_async().await {
-                    this.update_in(cx, |this, window, cx| {
-                        this.handle_auth(&req, window, cx);
-                    })
-                    .ok();
+                    this.handle_notifications(window, cx)
                 }
             }),
         );
 
         Self {
-            tasks,
+            pending_events: HashSet::default(),
+            tasks: smallvec![],
             _subscriptions: subscriptions,
         }
     }
 
-    // Handle nostr notifications
-    fn handle_notifications(
-        &mut self,
-        tx: flume::Sender<Arc<AuthRequest>>,
-        cx: &mut Context<Self>,
-    ) {
+    /// Handle nostr notifications
+    fn handle_notifications(&mut self, window: &mut Window, cx: &mut Context<Self>) {
         let nostr = NostrRegistry::global(cx);
         let client = nostr.read(cx).client();
 
-        let task = cx.background_spawn(async move {
+        // Channel for communication between nostr and gpui
+        let (tx, rx) = flume::bounded::<Signal>(256);
+
+        cx.background_spawn(async move {
             let mut notifications = client.notifications();
             let mut challenges: HashSet<Cow<'_, str>> = HashSet::default();
 
             while let Some(notification) = notifications.next().await {
-                match notification {
-                    ClientNotification::Message { relay_url, message } => {
-                        match message {
-                            RelayMessage::Auth { challenge } => {
-                                if challenges.insert(challenge.clone()) {
-                                    let request = AuthRequest::new(challenge, relay_url);
-                                    tx.send_async(Arc::new(request)).await.ok();
-                                }
-                            }
-                            RelayMessage::Ok {
-                                event_id, message, ..
-                            } => {
-                                let msg = MachineReadablePrefix::parse(&message);
-                                let mut tracker = tracker().write().await;
+                if let ClientNotification::Message { relay_url, message } = notification {
+                    match message {
+                        RelayMessage::Auth { challenge } => {
+                            if challenges.insert(challenge.clone()) {
+                                let request = AuthRequest::new(challenge, relay_url);
+                                let signal = Signal::Auth(Arc::new(request));
 
-                                // Handle authentication messages
-                                if let Some(MachineReadablePrefix::AuthRequired) = msg {
-                                    // Keep track of events that need to be resent after authentication
-                                    tracker.add_to_pending(event_id, relay_url);
-                                } else {
-                                    // Keep track of events sent by Coop
-                                    tracker.sent(event_id)
-                                }
+                                tx.send_async(signal).await.ok();
                             }
-                            _ => {}
                         }
+                        RelayMessage::Ok {
+                            event_id, message, ..
+                        } => {
+                            let msg = MachineReadablePrefix::parse(&message);
+
+                            // Handle authentication messages
+                            if let Some(MachineReadablePrefix::AuthRequired) = msg {
+                                let signal = Signal::Pending((event_id, relay_url));
+
+                                tx.send_async(signal).await.ok();
+                            }
+                        }
+                        _ => {}
                     }
-                    ClientNotification::Shutdown => break,
-                    _ => {}
                 }
             }
-        });
+        })
+        .detach();
 
-        self.tasks.push(task);
+        self.tasks.push(cx.spawn_in(window, async move |this, cx| {
+            while let Ok(signal) = rx.recv_async().await {
+                match signal {
+                    Signal::Auth(req) => {
+                        this.update_in(cx, |this, window, cx| {
+                            this.handle_auth(&req, window, cx);
+                        })
+                        .ok();
+                    }
+                    Signal::Pending((event_id, relay_url)) => {
+                        this.update_in(cx, |this, _window, cx| {
+                            this.insert_pending_event(event_id, relay_url, cx);
+                        })
+                        .ok();
+                    }
+                }
+            }
+        }));
     }
 
+    /// Insert a pending event waiting for resend after authentication
+    fn insert_pending_event(&mut self, id: EventId, relay: RelayUrl, cx: &mut Context<Self>) {
+        self.pending_events.insert((id, relay));
+        cx.notify();
+    }
+
+    /// Get all pending events for a specific relay,
+    fn get_pending_events(&self, relay: &RelayUrl, _cx: &App) -> Vec<EventId> {
+        let pending_events: Vec<EventId> = self
+            .pending_events
+            .iter()
+            .filter(|(_, pending_relay)| pending_relay == relay)
+            .map(|(id, _relay)| id)
+            .cloned()
+            .collect();
+
+        pending_events
+    }
+
+    /// Clear all pending events for a specific relay,
+    fn clear_pending_events(&mut self, relay: &RelayUrl, cx: &mut Context<Self>) {
+        self.pending_events
+            .retain(|(_, pending_relay)| pending_relay != relay);
+        cx.notify();
+    }
+
+    /// Handle authentication request
     fn handle_auth(&mut self, req: &Arc<AuthRequest>, window: &mut Window, cx: &mut Context<Self>) {
         let settings = AppSettings::global(cx);
         let trusted_relay = settings.read(cx).trusted_relay(req.url(), cx);
@@ -181,29 +214,25 @@ impl RelayAuth {
         }
     }
 
-    /// Respond to an authentication request.
-    fn response(&self, req: &Arc<AuthRequest>, window: &Window, cx: &Context<Self>) {
-        let settings = AppSettings::global(cx);
+    /// Send auth response and wait for confirmation
+    fn auth(&self, req: &Arc<AuthRequest>, cx: &App) -> Task<Result<(), Error>> {
         let nostr = NostrRegistry::global(cx);
         let client = nostr.read(cx).client();
-
         let req = req.clone();
-        let challenge = req.challenge().to_string();
-        let async_req = req.clone();
 
-        let task: Task<Result<(), Error>> = cx.background_spawn(async move {
+        // Get all pending events for the relay
+        let pending_events = self.get_pending_events(req.url(), cx);
+
+        cx.background_spawn(async move {
             // Construct event
-            let builder = EventBuilder::auth(async_req.challenge(), async_req.url().clone());
+            let builder = EventBuilder::auth(req.challenge(), req.url().clone());
             let event = client.sign_event_builder(builder).await?;
 
             // Get the event ID
             let id = event.id;
 
             // Get the relay
-            let relay = client
-                .relay(async_req.url())
-                .await?
-                .context("Relay not found")?;
+            let relay = client.relay(req.url()).await?.context("Relay not found")?;
 
             // Subscribe to notifications
             let mut notifications = relay.notifications();
@@ -219,17 +248,18 @@ impl RelayAuth {
                         message: RelayMessage::Ok { event_id, .. },
                     } => {
                         if id == event_id {
-                            // Re-subscribe to previous subscription
-                            // relay.resubscribe().await?;
+                            // Get all subscriptions
+                            let subscriptions = relay.subscriptions().await;
 
-                            // Get all pending events that need to be resent
-                            let mut tracker = tracker().write().await;
-                            let ids: Vec<EventId> = tracker.pending_resend(relay.url());
+                            // Re-subscribe to previous subscriptions
+                            for (id, filters) in subscriptions.into_iter() {
+                                relay.subscribe(filters).with_id(id).await?;
+                            }
 
-                            for id in ids.into_iter() {
+                            // Re-send pending events
+                            for id in pending_events {
                                 if let Some(event) = client.database().event_by_id(&id).await? {
-                                    let event_id = relay.send_event(&event).await?;
-                                    tracker.sent(event_id);
+                                    relay.send_event(&event).await?;
                                 }
                             }
 
@@ -242,22 +272,33 @@ impl RelayAuth {
             }
 
             Err(anyhow!("Authentication failed"))
-        });
+        })
+    }
+
+    /// Respond to an authentication request.
+    fn response(&self, req: &Arc<AuthRequest>, window: &Window, cx: &Context<Self>) {
+        let settings = AppSettings::global(cx);
+        let req = req.clone();
+        let challenge = req.challenge().to_string();
+
+        // Create a task for authentication
+        let task = self.auth(&req, cx);
 
         cx.spawn_in(window, async move |this, cx| {
             let result = task.await;
             let url = req.url();
 
-            this.update_in(cx, |_this, window, cx| {
+            this.update_in(cx, |this, window, cx| {
                 window.clear_notification(challenge, cx);
 
                 match result {
                     Ok(_) => {
+                        // Clear pending events for the authenticated relay
+                        this.clear_pending_events(url, cx);
                         // Save the authenticated relay to automatically authenticate future requests
                         settings.update(cx, |this, cx| {
                             this.add_trusted_relay(url, cx);
                         });
-
                         window.push_notification(format!("{} has been authenticated", url), cx);
                     }
                     Err(e) => {