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