From 7f978618e47981f2033cd7ca612c2811ff707e94 Mon Sep 17 00:00:00 2001
From: Mike Dilger <mike@mikedilger.com>
Date: Tue, 20 Feb 2024 09:28:10 +1300
Subject: [PATCH] Timeout idle connections with no subscriptions after 5
 seconds

---
 src/main.rs | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/src/main.rs b/src/main.rs
index 43a9928..d622e72 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -36,6 +36,7 @@ use std::time::Duration;
 use textnonce::TextNonce;
 use tokio::net::{TcpListener, TcpStream};
 use tokio::signal::unix::{signal, SignalKind};
+use tokio::time::Instant;
 use tungstenite::protocol::WebSocketConfig;
 use tungstenite::Message;
 
@@ -340,9 +341,25 @@ impl WebSocketService {
         let reply = NostrReply::Auth(self.challenge.clone());
         self.websocket.send(Message::text(reply.as_json())).await?;
 
+        let mut last_message_at = Instant::now();
+
         loop {
+            let interval = tokio::time::interval(Duration::from_secs(5));
+            tokio::pin!(interval);
+
             tokio::select! {
+                instant = interval.tick() => {
+                    // Drop them if they have no subscriptions
+                    if self.subscriptions.is_empty() {
+                        // And they are idle for 5 seconds with no subscriptions
+                        if last_message_at + Duration::from_secs(5) < instant {
+                            self.websocket.send(Message::Close(None)).await?;
+                            break;
+                        }
+                    }
+                }
                 message_option = self.websocket.next() => {
+                    last_message_at = Instant::now();
                     match message_option {
                         Some(message) => {
                             let message = message?;