From e9fd738213707807ada2d2a0f75e437b97d5eec8 Mon Sep 17 00:00:00 2001
From: Mike Dilger <mike@mikedilger.com>
Date: Tue, 21 Jan 2025 09:41:35 +1300
Subject: [PATCH] Improve negentropy error handling by using NEG-ERR when we
 can

---
 src/lib.rs   | 11 +++++++++--
 src/nostr.rs | 16 ++++++++++++----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 11d92a6..81fc8aa 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -227,6 +227,7 @@ async fn websocket_thread(peer: HashedPeer, websocket: HyperWebsocket, origin: S
                 user: None,
                 error_punishment: 0.0,
                 replied: false,
+                negentropy_sub: None,
             };
 
             // Increment connection count
@@ -362,6 +363,7 @@ struct WebSocketService {
     pub user: Option<Pubkey>,
     pub error_punishment: f32,
     pub replied: bool,
+    pub negentropy_sub: Option<String>,
 }
 
 impl WebSocketService {
@@ -544,8 +546,13 @@ impl WebSocketService {
                         }
                     }
                     if !self.replied {
-                        let reply = NostrReply::Notice(format!("error: {}", e.inner));
-                        self.send(Message::text(reply.as_json())).await?;
+                        if let Some(subid) = &self.negentropy_sub {
+                            let reply = NostrReply::NegErr(subid, format!("error: {e}"));
+                            self.send(Message::text(reply.as_json())).await?;
+                        } else {
+                            let reply = NostrReply::Notice(format!("error: {}", e.inner));
+                            self.send(Message::text(reply.as_json())).await?;
+                        }
                     }
                     if self.error_punishment >= 1.0 {
                         let reply = NostrReply::Notice("Closing due to error(s)".into());
diff --git a/src/nostr.rs b/src/nostr.rs
index f171476..b420660 100644
--- a/src/nostr.rs
+++ b/src/nostr.rs
@@ -468,8 +468,11 @@ impl WebSocketService {
             subid
         };
 
-        if ! GLOBALS.config.read().enable_negentropy {
-            let reply = NostrReply::NegErr(&subid, "blocked: Negentropy sync is disabled".to_owned());
+        self.negentropy_sub = Some(subid.clone());
+
+        if !GLOBALS.config.read().enable_negentropy {
+            let reply =
+                NostrReply::NegErr(&subid, "blocked: Negentropy sync is disabled".to_owned());
             self.send(Message::text(reply.as_json())).await?;
             return Ok(());
         }
@@ -600,8 +603,11 @@ impl WebSocketService {
             subid
         };
 
-        if ! GLOBALS.config.read().enable_negentropy {
-            let reply = NostrReply::NegErr(&subid, "blocked: Negentropy sync is disabled".to_owned());
+        self.negentropy_sub = Some(subid.clone());
+
+        if !GLOBALS.config.read().enable_negentropy {
+            let reply =
+                NostrReply::NegErr(&subid, "blocked: Negentropy sync is disabled".to_owned());
             self.send(Message::text(reply.as_json())).await?;
             return Ok(());
         }
@@ -677,6 +683,8 @@ impl WebSocketService {
             subid
         };
 
+        self.negentropy_sub = Some(subid.clone());
+
         // Close the subscription
         self.neg_subscriptions.remove(&subid);