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