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