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 recipientId = resolver.resolveRecipient(senderKey.recipientId);
  51                 if (recipientId == null) {
  52                     continue;
  53                 }
  54                 final var entry = new SenderKeySharedEntry(recipientId, senderKey.deviceId);
  55                 final var uuid = UuidUtil.parseOrNull(senderKey.distributionId);
  56                 if (uuid == null) {
  57                     logger.warn("Read invalid distribution id from storage {}, ignoring", senderKey.distributionId);
  58                     continue;
  59                 }
  60                 final var distributionId = DistributionId.from(uuid);
  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.debug("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<DistributionId, 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             return sharedSenderKeys.get(distributionId)
  93                     .stream()
  94                     .map(k -> new SignalProtocolAddress(addressResolver.resolveRecipientAddress(k.recipientId())
  95                             .getIdentifier(), k.deviceId()))
  96                     .collect(Collectors.toSet());
  97         }
  98     }
  99
 100     public void markSenderKeySharedWith(
 101             final DistributionId distributionId, final Collection<SignalProtocolAddress> addresses
 102     ) {
 103         final var newEntries = addresses.stream()
 104                 .map(a -> new SenderKeySharedEntry(resolveRecipient(a.getName()), a.getDeviceId()))
 105                 .collect(Collectors.toSet());
 106
 107         synchronized (sharedSenderKeys) {
 108             final var previousEntries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 109
 110             sharedSenderKeys.put(distributionId, new HashSet<>() {
 111                 {
 112                     addAll(previousEntries);
 113                     addAll(newEntries);
 114                 }
 115             });
 116             saveLocked();
 117         }
 118     }
 119
 120     public void clearSenderKeySharedWith(final Collection<SignalProtocolAddress> addresses) {
 121         final var entriesToDelete = addresses.stream()
 122                 .map(a -> new SenderKeySharedEntry(resolveRecipient(a.getName()), a.getDeviceId()))
 123                 .collect(Collectors.toSet());
 124
 125         synchronized (sharedSenderKeys) {
 126             for (final var distributionId : sharedSenderKeys.keySet()) {
 127                 final var entries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 128
 129                 sharedSenderKeys.put(distributionId, new HashSet<>(entries) {
 130                     {
 131                         removeAll(entriesToDelete);
 132                     }
 133                 });
 134             }
 135             saveLocked();
 136         }
 137     }
 138
 139     public void deleteAll() {
 140         synchronized (sharedSenderKeys) {
 141             sharedSenderKeys.clear();
 142             saveLocked();
 143         }
 144     }
 145
 146     public void deleteAllFor(final RecipientId recipientId) {
 147         synchronized (sharedSenderKeys) {
 148             for (final var distributionId : sharedSenderKeys.keySet()) {
 149                 final var entries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 150
 151                 sharedSenderKeys.put(distributionId, new HashSet<>(entries) {
 152                     {
 153                         removeIf(e -> e.recipientId().equals(recipientId));
 154                     }
 155                 });
 156             }
 157             saveLocked();
 158         }
 159     }
 160
 161     public void mergeRecipients(RecipientId recipientId, RecipientId toBeMergedRecipientId) {
 162         synchronized (sharedSenderKeys) {
 163             for (final var distributionId : sharedSenderKeys.keySet()) {
 164                 final var entries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 165
 166                 sharedSenderKeys.put(distributionId,
 167                         entries.stream()
 168                                 .map(e -> e.recipientId.equals(toBeMergedRecipientId) ? new SenderKeySharedEntry(
 169                                         recipientId,
 170                                         e.deviceId()) : e)
 171                                 .collect(Collectors.toSet()));
 172             }
 173             saveLocked();
 174         }
 175     }
 176
 177     /**
 178      * @param identifier can be either a serialized uuid or a e164 phone number
 179      */
 180     private RecipientId resolveRecipient(String identifier) {
 181         return resolver.resolveRecipient(identifier);
 182     }
 183
 184     private void saveLocked() {
 185         var storage = new Storage(sharedSenderKeys.entrySet().stream().flatMap(pair -> {
 186             final var sharedWith = pair.getValue();
 187             return sharedWith.stream()
 188                     .map(entry -> new Storage.SharedSenderKey(entry.recipientId().id(),
 189                             entry.deviceId(),
 190                             pair.getKey().asUuid().toString()));
 191         }).toList());
 192
 193         // Write to memory first to prevent corrupting the file in case of serialization errors
 194         try (var inMemoryOutput = new ByteArrayOutputStream()) {
 195             objectMapper.writeValue(inMemoryOutput, storage);
 196
 197             var input = new ByteArrayInputStream(inMemoryOutput.toByteArray());
 198             try (var outputStream = new FileOutputStream(file)) {
 199                 input.transferTo(outputStream);
 200             }
 201         } catch (Exception e) {
 202             logger.error("Error saving shared sender key store file: {}", e.getMessage());
 203         }
 204     }
 205
 206     private record Storage(List<SharedSenderKey> sharedSenderKeys) {
 207
 208         private record SharedSenderKey(long recipientId, int deviceId, String distributionId) {}
 209     }
 210
 211     private record SenderKeySharedEntry(RecipientId recipientId, int deviceId) {}
 212 }