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

   1 package org.asamk.signal.manager.storage.groups;
   2
   3 import org.asamk.signal.manager.api.GroupId;
   4 import org.asamk.signal.manager.api.GroupIdV1;
   5 import org.asamk.signal.manager.api.GroupIdV2;
   6 import org.asamk.signal.manager.groups.GroupUtils;
   7 import org.asamk.signal.manager.storage.Database;
   8 import org.asamk.signal.manager.storage.Utils;
   9 import org.asamk.signal.manager.storage.recipients.RecipientId;
  10 import org.asamk.signal.manager.storage.recipients.RecipientIdCreator;
  11 import org.asamk.signal.manager.storage.recipients.RecipientResolver;
  12 import org.signal.libsignal.zkgroup.InvalidInputException;
  13 import org.signal.libsignal.zkgroup.groups.GroupMasterKey;
  14 import org.signal.storageservice.protos.groups.local.DecryptedGroup;
  15 import org.slf4j.Logger;
  16 import org.slf4j.LoggerFactory;
  17 import org.whispersystems.signalservice.api.push.DistributionId;
  18 import org.whispersystems.signalservice.api.util.UuidUtil;
  19
  20 import java.io.IOException;
  21 import java.sql.Connection;
  22 import java.sql.ResultSet;
  23 import java.sql.SQLException;
  24 import java.sql.Types;
  25 import java.util.Arrays;
  26 import java.util.Collection;
  27 import java.util.List;
  28 import java.util.Objects;
  29 import java.util.Set;
  30 import java.util.stream.Collectors;
  31 import java.util.stream.Stream;
  32
  33 public class GroupStore {
  34
  35     private final static Logger logger = LoggerFactory.getLogger(GroupStore.class);
  36     private static final String TABLE_GROUP_V2 = "group_v2";
  37     private static final String TABLE_GROUP_V1 = "group_v1";
  38     private static final String TABLE_GROUP_V1_MEMBER = "group_v1_member";
  39
  40     private final Database database;
  41     private final RecipientResolver recipientResolver;
  42     private final RecipientIdCreator recipientIdCreator;
  43
  44     public static void createSql(Connection connection) throws SQLException {
  45         // When modifying the CREATE statement here, also add a migration in AccountDatabase.java
  46         try (final var statement = connection.createStatement()) {
  47             statement.executeUpdate("""
  48                                     CREATE TABLE group_v2 (
  49                                       _id INTEGER PRIMARY KEY,
  50                                       group_id BLOB UNIQUE NOT NULL,
  51                                       master_key BLOB NOT NULL,
  52                                       group_data BLOB,
  53                                       distribution_id BLOB UNIQUE NOT NULL,
  54                                       blocked INTEGER NOT NULL DEFAULT FALSE,
  55                                       permission_denied INTEGER NOT NULL DEFAULT FALSE
  56                                     ) STRICT;
  57                                     CREATE TABLE group_v1 (
  58                                       _id INTEGER PRIMARY KEY,
  59                                       group_id BLOB UNIQUE NOT NULL,
  60                                       group_id_v2 BLOB UNIQUE,
  61                                       name TEXT,
  62                                       color TEXT,
  63                                       expiration_time INTEGER NOT NULL DEFAULT 0,
  64                                       blocked INTEGER NOT NULL DEFAULT FALSE,
  65                                       archived INTEGER NOT NULL DEFAULT FALSE
  66                                     ) STRICT;
  67                                     CREATE TABLE group_v1_member (
  68                                       _id INTEGER PRIMARY KEY,
  69                                       group_id INTEGER NOT NULL REFERENCES group_v1 (_id) ON DELETE CASCADE,
  70                                       recipient_id INTEGER NOT NULL REFERENCES recipient (_id) ON DELETE CASCADE,
  71                                       UNIQUE(group_id, recipient_id)
  72                                     ) STRICT;
  73                                     """);
  74         }
  75     }
  76
  77     public GroupStore(
  78             final Database database,
  79             final RecipientResolver recipientResolver,
  80             final RecipientIdCreator recipientIdCreator
  81     ) {
  82         this.database = database;
  83         this.recipientResolver = recipientResolver;
  84         this.recipientIdCreator = recipientIdCreator;
  85     }
  86
  87     public void updateGroup(GroupInfo group) {
  88         try (final var connection = database.getConnection()) {
  89             connection.setAutoCommit(false);
  90             final Long internalId;
  91             final var sql = (
  92                     """
  93                     SELECT g._id
  94                     FROM %s g
  95                     WHERE g.group_id = ?
  96                     """
  97             ).formatted(group instanceof GroupInfoV1 ? TABLE_GROUP_V1 : TABLE_GROUP_V2);
  98             try (final var statement = connection.prepareStatement(sql)) {
  99                 statement.setBytes(1, group.getGroupId().serialize());
 100                 internalId = Utils.executeQueryForOptional(statement, res -> res.getLong("_id")).orElse(null);
 101             }
 102             insertOrReplaceGroup(connection, internalId, group);
 103             connection.commit();
 104         } catch (SQLException e) {
 105             throw new RuntimeException("Failed update recipient store", e);
 106         }
 107     }
 108
 109     public void deleteGroup(GroupId groupId) {
 110         if (groupId instanceof GroupIdV1 groupIdV1) {
 111             deleteGroup(groupIdV1);
 112         } else if (groupId instanceof GroupIdV2 groupIdV2) {
 113             deleteGroup(groupIdV2);
 114         }
 115     }
 116
 117     public void deleteGroup(GroupIdV1 groupIdV1) {
 118         final var sql = (
 119                 """
 120                 DELETE FROM %s
 121                 WHERE group_id = ?
 122                 """
 123         ).formatted(TABLE_GROUP_V1);
 124         try (final var connection = database.getConnection()) {
 125             try (final var statement = connection.prepareStatement(sql)) {
 126                 statement.setBytes(1, groupIdV1.serialize());
 127                 statement.executeUpdate();
 128             }
 129         } catch (SQLException e) {
 130             throw new RuntimeException("Failed update group store", e);
 131         }
 132     }
 133
 134     public void deleteGroup(GroupIdV2 groupIdV2) {
 135         final var sql = (
 136                 """
 137                 DELETE FROM %s
 138                 WHERE group_id = ?
 139                 """
 140         ).formatted(TABLE_GROUP_V2);
 141         try (final var connection = database.getConnection()) {
 142             try (final var statement = connection.prepareStatement(sql)) {
 143                 statement.setBytes(1, groupIdV2.serialize());
 144                 statement.executeUpdate();
 145             }
 146         } catch (SQLException e) {
 147             throw new RuntimeException("Failed update group store", e);
 148         }
 149     }
 150
 151     public GroupInfo getGroup(GroupId groupId) {
 152         try (final var connection = database.getConnection()) {
 153             if (groupId instanceof GroupIdV1 groupIdV1) {
 154                 final var group = getGroup(connection, groupIdV1);
 155                 if (group != null) {
 156                     return group;
 157                 }
 158                 return getGroupV2ByV1Id(connection, groupIdV1);
 159             } else if (groupId instanceof GroupIdV2 groupIdV2) {
 160                 final var group = getGroup(connection, groupIdV2);
 161                 if (group != null) {
 162                     return group;
 163                 }
 164                 return getGroupV1ByV2Id(connection, groupIdV2);
 165             }
 166         } catch (SQLException e) {
 167             throw new RuntimeException("Failed read from group store", e);
 168         }
 169         throw new AssertionError("Invalid group id type");
 170     }
 171
 172     public GroupInfoV1 getOrCreateGroupV1(GroupIdV1 groupId) {
 173         try (final var connection = database.getConnection()) {
 174             var group = getGroup(connection, groupId);
 175
 176             if (group != null) {
 177                 return group;
 178             }
 179
 180             if (getGroupV2ByV1Id(connection, groupId) == null) {
 181                 return new GroupInfoV1(groupId);
 182             }
 183
 184             return null;
 185         } catch (SQLException e) {
 186             throw new RuntimeException("Failed read from group store", e);
 187         }
 188     }
 189
 190     public List<GroupInfo> getGroups() {
 191         return Stream.concat(getGroupsV2().stream(), getGroupsV1().stream()).toList();
 192     }
 193
 194     public void mergeRecipients(
 195             final Connection connection, final RecipientId recipientId, final RecipientId toBeMergedRecipientId
 196     ) throws SQLException {
 197         final var sql = (
 198                 """
 199                 UPDATE OR REPLACE %s
 200                 SET recipient_id = ?
 201                 WHERE recipient_id = ?
 202                 """
 203         ).formatted(TABLE_GROUP_V1_MEMBER);
 204         try (final var statement = connection.prepareStatement(sql)) {
 205             statement.setLong(1, recipientId.id());
 206             statement.setLong(2, toBeMergedRecipientId.id());
 207             final var updatedRows = statement.executeUpdate();
 208             if (updatedRows > 0) {
 209                 logger.info("Updated {} group members when merging recipients", updatedRows);
 210             }
 211         }
 212     }
 213
 214     void addLegacyGroups(final Collection<GroupInfo> groups) {
 215         logger.debug("Migrating legacy groups to database");
 216         long start = System.nanoTime();
 217         try (final var connection = database.getConnection()) {
 218             connection.setAutoCommit(false);
 219             for (final var group : groups) {
 220                 insertOrReplaceGroup(connection, null, group);
 221             }
 222             connection.commit();
 223         } catch (SQLException e) {
 224             throw new RuntimeException("Failed update group store", e);
 225         }
 226         logger.debug("Complete groups migration took {}ms", (System.nanoTime() - start) / 1000000);
 227     }
 228
 229     private void insertOrReplaceGroup(
 230             final Connection connection, Long internalId, final GroupInfo group
 231     ) throws SQLException {
 232         if (group instanceof GroupInfoV1 groupV1) {
 233             if (internalId != null) {
 234                 final var sqlDeleteMembers = "DELETE FROM %s where group_id = ?".formatted(TABLE_GROUP_V1_MEMBER);
 235                 try (final var statement = connection.prepareStatement(sqlDeleteMembers)) {
 236                     statement.setLong(1, internalId);
 237                     statement.executeUpdate();
 238                 }
 239             }
 240             final var sql = """
 241                             INSERT OR REPLACE INTO %s (_id, group_id, group_id_v2, name, color, expiration_time, blocked, archived)
 242                             VALUES (?, ?, ?, ?, ?, ?, ?, ?)
 243                             RETURNING _id
 244                             """.formatted(TABLE_GROUP_V1);
 245             try (final var statement = connection.prepareStatement(sql)) {
 246                 if (internalId == null) {
 247                     statement.setNull(1, Types.NUMERIC);
 248                 } else {
 249                     statement.setLong(1, internalId);
 250                 }
 251                 statement.setBytes(2, groupV1.getGroupId().serialize());
 252                 statement.setBytes(3, groupV1.getExpectedV2Id().serialize());
 253                 statement.setString(4, groupV1.getTitle());
 254                 statement.setString(5, groupV1.color);
 255                 statement.setLong(6, groupV1.getMessageExpirationTimer());
 256                 statement.setBoolean(7, groupV1.isBlocked());
 257                 statement.setBoolean(8, groupV1.archived);
 258                 final var generatedKey = Utils.executeQueryForOptional(statement, Utils::getIdMapper);
 259
 260                 if (internalId == null) {
 261                     if (generatedKey.isPresent()) {
 262                         internalId = generatedKey.get();
 263                     } else {
 264                         throw new RuntimeException("Failed to add new group to database");
 265                     }
 266                 }
 267             }
 268             final var sqlInsertMember = """
 269                                         INSERT OR REPLACE INTO %s (group_id, recipient_id)
 270                                         VALUES (?, ?)
 271                                         """.formatted(TABLE_GROUP_V1_MEMBER);
 272             try (final var statement = connection.prepareStatement(sqlInsertMember)) {
 273                 for (final var recipient : groupV1.getMembers()) {
 274                     statement.setLong(1, internalId);
 275                     statement.setLong(2, recipient.id());
 276                     statement.executeUpdate();
 277                 }
 278             }
 279         } else if (group instanceof GroupInfoV2 groupV2) {
 280             final var sql = (
 281                     """
 282                     INSERT OR REPLACE INTO %s (_id, group_id, master_key, group_data, distribution_id, blocked, distribution_id)
 283                     VALUES (?, ?, ?, ?, ?, ?, ?)
 284                     """
 285             ).formatted(TABLE_GROUP_V2);
 286             try (final var statement = connection.prepareStatement(sql)) {
 287                 if (internalId == null) {
 288                     statement.setNull(1, Types.NUMERIC);
 289                 } else {
 290                     statement.setLong(1, internalId);
 291                 }
 292                 statement.setBytes(2, groupV2.getGroupId().serialize());
 293                 statement.setBytes(3, groupV2.getMasterKey().serialize());
 294                 if (groupV2.getGroup() == null) {
 295                     statement.setNull(4, Types.NUMERIC);
 296                 } else {
 297                     statement.setBytes(4, groupV2.getGroup().encode());
 298                 }
 299                 statement.setBytes(5, UuidUtil.toByteArray(groupV2.getDistributionId().asUuid()));
 300                 statement.setBoolean(6, groupV2.isBlocked());
 301                 statement.setBoolean(7, groupV2.isPermissionDenied());
 302                 statement.executeUpdate();
 303             }
 304         } else {
 305             throw new AssertionError("Invalid group id type");
 306         }
 307     }
 308
 309     private List<GroupInfoV2> getGroupsV2() {
 310         final var sql = (
 311                 """
 312                 SELECT g.group_id, g.master_key, g.group_data, g.distribution_id, g.blocked, g.permission_denied
 313                 FROM %s g
 314                 """
 315         ).formatted(TABLE_GROUP_V2);
 316         try (final var connection = database.getConnection()) {
 317             try (final var statement = connection.prepareStatement(sql)) {
 318                 return Utils.executeQueryForStream(statement, this::getGroupInfoV2FromResultSet)
 319                         .filter(Objects::nonNull)
 320                         .toList();
 321             }
 322         } catch (SQLException e) {
 323             throw new RuntimeException("Failed read from group store", e);
 324         }
 325     }
 326
 327     private GroupInfoV2 getGroup(Connection connection, GroupIdV2 groupIdV2) throws SQLException {
 328         final var sql = (
 329                 """
 330                 SELECT g.group_id, g.master_key, g.group_data, g.distribution_id, g.blocked, g.permission_denied
 331                 FROM %s g
 332                 WHERE g.group_id = ?
 333                 """
 334         ).formatted(TABLE_GROUP_V2);
 335         try (final var statement = connection.prepareStatement(sql)) {
 336             statement.setBytes(1, groupIdV2.serialize());
 337             return Utils.executeQueryForOptional(statement, this::getGroupInfoV2FromResultSet).orElse(null);
 338         }
 339     }
 340
 341     private GroupInfoV2 getGroupInfoV2FromResultSet(ResultSet resultSet) throws SQLException {
 342         try {
 343             final var groupId = resultSet.getBytes("group_id");
 344             final var masterKey = resultSet.getBytes("master_key");
 345             final var groupData = resultSet.getBytes("group_data");
 346             final var distributionId = resultSet.getBytes("distribution_id");
 347             final var blocked = resultSet.getBoolean("blocked");
 348             final var permissionDenied = resultSet.getBoolean("permission_denied");
 349             return new GroupInfoV2(GroupId.v2(groupId),
 350                     new GroupMasterKey(masterKey),
 351                     groupData == null ? null : DecryptedGroup.ADAPTER.decode(groupData),
 352                     DistributionId.from(UuidUtil.parseOrThrow(distributionId)),
 353                     blocked,
 354                     permissionDenied,
 355                     recipientResolver);
 356         } catch (InvalidInputException | IOException e) {
 357             return null;
 358         }
 359     }
 360
 361     private List<GroupInfoV1> getGroupsV1() {
 362         final var sql = (
 363                 """
 364                 SELECT g.group_id, g.group_id_v2, g.name, g.color, (select group_concat(gm.recipient_id) from %s gm where gm.group_id = g._id) as members, g.expiration_time, g.blocked, g.archived
 365                 FROM %s g
 366                 """
 367         ).formatted(TABLE_GROUP_V1_MEMBER, TABLE_GROUP_V1);
 368         try (final var connection = database.getConnection()) {
 369             try (final var statement = connection.prepareStatement(sql)) {
 370                 return Utils.executeQueryForStream(statement, this::getGroupInfoV1FromResultSet)
 371                         .filter(Objects::nonNull)
 372                         .toList();
 373             }
 374         } catch (SQLException e) {
 375             throw new RuntimeException("Failed read from group store", e);
 376         }
 377     }
 378
 379     private GroupInfoV1 getGroup(Connection connection, GroupIdV1 groupIdV1) throws SQLException {
 380         final var sql = (
 381                 """
 382                 SELECT g.group_id, g.group_id_v2, g.name, g.color, (select group_concat(gm.recipient_id) from %s gm where gm.group_id = g._id) as members, g.expiration_time, g.blocked, g.archived
 383                 FROM %s g
 384                 WHERE g.group_id = ?
 385                 """
 386         ).formatted(TABLE_GROUP_V1_MEMBER, TABLE_GROUP_V1);
 387         try (final var statement = connection.prepareStatement(sql)) {
 388             statement.setBytes(1, groupIdV1.serialize());
 389             return Utils.executeQueryForOptional(statement, this::getGroupInfoV1FromResultSet).orElse(null);
 390         }
 391     }
 392
 393     private GroupInfoV1 getGroupInfoV1FromResultSet(ResultSet resultSet) throws SQLException {
 394         final var groupId = resultSet.getBytes("group_id");
 395         final var groupIdV2 = resultSet.getBytes("group_id_v2");
 396         final var name = resultSet.getString("name");
 397         final var color = resultSet.getString("color");
 398         final var membersString = resultSet.getString("members");
 399         final var members = membersString == null
 400                 ? Set.<RecipientId>of()
 401                 : Arrays.stream(membersString.split(","))
 402                         .map(Integer::valueOf)
 403                         .map(recipientIdCreator::create)
 404                         .collect(Collectors.toSet());
 405         final var expirationTime = resultSet.getInt("expiration_time");
 406         final var blocked = resultSet.getBoolean("blocked");
 407         final var archived = resultSet.getBoolean("archived");
 408         return new GroupInfoV1(GroupId.v1(groupId),
 409                 groupIdV2 == null ? null : GroupId.v2(groupIdV2),
 410                 name,
 411                 members,
 412                 color,
 413                 expirationTime,
 414                 blocked,
 415                 archived);
 416     }
 417
 418     private GroupInfoV2 getGroupV2ByV1Id(final Connection connection, final GroupIdV1 groupId) throws SQLException {
 419         return getGroup(connection, GroupUtils.getGroupIdV2(groupId));
 420     }
 421
 422     private GroupInfoV1 getGroupV1ByV2Id(Connection connection, GroupIdV2 groupIdV2) throws SQLException {
 423         final var sql = (
 424                 """
 425                 SELECT g.group_id, g.group_id_v2, g.name, g.color, (select group_concat(gm.recipient_id) from %s gm where gm.group_id = g._id) as members, g.expiration_time, g.blocked, g.archived
 426                 FROM %s g
 427                 WHERE g.group_id_v2 = ?
 428                 """
 429         ).formatted(TABLE_GROUP_V1_MEMBER, TABLE_GROUP_V1);
 430         try (final var statement = connection.prepareStatement(sql)) {
 431             statement.setBytes(1, groupIdV2.serialize());
 432             return Utils.executeQueryForOptional(statement, this::getGroupInfoV1FromResultSet).orElse(null);
 433         }
 434     }
 435 }