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     ) {
  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         } catch (IOException e) {
  74             logger.warn("Failed to load shared sender key store", e);
  75             throw new RuntimeException(e);
  76         }
  77     }
  78
  79     private SenderKeySharedStore(
  80             final Map<UUID, Set<SenderKeySharedEntry>> sharedSenderKeys,
  81             final ObjectMapper objectMapper,
  82             final File file,
  83             final RecipientAddressResolver addressResolver,
  84             final RecipientResolver resolver
  85     ) {
  86         this.sharedSenderKeys = sharedSenderKeys;
  87         this.objectMapper = objectMapper;
  88         this.file = file;
  89         this.addressResolver = addressResolver;
  90         this.resolver = resolver;
  91     }
  92
  93     public Set<SignalProtocolAddress> getSenderKeySharedWith(final DistributionId distributionId) {
  94         synchronized (sharedSenderKeys) {
  95             final var addresses = sharedSenderKeys.get(distributionId.asUuid());
  96             if (addresses == null) {
  97                 return Set.of();
  98             }
  99             return addresses.stream()
 100                     .map(k -> new SignalProtocolAddress(addressResolver.resolveRecipientAddress(k.recipientId())
 101                             .getIdentifier(), k.deviceId()))
 102                     .collect(Collectors.toSet());
 103         }
 104     }
 105
 106     public void markSenderKeySharedWith(
 107             final DistributionId distributionId, final Collection<SignalProtocolAddress> addresses
 108     ) {
 109         final var newEntries = addresses.stream()
 110                 .map(a -> new SenderKeySharedEntry(resolveRecipient(a.getName()), a.getDeviceId()))
 111                 .collect(Collectors.toSet());
 112
 113         synchronized (sharedSenderKeys) {
 114             final var previousEntries = sharedSenderKeys.getOrDefault(distributionId.asUuid(), Set.of());
 115
 116             sharedSenderKeys.put(distributionId.asUuid(), new HashSet<>() {
 117                 {
 118                     addAll(previousEntries);
 119                     addAll(newEntries);
 120                 }
 121             });
 122             saveLocked();
 123         }
 124     }
 125
 126     public void clearSenderKeySharedWith(final Collection<SignalProtocolAddress> addresses) {
 127         final var entriesToDelete = addresses.stream()
 128                 .map(a -> new SenderKeySharedEntry(resolveRecipient(a.getName()), a.getDeviceId()))
 129                 .collect(Collectors.toSet());
 130
 131         synchronized (sharedSenderKeys) {
 132             for (final var distributionId : sharedSenderKeys.keySet()) {
 133                 final var entries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 134
 135                 sharedSenderKeys.put(distributionId, new HashSet<>(entries) {
 136                     {
 137                         removeAll(entriesToDelete);
 138                     }
 139                 });
 140             }
 141             saveLocked();
 142         }
 143     }
 144
 145     public void deleteAll() {
 146         synchronized (sharedSenderKeys) {
 147             sharedSenderKeys.clear();
 148             saveLocked();
 149         }
 150     }
 151
 152     public void deleteAllFor(final RecipientId recipientId) {
 153         synchronized (sharedSenderKeys) {
 154             for (final var distributionId : sharedSenderKeys.keySet()) {
 155                 final var entries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 156
 157                 sharedSenderKeys.put(distributionId, new HashSet<>(entries) {
 158                     {
 159                         removeIf(e -> e.recipientId().equals(recipientId));
 160                     }
 161                 });
 162             }
 163             saveLocked();
 164         }
 165     }
 166
 167     public void deleteAllFor(final DistributionId distributionId) {
 168         synchronized (sharedSenderKeys) {
 169             if (sharedSenderKeys.remove(distributionId.asUuid()) != null) {
 170                 saveLocked();
 171             }
 172         }
 173     }
 174
 175     public void mergeRecipients(RecipientId recipientId, RecipientId toBeMergedRecipientId) {
 176         synchronized (sharedSenderKeys) {
 177             for (final var distributionId : sharedSenderKeys.keySet()) {
 178                 final var entries = sharedSenderKeys.getOrDefault(distributionId, Set.of());
 179
 180                 sharedSenderKeys.put(distributionId,
 181                         entries.stream()
 182                                 .map(e -> e.recipientId.equals(toBeMergedRecipientId) ? new SenderKeySharedEntry(
 183                                         recipientId,
 184                                         e.deviceId()) : e)
 185                                 .collect(Collectors.toSet()));
 186             }
 187             saveLocked();
 188         }
 189     }
 190
 191     /**
 192      * @param identifier can be either a serialized uuid or a e164 phone number
 193      */
 194     private RecipientId resolveRecipient(String identifier) {
 195         return resolver.resolveRecipient(identifier);
 196     }
 197
 198     private void saveLocked() {
 199         var storage = new Storage(sharedSenderKeys.entrySet().stream().flatMap(pair -> {
 200             final var sharedWith = pair.getValue();
 201             return sharedWith.stream()
 202                     .map(entry -> new Storage.SharedSenderKey(entry.recipientId().id(),
 203                             entry.deviceId(),
 204                             pair.getKey().toString()));
 205         }).toList());
 206
 207         // Write to memory first to prevent corrupting the file in case of serialization errors
 208         try (var inMemoryOutput = new ByteArrayOutputStream()) {
 209             objectMapper.writeValue(inMemoryOutput, storage);
 210
 211             var input = new ByteArrayInputStream(inMemoryOutput.toByteArray());
 212             try (var outputStream = new FileOutputStream(file)) {
 213                 input.transferTo(outputStream);
 214             }
 215         } catch (Exception e) {
 216             logger.error("Error saving shared sender key store file: {}", e.getMessage());
 217         }
 218     }
 219
 220     private record Storage(List<SharedSenderKey> sharedSenderKeys) {
 221
 222         private record SharedSenderKey(long recipientId, int deviceId, String distributionId) {}
 223     }
 224
 225     private record SenderKeySharedEntry(RecipientId recipientId, int deviceId) {}
 226 }