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