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