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 static final 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             return getGroup(connection, groupId);
 154         } catch (SQLException e) {
 155             throw new RuntimeException("Failed read from group store", e);
 156         }
 157     }
 158
 159     public GroupInfo getGroup(final Connection connection, final GroupId groupId) throws SQLException {
 160         switch (groupId) {
 161             case GroupIdV1 groupIdV1 -> {
 162                 final var group = getGroup(connection, groupIdV1);
 163                 if (group != null) {
 164                     return group;
 165                 }
 166                 return getGroupV2ByV1Id(connection, groupIdV1);
 167             }
 168             case GroupIdV2 groupIdV2 -> {
 169                 final var group = getGroup(connection, groupIdV2);
 170                 if (group != null) {
 171                     return group;
 172                 }
 173                 return getGroupV1ByV2Id(connection, groupIdV2);
 174             }
 175         }
 176     }
 177
 178     public GroupInfoV1 getOrCreateGroupV1(GroupIdV1 groupId) {
 179         try (final var connection = database.getConnection()) {
 180             var group = getGroup(connection, groupId);
 181
 182             if (group != null) {
 183                 return group;
 184             }
 185
 186             if (getGroupV2ByV1Id(connection, groupId) == null) {
 187                 return new GroupInfoV1(groupId);
 188             }
 189
 190             return null;
 191         } catch (SQLException e) {
 192             throw new RuntimeException("Failed read from group store", e);
 193         }
 194     }
 195
 196     public List<GroupInfo> getGroups() {
 197         return Stream.concat(getGroupsV2().stream(), getGroupsV1().stream()).toList();
 198     }
 199
 200     public void mergeRecipients(
 201             final Connection connection, final RecipientId recipientId, final RecipientId toBeMergedRecipientId
 202     ) throws SQLException {
 203         final var sql = (
 204                 """
 205                 UPDATE OR REPLACE %s
 206                 SET recipient_id = ?
 207                 WHERE recipient_id = ?
 208                 """
 209         ).formatted(TABLE_GROUP_V1_MEMBER);
 210         try (final var statement = connection.prepareStatement(sql)) {
 211             statement.setLong(1, recipientId.id());
 212             statement.setLong(2, toBeMergedRecipientId.id());
 213             final var updatedRows = statement.executeUpdate();
 214             if (updatedRows > 0) {
 215                 logger.info("Updated {} group members when merging recipients", updatedRows);
 216             }
 217         }
 218     }
 219
 220     void addLegacyGroups(final Collection<GroupInfo> groups) {
 221         logger.debug("Migrating legacy groups to database");
 222         long start = System.nanoTime();
 223         try (final var connection = database.getConnection()) {
 224             connection.setAutoCommit(false);
 225             for (final var group : groups) {
 226                 insertOrReplaceGroup(connection, null, group);
 227             }
 228             connection.commit();
 229         } catch (SQLException e) {
 230             throw new RuntimeException("Failed update group store", e);
 231         }
 232         logger.debug("Complete groups migration took {}ms", (System.nanoTime() - start) / 1000000);
 233     }
 234
 235     private void insertOrReplaceGroup(
 236             final Connection connection, Long internalId, final GroupInfo group
 237     ) throws SQLException {
 238         if (group instanceof GroupInfoV1 groupV1) {
 239             if (internalId != null) {
 240                 final var sqlDeleteMembers = "DELETE FROM %s where group_id = ?".formatted(TABLE_GROUP_V1_MEMBER);
 241                 try (final var statement = connection.prepareStatement(sqlDeleteMembers)) {
 242                     statement.setLong(1, internalId);
 243                     statement.executeUpdate();
 244                 }
 245             }
 246             final var sql = """
 247                             INSERT OR REPLACE INTO %s (_id, group_id, group_id_v2, name, color, expiration_time, blocked, archived)
 248                             VALUES (?, ?, ?, ?, ?, ?, ?, ?)
 249                             RETURNING _id
 250                             """.formatted(TABLE_GROUP_V1);
 251             try (final var statement = connection.prepareStatement(sql)) {
 252                 if (internalId == null) {
 253                     statement.setNull(1, Types.NUMERIC);
 254                 } else {
 255                     statement.setLong(1, internalId);
 256                 }
 257                 statement.setBytes(2, groupV1.getGroupId().serialize());
 258                 statement.setBytes(3, groupV1.getExpectedV2Id().serialize());
 259                 statement.setString(4, groupV1.getTitle());
 260                 statement.setString(5, groupV1.color);
 261                 statement.setLong(6, groupV1.getMessageExpirationTimer());
 262                 statement.setBoolean(7, groupV1.isBlocked());
 263                 statement.setBoolean(8, groupV1.archived);
 264                 final var generatedKey = Utils.executeQueryForOptional(statement, Utils::getIdMapper);
 265
 266                 if (internalId == null) {
 267                     if (generatedKey.isPresent()) {
 268                         internalId = generatedKey.get();
 269                     } else {
 270                         throw new RuntimeException("Failed to add new group to database");
 271                     }
 272                 }
 273             }
 274             final var sqlInsertMember = """
 275                                         INSERT OR REPLACE INTO %s (group_id, recipient_id)
 276                                         VALUES (?, ?)
 277                                         """.formatted(TABLE_GROUP_V1_MEMBER);
 278             try (final var statement = connection.prepareStatement(sqlInsertMember)) {
 279                 for (final var recipient : groupV1.getMembers()) {
 280                     statement.setLong(1, internalId);
 281                     statement.setLong(2, recipient.id());
 282                     statement.executeUpdate();
 283                 }
 284             }
 285         } else if (group instanceof GroupInfoV2 groupV2) {
 286             final var sql = (
 287                     """
 288                     INSERT OR REPLACE INTO %s (_id, group_id, master_key, group_data, distribution_id, blocked, distribution_id)
 289                     VALUES (?, ?, ?, ?, ?, ?, ?)
 290                     """
 291             ).formatted(TABLE_GROUP_V2);
 292             try (final var statement = connection.prepareStatement(sql)) {
 293                 if (internalId == null) {
 294                     statement.setNull(1, Types.NUMERIC);
 295                 } else {
 296                     statement.setLong(1, internalId);
 297                 }
 298                 statement.setBytes(2, groupV2.getGroupId().serialize());
 299                 statement.setBytes(3, groupV2.getMasterKey().serialize());
 300                 if (groupV2.getGroup() == null) {
 301                     statement.setNull(4, Types.NUMERIC);
 302                 } else {
 303                     statement.setBytes(4, groupV2.getGroup().encode());
 304                 }
 305                 statement.setBytes(5, UuidUtil.toByteArray(groupV2.getDistributionId().asUuid()));
 306                 statement.setBoolean(6, groupV2.isBlocked());
 307                 statement.setBoolean(7, groupV2.isPermissionDenied());
 308                 statement.executeUpdate();
 309             }
 310         } else {
 311             throw new AssertionError("Invalid group id type");
 312         }
 313     }
 314
 315     private List<GroupInfoV2> getGroupsV2() {
 316         final var sql = (
 317                 """
 318                 SELECT g.group_id, g.master_key, g.group_data, g.distribution_id, g.blocked, g.permission_denied
 319                 FROM %s g
 320                 """
 321         ).formatted(TABLE_GROUP_V2);
 322         try (final var connection = database.getConnection()) {
 323             try (final var statement = connection.prepareStatement(sql)) {
 324                 return Utils.executeQueryForStream(statement, this::getGroupInfoV2FromResultSet)
 325                         .filter(Objects::nonNull)
 326                         .toList();
 327             }
 328         } catch (SQLException e) {
 329             throw new RuntimeException("Failed read from group store", e);
 330         }
 331     }
 332
 333     private GroupInfoV2 getGroup(Connection connection, GroupIdV2 groupIdV2) throws SQLException {
 334         final var sql = (
 335                 """
 336                 SELECT g.group_id, g.master_key, g.group_data, g.distribution_id, g.blocked, g.permission_denied
 337                 FROM %s g
 338                 WHERE g.group_id = ?
 339                 """
 340         ).formatted(TABLE_GROUP_V2);
 341         try (final var statement = connection.prepareStatement(sql)) {
 342             statement.setBytes(1, groupIdV2.serialize());
 343             return Utils.executeQueryForOptional(statement, this::getGroupInfoV2FromResultSet).orElse(null);
 344         }
 345     }
 346
 347     private GroupInfoV2 getGroupInfoV2FromResultSet(ResultSet resultSet) throws SQLException {
 348         try {
 349             final var groupId = resultSet.getBytes("group_id");
 350             final var masterKey = resultSet.getBytes("master_key");
 351             final var groupData = resultSet.getBytes("group_data");
 352             final var distributionId = resultSet.getBytes("distribution_id");
 353             final var blocked = resultSet.getBoolean("blocked");
 354             final var permissionDenied = resultSet.getBoolean("permission_denied");
 355             return new GroupInfoV2(GroupId.v2(groupId),
 356                     new GroupMasterKey(masterKey),
 357                     groupData == null ? null : DecryptedGroup.ADAPTER.decode(groupData),
 358                     DistributionId.from(UuidUtil.parseOrThrow(distributionId)),
 359                     blocked,
 360                     permissionDenied,
 361                     recipientResolver);
 362         } catch (InvalidInputException | IOException e) {
 363             return null;
 364         }
 365     }
 366
 367     private List<GroupInfoV1> getGroupsV1() {
 368         final var sql = (
 369                 """
 370                 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
 371                 FROM %s g
 372                 """
 373         ).formatted(TABLE_GROUP_V1_MEMBER, TABLE_GROUP_V1);
 374         try (final var connection = database.getConnection()) {
 375             try (final var statement = connection.prepareStatement(sql)) {
 376                 return Utils.executeQueryForStream(statement, this::getGroupInfoV1FromResultSet)
 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 GroupInfoV1 getGroup(Connection connection, GroupIdV1 groupIdV1) throws SQLException {
 386         final var sql = (
 387                 """
 388                 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
 389                 FROM %s g
 390                 WHERE g.group_id = ?
 391                 """
 392         ).formatted(TABLE_GROUP_V1_MEMBER, TABLE_GROUP_V1);
 393         try (final var statement = connection.prepareStatement(sql)) {
 394             statement.setBytes(1, groupIdV1.serialize());
 395             return Utils.executeQueryForOptional(statement, this::getGroupInfoV1FromResultSet).orElse(null);
 396         }
 397     }
 398
 399     private GroupInfoV1 getGroupInfoV1FromResultSet(ResultSet resultSet) throws SQLException {
 400         final var groupId = resultSet.getBytes("group_id");
 401         final var groupIdV2 = resultSet.getBytes("group_id_v2");
 402         final var name = resultSet.getString("name");
 403         final var color = resultSet.getString("color");
 404         final var membersString = resultSet.getString("members");
 405         final var members = membersString == null
 406                 ? Set.<RecipientId>of()
 407                 : Arrays.stream(membersString.split(","))
 408                         .map(Integer::valueOf)
 409                         .map(recipientIdCreator::create)
 410                         .collect(Collectors.toSet());
 411         final var expirationTime = resultSet.getInt("expiration_time");
 412         final var blocked = resultSet.getBoolean("blocked");
 413         final var archived = resultSet.getBoolean("archived");
 414         return new GroupInfoV1(GroupId.v1(groupId),
 415                 groupIdV2 == null ? null : GroupId.v2(groupIdV2),
 416                 name,
 417                 members,
 418                 color,
 419                 expirationTime,
 420                 blocked,
 421                 archived);
 422     }
 423
 424     private GroupInfoV2 getGroupV2ByV1Id(final Connection connection, final GroupIdV1 groupId) throws SQLException {
 425         return getGroup(connection, GroupUtils.getGroupIdV2(groupId));
 426     }
 427
 428     private GroupInfoV1 getGroupV1ByV2Id(Connection connection, GroupIdV2 groupIdV2) throws SQLException {
 429         final var sql = (
 430                 """
 431                 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
 432                 FROM %s g
 433                 WHERE g.group_id_v2 = ?
 434                 """
 435         ).formatted(TABLE_GROUP_V1_MEMBER, TABLE_GROUP_V1);
 436         try (final var statement = connection.prepareStatement(sql)) {
 437             statement.setBytes(1, groupIdV2.serialize());
 438             return Utils.executeQueryForOptional(statement, this::getGroupInfoV1FromResultSet).orElse(null);
 439         }
 440     }
 441 }