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