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.ArrayNode;
   9 import com.fasterxml.jackson.databind.node.ContainerNode;
  10 import com.fasterxml.jackson.databind.node.IntNode;
  11 import com.fasterxml.jackson.databind.node.ObjectNode;
  12
  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.json.JsonReceiveMessageHandler;
  23 import org.asamk.signal.manager.Manager;
  24 import org.asamk.signal.manager.MultiAccountManager;
  25 import org.asamk.signal.manager.RegistrationManager;
  26 import org.asamk.signal.manager.api.Pair;
  27 import org.asamk.signal.output.JsonWriter;
  28 import org.asamk.signal.util.Util;
  29 import org.slf4j.Logger;
  30 import org.slf4j.LoggerFactory;
  31
  32 import java.io.IOException;
  33 import java.nio.channels.OverlappingFileLockException;
  34 import java.util.HashMap;
  35 import java.util.List;
  36 import java.util.Map;
  37 import java.util.concurrent.atomic.AtomicInteger;
  38 import java.util.function.Supplier;
  39
  40 public class SignalJsonRpcDispatcherHandler {
  41
  42     private final static Logger logger = LoggerFactory.getLogger(SignalJsonRpcDispatcherHandler.class);
  43
  44     private static final int USER_ERROR = -1;
  45     private static final int IO_ERROR = -3;
  46     private static final int UNTRUSTED_KEY_ERROR = -4;
  47
  48     private final ObjectMapper objectMapper;
  49     private final JsonRpcSender jsonRpcSender;
  50     private final JsonRpcReader jsonRpcReader;
  51     private final boolean noReceiveOnStart;
  52
  53     private MultiAccountManager c;
  54     private final Map<Integer, List<Pair<Manager, Manager.ReceiveMessageHandler>>> receiveHandlers = new HashMap<>();
  55
  56     private Manager m;
  57
  58     public SignalJsonRpcDispatcherHandler(
  59             final JsonWriter jsonWriter, final Supplier<String> lineSupplier, final boolean noReceiveOnStart
  60     ) {
  61         this.noReceiveOnStart = noReceiveOnStart;
  62         this.objectMapper = Util.createJsonObjectMapper();
  63         this.jsonRpcSender = new JsonRpcSender(jsonWriter);
  64         this.jsonRpcReader = new JsonRpcReader(jsonRpcSender, lineSupplier);
  65     }
  66
  67     public void handleConnection(final MultiAccountManager c) {
  68         this.c = c;
  69
  70         if (!noReceiveOnStart) {
  71             this.subscribeReceive(c.getManagers());
  72             c.addOnManagerAddedHandler(this::subscribeReceive);
  73             c.addOnManagerRemovedHandler(this::unsubscribeReceive);
  74         }
  75
  76         handleConnection();
  77     }
  78
  79     public void handleConnection(final Manager m) {
  80         this.m = m;
  81
  82         if (!noReceiveOnStart) {
  83             subscribeReceive(m);
  84         }
  85
  86         final var currentThread = Thread.currentThread();
  87         m.addClosedListener(currentThread::interrupt);
  88
  89         handleConnection();
  90     }
  91
  92     private static final AtomicInteger nextSubscriptionId = new AtomicInteger(0);
  93
  94     private int subscribeReceive(final Manager manager) {
  95         return subscribeReceive(List.of(manager));
  96     }
  97
  98     private int subscribeReceive(final List<Manager> managers) {
  99         final var subscriptionId = nextSubscriptionId.getAndIncrement();
 100         final var handlers = managers.stream().map(m -> {
 101             final var receiveMessageHandler = new JsonReceiveMessageHandler(m, s -> {
 102                 final ContainerNode<?> params = objectMapper.valueToTree(s);
 103                 ((ObjectNode) params).set("subscription", IntNode.valueOf(subscriptionId));
 104                 jsonRpcSender.sendRequest(JsonRpcRequest.forNotification("receive", params, null));
 105             });
 106             m.addReceiveHandler(receiveMessageHandler);
 107             return new Pair<>(m, (Manager.ReceiveMessageHandler) receiveMessageHandler);
 108         }).toList();
 109         receiveHandlers.put(subscriptionId, handlers);
 110
 111         return subscriptionId;
 112     }
 113
 114     private boolean unsubscribeReceive(final int subscriptionId) {
 115         final var handlers = receiveHandlers.remove(subscriptionId);
 116         if (handlers == null) {
 117             return false;
 118         }
 119         for (final var pair : handlers) {
 120             unsubscribeReceiveHandler(pair);
 121         }
 122         return true;
 123     }
 124
 125     private void unsubscribeReceive(final Manager m) {
 126         final var subscriptionId = receiveHandlers.entrySet()
 127                 .stream()
 128                 .filter(e -> e.getValue().size() == 1 && e.getValue().get(0).first().equals(m))
 129                 .map(Map.Entry::getKey)
 130                 .findFirst();
 131         subscriptionId.ifPresent(this::unsubscribeReceive);
 132     }
 133
 134     private void handleConnection() {
 135         try {
 136             jsonRpcReader.readMessages((method, params) -> handleRequest(objectMapper, method, params),
 137                     response -> logger.debug("Received unexpected response for id {}", response.getId()));
 138         } finally {
 139             receiveHandlers.forEach((_subscriptionId, handlers) -> handlers.forEach(this::unsubscribeReceiveHandler));
 140             receiveHandlers.clear();
 141         }
 142     }
 143
 144     private void unsubscribeReceiveHandler(final Pair<Manager, Manager.ReceiveMessageHandler> pair) {
 145         final var m = pair.first();
 146         final var handler = pair.second();
 147         m.removeReceiveHandler(handler);
 148     }
 149
 150     private JsonNode handleRequest(
 151             final ObjectMapper objectMapper, final String method, ContainerNode<?> params
 152     ) throws JsonRpcException {
 153         var command = getCommand(method);
 154         if (c != null) {
 155             if (command instanceof JsonRpcSingleCommand<?> jsonRpcCommand) {
 156                 final var manager = getManagerFromParams(params);
 157                 if (manager != null) {
 158                     return runCommand(objectMapper, params, new CommandRunnerImpl<>(manager, jsonRpcCommand));
 159                 }
 160             }
 161             if (command instanceof JsonRpcMultiCommand<?> jsonRpcCommand) {
 162                 return runCommand(objectMapper, params, new MultiCommandRunnerImpl<>(c, jsonRpcCommand));
 163             }
 164             if (command instanceof JsonRpcRegistrationCommand<?> jsonRpcCommand) {
 165                 try (var manager = getRegistrationManagerFromParams(params)) {
 166                     if (manager != null) {
 167                         return runCommand(objectMapper,
 168                                 params,
 169                                 new RegistrationCommandRunnerImpl<>(manager, c, jsonRpcCommand));
 170                     } else {
 171                         throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INVALID_PARAMS,
 172                                 "Method requires valid account parameter",
 173                                 null));
 174                     }
 175                 } catch (IOException e) {
 176                     logger.warn("Failed to close registration manager", e);
 177                 }
 178             }
 179         }
 180         if (command instanceof JsonRpcSingleCommand<?> jsonRpcCommand) {
 181             if (m != null) {
 182                 return runCommand(objectMapper, params, new CommandRunnerImpl<>(m, jsonRpcCommand));
 183             }
 184
 185             final var manager = getManagerFromParams(params);
 186             if (manager != null) {
 187                 return runCommand(objectMapper, params, new CommandRunnerImpl<>(manager, jsonRpcCommand));
 188             } else {
 189                 throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INVALID_PARAMS,
 190                         "Method requires valid account parameter",
 191                         null));
 192             }
 193         }
 194
 195         throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.METHOD_NOT_FOUND,
 196                 "Method not implemented",
 197                 null));
 198     }
 199
 200     private Manager getManagerFromParams(final ContainerNode<?> params) throws JsonRpcException {
 201         if (params != null && params.hasNonNull("account")) {
 202             final var manager = c.getManager(params.get("account").asText());
 203             ((ObjectNode) params).remove("account");
 204             if (manager == null) {
 205                 throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INVALID_PARAMS,
 206                         "Specified account does not exist",
 207                         null));
 208             }
 209             return manager;
 210         }
 211         return null;
 212     }
 213
 214     private RegistrationManager getRegistrationManagerFromParams(final ContainerNode<?> params) {
 215         if (params != null && params.has("account")) {
 216             try {
 217                 final var registrationManager = c.getNewRegistrationManager(params.get("account").asText());
 218                 ((ObjectNode) params).remove("account");
 219                 return registrationManager;
 220             } catch (OverlappingFileLockException e) {
 221                 logger.warn("Account is already in use");
 222                 return null;
 223             } catch (IOException | IllegalStateException e) {
 224                 logger.warn("Failed to load registration manager", e);
 225                 return null;
 226             }
 227         }
 228         return null;
 229     }
 230
 231     private Command getCommand(final String method) {
 232         if ("subscribeReceive".equals(method)) {
 233             return new SubscribeReceiveCommand();
 234         }
 235         if ("unsubscribeReceive".equals(method)) {
 236             return new UnsubscribeReceiveCommand();
 237         }
 238         return Commands.getCommand(method);
 239     }
 240
 241     private record CommandRunnerImpl<T>(Manager m, JsonRpcSingleCommand<T> command) implements CommandRunner<T> {
 242
 243         @Override
 244         public void handleCommand(final T request, final JsonWriter jsonWriter) throws CommandException {
 245             command.handleCommand(request, m, jsonWriter);
 246         }
 247
 248         @Override
 249         public TypeReference<T> getRequestType() {
 250             return command.getRequestType();
 251         }
 252     }
 253
 254     private record RegistrationCommandRunnerImpl<T>(
 255             RegistrationManager m, MultiAccountManager c, JsonRpcRegistrationCommand<T> command
 256     ) implements CommandRunner<T> {
 257
 258         @Override
 259         public void handleCommand(final T request, final JsonWriter jsonWriter) throws CommandException {
 260             command.handleCommand(request, m, jsonWriter);
 261         }
 262
 263         @Override
 264         public TypeReference<T> getRequestType() {
 265             return command.getRequestType();
 266         }
 267     }
 268
 269     private record MultiCommandRunnerImpl<T>(
 270             MultiAccountManager c, JsonRpcMultiCommand<T> command
 271     ) implements CommandRunner<T> {
 272
 273         @Override
 274         public void handleCommand(final T request, final JsonWriter jsonWriter) throws CommandException {
 275             command.handleCommand(request, c, jsonWriter);
 276         }
 277
 278         @Override
 279         public TypeReference<T> getRequestType() {
 280             return command.getRequestType();
 281         }
 282     }
 283
 284     interface CommandRunner<T> {
 285
 286         void handleCommand(T request, JsonWriter jsonWriter) throws CommandException;
 287
 288         TypeReference<T> getRequestType();
 289     }
 290
 291     private JsonNode runCommand(
 292             final ObjectMapper objectMapper, final ContainerNode<?> params, final CommandRunner<?> command
 293     ) throws JsonRpcException {
 294         final Object[] result = {null};
 295         final JsonWriter commandJsonWriter = s -> {
 296             if (result[0] != null) {
 297                 throw new AssertionError("Command may only write one json result");
 298             }
 299
 300             result[0] = s;
 301         };
 302
 303         try {
 304             parseParamsAndRunCommand(objectMapper, params, commandJsonWriter, command);
 305         } catch (JsonMappingException e) {
 306             throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INVALID_REQUEST,
 307                     e.getMessage(),
 308                     null));
 309         } catch (UserErrorException e) {
 310             throw new JsonRpcException(new JsonRpcResponse.Error(USER_ERROR,
 311                     e.getMessage(),
 312                     getErrorDataNode(objectMapper, result)));
 313         } catch (IOErrorException e) {
 314             throw new JsonRpcException(new JsonRpcResponse.Error(IO_ERROR,
 315                     e.getMessage(),
 316                     getErrorDataNode(objectMapper, result)));
 317         } catch (UntrustedKeyErrorException e) {
 318             throw new JsonRpcException(new JsonRpcResponse.Error(UNTRUSTED_KEY_ERROR,
 319                     e.getMessage(),
 320                     getErrorDataNode(objectMapper, result)));
 321         } catch (Throwable e) {
 322             logger.error("Command execution failed", e);
 323             throw new JsonRpcException(new JsonRpcResponse.Error(JsonRpcResponse.Error.INTERNAL_ERROR,
 324                     e.getMessage(),
 325                     getErrorDataNode(objectMapper, result)));
 326         }
 327
 328         Object output = result[0] == null ? Map.of() : result[0];
 329         return objectMapper.valueToTree(output);
 330     }
 331
 332     private JsonNode getErrorDataNode(final ObjectMapper objectMapper, final Object[] result) {
 333         if (result[0] == null) {
 334             return null;
 335         }
 336         return objectMapper.valueToTree(Map.of("response", result[0]));
 337     }
 338
 339     private <T> void parseParamsAndRunCommand(
 340             final ObjectMapper objectMapper,
 341             final TreeNode params,
 342             final JsonWriter jsonWriter,
 343             final CommandRunner<T> command
 344     ) throws CommandException, JsonMappingException {
 345         T requestParams = null;
 346         final var requestType = command.getRequestType();
 347         if (params != null && requestType != null) {
 348             try {
 349                 requestParams = objectMapper.readValue(objectMapper.treeAsTokens(params), requestType);
 350             } catch (JsonMappingException e) {
 351                 throw e;
 352             } catch (IOException e) {
 353                 throw new AssertionError(e);
 354             }
 355         }
 356         command.handleCommand(requestParams, jsonWriter);
 357     }
 358
 359     private class SubscribeReceiveCommand implements JsonRpcSingleCommand<Void>, JsonRpcMultiCommand<Void> {
 360
 361         @Override
 362         public String getName() {
 363             return "subscribeReceive";
 364         }
 365
 366         @Override
 367         public void handleCommand(
 368                 final Void request, final Manager m, final JsonWriter jsonWriter
 369         ) throws CommandException {
 370             final var subscriptionId = subscribeReceive(m);
 371             jsonWriter.write(subscriptionId);
 372         }
 373
 374         @Override
 375         public void handleCommand(
 376                 final Void request, final MultiAccountManager c, final JsonWriter jsonWriter
 377         ) throws CommandException {
 378             final var subscriptionId = subscribeReceive(c.getManagers());
 379             jsonWriter.write(subscriptionId);
 380         }
 381     }
 382
 383     private class UnsubscribeReceiveCommand implements JsonRpcSingleCommand<JsonNode>, JsonRpcMultiCommand<JsonNode> {
 384
 385         @Override
 386         public String getName() {
 387             return "unsubscribeReceive";
 388         }
 389
 390         @Override
 391         public TypeReference<JsonNode> getRequestType() {
 392             return new TypeReference<>() {};
 393         }
 394
 395         @Override
 396         public void handleCommand(
 397                 final JsonNode request, final Manager m, final JsonWriter jsonWriter
 398         ) throws CommandException {
 399             final var subscriptionId = getSubscriptionId(request);
 400             if (subscriptionId == null) {
 401                 unsubscribeReceive(m);
 402             } else {
 403                 if (!unsubscribeReceive(subscriptionId)) {
 404                     throw new UserErrorException("Unknown subscription id");
 405                 }
 406             }
 407         }
 408
 409         @Override
 410         public void handleCommand(
 411                 final JsonNode request, final MultiAccountManager c, final JsonWriter jsonWriter
 412         ) throws CommandException {
 413             final var subscriptionId = getSubscriptionId(request);
 414             if (subscriptionId == null) {
 415                 throw new UserErrorException("Missing subscription parameter with subscription id");
 416             } else {
 417                 if (!unsubscribeReceive(subscriptionId)) {
 418                     throw new UserErrorException("Unknown subscription id");
 419                 }
 420             }
 421         }
 422
 423         private Integer getSubscriptionId(final JsonNode request) {
 424             if (request instanceof ArrayNode req) {
 425                 return req.get(0).asInt();
 426             } else if (request instanceof ObjectNode req) {
 427                 return req.get("subscription").asInt();
 428             } else {
 429                 return null;
 430             }
 431         }
 432     }
 433 }