Adapt receive subscription notification to have payload in result field

[signal-cli] / src / main / java / org / asamk / signal / jsonrpc / SignalJsonRpcDispatcherHandler.java

diff --git a/src/main/java/org/asamk/signal/jsonrpc/SignalJsonRpcDispatcherHandler.java b/src/main/java/org/asamk/signal/jsonrpc/SignalJsonRpcDispatcherHandler.java
index 9085609ca733c40246d9abe7c35049d6964989af..aea332fd47e384c8f334a48dc2cd8244503608b5 100644 (file)
--- a/src/main/java/org/asamk/signal/jsonrpc/SignalJsonRpcDispatcherHandler.java
+++ b/src/main/java/org/asamk/signal/jsonrpc/SignalJsonRpcDispatcherHandler.java
@@ -1,138 +1,228 @@
 package org.asamk.signal.jsonrpc;
 
 package org.asamk.signal.jsonrpc;
 

-import com.fasterxml.jackson.core.TreeNode;
-import com.fasterxml.jackson.databind.JsonMappingException;
+import com.fasterxml.jackson.core.type.TypeReference;

 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;

-import com.fasterxml.jackson.databind.node.ContainerNode;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.IntNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;

 
 

-import org.asamk.signal.JsonReceiveMessageHandler;
-import org.asamk.signal.JsonWriter;
-import org.asamk.signal.OutputWriter;
+import org.asamk.signal.commands.Command;

 import org.asamk.signal.commands.Commands;
 import org.asamk.signal.commands.Commands;

-import org.asamk.signal.commands.JsonRpcCommand;
+import org.asamk.signal.commands.JsonRpcMultiCommand;
+import org.asamk.signal.commands.JsonRpcSingleCommand;

 import org.asamk.signal.commands.exceptions.CommandException;
 import org.asamk.signal.commands.exceptions.CommandException;

-import org.asamk.signal.commands.exceptions.IOErrorException;
-import org.asamk.signal.commands.exceptions.UntrustedKeyErrorException;

 import org.asamk.signal.commands.exceptions.UserErrorException;
 import org.asamk.signal.commands.exceptions.UserErrorException;

+import org.asamk.signal.json.JsonReceiveMessageHandler;

 import org.asamk.signal.manager.Manager;
 import org.asamk.signal.manager.Manager;

+import org.asamk.signal.manager.MultiAccountManager;
+import org.asamk.signal.manager.api.Pair;
+import org.asamk.signal.output.JsonWriter;

 import org.asamk.signal.util.Util;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import org.asamk.signal.util.Util;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 

-import java.io.IOException;
+import java.nio.channels.ClosedChannelException;
+import java.util.HashMap;
+import java.util.List;

 import java.util.Map;
 import java.util.Map;

+import java.util.concurrent.atomic.AtomicInteger;

 import java.util.function.Supplier;
 
 public class SignalJsonRpcDispatcherHandler {
 
     private final static Logger logger = LoggerFactory.getLogger(SignalJsonRpcDispatcherHandler.class);
 
 import java.util.function.Supplier;
 
 public class SignalJsonRpcDispatcherHandler {
 
     private final static Logger logger = LoggerFactory.getLogger(SignalJsonRpcDispatcherHandler.class);
 

-    private static final int USER_ERROR = -1;
-    private static final int IO_ERROR = -3;
-    private static final int UNTRUSTED_KEY_ERROR = -4;
+    private final ObjectMapper objectMapper;
+    private final JsonRpcSender jsonRpcSender;
+    private final JsonRpcReader jsonRpcReader;
+    private final boolean noReceiveOnStart;

 
 

-    private final Manager m;
-    private final JsonWriter outputWriter;
-    private final Supplier<String> lineSupplier;
+    private final Map<Integer, List<Pair<Manager, Manager.ReceiveMessageHandler>>> receiveHandlers = new HashMap<>();
+    private SignalJsonRpcCommandHandler commandHandler;

 
     public SignalJsonRpcDispatcherHandler(
 
     public SignalJsonRpcDispatcherHandler(

-            final Manager m, final JsonWriter outputWriter, final Supplier<String> lineSupplier
+            final JsonWriter jsonWriter, final Supplier<String> lineSupplier, final boolean noReceiveOnStart

     ) {
     ) {

-        this.m = m;
-        this.outputWriter = outputWriter;
-        this.lineSupplier = lineSupplier;
+        this.noReceiveOnStart = noReceiveOnStart;
+        this.objectMapper = Util.createJsonObjectMapper();
+        this.jsonRpcSender = new JsonRpcSender(jsonWriter);
+        this.jsonRpcReader = new JsonRpcReader(jsonRpcSender, lineSupplier);

     }
 
     }
 

-    public void handleConnection() {
-        final var objectMapper = Util.createJsonObjectMapper();
-        final var jsonRpcSender = new JsonRpcSender(outputWriter);
+    public void handleConnection(final MultiAccountManager c) {
+        this.commandHandler = new SignalJsonRpcCommandHandler(c, this::getCommand);

 
 

-        final var receiveMessageHandler = new JsonReceiveMessageHandler(m,
-                s -> jsonRpcSender.sendRequest(JsonRpcRequest.forNotification("receive",
-                        objectMapper.valueToTree(s),
-                        null)));
-        try {
-            m.addReceiveHandler(receiveMessageHandler);
+        if (!noReceiveOnStart) {
+            this.subscribeReceive(c.getManagers());
+            c.addOnManagerAddedHandler(this::subscribeReceive);
+            c.addOnManagerRemovedHandler(this::unsubscribeReceive);
+        }
+
+        handleConnection();
+    }
+
+    public void handleConnection(final Manager m) {
+        this.commandHandler = new SignalJsonRpcCommandHandler(m, this::getCommand);
+
+        if (!noReceiveOnStart) {
+            subscribeReceive(m);
+        }
+
+        final var currentThread = Thread.currentThread();
+        m.addClosedListener(currentThread::interrupt);
+
+        handleConnection();
+    }
+
+    private static final AtomicInteger nextSubscriptionId = new AtomicInteger(0);

 
 

-            // Maybe this should be handled inside the Manager
-            while (!m.hasCaughtUpWithOldMessages()) {
+    private int subscribeReceive(final Manager manager) {
+        return subscribeReceive(List.of(manager));
+    }
+
+    private int subscribeReceive(final List<Manager> managers) {
+        final var subscriptionId = nextSubscriptionId.getAndIncrement();
+        final var handlers = managers.stream().map(m -> {
+            final var receiveMessageHandler = new JsonReceiveMessageHandler(m, s -> {
+                final var params = new ObjectNode(objectMapper.getNodeFactory());
+                params.set("subscription", IntNode.valueOf(subscriptionId));
+                params.set("result", objectMapper.valueToTree(s));
+                final var jsonRpcRequest = JsonRpcRequest.forNotification("receive", params, null);

                 try {
                 try {

-                    synchronized (m) {
-                        m.wait();
+                    jsonRpcSender.sendRequest(jsonRpcRequest);
+                } catch (AssertionError e) {
+                    if (e.getCause() instanceof ClosedChannelException) {
+                        unsubscribeReceive(subscriptionId);

                     }
                     }

-                } catch (InterruptedException ignored) {

                 }
                 }

-            }
+            });
+            m.addReceiveHandler(receiveMessageHandler);
+            return new Pair<>(m, (Manager.ReceiveMessageHandler) receiveMessageHandler);
+        }).toList();
+        receiveHandlers.put(subscriptionId, handlers);
+
+        return subscriptionId;
+    }
+
+    private boolean unsubscribeReceive(final int subscriptionId) {
+        final var handlers = receiveHandlers.remove(subscriptionId);
+        if (handlers == null) {
+            return false;
+        }
+        for (final var pair : handlers) {
+            unsubscribeReceiveHandler(pair);
+        }
+        return true;
+    }
+
+    private void unsubscribeReceive(final Manager m) {
+        final var subscriptionId = receiveHandlers.entrySet()
+                .stream()
+                .filter(e -> e.getValue().size() == 1 && e.getValue().get(0).first().equals(m))
+                .map(Map.Entry::getKey)
+                .findFirst();
+        subscriptionId.ifPresent(this::unsubscribeReceive);
+    }

 
 

-            final var jsonRpcReader = new JsonRpcReader(jsonRpcSender, lineSupplier);
-            jsonRpcReader.readRequests((method, params) -> handleRequest(m, objectMapper, method, params),
+    private void handleConnection() {
+        try {
+            jsonRpcReader.readMessages((method, params) -> commandHandler.handleRequest(objectMapper, method, params),

                     response -> logger.debug("Received unexpected response for id {}", response.getId()));
         } finally {
                     response -> logger.debug("Received unexpected response for id {}", response.getId()));
         } finally {

-            m.removeReceiveHandler(receiveMessageHandler);
+            receiveHandlers.forEach((_subscriptionId, handlers) -> handlers.forEach(this::unsubscribeReceiveHandler));
+            receiveHandlers.clear();

         }
     }
 
         }
     }
 

-    private JsonNode handleRequest(
-            final Manager m, final ObjectMapper objectMapper, final String method, ContainerNode<?> params
-    ) throws JsonRpcException {
-        final Object[] result = {null};
-        final JsonWriter commandOutputWriter = s -> {
-            if (result[0] != null) {
-                throw new AssertionError("Command may only write one json result");
-            }
+    private void unsubscribeReceiveHandler(final Pair<Manager, Manager.ReceiveMessageHandler> pair) {
+        final var m = pair.first();
+        final var handler = pair.second();
+        m.removeReceiveHandler(handler);
+    }
+
+    private Command getCommand(final String method) {
+        if ("subscribeReceive".equals(method)) {
+            return new SubscribeReceiveCommand();
+        }
+        if ("unsubscribeReceive".equals(method)) {
+            return new UnsubscribeReceiveCommand();
+        }
+        return Commands.getCommand(method);
+    }

 
 

-            result[0] = s;
-        };
+    private class SubscribeReceiveCommand implements JsonRpcSingleCommand<Void>, JsonRpcMultiCommand<Void> {

 
 

-        var command = Commands.getCommand(method);
-        if (!(command instanceof JsonRpcCommand)) {
-            throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.METHOD_NOT_FOUND,
-                    "Method not implemented",
-                    null));
+        @Override
+        public String getName() {
+            return "subscribeReceive";

         }
 
         }
 

-        try {
-            parseParamsAndRunCommand(m, objectMapper, params, commandOutputWriter, (JsonRpcCommand<?>) command);
-        } catch (JsonMappingException e) {
-            throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INVALID_REQUEST,
-                    e.getMessage(),
-                    null));
-        } catch (UserErrorException e) {
-            throw new JsonRpcException(new JsonRpcResponse.Error(USER_ERROR, e.getMessage(), null));
-        } catch (IOErrorException e) {
-            throw new JsonRpcException(new JsonRpcResponse.Error(IO_ERROR, e.getMessage(), null));
-        } catch (UntrustedKeyErrorException e) {
-            throw new JsonRpcException(new JsonRpcResponse.Error(UNTRUSTED_KEY_ERROR, e.getMessage(), null));
-        } catch (Throwable e) {
-            logger.error("Command execution failed", e);
-            throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INTERNAL_ERROR,
-                    e.getMessage(),
-                    null));
+        @Override
+        public void handleCommand(
+                final Void request, final Manager m, final JsonWriter jsonWriter
+        ) throws CommandException {
+            final var subscriptionId = subscribeReceive(m);
+            jsonWriter.write(subscriptionId);

         }
 
         }
 

-        Object output = result[0] == null ? Map.of() : result[0];
-        return objectMapper.valueToTree(output);
+        @Override
+        public void handleCommand(
+                final Void request, final MultiAccountManager c, final JsonWriter jsonWriter
+        ) throws CommandException {
+            final var subscriptionId = subscribeReceive(c.getManagers());
+            jsonWriter.write(subscriptionId);
+        }

     }
 
     }
 

-    private <T> void parseParamsAndRunCommand(
-            final Manager m,
-            final ObjectMapper objectMapper,
-            final TreeNode params,
-            final OutputWriter outputWriter,
-            final JsonRpcCommand<T> command
-    ) throws CommandException, JsonMappingException {
-        T requestParams = null;
-        final var requestType = command.getRequestType();
-        if (params != null && requestType != null) {
-            try {
-                requestParams = objectMapper.readValue(objectMapper.treeAsTokens(params), requestType);
-            } catch (JsonMappingException e) {
-                throw e;
-            } catch (IOException e) {
-                throw new AssertionError(e);
+    private class UnsubscribeReceiveCommand implements JsonRpcSingleCommand<JsonNode>, JsonRpcMultiCommand<JsonNode> {
+
+        @Override
+        public String getName() {
+            return "unsubscribeReceive";
+        }
+
+        @Override
+        public TypeReference<JsonNode> getRequestType() {
+            return new TypeReference<>() {};
+        }
+
+        @Override
+        public void handleCommand(
+                final JsonNode request, final Manager m, final JsonWriter jsonWriter
+        ) throws CommandException {
+            final var subscriptionId = getSubscriptionId(request);
+            if (subscriptionId == null) {
+                unsubscribeReceive(m);
+            } else {
+                if (!unsubscribeReceive(subscriptionId)) {
+                    throw new UserErrorException("Unknown subscription id");
+                }
+            }
+        }
+
+        @Override
+        public void handleCommand(
+                final JsonNode request, final MultiAccountManager c, final JsonWriter jsonWriter
+        ) throws CommandException {
+            final var subscriptionId = getSubscriptionId(request);
+            if (subscriptionId == null) {
+                throw new UserErrorException("Missing subscription parameter with subscription id");
+            } else {
+                if (!unsubscribeReceive(subscriptionId)) {
+                    throw new UserErrorException("Unknown subscription id");
+                }
+            }
+        }
+
+        private Integer getSubscriptionId(final JsonNode request) {
+            if (request instanceof ArrayNode req) {
+                return req.get(0).asInt();
+            } else if (request instanceof ObjectNode req) {
+                return req.get("subscription").asInt();
+            } else {
+                return null;

             }
         }
             }
         }

-        command.handleCommand(requestParams, m, outputWriter);

     }
 }
     }
 }