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

   1 package org.asamk.signal.manager.storage.senderKeys;
   2
   3 import org.asamk.signal.manager.storage.recipients.RecipientId;
   4 import org.asamk.signal.manager.storage.recipients.RecipientResolver;
   5 import org.asamk.signal.manager.util.IOUtils;
   6 import org.slf4j.Logger;
   7 import org.slf4j.LoggerFactory;
   8 import org.whispersystems.libsignal.SignalProtocolAddress;
   9 import org.whispersystems.libsignal.groups.state.SenderKeyRecord;
  10
  11 import java.io.File;
  12 import java.io.FileInputStream;
  13 import java.io.FileOutputStream;
  14 import java.io.IOException;
  15 import java.nio.file.Files;
  16 import java.util.Arrays;
  17 import java.util.HashMap;
  18 import java.util.List;
  19 import java.util.Map;
  20 import java.util.Objects;
  21 import java.util.UUID;
  22 import java.util.regex.Matcher;
  23 import java.util.regex.Pattern;
  24
  25 public class SenderKeyRecordStore implements org.whispersystems.libsignal.groups.state.SenderKeyStore {
  26
  27     private final static Logger logger = LoggerFactory.getLogger(SenderKeyRecordStore.class);
  28
  29     private final Map<Key, SenderKeyRecord> cachedSenderKeys = new HashMap<>();
  30
  31     private final File senderKeysPath;
  32
  33     private final RecipientResolver resolver;
  34
  35     public SenderKeyRecordStore(
  36             final File senderKeysPath, final RecipientResolver resolver
  37     ) {
  38         this.senderKeysPath = senderKeysPath;
  39         this.resolver = resolver;
  40     }
  41
  42     @Override
  43     public SenderKeyRecord loadSenderKey(final SignalProtocolAddress address, final UUID distributionId) {
  44         final var key = getKey(address, distributionId);
  45
  46         synchronized (cachedSenderKeys) {
  47             return loadSenderKeyLocked(key);
  48         }
  49     }
  50
  51     @Override
  52     public void storeSenderKey(
  53             final SignalProtocolAddress address, final UUID distributionId, final SenderKeyRecord record
  54     ) {
  55         final var key = getKey(address, distributionId);
  56
  57         synchronized (cachedSenderKeys) {
  58             storeSenderKeyLocked(key, record);
  59         }
  60     }
  61
  62     public void deleteAll() {
  63         synchronized (cachedSenderKeys) {
  64             cachedSenderKeys.clear();
  65             final var files = senderKeysPath.listFiles((_file, s) -> senderKeyFileNamePattern.matcher(s).matches());
  66             if (files == null) {
  67                 return;
  68             }
  69
  70             for (final var file : files) {
  71                 try {
  72                     Files.delete(file.toPath());
  73                 } catch (IOException e) {
  74                     logger.error("Failed to delete sender key file {}: {}", file, e.getMessage());
  75                 }
  76             }
  77         }
  78     }
  79
  80     public void deleteAllFor(final RecipientId recipientId) {
  81         synchronized (cachedSenderKeys) {
  82             cachedSenderKeys.clear();
  83             final var keys = getKeysLocked(recipientId);
  84             for (var key : keys) {
  85                 deleteSenderKeyLocked(key);
  86             }
  87         }
  88     }
  89
  90     public void mergeRecipients(RecipientId recipientId, RecipientId toBeMergedRecipientId) {
  91         synchronized (cachedSenderKeys) {
  92             final var keys = getKeysLocked(toBeMergedRecipientId);
  93             final var otherHasSenderKeys = keys.size() > 0;
  94             if (!otherHasSenderKeys) {
  95                 return;
  96             }
  97
  98             logger.debug("To be merged recipient had sender keys, re-assigning to the new recipient.");
  99             for (var key : keys) {
 100                 final var toBeMergedSenderKey = loadSenderKeyLocked(key);
 101                 deleteSenderKeyLocked(key);
 102                 if (toBeMergedSenderKey == null) {
 103                     continue;
 104                 }
 105
 106                 final var newKey = new Key(recipientId, key.deviceId(), key.distributionId);
 107                 final var senderKeyRecord = loadSenderKeyLocked(newKey);
 108                 if (senderKeyRecord != null) {
 109                     continue;
 110                 }
 111                 storeSenderKeyLocked(newKey, toBeMergedSenderKey);
 112             }
 113         }
 114     }
 115
 116     /**
 117      * @param identifier can be either a serialized uuid or a e164 phone number
 118      */
 119     private RecipientId resolveRecipient(String identifier) {
 120         return resolver.resolveRecipient(identifier);
 121     }
 122
 123     private Key getKey(final SignalProtocolAddress address, final UUID distributionId) {
 124         final var recipientId = resolveRecipient(address.getName());
 125         return new Key(recipientId, address.getDeviceId(), distributionId);
 126     }
 127
 128     private List<Key> getKeysLocked(RecipientId recipientId) {
 129         final var files = senderKeysPath.listFiles((_file, s) -> s.startsWith(recipientId.id() + "_"));
 130         if (files == null) {
 131             return List.of();
 132         }
 133         return parseFileNames(files);
 134     }
 135
 136     final Pattern senderKeyFileNamePattern = Pattern.compile("([0-9]+)_([0-9]+)_([0-9a-z\\-]+)");
 137
 138     private List<Key> parseFileNames(final File[] files) {
 139         return Arrays.stream(files)
 140                 .map(f -> senderKeyFileNamePattern.matcher(f.getName()))
 141                 .filter(Matcher::matches)
 142                 .map(matcher -> {
 143                     final var recipientId = resolver.resolveRecipient(Long.parseLong(matcher.group(1)));
 144                     if (recipientId == null) {
 145                         return null;
 146                     }
 147                     return new Key(recipientId, Integer.parseInt(matcher.group(2)), UUID.fromString(matcher.group(3)));
 148                 })
 149                 .filter(Objects::nonNull)
 150                 .toList();
 151     }
 152
 153     private File getSenderKeyFile(Key key) {
 154         try {
 155             IOUtils.createPrivateDirectories(senderKeysPath);
 156         } catch (IOException e) {
 157             throw new AssertionError("Failed to create sender keys path", e);
 158         }
 159         return new File(senderKeysPath,
 160                 key.recipientId().id() + "_" + key.deviceId() + "_" + key.distributionId.toString());
 161     }
 162
 163     private SenderKeyRecord loadSenderKeyLocked(final Key key) {
 164         {
 165             final var senderKeyRecord = cachedSenderKeys.get(key);
 166             if (senderKeyRecord != null) {
 167                 return senderKeyRecord;
 168             }
 169         }
 170
 171         final var file = getSenderKeyFile(key);
 172         if (!file.exists()) {
 173             return null;
 174         }
 175         try (var inputStream = new FileInputStream(file)) {
 176             final var senderKeyRecord = new SenderKeyRecord(inputStream.readAllBytes());
 177             cachedSenderKeys.put(key, senderKeyRecord);
 178             return senderKeyRecord;
 179         } catch (IOException e) {
 180             logger.warn("Failed to load sender key, resetting sender key: {}", e.getMessage());
 181             return null;
 182         }
 183     }
 184
 185     private void storeSenderKeyLocked(final Key key, final SenderKeyRecord senderKeyRecord) {
 186         cachedSenderKeys.put(key, senderKeyRecord);
 187
 188         final var file = getSenderKeyFile(key);
 189         try {
 190             try (var outputStream = new FileOutputStream(file)) {
 191                 outputStream.write(senderKeyRecord.serialize());
 192             }
 193         } catch (IOException e) {
 194             logger.warn("Failed to store sender key, trying to delete file and retry: {}", e.getMessage());
 195             try {
 196                 Files.delete(file.toPath());
 197                 try (var outputStream = new FileOutputStream(file)) {
 198                     outputStream.write(senderKeyRecord.serialize());
 199                 }
 200             } catch (IOException e2) {
 201                 logger.error("Failed to store sender key file {}: {}", file, e2.getMessage());
 202             }
 203         }
 204     }
 205
 206     private void deleteSenderKeyLocked(final Key key) {
 207         cachedSenderKeys.remove(key);
 208
 209         final var file = getSenderKeyFile(key);
 210         if (!file.exists()) {
 211             return;
 212         }
 213         try {
 214             Files.delete(file.toPath());
 215         } catch (IOException e) {
 216             logger.error("Failed to delete sender key file {}: {}", file, e.getMessage());
 217         }
 218     }
 219
 220     private record Key(RecipientId recipientId, int deviceId, UUID distributionId) {
 221
 222     }
 223 }