Skip to content

Commit

Permalink
Refactor group v2 migration
Browse files Browse the repository at this point in the history
  • Loading branch information
@AsamK
AsamK committed Nov 12, 2023
1 parent f06eeb0 commit 0d60c4d
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 43 deletions.
19 changes: 2 additions & 17 deletions lib/src/main/java/org/asamk/signal/manager/helper/GroupHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,23 +107,8 @@ public GroupInfoV2 getOrMigrateGroup(
) {
final var groupSecretParams = GroupSecretParams.deriveFromMasterKey(groupMasterKey);

var groupId = GroupUtils.getGroupIdV2(groupSecretParams);
var groupInfo = getGroup(groupId);
final GroupInfoV2 groupInfoV2;
if (groupInfo instanceof GroupInfoV1) {
// Received a v2 group message for a v1 group, we need to locally migrate the group
account.getGroupStore().deleteGroup(((GroupInfoV1) groupInfo).getGroupId());
groupInfoV2 = new GroupInfoV2(groupId, groupMasterKey, account.getRecipientResolver());
groupInfoV2.setBlocked(groupInfo.isBlocked());
account.getGroupStore().updateGroup(groupInfoV2);
logger.info("Locally migrated group {} to group v2, id: {}",
groupInfo.getGroupId().toBase64(),
groupInfoV2.getGroupId().toBase64());
} else if (groupInfo instanceof GroupInfoV2) {
groupInfoV2 = (GroupInfoV2) groupInfo;
} else {
groupInfoV2 = new GroupInfoV2(groupId, groupMasterKey, account.getRecipientResolver());
}
final var groupId = GroupUtils.getGroupIdV2(groupSecretParams);
final var groupInfoV2 = account.getGroupStore().getGroupOrPartialMigrate(groupMasterKey, groupId);

if (groupInfoV2.getGroup() == null || groupInfoV2.getGroup().revision < revision) {
DecryptedGroup group = null;
Expand Down
104 changes: 78 additions & 26 deletions lib/src/main/java/org/asamk/signal/manager/storage/groups/GroupStore.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import org.asamk.signal.manager.storage.recipients.RecipientResolver;
import org.signal.libsignal.zkgroup.InvalidInputException;
import org.signal.libsignal.zkgroup.groups.GroupMasterKey;
import org.signal.libsignal.zkgroup.groups.GroupSecretParams;
import org.signal.storageservice.protos.groups.local.DecryptedGroup;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand Down Expand Up @@ -87,25 +88,29 @@ public GroupStore(
public void updateGroup(GroupInfo group) {
try (final var connection = database.getConnection()) {
connection.setAutoCommit(false);
final Long internalId;
final var sql = (
"""
SELECT g._id
FROM %s g
WHERE g.group_id = ?
"""
).formatted(group instanceof GroupInfoV1 ? TABLE_GROUP_V1 : TABLE_GROUP_V2);
try (final var statement = connection.prepareStatement(sql)) {
statement.setBytes(1, group.getGroupId().serialize());
internalId = Utils.executeQueryForOptional(statement, res -> res.getLong("_id")).orElse(null);
}
insertOrReplaceGroup(connection, internalId, group);
updateGroup(connection, group);
connection.commit();
} catch (SQLException e) {
throw new RuntimeException("Failed update recipient store", e);
}
}

public void updateGroup(final Connection connection, final GroupInfo group) throws SQLException {
final Long internalId;
final var sql = (
"""
SELECT g._id
FROM %s g
WHERE g.group_id = ?
"""
).formatted(group instanceof GroupInfoV1 ? TABLE_GROUP_V1 : TABLE_GROUP_V2);
try (final var statement = connection.prepareStatement(sql)) {
statement.setBytes(1, group.getGroupId().serialize());
internalId = Utils.executeQueryForOptional(statement, res -> res.getLong("_id")).orElse(null);
}
insertOrReplaceGroup(connection, internalId, group);
}

public void deleteGroup(GroupId groupId) {
if (groupId instanceof GroupIdV1 groupIdV1) {
deleteGroup(groupIdV1);
Expand All @@ -115,30 +120,34 @@ public void deleteGroup(GroupId groupId) {
}

public void deleteGroup(GroupIdV1 groupIdV1) {
final var sql = (
"""
DELETE FROM %s
WHERE group_id = ?
"""
).formatted(TABLE_GROUP_V1);
try (final var connection = database.getConnection()) {
try (final var statement = connection.prepareStatement(sql)) {
statement.setBytes(1, groupIdV1.serialize());
statement.executeUpdate();
}
deleteGroup(connection, groupIdV1);
} catch (SQLException e) {
throw new RuntimeException("Failed update group store", e);
}
}

public void deleteGroup(GroupIdV2 groupIdV2) {
private void deleteGroup(final Connection connection, final GroupIdV1 groupIdV1) throws SQLException {
final var sql = (
"""
DELETE FROM %s
WHERE group_id = ?
"""
).formatted(TABLE_GROUP_V2);
).formatted(TABLE_GROUP_V1);
try (final var statement = connection.prepareStatement(sql)) {
statement.setBytes(1, groupIdV1.serialize());
statement.executeUpdate();
}
}

public void deleteGroup(GroupIdV2 groupIdV2) {
try (final var connection = database.getConnection()) {
final var sql = (
"""
DELETE FROM %s
WHERE group_id = ?
"""
).formatted(TABLE_GROUP_V2);
try (final var statement = connection.prepareStatement(sql)) {
statement.setBytes(1, groupIdV2.serialize());
statement.executeUpdate();
Expand Down Expand Up @@ -193,6 +202,49 @@ public GroupInfoV1 getOrCreateGroupV1(GroupIdV1 groupId) {
}
}

public GroupInfoV2 getGroupOrPartialMigrate(
Connection connection, final GroupMasterKey groupMasterKey
) throws SQLException {
final var groupSecretParams = GroupSecretParams.deriveFromMasterKey(groupMasterKey);
final var groupId = GroupUtils.getGroupIdV2(groupSecretParams);

return getGroupOrPartialMigrate(connection, groupMasterKey, groupId);
}

public GroupInfoV2 getGroupOrPartialMigrate(
final GroupMasterKey groupMasterKey, final GroupIdV2 groupId
) {
try (final var connection = database.getConnection()) {
return getGroupOrPartialMigrate(connection, groupMasterKey, groupId);
} catch (SQLException e) {
throw new RuntimeException("Failed read from group store", e);
}
}

private GroupInfoV2 getGroupOrPartialMigrate(
Connection connection, final GroupMasterKey groupMasterKey, final GroupIdV2 groupId
) throws SQLException {
switch (getGroup(groupId)) {
case GroupInfoV1 groupInfoV1 -> {
// Received a v2 group message for a v1 group, we need to locally migrate the group
deleteGroup(connection, groupInfoV1.getGroupId());
final var groupInfoV2 = new GroupInfoV2(groupId, groupMasterKey, recipientResolver);
groupInfoV2.setBlocked(groupInfoV1.isBlocked());
updateGroup(connection, groupInfoV2);
logger.debug("Locally migrated group {} to group v2, id: {}",
groupInfoV1.getGroupId().toBase64(),
groupInfoV2.getGroupId().toBase64());
return groupInfoV2;
}
case GroupInfoV2 groupInfoV2 -> {
return groupInfoV2;
}
case null -> {
return new GroupInfoV2(groupId, groupMasterKey, recipientResolver);
}
}
}

public List<GroupInfo> getGroups() {
return Stream.concat(getGroupsV2().stream(), getGroupsV1().stream()).toList();
}
Expand All @@ -212,7 +264,7 @@ public void mergeRecipients(
statement.setLong(2, toBeMergedRecipientId.id());
final var updatedRows = statement.executeUpdate();
if (updatedRows > 0) {
logger.info("Updated {} group members when merging recipients", updatedRows);
logger.debug("Updated {} group members when merging recipients", updatedRows);
}
}
}
Expand Down

0 comments on commit 0d60c4d

Please sign in to comment.