nmode's Git Repositories - signal-cli/blob - lib/src/main/java/org/asamk/signal/manager/storage/senderKeys/SenderKeySharedStore.java

   1 package org.asamk.signal.manager.storage.senderKeys;
   2
   3 import com.fasterxml.jackson.databind.ObjectMapper;
   4
   5 import org.asamk.signal.manager.helper.RecipientAddressResolver;
   6 import org.asamk.signal.manager.storage.Utils;
   7 import org.asamk.signal.manager.storage.recipients.RecipientId;
   8 import org.asamk.signal.manager.storage.recipients.RecipientResolver;
   9 import org.slf4j.Logger;
  10 import org.slf4j.LoggerFactory;
  11 import org.whispersystems.libsignal.SignalProtocolAddress;
  12 import org.whispersystems.signalservice.api.push.DistributionId;
  13 import org.whispersystems.signalservice.api.util.UuidUtil;
  14
  15 import java.io.ByteArrayInputStream;
  16 import java.io.ByteArrayOutputStream;
  17 import java.io.File;
  18 import java.io.FileInputStream;
  19 import java.io.FileNotFoundException;
  20 import java.io.FileOutputStream;
  21 import java.io.IOException;
  22 import java.util.Collection;
  23 import java.util.HashMap;
  24 import java.util.HashSet;
  25 import java.util.List;
  26 import java.util.Map;
  27 import java.util.Set;
  28 import java.util.stream.Collectors;
  29
  30 public class SenderKeySharedStore {
  31
  32     private final static Logger logger = LoggerFactory.getLogger(SenderKeySharedStore.class);
  33
  34     private final Map<DistributionId, Set<SenderKeySharedEntry>> sharedSenderKeys;
  35
  36     private final ObjectMapper objectMapper;
  37     private final File file;
  38
  39     private final RecipientResolver resolver;
  40     private final RecipientAddressResolver addressResolver;
  41
  42     public static SenderKeySharedStore load(
  43             final File file, final RecipientAddressResolver addressResolver, final RecipientResolver resolver
  44     ) throws IOException {
  45         final var objectMapper = Utils.createStorageObjectMapper();
  46         try (var inputStream = new FileInputStream(file)) {
  47             final var storage = objectMapper.readValue(inputStream, Storage.class);
  48             final var sharedSenderKeys = new HashMap<DistributionId, Set<SenderKeySharedEntry>>();
  49             for (final var senderKey : storage.sharedSenderKeys) {
  50                 final var entry = new SenderKeySharedEntry(RecipientId.of(senderKey.recipientId), senderKey.deviceId);
  51                 final var uuid = UuidUtil.parseOrNull(senderKey.distributionId);
  52                 if (uuid == null) {
  53                     logger.warn("Read invalid distribution id from storage {}, ignoring", senderKey.distributionId);
  54                     continue;
  55                 }
  56                 final var distributionId = DistributionId.from(uuid);
  57                 var entries = sharedSenderKeys.get(distributionId);
  58                 if (entries == null) {
  59                     entries = new HashSet<>();
  60                 }
  61                 entries.add(entry);
  62                 sharedSenderKeys.put(distributionId, entries);
  63             }
  64
  65             return new SenderKeySharedStore(sharedSenderKeys, objectMapper, file, addressResolver, resolver);
  66         } catch (FileNotFoundException e) {
  67             logger.debug("Creating new shared sender key store.");
  68             return new SenderKeySharedStore(new HashMap<>(), objectMapper, file, addressResolver, resolver);
  69         }
  70     }
  71
  72     private SenderKeySharedStore(
  73             final Map<DistributionId, Set<SenderKeySharedEntry>> sharedSenderKeys,
  74             final ObjectMapper objectMapper,
  75             final File file,
  76             final RecipientAddressResolver addressResolver,
  77             final RecipientResolver resolver
  78     ) {
  79         this.sharedSenderKeys = sharedSenderKeys;
  80         this.objectMapper = objectMapper;
  81         this.file = file;
  82         this.addressResolver = addressResolver;
  83         this.resolver = resolver;
  84     }
  85
  86     public Set<SignalProtocolAddress> getSenderKeySharedWith(final DistributionId distributionId) {
  87         synchronized (sharedSenderKeys) {
  88             return sharedSenderKeys.get(distributionId)
  89                     .stream()
  90                     .map(k -> new SignalProtocolAddress(addressResolver.resolveRecipientAddress(k.recipientId())
  91                             .getIdentifier(), k.deviceId()))
  92                     .collect(Collectors.toSet());
  93         }
  94     }
  95
  96     public void markSenderKeySharedWith(
  97             final DistributionId distributionId, final Collection<SignalProtocolAddress> addresses
  98     ) {
  99         final var newEntries = addresses.stream()
 100                 .map(a -> new SenderKeySharedEntry(resolveRecipient(a.getName()), a.getDeviceId()))
 101                 .collect(Collectors.toSet());
 102
 103         synchronized (sharedSenderKeys) {
 104             final var previousEntries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 105
 106             sharedSenderKeys.put(distributionId, new HashSet<>() {
 107                 {
 108                     addAll(previousEntries);
 109                     addAll(newEntries);
 110                 }
 111             });
 112             saveLocked();
 113         }
 114     }
 115
 116     public void clearSenderKeySharedWith(final Collection<SignalProtocolAddress> addresses) {
 117         final var entriesToDelete = addresses.stream()
 118                 .map(a -> new SenderKeySharedEntry(resolveRecipient(a.getName()), a.getDeviceId()))
 119                 .collect(Collectors.toSet());
 120
 121         synchronized (sharedSenderKeys) {
 122             for (final var distributionId : sharedSenderKeys.keySet()) {
 123                 final var entries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 124
 125                 sharedSenderKeys.put(distributionId, new HashSet<>(entries) {
 126                     {
 127                         removeAll(entriesToDelete);
 128                     }
 129                 });
 130             }
 131             saveLocked();
 132         }
 133     }
 134
 135     public void deleteAll() {
 136         synchronized (sharedSenderKeys) {
 137             sharedSenderKeys.clear();
 138             saveLocked();
 139         }
 140     }
 141
 142     public void deleteAllFor(final RecipientId recipientId) {
 143         synchronized (sharedSenderKeys) {
 144             for (final var distributionId : sharedSenderKeys.keySet()) {
 145                 final var entries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 146
 147                 sharedSenderKeys.put(distributionId, new HashSet<>(entries) {
 148                     {
 149                         removeIf(e -> e.recipientId().equals(recipientId));
 150                     }
 151                 });
 152             }
 153             saveLocked();
 154         }
 155     }
 156
 157     public void mergeRecipients(RecipientId recipientId, RecipientId toBeMergedRecipientId) {
 158         synchronized (sharedSenderKeys) {
 159             for (final var distributionId : sharedSenderKeys.keySet()) {
 160                 final var entries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 161
 162                 sharedSenderKeys.put(distributionId,
 163                         entries.stream()
 164                                 .map(e -> e.recipientId.equals(toBeMergedRecipientId) ? new SenderKeySharedEntry(
 165                                         recipientId,
 166                                         e.deviceId()) : e)
 167                                 .collect(Collectors.toSet()));
 168             }
 169             saveLocked();
 170         }
 171     }
 172
 173     /**
 174      * @param identifier can be either a serialized uuid or a e164 phone number
 175      */
 176     private RecipientId resolveRecipient(String identifier) {
 177         return resolver.resolveRecipient(identifier);
 178     }
 179
 180     private void saveLocked() {
 181         var storage = new Storage(sharedSenderKeys.entrySet().stream().flatMap(pair -> {
 182             final var sharedWith = pair.getValue();
 183             return sharedWith.stream()
 184                     .map(entry -> new Storage.SharedSenderKey(entry.recipientId().id(),
 185                             entry.deviceId(),
 186                             pair.getKey().asUuid().toString()));
 187         }).collect(Collectors.toList()));
 188
 189         // Write to memory first to prevent corrupting the file in case of serialization errors
 190         try (var inMemoryOutput = new ByteArrayOutputStream()) {
 191             objectMapper.writeValue(inMemoryOutput, storage);
 192
 193             var input = new ByteArrayInputStream(inMemoryOutput.toByteArray());
 194             try (var outputStream = new FileOutputStream(file)) {
 195                 input.transferTo(outputStream);
 196             }
 197         } catch (Exception e) {
 198             logger.error("Error saving shared sender key store file: {}", e.getMessage());
 199         }
 200     }
 201
 202     private record Storage(List<SharedSenderKey> sharedSenderKeys) {
 203
 204         private record SharedSenderKey(long recipientId, int deviceId, String distributionId) {}
 205     }
 206
 207     private record SenderKeySharedEntry(RecipientId recipientId, int deviceId) {}
 208 }