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

   1 package org.asamk.signal.manager.storage.sendLog;
   2
   3 import org.asamk.signal.manager.groups.GroupId;
   4 import org.asamk.signal.manager.groups.GroupUtils;
   5 import org.asamk.signal.manager.storage.Database;
   6 import org.asamk.signal.manager.storage.recipients.RecipientId;
   7 import org.asamk.signal.manager.storage.recipients.RecipientResolver;
   8 import org.signal.libsignal.zkgroup.InvalidInputException;
   9 import org.signal.libsignal.zkgroup.groups.GroupMasterKey;
  10 import org.slf4j.Logger;
  11 import org.slf4j.LoggerFactory;
  12 import org.whispersystems.signalservice.api.crypto.ContentHint;
  13 import org.whispersystems.signalservice.api.messages.SendMessageResult;
  14 import org.whispersystems.signalservice.internal.push.SignalServiceProtos;
  15
  16 import java.io.IOException;
  17 import java.sql.Connection;
  18 import java.sql.PreparedStatement;
  19 import java.sql.ResultSet;
  20 import java.sql.SQLException;
  21 import java.time.Duration;
  22 import java.util.List;
  23 import java.util.Objects;
  24 import java.util.Optional;
  25 import java.util.Spliterator;
  26 import java.util.Spliterators;
  27 import java.util.function.Consumer;
  28 import java.util.stream.Stream;
  29 import java.util.stream.StreamSupport;
  30
  31 public class MessageSendLogStore implements AutoCloseable {
  32
  33     private static final Logger logger = LoggerFactory.getLogger(MessageSendLogStore.class);
  34
  35     private static final String TABLE_MESSAGE_SEND_LOG = "message_send_log";
  36     private static final String TABLE_MESSAGE_SEND_LOG_CONTENT = "message_send_log_content";
  37
  38     private static final Duration LOG_DURATION = Duration.ofDays(1);
  39
  40     private final RecipientResolver recipientResolver;
  41     private final Database database;
  42     private final Thread cleanupThread;
  43
  44     public MessageSendLogStore(
  45             final RecipientResolver recipientResolver, final Database database
  46     ) {
  47         this.recipientResolver = recipientResolver;
  48         this.database = database;
  49         this.cleanupThread = new Thread(() -> {
  50             try {
  51                 final var interval = Duration.ofHours(1).toMillis();
  52                 while (!Thread.interrupted()) {
  53                     try (final var connection = database.getConnection()) {
  54                         deleteOutdatedEntries(connection);
  55                     } catch (SQLException e) {
  56                         logger.warn("Deleting outdated entries failed");
  57                         break;
  58                     }
  59                     Thread.sleep(interval);
  60                 }
  61             } catch (InterruptedException e) {
  62                 logger.debug("Stopping msl cleanup thread");
  63             }
  64         });
  65         cleanupThread.setName("msl-cleanup");
  66         cleanupThread.setDaemon(true);
  67         cleanupThread.start();
  68     }
  69
  70     public static void createSql(Connection connection) throws SQLException {
  71         try (final var statement = connection.createStatement()) {
  72             statement.executeUpdate("""
  73                                     CREATE TABLE message_send_log (
  74                                       _id INTEGER PRIMARY KEY,
  75                                       content_id INTEGER NOT NULL REFERENCES message_send_log_content (_id) ON DELETE CASCADE,
  76                                       recipient_id INTEGER NOT NULL,
  77                                       device_id INTEGER NOT NULL
  78                                     );
  79                                     CREATE TABLE message_send_log_content (
  80                                       _id INTEGER PRIMARY KEY,
  81                                       group_id BLOB,
  82                                       timestamp INTEGER NOT NULL,
  83                                       content BLOB NOT NULL,
  84                                       content_hint INTEGER NOT NULL
  85                                     );
  86                                     CREATE INDEX mslc_timestamp_index ON message_send_log_content (timestamp);
  87                                     CREATE INDEX msl_recipient_index ON message_send_log (recipient_id, device_id, content_id);
  88                                     CREATE INDEX msl_content_index ON message_send_log (content_id);
  89                                     """);
  90         }
  91     }
  92
  93     public List<MessageSendLogEntry> findMessages(
  94             final RecipientId recipientId, final int deviceId, final long timestamp, final boolean isSenderKey
  95     ) {
  96         final var sql = """
  97                         SELECT group_id, content, content_hint
  98                         FROM %s l
  99                              INNER JOIN %s lc ON l.content_id = lc._id
 100                         WHERE l.recipient_id = ? AND l.device_id = ? AND lc.timestamp = ?
 101                         """.formatted(TABLE_MESSAGE_SEND_LOG, TABLE_MESSAGE_SEND_LOG_CONTENT);
 102         try (final var connection = database.getConnection()) {
 103             deleteOutdatedEntries(connection);
 104
 105             try (final var statement = connection.prepareStatement(sql)) {
 106                 statement.setLong(1, recipientId.id());
 107                 statement.setInt(2, deviceId);
 108                 statement.setLong(3, timestamp);
 109                 try (var result = executeQueryForStream(statement, resultSet -> {
 110                     final var groupId = Optional.ofNullable(resultSet.getBytes("group_id"))
 111                             .map(GroupId::unknownVersion);
 112                     final SignalServiceProtos.Content content;
 113                     try {
 114                         content = SignalServiceProtos.Content.parseFrom(resultSet.getBinaryStream("content"));
 115                     } catch (IOException e) {
 116                         logger.warn("Failed to parse content from message send log", e);
 117                         return null;
 118                     }
 119                     final var contentHint = ContentHint.fromType(resultSet.getInt("content_hint"));
 120                     return new MessageSendLogEntry(groupId, content, contentHint);
 121                 })) {
 122                     return result.filter(Objects::nonNull)
 123                             .filter(e -> !isSenderKey || e.groupId().isPresent())
 124                             .toList();
 125                 }
 126             }
 127         } catch (SQLException e) {
 128             logger.warn("Failed read from message send log", e);
 129             return List.of();
 130         }
 131     }
 132
 133     public long insertIfPossible(
 134             long sentTimestamp, SendMessageResult sendMessageResult, ContentHint contentHint
 135     ) {
 136         final RecipientDevices recipientDevice = getRecipientDevices(sendMessageResult);
 137         if (recipientDevice == null) {
 138             return -1;
 139         }
 140
 141         return insert(List.of(recipientDevice),
 142                 sentTimestamp,
 143                 sendMessageResult.getSuccess().getContent().get(),
 144                 contentHint);
 145     }
 146
 147     public long insertIfPossible(
 148             long sentTimestamp, List<SendMessageResult> sendMessageResults, ContentHint contentHint
 149     ) {
 150         final var recipientDevices = sendMessageResults.stream()
 151                 .map(this::getRecipientDevices)
 152                 .filter(Objects::nonNull)
 153                 .toList();
 154         if (recipientDevices.isEmpty()) {
 155             return -1;
 156         }
 157
 158         final var content = sendMessageResults.stream()
 159                 .filter(r -> r.isSuccess() && r.getSuccess().getContent().isPresent())
 160                 .map(r -> r.getSuccess().getContent().get())
 161                 .findFirst()
 162                 .get();
 163
 164         return insert(recipientDevices, sentTimestamp, content, contentHint);
 165     }
 166
 167     public void addRecipientToExistingEntryIfPossible(final long contentId, final SendMessageResult sendMessageResult) {
 168         final RecipientDevices recipientDevice = getRecipientDevices(sendMessageResult);
 169         if (recipientDevice == null) {
 170             return;
 171         }
 172
 173         insertRecipientsForExistingContent(contentId, List.of(recipientDevice));
 174     }
 175
 176     public void addRecipientToExistingEntryIfPossible(
 177             final long contentId, final List<SendMessageResult> sendMessageResults
 178     ) {
 179         final var recipientDevices = sendMessageResults.stream()
 180                 .map(this::getRecipientDevices)
 181                 .filter(Objects::nonNull)
 182                 .toList();
 183         if (recipientDevices.isEmpty()) {
 184             return;
 185         }
 186
 187         insertRecipientsForExistingContent(contentId, recipientDevices);
 188     }
 189
 190     public void deleteEntryForGroup(long sentTimestamp, GroupId groupId) {
 191         final var sql = """
 192                         DELETE FROM %s AS lc
 193                         WHERE lc.timestamp = ? AND lc.group_id = ?
 194                         """.formatted(TABLE_MESSAGE_SEND_LOG_CONTENT);
 195         try (final var connection = database.getConnection()) {
 196             try (final var statement = connection.prepareStatement(sql)) {
 197                 statement.setLong(1, sentTimestamp);
 198                 statement.setBytes(2, groupId.serialize());
 199                 statement.executeUpdate();
 200             }
 201         } catch (SQLException e) {
 202             logger.warn("Failed delete from message send log", e);
 203         }
 204     }
 205
 206     public void deleteEntryForRecipientNonGroup(long sentTimestamp, RecipientId recipientId) {
 207         final var sql = """
 208                         DELETE FROM %s AS lc
 209                         WHERE lc.timestamp = ? AND lc.group_id IS NULL AND lc._id IN (SELECT content_id FROM %s l WHERE l.recipient_id = ?)
 210                         """.formatted(TABLE_MESSAGE_SEND_LOG_CONTENT, TABLE_MESSAGE_SEND_LOG);
 211         try (final var connection = database.getConnection()) {
 212             connection.setAutoCommit(false);
 213             try (final var statement = connection.prepareStatement(sql)) {
 214                 statement.setLong(1, sentTimestamp);
 215                 statement.setLong(2, recipientId.id());
 216                 statement.executeUpdate();
 217             }
 218
 219             deleteOrphanedLogContents(connection);
 220             connection.commit();
 221         } catch (SQLException e) {
 222             logger.warn("Failed delete from message send log", e);
 223         }
 224     }
 225
 226     public void deleteEntryForRecipient(long sentTimestamp, RecipientId recipientId, int deviceId) {
 227         deleteEntriesForRecipient(List.of(sentTimestamp), recipientId, deviceId);
 228     }
 229
 230     public void deleteEntriesForRecipient(List<Long> sentTimestamps, RecipientId recipientId, int deviceId) {
 231         final var sql = """
 232                         DELETE FROM %s AS l
 233                         WHERE l.content_id IN (SELECT _id FROM %s lc WHERE lc.timestamp = ?) AND l.recipient_id = ? AND l.device_id = ?
 234                         """.formatted(TABLE_MESSAGE_SEND_LOG, TABLE_MESSAGE_SEND_LOG_CONTENT);
 235         try (final var connection = database.getConnection()) {
 236             connection.setAutoCommit(false);
 237             try (final var statement = connection.prepareStatement(sql)) {
 238                 for (final var sentTimestamp : sentTimestamps) {
 239                     statement.setLong(1, sentTimestamp);
 240                     statement.setLong(2, recipientId.id());
 241                     statement.setInt(3, deviceId);
 242                     statement.executeUpdate();
 243                 }
 244             }
 245
 246             deleteOrphanedLogContents(connection);
 247             connection.commit();
 248         } catch (SQLException e) {
 249             logger.warn("Failed delete from message send log", e);
 250         }
 251     }
 252
 253     @Override
 254     public void close() {
 255         cleanupThread.interrupt();
 256         try {
 257             cleanupThread.join();
 258         } catch (InterruptedException ignored) {
 259         }
 260     }
 261
 262     private RecipientDevices getRecipientDevices(final SendMessageResult sendMessageResult) {
 263         if (sendMessageResult.isSuccess() && sendMessageResult.getSuccess().getContent().isPresent()) {
 264             final var recipientId = recipientResolver.resolveRecipient(sendMessageResult.getAddress());
 265             return new RecipientDevices(recipientId, sendMessageResult.getSuccess().getDevices());
 266         } else {
 267             return null;
 268         }
 269     }
 270
 271     private long insert(
 272             final List<RecipientDevices> recipientDevices,
 273             final long sentTimestamp,
 274             final SignalServiceProtos.Content content,
 275             final ContentHint contentHint
 276     ) {
 277         byte[] groupId = getGroupId(content);
 278
 279         final var sql = """
 280                         INSERT INTO %s (timestamp, group_id, content, content_hint)
 281                         VALUES (?,?,?,?)
 282                         """.formatted(TABLE_MESSAGE_SEND_LOG_CONTENT);
 283         try (final var connection = database.getConnection()) {
 284             connection.setAutoCommit(false);
 285             final long contentId;
 286             try (final var statement = connection.prepareStatement(sql)) {
 287                 statement.setLong(1, sentTimestamp);
 288                 statement.setBytes(2, groupId);
 289                 statement.setBytes(3, content.toByteArray());
 290                 statement.setInt(4, contentHint.getType());
 291                 statement.executeUpdate();
 292                 final var generatedKeys = statement.getGeneratedKeys();
 293                 if (generatedKeys.next()) {
 294                     contentId = generatedKeys.getLong(1);
 295                 } else {
 296                     contentId = -1;
 297                 }
 298             }
 299             if (contentId == -1) {
 300                 logger.warn("Failed to insert message send log content");
 301                 return -1;
 302             }
 303             insertRecipientsForExistingContent(contentId, recipientDevices, connection);
 304
 305             connection.commit();
 306             return contentId;
 307         } catch (SQLException e) {
 308             logger.warn("Failed to insert into message send log", e);
 309             return -1;
 310         }
 311     }
 312
 313     private byte[] getGroupId(final SignalServiceProtos.Content content) {
 314         try {
 315             return !content.hasDataMessage()
 316                     ? null
 317                     : content.getDataMessage().hasGroup()
 318                             ? content.getDataMessage().getGroup().getId().toByteArray()
 319                             : content.getDataMessage().hasGroupV2()
 320                                     ? GroupUtils.getGroupIdV2(new GroupMasterKey(content.getDataMessage()
 321                                     .getGroupV2()
 322                                     .getMasterKey()
 323                                     .toByteArray())).serialize()
 324                                     : null;
 325         } catch (InvalidInputException e) {
 326             logger.warn("Failed to parse groupId id from content");
 327             return null;
 328         }
 329     }
 330
 331     private void insertRecipientsForExistingContent(
 332             final long contentId, final List<RecipientDevices> recipientDevices
 333     ) {
 334         try (final var connection = database.getConnection()) {
 335             connection.setAutoCommit(false);
 336             insertRecipientsForExistingContent(contentId, recipientDevices, connection);
 337             connection.commit();
 338         } catch (SQLException e) {
 339             logger.warn("Failed to append recipients to message send log", e);
 340         }
 341     }
 342
 343     private void insertRecipientsForExistingContent(
 344             final long contentId, final List<RecipientDevices> recipientDevices, final Connection connection
 345     ) throws SQLException {
 346         final var sql = """
 347                         INSERT INTO %s (recipient_id, device_id, content_id)
 348                         VALUES (?,?,?)
 349                         """.formatted(TABLE_MESSAGE_SEND_LOG);
 350         try (final var statement = connection.prepareStatement(sql)) {
 351             for (final var recipientDevice : recipientDevices) {
 352                 for (final var deviceId : recipientDevice.deviceIds()) {
 353                     statement.setLong(1, recipientDevice.recipientId().id());
 354                     statement.setInt(2, deviceId);
 355                     statement.setLong(3, contentId);
 356                     statement.executeUpdate();
 357                 }
 358             }
 359         }
 360     }
 361
 362     private void deleteOutdatedEntries(final Connection connection) throws SQLException {
 363         final var sql = """
 364                         DELETE FROM %s
 365                         WHERE timestamp < ?
 366                         """.formatted(TABLE_MESSAGE_SEND_LOG_CONTENT);
 367         try (final var statement = connection.prepareStatement(sql)) {
 368             statement.setLong(1, System.currentTimeMillis() - LOG_DURATION.toMillis());
 369             final var rowCount = statement.executeUpdate();
 370             if (rowCount > 0) {
 371                 logger.debug("Removed {} outdated entries from the message send log", rowCount);
 372             } else {
 373                 logger.trace("No outdated entries to be removed from message send log.");
 374             }
 375         }
 376     }
 377
 378     private void deleteOrphanedLogContents(final Connection connection) throws SQLException {
 379         final var sql = """
 380                         DELETE FROM %s
 381                         WHERE _id NOT IN (SELECT content_id FROM %s)
 382                         """.formatted(TABLE_MESSAGE_SEND_LOG_CONTENT, TABLE_MESSAGE_SEND_LOG);
 383         try (final var statement = connection.prepareStatement(sql)) {
 384             statement.executeUpdate();
 385         }
 386     }
 387
 388     private <T> Stream<T> executeQueryForStream(
 389             PreparedStatement statement, ResultSetMapper<T> mapper
 390     ) throws SQLException {
 391         final var resultSet = statement.executeQuery();
 392
 393         return StreamSupport.stream(new Spliterators.AbstractSpliterator<>(Long.MAX_VALUE, Spliterator.ORDERED) {
 394             @Override
 395             public boolean tryAdvance(final Consumer<? super T> consumer) {
 396                 try {
 397                     if (!resultSet.next()) {
 398                         return false;
 399                     }
 400                     consumer.accept(mapper.apply(resultSet));
 401                     return true;
 402                 } catch (SQLException e) {
 403                     logger.warn("Failed to read from database result", e);
 404                     throw new RuntimeException(e);
 405                 }
 406             }
 407         }, false);
 408     }
 409
 410     private interface ResultSetMapper<T> {
 411
 412         T apply(ResultSet resultSet) throws SQLException;
 413     }
 414
 415     private record RecipientDevices(RecipientId recipientId, List<Integer> deviceIds) {}
 416 }