nmode's Git Repositories - signal-cli/blob - src/main/java/org/asamk/signal/jsonrpc/SignalJsonRpcDispatcherHandler.java

   1 package org.asamk.signal.jsonrpc;
   2
   3 import com.fasterxml.jackson.core.TreeNode;
   4 import com.fasterxml.jackson.core.type.TypeReference;
   5 import com.fasterxml.jackson.databind.JsonMappingException;
   6 import com.fasterxml.jackson.databind.JsonNode;
   7 import com.fasterxml.jackson.databind.ObjectMapper;
   8 import com.fasterxml.jackson.databind.node.ContainerNode;
   9 import com.fasterxml.jackson.databind.node.ObjectNode;
  10
  11 import org.asamk.signal.JsonReceiveMessageHandler;
  12 import org.asamk.signal.JsonWriter;
  13 import org.asamk.signal.commands.Command;
  14 import org.asamk.signal.commands.Commands;
  15 import org.asamk.signal.commands.JsonRpcMultiCommand;
  16 import org.asamk.signal.commands.JsonRpcRegistrationCommand;
  17 import org.asamk.signal.commands.JsonRpcSingleCommand;
  18 import org.asamk.signal.commands.exceptions.CommandException;
  19 import org.asamk.signal.commands.exceptions.IOErrorException;
  20 import org.asamk.signal.commands.exceptions.UntrustedKeyErrorException;
  21 import org.asamk.signal.commands.exceptions.UserErrorException;
  22 import org.asamk.signal.manager.Manager;
  23 import org.asamk.signal.manager.MultiAccountManager;
  24 import org.asamk.signal.manager.RegistrationManager;
  25 import org.asamk.signal.util.Util;
  26 import org.slf4j.Logger;
  27 import org.slf4j.LoggerFactory;
  28
  29 import java.io.IOException;
  30 import java.util.HashMap;
  31 import java.util.Map;
  32 import java.util.Objects;
  33 import java.util.function.Supplier;
  34
  35 public class SignalJsonRpcDispatcherHandler {
  36
  37     private final static Logger logger = LoggerFactory.getLogger(SignalJsonRpcDispatcherHandler.class);
  38
  39     private static final int USER_ERROR = -1;
  40     private static final int IO_ERROR = -3;
  41     private static final int UNTRUSTED_KEY_ERROR = -4;
  42
  43     private final ObjectMapper objectMapper;
  44     private final JsonRpcSender jsonRpcSender;
  45     private final JsonRpcReader jsonRpcReader;
  46     private final boolean noReceiveOnStart;
  47
  48     private MultiAccountManager c;
  49     private final Map<Manager, Manager.ReceiveMessageHandler> receiveHandlers = new HashMap<>();
  50
  51     private Manager m;
  52
  53     public SignalJsonRpcDispatcherHandler(
  54             final JsonWriter jsonWriter, final Supplier<String> lineSupplier, final boolean noReceiveOnStart
  55     ) {
  56         this.noReceiveOnStart = noReceiveOnStart;
  57         this.objectMapper = Util.createJsonObjectMapper();
  58         this.jsonRpcSender = new JsonRpcSender(jsonWriter);
  59         this.jsonRpcReader = new JsonRpcReader(jsonRpcSender, lineSupplier);
  60     }
  61
  62     public void handleConnection(final MultiAccountManager c) {
  63         this.c = c;
  64
  65         if (!noReceiveOnStart) {
  66             c.getAccountNumbers().stream().map(c::getManager).filter(Objects::nonNull).forEach(this::subscribeReceive);
  67         }
  68
  69         handleConnection();
  70     }
  71
  72     public void handleConnection(final Manager m) {
  73         this.m = m;
  74
  75         if (!noReceiveOnStart) {
  76             subscribeReceive(m);
  77         }
  78
  79         handleConnection();
  80     }
  81
  82     private void subscribeReceive(final Manager m) {
  83         if (receiveHandlers.containsKey(m)) {
  84             return;
  85         }
  86
  87         final var receiveMessageHandler = new JsonReceiveMessageHandler(m,
  88                 s -> jsonRpcSender.sendRequest(JsonRpcRequest.forNotification("receive",
  89                         objectMapper.valueToTree(s),
  90                         null)));
  91         m.addReceiveHandler(receiveMessageHandler);
  92         receiveHandlers.put(m, receiveMessageHandler);
  93
  94         while (!m.hasCaughtUpWithOldMessages()) {
  95             try {
  96                 synchronized (m) {
  97                     m.wait();
  98                 }
  99             } catch (InterruptedException ignored) {
 100             }
 101         }
 102     }
 103
 104     void unsubscribeReceive(final Manager m) {
 105         final var receiveMessageHandler = receiveHandlers.remove(m);
 106         if (receiveMessageHandler != null) {
 107             m.removeReceiveHandler(receiveMessageHandler);
 108         }
 109     }
 110
 111     private void handleConnection() {
 112         try {
 113             jsonRpcReader.readMessages((method, params) -> handleRequest(objectMapper, method, params),
 114                     response -> logger.debug("Received unexpected response for id {}", response.getId()));
 115         } finally {
 116             receiveHandlers.forEach(Manager::removeReceiveHandler);
 117             receiveHandlers.clear();
 118         }
 119     }
 120
 121     private JsonNode handleRequest(
 122             final ObjectMapper objectMapper, final String method, ContainerNode<?> params
 123     ) throws JsonRpcException {
 124         var command = getCommand(method);
 125         if (c != null) {
 126             if (command instanceof JsonRpcMultiCommand<?> jsonRpcCommand) {
 127                 return runCommand(objectMapper, params, new MultiCommandRunnerImpl<>(c, jsonRpcCommand));
 128             }
 129             if (command instanceof JsonRpcRegistrationCommand<?> jsonRpcCommand) {
 130                 try (var manager = getRegistrationManagerFromParams(params)) {
 131                     if (manager != null) {
 132                         return runCommand(objectMapper,
 133                                 params,
 134                                 new RegistrationCommandRunnerImpl<>(manager, c, jsonRpcCommand));
 135                     } else {
 136                         throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INVALID_PARAMS,
 137                                 "Method requires valid account parameter",
 138                                 null));
 139                     }
 140                 } catch (IOException e) {
 141                     logger.warn("Failed to close registration manager", e);
 142                 }
 143             }
 144         }
 145         if (command instanceof JsonRpcSingleCommand<?> jsonRpcCommand) {
 146             if (m != null) {
 147                 return runCommand(objectMapper, params, new CommandRunnerImpl<>(m, jsonRpcCommand));
 148             }
 149
 150             final var manager = getManagerFromParams(params);
 151             if (manager != null) {
 152                 return runCommand(objectMapper, params, new CommandRunnerImpl<>(manager, jsonRpcCommand));
 153             } else {
 154                 throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INVALID_PARAMS,
 155                         "Method requires valid account parameter",
 156                         null));
 157             }
 158         }
 159
 160         throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.METHOD_NOT_FOUND,
 161                 "Method not implemented",
 162                 null));
 163     }
 164
 165     private Manager getManagerFromParams(final ContainerNode<?> params) {
 166         if (params != null && params.has("account")) {
 167             final var manager = c.getManager(params.get("account").asText());
 168             ((ObjectNode) params).remove("account");
 169             return manager;
 170         }
 171         return null;
 172     }
 173
 174     private RegistrationManager getRegistrationManagerFromParams(final ContainerNode<?> params) {
 175         if (params != null && params.has("account")) {
 176             try {
 177                 final var registrationManager = c.getNewRegistrationManager(params.get("account").asText());
 178                 ((ObjectNode) params).remove("account");
 179                 return registrationManager;
 180             } catch (IOException | IllegalStateException e) {
 181                 logger.warn("Failed to load registration manager", e);
 182                 return null;
 183             }
 184         }
 185         return null;
 186     }
 187
 188     private Command getCommand(final String method) {
 189         if ("subscribeReceive".equals(method)) {
 190             return new SubscribeReceiveCommand();
 191         }
 192         if ("unsubscribeReceive".equals(method)) {
 193             return new UnsubscribeReceiveCommand();
 194         }
 195         return Commands.getCommand(method);
 196     }
 197
 198     private record CommandRunnerImpl<T>(Manager m, JsonRpcSingleCommand<T> command) implements CommandRunner<T> {
 199
 200         @Override
 201         public void handleCommand(final T request, final JsonWriter jsonWriter) throws CommandException {
 202             command.handleCommand(request, m, jsonWriter);
 203         }
 204
 205         @Override
 206         public TypeReference<T> getRequestType() {
 207             return command.getRequestType();
 208         }
 209     }
 210
 211     private record RegistrationCommandRunnerImpl<T>(
 212             RegistrationManager m, MultiAccountManager c, JsonRpcRegistrationCommand<T> command
 213     ) implements CommandRunner<T> {
 214
 215         @Override
 216         public void handleCommand(final T request, final JsonWriter jsonWriter) throws CommandException {
 217             command.handleCommand(request, m, jsonWriter);
 218         }
 219
 220         @Override
 221         public TypeReference<T> getRequestType() {
 222             return command.getRequestType();
 223         }
 224     }
 225
 226     private record MultiCommandRunnerImpl<T>(
 227             MultiAccountManager c, JsonRpcMultiCommand<T> command
 228     ) implements CommandRunner<T> {
 229
 230         @Override
 231         public void handleCommand(final T request, final JsonWriter jsonWriter) throws CommandException {
 232             command.handleCommand(request, c, jsonWriter);
 233         }
 234
 235         @Override
 236         public TypeReference<T> getRequestType() {
 237             return command.getRequestType();
 238         }
 239     }
 240
 241     interface CommandRunner<T> {
 242
 243         void handleCommand(T request, JsonWriter jsonWriter) throws CommandException;
 244
 245         TypeReference<T> getRequestType();
 246     }
 247
 248     private JsonNode runCommand(
 249             final ObjectMapper objectMapper, final ContainerNode<?> params, final CommandRunner<?> command
 250     ) throws JsonRpcException {
 251         final Object[] result = {null};
 252         final JsonWriter commandJsonWriter = s -> {
 253             if (result[0] != null) {
 254                 throw new AssertionError("Command may only write one json result");
 255             }
 256
 257             result[0] = s;
 258         };
 259
 260         try {
 261             parseParamsAndRunCommand(objectMapper, params, commandJsonWriter, command);
 262         } catch (JsonMappingException e) {
 263             throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INVALID_REQUEST,
 264                     e.getMessage(),
 265                     null));
 266         } catch (UserErrorException e) {
 267             throw new JsonRpcException(new JsonRpcResponse.Error(USER_ERROR, e.getMessage(), null));
 268         } catch (IOErrorException e) {
 269             throw new JsonRpcException(new JsonRpcResponse.Error(IO_ERROR, e.getMessage(), null));
 270         } catch (UntrustedKeyErrorException e) {
 271             throw new JsonRpcException(new JsonRpcResponse.Error(UNTRUSTED_KEY_ERROR, e.getMessage(), null));
 272         } catch (Throwable e) {
 273             logger.error("Command execution failed", e);
 274             throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INTERNAL_ERROR,
 275                     e.getMessage(),
 276                     null));
 277         }
 278
 279         Object output = result[0] == null ? Map.of() : result[0];
 280         return objectMapper.valueToTree(output);
 281     }
 282
 283     private <T> void parseParamsAndRunCommand(
 284             final ObjectMapper objectMapper,
 285             final TreeNode params,
 286             final JsonWriter jsonWriter,
 287             final CommandRunner<T> command
 288     ) throws CommandException, JsonMappingException {
 289         T requestParams = null;
 290         final var requestType = command.getRequestType();
 291         if (params != null && requestType != null) {
 292             try {
 293                 requestParams = objectMapper.readValue(objectMapper.treeAsTokens(params), requestType);
 294             } catch (JsonMappingException e) {
 295                 throw e;
 296             } catch (IOException e) {
 297                 throw new AssertionError(e);
 298             }
 299         }
 300         command.handleCommand(requestParams, jsonWriter);
 301     }
 302
 303     private class SubscribeReceiveCommand implements JsonRpcSingleCommand<Void> {
 304
 305         @Override
 306         public String getName() {
 307             return "subscribeReceive";
 308         }
 309
 310         @Override
 311         public void handleCommand(
 312                 final Void request, final Manager m, final JsonWriter jsonWriter
 313         ) throws CommandException {
 314             subscribeReceive(m);
 315         }
 316     }
 317
 318     private class UnsubscribeReceiveCommand implements JsonRpcSingleCommand<Void> {
 319
 320         @Override
 321         public String getName() {
 322             return "unsubscribeReceive";
 323         }
 324
 325         @Override
 326         public void handleCommand(
 327                 final Void request, final Manager m, final JsonWriter jsonWriter
 328         ) throws CommandException {
 329             unsubscribeReceive(m);
 330         }
 331     }
 332 }