Tie unauthenticated WebSocket keep-alive to active client connections

Keep the unidentified socket alive while a JSON-RPC connection is open (including stdio mode) and while a D-Bus object is exported, instead of for the lifetime of the receive loop. The receive loop does not use the unauthenticated socket, so keeping it alive there was semantically wrong.

This also covers --receive-mode=manual, where no receive loop runs butclients still send messages.

lib/src/main/java/org/asamk/signal/manager/Manager.java

@@ -439,6 +439,10 @@ public interface Manager extends Closeable {
 
     void addClosedListener(Runnable listener);
 
+    void addUnidentifiedKeepAlive(String token);
+
+    void removeUnidentifiedKeepAlive(String token);
+
     InputStream retrieveAttachment(final String id) throws IOException;
 
     InputStream retrieveContactAvatar(final RecipientIdentifier.Single recipient) throws IOException, UnregisteredRecipientException;

lib/src/main/java/org/asamk/signal/manager/helper/ReceiveHelper.java

@@ -102,9 +102,6 @@ public class ReceiveHelper {
         signalWebSocket.connect();
         signalWebSocket.registerKeepAliveToken("receive");
 
-        final var unauthenticatedSignalWebSocket = dependencies.getUnauthenticatedSignalWebSocket();
-        unauthenticatedSignalWebSocket.registerKeepAliveToken("receive");
-
         try {
             receiveMessagesInternal(signalWebSocket, timeout, maxMessages, handler, queuedActions);
         } finally {
@@ -113,7 +110,6 @@ public class ReceiveHelper {
             queuedActions.clear();
             signalWebSocket.removeKeepAliveToken("receive");
             signalWebSocket.disconnect();
-            unauthenticatedSignalWebSocket.removeKeepAliveToken("receive");
             webSocketStateDisposable.dispose();
             shouldStop = false;
         }

lib/src/main/java/org/asamk/signal/manager/internal/ManagerImpl.java

@@ -1826,6 +1826,16 @@ public class ManagerImpl implements Manager {
         }
     }
 
+    @Override
+    public void addUnidentifiedKeepAlive(final String token) {
+        dependencies.getUnauthenticatedSignalWebSocket().registerKeepAliveToken(token);
+    }
+
+    @Override
+    public void removeUnidentifiedKeepAlive(final String token) {
+        dependencies.getUnauthenticatedSignalWebSocket().removeKeepAliveToken(token);
+    }
+
     @Override
     public void addCallEventListener(final CallEventListener listener) {
         context.getCallManager().addCallEventListener(listener);

src/main/java/org/asamk/signal/dbus/DbusManagerImpl.java

@@ -951,6 +951,14 @@ public class DbusManagerImpl implements Manager {
         }
     }
 
+    @Override
+    public void addUnidentifiedKeepAlive(final String token) {
+    }
+
+    @Override
+    public void removeUnidentifiedKeepAlive(final String token) {
+    }
+
     @Override
     public void addCallEventListener(final CallEventListener listener) {
         // Not supported over DBus

src/main/java/org/asamk/signal/dbus/DbusSignalImpl.java

@@ -100,6 +100,7 @@ public class DbusSignalImpl implements Signal, AutoCloseable {
 
     public void initObjects() {
         exportObjects();
+        m.addUnidentifiedKeepAlive("dbus");
         if (!noReceiveOnStart) {
             subscribeReceive();
         }
@@ -116,6 +117,7 @@ public class DbusSignalImpl implements Signal, AutoCloseable {
 
     @Override
     public void close() {
+        m.removeUnidentifiedKeepAlive("dbus");
         if (dbusMessageHandler != null) {
             m.removeReceiveHandler(dbusMessageHandler);
             dbusMessageHandler = null;

src/main/java/org/asamk/signal/jsonrpc/SignalJsonRpcDispatcherHandler.java

@@ -30,6 +30,7 @@ import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.UUID;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
@@ -44,8 +45,11 @@ public class SignalJsonRpcDispatcherHandler {
     private final JsonRpcReader jsonRpcReader;
     private final boolean noReceiveOnStart;
 
-    private final Map<Integer, ArrayList<Pair<Manager, Manager.ReceiveMessageHandler>>> receiveHandlers = new HashMap<>();
-    private final Map<Integer, ArrayList<Pair<Manager, Manager.CallEventListener>>> callEventHandlers = new HashMap<>();
+    private final Map<Integer, List<Pair<Manager, Manager.ReceiveMessageHandler>>> receiveHandlers = new HashMap<>();
+    private final Map<Integer, List<Pair<Manager, Manager.CallEventListener>>> callEventHandlers = new HashMap<>();
+    private final String connectionKeepAliveToken = "jsonrpc-" + UUID.randomUUID();
+    private final List<Manager> keepAliveManagers = new ArrayList<>();
+    private boolean connectionActive = true;
     private SignalJsonRpcCommandHandler commandHandler;
 
     public SignalJsonRpcDispatcherHandler(
@@ -68,6 +72,8 @@ public class SignalJsonRpcDispatcherHandler {
                 createReceiveHandler(m, subscriptionId, false)));
         final Consumer<Manager> onManagerAddedCallEventSubscriptionsHandler = m -> callEventHandlers.forEach((subscriptionId, handlers) -> handlers.add(
                 createCallEventHandler(m, subscriptionId)));
+        final Consumer<Manager> onManagerAddedKeepAliveHandler = this::registerKeepAlive;
+        final Consumer<Manager> onManagerRemovedKeepAliveHandler = this::unregisterKeepAlive;
 
         if (!noReceiveOnStart) {
             this.subscribeReceive(c.getManagers(), true);
@@ -76,6 +82,9 @@ public class SignalJsonRpcDispatcherHandler {
         }
         c.addOnManagerAddedHandler(onManagerAddedReceiveSubscriptionsHandler);
         c.addOnManagerAddedHandler(onManagerAddedCallEventSubscriptionsHandler);
+        c.getManagers().forEach(this::registerKeepAlive);
+        c.addOnManagerAddedHandler(onManagerAddedKeepAliveHandler);
+        c.addOnManagerRemovedHandler(onManagerRemovedKeepAliveHandler);
 
         try {
             handleConnection();
@@ -86,6 +95,8 @@ public class SignalJsonRpcDispatcherHandler {
             }
             c.removeOnManagerAddedHandler(onManagerAddedReceiveSubscriptionsHandler);
             c.removeOnManagerAddedHandler(onManagerAddedCallEventSubscriptionsHandler);
+            c.removeOnManagerAddedHandler(onManagerAddedKeepAliveHandler);
+            c.removeOnManagerRemovedHandler(onManagerRemovedKeepAliveHandler);
         }
     }
 
@@ -99,6 +110,8 @@ public class SignalJsonRpcDispatcherHandler {
         final var currentThread = Thread.currentThread();
         m.addClosedListener(currentThread::interrupt);
 
+        registerKeepAlive(m);
+
         handleConnection();
     }
 
@@ -219,14 +232,29 @@ public class SignalJsonRpcDispatcherHandler {
         subscriptionId.ifPresent(this::unsubscribeReceive);
     }
 
+    private void registerKeepAlive(final Manager m) {
+        if (!connectionActive) return;
+        m.addUnidentifiedKeepAlive(connectionKeepAliveToken);
+        keepAliveManagers.add(m);
+    }
+
+    private void unregisterKeepAlive(final Manager m) {
+        if (!connectionActive) return;
+        m.removeUnidentifiedKeepAlive(connectionKeepAliveToken);
+        keepAliveManagers.remove(m);
+    }
+
     private void handleConnection() {
         try {
             jsonRpcReader.readMessages((method, params) -> commandHandler.handleRequest(objectMapper, method, params),
                     response -> logger.debug("Received unexpected response for id {}", response.getId()));
         } finally {
+            connectionActive = false;
             receiveHandlers.forEach((_subscriptionId, handlers) -> handlers.forEach(this::unsubscribeReceiveHandler));
             receiveHandlers.clear();
             unsubscribeAllCallEvents();
+            keepAliveManagers.forEach(m -> m.removeUnidentifiedKeepAlive(connectionKeepAliveToken));
+            keepAliveManagers.clear();
         }
     }
 

src/test/java/org/asamk/signal/dbus/NumberlessDbusTest.java

@@ -379,6 +379,7 @@ class NumberlessDbusTest {
                             case "isContactBlocked" -> false;
                             case "getContactOrProfileName" -> "Test contact";
                             case "addAddressChangedListener" -> null;
+                            case "addUnidentifiedKeepAlive", "removeUnidentifiedKeepAlive" -> null;
                             case "addClosedListener" -> {
                                 closedListeners.add((Runnable) args[0]);
                                 yield null;

src/test/java/org/asamk/signal/jsonrpc/SubscribeCallEventsTest.java

@@ -69,6 +69,7 @@ class SubscribeCallEventsTest {
         final List<Consumer<Manager>> addedHandlers = new ArrayList<>();
         final List<Consumer<Manager>> removedHandlers = new ArrayList<>();
         int addedHandlerRegistrationCount;
+        int removedHandlerRegistrationCount;
 
         StubMultiAccountManager(List<Manager> managers) {
             this.managers = new ArrayList<>(managers);
@@ -92,6 +93,7 @@ class SubscribeCallEventsTest {
 
         @Override
         public void addOnManagerRemovedHandler(Consumer<Manager> handler) {
+            removedHandlerRegistrationCount++;
             removedHandlers.add(handler);
         }
 
@@ -237,9 +239,11 @@ class SubscribeCallEventsTest {
 
         assertEquals(1, manager1.state().addCallEventListenerCount.get(), "manager1 should have one listener");
         assertEquals(1, manager2.state().addCallEventListenerCount.get(), "manager2 should have one listener");
-        assertEquals(2, multi.addedHandlerRegistrationCount, "should register onManagerAdded handlers");
+        // Registers onManagerAdded handlers for receive subscriptions, call events, and keep-alive
+        assertEquals(3, multi.addedHandlerRegistrationCount, "should register onManagerAdded handlers");
         assertEquals(0, multi.addedHandlers.size(), "should remove onManagerAdded handlers after connection closes");
-        assertEquals(0, multi.removedHandlers.size(), "should not register onManagerRemoved handlers in manual mode");
+        assertEquals(1, multi.removedHandlerRegistrationCount, "should register keep-alive onManagerRemoved handler");
+        assertEquals(0, multi.removedHandlers.size(), "should remove onManagerRemoved handler after close");
     }
 
     @Test
@@ -287,8 +291,9 @@ class SubscribeCallEventsTest {
 
         handler.handleConnection(multi);
 
-        assertEquals(3, multi.addedHandlerRegistrationCount);
+        assertEquals(4, multi.addedHandlerRegistrationCount);
         assertEquals(0, multi.addedHandlers.size(), "should remove all onManagerAdded handlers after close");
-        assertEquals(0, multi.removedHandlers.size(), "should remove onManagerRemoved handler after close");
+        assertEquals(2, multi.removedHandlerRegistrationCount);
+        assertEquals(0, multi.removedHandlers.size(), "should remove onManagerRemoved handlers after close");
     }
 }