Fix discrepancy in message counting between export and import backups.

This commit is contained in:
Cody Henthorne 2021-03-17 10:22:53 -04:00 committed by Greyson Parrelli
parent cb6e3ade15
commit 9366596f5f
2 changed files with 24 additions and 9 deletions

View file

@ -10,7 +10,6 @@ import androidx.annotation.Nullable;
import androidx.annotation.RequiresApi; import androidx.annotation.RequiresApi;
import androidx.documentfile.provider.DocumentFile; import androidx.documentfile.provider.DocumentFile;
import com.annimon.stream.function.Consumer;
import com.annimon.stream.function.Predicate; import com.annimon.stream.function.Predicate;
import com.google.protobuf.ByteString; import com.google.protobuf.ByteString;
@ -127,8 +126,10 @@ public class FullBackupExporter extends FullBackupBase {
try { try {
outputStream.writeDatabaseVersion(input.getVersion()); outputStream.writeDatabaseVersion(input.getVersion());
count++;
List<String> tables = exportSchema(input, outputStream); List<String> tables = exportSchema(input, outputStream);
count += tables.size() * 3;
Stopwatch stopwatch = new Stopwatch("Backup"); Stopwatch stopwatch = new Stopwatch("Backup");
@ -141,9 +142,9 @@ public class FullBackupExporter extends FullBackupBase {
} else if (table.equals(GroupReceiptDatabase.TABLE_NAME)) { } else if (table.equals(GroupReceiptDatabase.TABLE_NAME)) {
count = exportTable(table, input, outputStream, cursor -> isForNonExpiringMessage(input, cursor.getLong(cursor.getColumnIndexOrThrow(GroupReceiptDatabase.MMS_ID))), null, count, cancellationSignal); count = exportTable(table, input, outputStream, cursor -> isForNonExpiringMessage(input, cursor.getLong(cursor.getColumnIndexOrThrow(GroupReceiptDatabase.MMS_ID))), null, count, cancellationSignal);
} else if (table.equals(AttachmentDatabase.TABLE_NAME)) { } else if (table.equals(AttachmentDatabase.TABLE_NAME)) {
count = exportTable(table, input, outputStream, cursor -> isForNonExpiringMessage(input, cursor.getLong(cursor.getColumnIndexOrThrow(AttachmentDatabase.MMS_ID))), cursor -> exportAttachment(attachmentSecret, cursor, outputStream), count, cancellationSignal); count = exportTable(table, input, outputStream, cursor -> isForNonExpiringMessage(input, cursor.getLong(cursor.getColumnIndexOrThrow(AttachmentDatabase.MMS_ID))), (cursor, innerCount) -> exportAttachment(attachmentSecret, cursor, outputStream, innerCount), count, cancellationSignal);
} else if (table.equals(StickerDatabase.TABLE_NAME)) { } else if (table.equals(StickerDatabase.TABLE_NAME)) {
count = exportTable(table, input, outputStream, cursor -> true, cursor -> exportSticker(attachmentSecret, cursor, outputStream), count, cancellationSignal); count = exportTable(table, input, outputStream, cursor -> true, (cursor, innerCount) -> exportSticker(attachmentSecret, cursor, outputStream, innerCount), count, cancellationSignal);
} else if (!BLACKLISTED_TABLES.contains(table) && !table.startsWith("sqlite_")) { } else if (!BLACKLISTED_TABLES.contains(table) && !table.startsWith("sqlite_")) {
count = exportTable(table, input, outputStream, null, null, count, cancellationSignal); count = exportTable(table, input, outputStream, null, null, count, cancellationSignal);
} }
@ -224,7 +225,7 @@ public class FullBackupExporter extends FullBackupBase {
@NonNull SQLiteDatabase input, @NonNull SQLiteDatabase input,
@NonNull BackupFrameOutputStream outputStream, @NonNull BackupFrameOutputStream outputStream,
@Nullable Predicate<Cursor> predicate, @Nullable Predicate<Cursor> predicate,
@Nullable Consumer<Cursor> postProcess, @Nullable PostProcessor postProcess,
int count, int count,
@NonNull BackupCancellationSignal cancellationSignal) @NonNull BackupCancellationSignal cancellationSignal)
throws IOException throws IOException
@ -234,7 +235,6 @@ public class FullBackupExporter extends FullBackupBase {
try (Cursor cursor = input.rawQuery("SELECT * FROM " + table, null)) { try (Cursor cursor = input.rawQuery("SELECT * FROM " + table, null)) {
while (cursor != null && cursor.moveToNext()) { while (cursor != null && cursor.moveToNext()) {
throwIfCanceled(cancellationSignal); throwIfCanceled(cancellationSignal);
EventBus.getDefault().post(new BackupEvent(BackupEvent.Type.PROGRESS, ++count));
if (predicate == null || predicate.test(cursor)) { if (predicate == null || predicate.test(cursor)) {
StringBuilder statement = new StringBuilder(template); StringBuilder statement = new StringBuilder(template);
@ -266,9 +266,12 @@ public class FullBackupExporter extends FullBackupBase {
statement.append(')'); statement.append(')');
EventBus.getDefault().post(new BackupEvent(BackupEvent.Type.PROGRESS, ++count));
outputStream.write(statementBuilder.setStatement(statement.toString()).build()); outputStream.write(statementBuilder.setStatement(statement.toString()).build());
if (postProcess != null) postProcess.accept(cursor); if (postProcess != null) {
count = postProcess.postProcess(cursor, count);
}
} }
} }
} }
@ -276,7 +279,7 @@ public class FullBackupExporter extends FullBackupBase {
return count; return count;
} }
private static void exportAttachment(@NonNull AttachmentSecret attachmentSecret, @NonNull Cursor cursor, @NonNull BackupFrameOutputStream outputStream) { private static int exportAttachment(@NonNull AttachmentSecret attachmentSecret, @NonNull Cursor cursor, @NonNull BackupFrameOutputStream outputStream, int count) {
try { try {
long rowId = cursor.getLong(cursor.getColumnIndexOrThrow(AttachmentDatabase.ROW_ID)); long rowId = cursor.getLong(cursor.getColumnIndexOrThrow(AttachmentDatabase.ROW_ID));
long uniqueId = cursor.getLong(cursor.getColumnIndexOrThrow(AttachmentDatabase.UNIQUE_ID)); long uniqueId = cursor.getLong(cursor.getColumnIndexOrThrow(AttachmentDatabase.UNIQUE_ID));
@ -301,14 +304,17 @@ public class FullBackupExporter extends FullBackupBase {
if (random != null && random.length == 32) inputStream = ModernDecryptingPartInputStream.createFor(attachmentSecret, random, new File(data), 0); if (random != null && random.length == 32) inputStream = ModernDecryptingPartInputStream.createFor(attachmentSecret, random, new File(data), 0);
else inputStream = ClassicDecryptingPartInputStream.createFor(attachmentSecret, new File(data)); else inputStream = ClassicDecryptingPartInputStream.createFor(attachmentSecret, new File(data));
EventBus.getDefault().post(new BackupEvent(BackupEvent.Type.PROGRESS, ++count));
outputStream.write(new AttachmentId(rowId, uniqueId), inputStream, size); outputStream.write(new AttachmentId(rowId, uniqueId), inputStream, size);
} }
} catch (IOException e) { } catch (IOException e) {
Log.w(TAG, e); Log.w(TAG, e);
} }
return count;
} }
private static void exportSticker(@NonNull AttachmentSecret attachmentSecret, @NonNull Cursor cursor, @NonNull BackupFrameOutputStream outputStream) { private static int exportSticker(@NonNull AttachmentSecret attachmentSecret, @NonNull Cursor cursor, @NonNull BackupFrameOutputStream outputStream, int count) {
try { try {
long rowId = cursor.getLong(cursor.getColumnIndexOrThrow(StickerDatabase._ID)); long rowId = cursor.getLong(cursor.getColumnIndexOrThrow(StickerDatabase._ID));
long size = cursor.getLong(cursor.getColumnIndexOrThrow(StickerDatabase.FILE_LENGTH)); long size = cursor.getLong(cursor.getColumnIndexOrThrow(StickerDatabase.FILE_LENGTH));
@ -317,12 +323,15 @@ public class FullBackupExporter extends FullBackupBase {
byte[] random = cursor.getBlob(cursor.getColumnIndexOrThrow(StickerDatabase.FILE_RANDOM)); byte[] random = cursor.getBlob(cursor.getColumnIndexOrThrow(StickerDatabase.FILE_RANDOM));
if (!TextUtils.isEmpty(data) && size > 0) { if (!TextUtils.isEmpty(data) && size > 0) {
EventBus.getDefault().post(new BackupEvent(BackupEvent.Type.PROGRESS, ++count));
InputStream inputStream = ModernDecryptingPartInputStream.createFor(attachmentSecret, random, new File(data), 0); InputStream inputStream = ModernDecryptingPartInputStream.createFor(attachmentSecret, random, new File(data), 0);
outputStream.writeSticker(rowId, inputStream, size); outputStream.writeSticker(rowId, inputStream, size);
} }
} catch (IOException e) { } catch (IOException e) {
Log.w(TAG, e); Log.w(TAG, e);
} }
return count;
} }
private static long calculateVeryOldStreamLength(@NonNull AttachmentSecret attachmentSecret, @Nullable byte[] random, @NonNull String data) throws IOException { private static long calculateVeryOldStreamLength(@NonNull AttachmentSecret attachmentSecret, @Nullable byte[] random, @NonNull String data) throws IOException {
@ -528,6 +537,10 @@ public class FullBackupExporter extends FullBackupBase {
} }
} }
public interface PostProcessor {
int postProcess(@NonNull Cursor cursor, int count);
}
public interface BackupCancellationSignal { public interface BackupCancellationSignal {
boolean isCanceled(); boolean isCanceled();
} }

View file

@ -88,7 +88,8 @@ public class FullBackupImporter extends FullBackupBase {
BackupFrame frame; BackupFrame frame;
while (!(frame = inputStream.readFrame()).getEnd()) { while (!(frame = inputStream.readFrame()).getEnd()) {
if (count++ % 100 == 0) EventBus.getDefault().post(new BackupEvent(BackupEvent.Type.PROGRESS, count)); if (count % 100 == 0) EventBus.getDefault().post(new BackupEvent(BackupEvent.Type.PROGRESS, count));
count++;
if (frame.hasVersion()) processVersion(db, frame.getVersion()); if (frame.hasVersion()) processVersion(db, frame.getVersion());
else if (frame.hasStatement()) processStatement(db, frame.getStatement()); else if (frame.hasStatement()) processStatement(db, frame.getStatement());
@ -96,6 +97,7 @@ public class FullBackupImporter extends FullBackupBase {
else if (frame.hasAttachment()) processAttachment(context, attachmentSecret, db, frame.getAttachment(), inputStream); else if (frame.hasAttachment()) processAttachment(context, attachmentSecret, db, frame.getAttachment(), inputStream);
else if (frame.hasSticker()) processSticker(context, attachmentSecret, db, frame.getSticker(), inputStream); else if (frame.hasSticker()) processSticker(context, attachmentSecret, db, frame.getSticker(), inputStream);
else if (frame.hasAvatar()) processAvatar(context, db, frame.getAvatar(), inputStream); else if (frame.hasAvatar()) processAvatar(context, db, frame.getAvatar(), inputStream);
else count--;
} }
db.setTransactionSuccessful(); db.setTransactionSuccessful();