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