Refactor AttachmentCipherInputStream.

This commit is contained in:
Greyson Parrelli
2025-06-18 10:35:45 -04:00
committed by Michelle Tang
parent 9798f5cc7c
commit 4f6a5de227
10 changed files with 572 additions and 403 deletions

View File

@@ -252,7 +252,6 @@ class RestoreAttachmentJob private constructor(
pointer,
attachmentFile,
maxReceiveSize,
false,
progressListener
)
} else {

View File

@@ -132,14 +132,13 @@ class RestoreAttachmentThumbnailJob private constructor(
Log.i(TAG, "Downloading thumbnail for $attachmentId")
val downloadResult = AppDependencies.signalServiceMessageReceiver
.retrieveArchivedAttachment(
.retrieveArchivedThumbnail(
SignalStore.backup.mediaRootBackupKey.deriveMediaSecrets(attachment.requireThumbnailMediaName()),
cdnCredentials,
thumbnailTransferFile,
pointer,
thumbnailFile,
maxThumbnailSize,
true,
progressListener
)

View File

@@ -149,7 +149,15 @@ class RestoreLocalAttachmentJob private constructor(
try {
val iv = ByteArray(16)
streamSupplier.openStream().use { StreamUtil.readFully(it, iv) }
AttachmentCipherInputStream.createForAttachment(streamSupplier, size, attachment.size, combinedKey, attachment.remoteDigest, null, 0, false).use { input ->
AttachmentCipherInputStream.createForAttachment(
streamSupplier = streamSupplier,
streamLength = size,
plaintextLength = attachment.size,
combinedKeyMaterial = combinedKey,
digest = attachment.remoteDigest,
incrementalDigest = null,
incrementalMacChunkSize = 0
).use { input ->
SignalDatabase.attachments.finalizeAttachmentAfterDownload(attachment.mmsId, attachment.attachmentId, input, iv)
}
} catch (e: InvalidMessageException) {

View File

@@ -11,6 +11,7 @@ import org.signal.libsignal.protocol.InvalidMessageException;
import org.whispersystems.signalservice.api.crypto.AttachmentCipherInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Optional;
@@ -37,7 +38,12 @@ class AttachmentStreamLocalUriFetcher implements DataFetcher<InputStream> {
public void loadData(@NonNull Priority priority, @NonNull DataCallback<? super InputStream> callback) {
try {
if (!digest.isPresent()) throw new InvalidMessageException("No attachment digest!");
is = AttachmentCipherInputStream.createForAttachment(attachment, plaintextLength, key, digest.get(), null, 0);
is = AttachmentCipherInputStream.createForAttachment(attachment,
plaintextLength,
key,
digest.get(),
null,
0);
callback.onDataReady(is);
} catch (IOException | InvalidMessageException e) {
callback.onLoadFailed(e);

View File

@@ -71,7 +71,8 @@ class PartDataSource implements DataSource {
final boolean hasData = attachment.hasData;
if (inProgress && !hasData && hasIncrementalDigest && attachmentKey != null) {
final byte[] decode = Base64.decode(attachmentKey);
final byte[] decodedKey = Base64.decode(attachmentKey);
if (attachment.transferState == AttachmentTable.TRANSFER_RESTORE_IN_PROGRESS && attachment.archiveTransferState == AttachmentTable.ArchiveTransferState.FINISHED) {
final File archiveFile = attachmentDatabase.getOrCreateArchiveTransferFile(attachment.attachmentId);
try {
@@ -81,7 +82,11 @@ class PartDataSource implements DataSource {
MediaRootBackupKey.MediaKeyMaterial mediaKeyMaterial = SignalStore.backup().getMediaRootBackupKey().deriveMediaSecretsFromMediaId(mediaId);
long originalCipherLength = AttachmentCipherStreamUtil.getCiphertextLength(PaddingInputStream.getPaddedSize(attachment.size));
this.inputStream = AttachmentCipherInputStream.createStreamingForArchivedAttachment(mediaKeyMaterial, archiveFile, originalCipherLength, attachment.size, attachment.remoteDigest, decode, attachment.getIncrementalDigest(), attachment.incrementalMacChunkSize);
if (attachment.remoteDigest == null) {
throw new InvalidMessageException("Missing digest!");
}
this.inputStream = AttachmentCipherInputStream.createForArchivedMediaOuterAndInnerLayers(mediaKeyMaterial, archiveFile, originalCipherLength, attachment.size, decodedKey, attachment.remoteDigest, attachment.getIncrementalDigest(), attachment.incrementalMacChunkSize);
} catch (InvalidMessageException e) {
throw new IOException("Error decrypting attachment stream!", e);
}
@@ -95,7 +100,7 @@ class PartDataSource implements DataSource {
throw new InvalidMessageException("Missing digest!");
}
this.inputStream = AttachmentCipherInputStream.createForAttachment(streamSupplier, streamLength, attachment.size, decode, attachment.remoteDigest, attachment.getIncrementalDigest(), attachment.incrementalMacChunkSize, false);
this.inputStream = AttachmentCipherInputStream.createForAttachment(streamSupplier, streamLength, attachment.size, decodedKey, attachment.remoteDigest, attachment.getIncrementalDigest(), attachment.incrementalMacChunkSize);
} catch (InvalidMessageException e) {
throw new IOException("Error decrypting attachment stream!", e);
}