Use streaming decryption for profile avatars

This commit is contained in:
Max Radermacher 2024-04-05 20:08:42 -05:00 committed by GitHub
parent 9e2d753bc0
commit 269e2cd972
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 156 additions and 104 deletions

View File

@ -132,9 +132,6 @@ typedef void (^ProfileManagerFailureBlock)(NSError *error);
#pragma mark -
// This method is only exposed for usage by the Swift extensions.
- (NSString *)generateAvatarFilename;
- (NSString *)groupKeyForGroupId:(NSData *)groupId;
#pragma mark - Internals exposed for Swift

View File

@ -1190,45 +1190,6 @@ NSString *const kNSNotificationKey_UserProfileWriter = @"kNSNotificationKey_User
return [self.modelReadCaches.userProfileReadCache getUserProfileWithAddress:address transaction:transaction];
}
- (nullable NSURL *)writeAvatarDataToFile:(NSData *)avatarData
{
OWSAssertDebug(avatarData.length > 0);
if (![avatarData ows_isValidImage]) {
OWSFailDebug(@"Invalid avatar format");
return nil;
}
NSString *filename = [self generateAvatarFilename];
NSString *avatarPath = [OWSUserProfile profileAvatarFilePathFor:filename];
NSURL *avatarUrl = [NSURL fileURLWithPath:avatarPath];
if (!avatarUrl) {
OWSFailDebug(@"Invalid URL for avatarPath %@", avatarPath);
return nil;
}
NSError *error = nil;
BOOL success = [avatarData writeToURL:avatarUrl options:NSDataWritingAtomic error:&error];
if (!success || error) {
OWSFailDebug(@"Failed write to url %@: %@", avatarUrl, error);
return nil;
}
// We were double checking that a UIImage could be instantiated from this file before recording the
// avatar to the profile. That behavior is preserved here:
UIImage *_Nullable avatarImage = [UIImage imageWithContentsOfFile:avatarUrl.path];
if (avatarImage) {
return avatarUrl;
} else {
OWSFailDebug(@"Failed to open avatar image written to disk");
return nil;
}
}
- (NSString *)generateAvatarFilename
{
return [[NSUUID UUID].UUIDString stringByAppendingPathExtension:@"jpg"];
}
#pragma mark - Avatar Disk Cache
- (nullable NSData *)loadProfileAvatarDataWithFilename:(NSString *)filename

View File

@ -9,7 +9,6 @@ import SignalCoreKit
public class OWSProfileManagerSwiftValues {
fileprivate let pendingUpdateRequests = AtomicValue<[OWSProfileManager.ProfileUpdateRequest]>([], lock: .init())
fileprivate let avatarDownloadRequests = AtomicValue<[OWSProfileManager.AvatarDownloadKey: Promise<Data>]>([:], lock: .init())
public init() {
}
@ -1031,7 +1030,7 @@ extension OWSProfileManager: ProfileManager, Dependencies {
}
private func writeProfileAvatarToDisk(avatarData: Data) throws -> String {
let filename = self.generateAvatarFilename()
let filename = OWSUserProfile.generateAvatarFilename()
let filePath = OWSUserProfile.profileAvatarFilePath(for: filename)
do {
try avatarData.write(to: URL(fileURLWithPath: filePath), options: [.atomic])
@ -1355,18 +1354,12 @@ extension OWSProfileManager {
}
let shouldRetry: Bool
do {
let filename = generateAvatarFilename()
let filePath = OWSUserProfile.profileAvatarFilePath(for: filename)
let avatarData = try await downloadAndDecryptAvatar(avatarUrlPath: avatarUrlPath, profileKey: profileKey).awaitable()
try avatarData.write(to: URL(fileURLWithPath: filePath), options: [.atomic])
let temporaryFileUrl = try await downloadAndDecryptAvatar(avatarUrlPath: avatarUrlPath, profileKey: profileKey)
var didConsumeFilePath = false
defer {
if !didConsumeFilePath { try? FileManager.default.removeItem(atPath: filePath) }
if !didConsumeFilePath { try? FileManager.default.removeItem(at: temporaryFileUrl) }
}
guard UIImage(contentsOfFile: filePath) != nil else {
throw OWSGenericError("Avatar image can't be decoded.")
}
(shouldRetry, didConsumeFilePath) = await databaseStorage.awaitableWrite { tx in
(shouldRetry, didConsumeFilePath) = try await databaseStorage.awaitableWrite { tx in
let newProfile = OWSUserProfile.getOrBuildUserProfile(for: internalAddress, authedAccount: authedAccount, transaction: tx)
guard newProfile.avatarFileName == nil else {
return (false, false)
@ -1375,8 +1368,9 @@ extension OWSProfileManager {
guard newProfile.profileKey == profileKey, newProfile.avatarUrlPath == avatarUrlPath else {
return (true, false)
}
let avatarFilename = try OWSUserProfile.consumeTemporaryAvatarFileUrl(.setTo(temporaryFileUrl), tx: tx)
newProfile.update(
avatarFileName: .setTo(filename),
avatarFileName: avatarFilename,
userProfileWriter: .avatarDownload,
authedAccount: authedAccount,
transaction: tx,
@ -1390,59 +1384,36 @@ extension OWSProfileManager {
}
}
public func downloadAndDecryptAvatar(avatarUrlPath: String, profileKey: OWSAES256Key) async throws -> Data {
try await downloadAndDecryptAvatar(avatarUrlPath: avatarUrlPath, profileKey: profileKey).awaitable()
}
public func downloadAndDecryptAvatar(avatarUrlPath: String, profileKey: OWSAES256Key) async throws -> URL {
let backgroundTask = OWSBackgroundTask(label: "\(#function)")
defer { backgroundTask.end() }
fileprivate func downloadAndDecryptAvatar(avatarUrlPath: String, profileKey: OWSAES256Key) -> Promise<Data> {
let key = AvatarDownloadKey(avatarUrlPath: avatarUrlPath, profileKey: profileKey.keyData)
let (promise, futureToStart): (Promise<Data>, Future<Data>?) = swiftValues.avatarDownloadRequests.update {
if let existingPromise = $0[key] {
return (existingPromise, nil)
}
let (promise, futureToStart) = Promise<Data>.pending()
$0[key] = promise.ensure(on: DispatchQueue.global()) { [weak self] in
self?.swiftValues.avatarDownloadRequests.update { _ = $0.removeValue(forKey: key) }
}
return (promise, futureToStart)
}
if let futureToStart {
futureToStart.resolve(on: SyncScheduler(), with: Promise.wrapAsync {
let backgroundTask = OWSBackgroundTask(label: "\(#function)")
defer { backgroundTask.end() }
return try await Self._downloadAndDecryptAvatar(
avatarUrlPath: avatarUrlPath,
profileKey: profileKey,
remainingRetries: 3
)
})
}
return promise
return try await Self._downloadAndDecryptAvatar(
avatarUrlPath: avatarUrlPath,
profileKey: profileKey,
remainingRetries: 3
)
}
private static func _downloadAndDecryptAvatar(
avatarUrlPath: String,
profileKey: OWSAES256Key,
remainingRetries: Int
) async throws -> Data {
) async throws -> URL {
assert(!avatarUrlPath.isEmpty)
do {
Logger.info("")
let urlSession = self.avatarUrlSession
let response = try await urlSession.downloadTaskPromise(avatarUrlPath, method: .get).awaitable()
let encryptedData: Data
do {
encryptedData = try Data(contentsOf: response.downloadUrl)
} catch {
owsFailDebug("Could not load data failed: \(error)")
// Fail immediately; do not retry.
throw SSKUnretryableError.couldNotLoadFileData
let decryptedFileUrl = OWSFileSystem.temporaryFileUrl(isAvailableWhileDeviceLocked: true)
try decryptAvatar(at: response.downloadUrl, to: decryptedFileUrl, profileKey: profileKey)
guard NSData.ows_isValidImage(at: decryptedFileUrl, mimeType: nil) else {
throw OWSGenericError("Couldn't validate avatar")
}
return try OWSUserProfile.decrypt(profileData: encryptedData, profileKey: profileKey)
guard UIImage(contentsOfFile: decryptedFileUrl.path) != nil else {
throw OWSGenericError("Couldn't decode image")
}
return decryptedFileUrl
} catch where error.isNetworkFailureOrTimeout && remainingRetries > 0 {
return try await _downloadAndDecryptAvatar(
avatarUrlPath: avatarUrlPath,
@ -1451,4 +1422,89 @@ extension OWSProfileManager {
)
}
}
private static func decryptAvatar(
at encryptedFileUrl: URL,
to decryptedFileUrl: URL,
profileKey: OWSAES256Key
) throws {
let readHandle = try FileHandle(forReadingFrom: encryptedFileUrl)
defer {
try? readHandle.close()
}
guard
FileManager.default.createFile(
atPath: decryptedFileUrl.path,
contents: nil,
attributes: [.protectionKey: FileProtectionType.completeUntilFirstUserAuthentication]
)
else {
throw OWSGenericError("Couldn't create temporary file")
}
let writeHandle = try FileHandle(forWritingTo: decryptedFileUrl)
defer {
try? writeHandle.close()
}
let concatenatedLength = Int(try readHandle._seekToEnd())
try readHandle._seek(toOffset: 0)
let nonceLength = Aes256GcmEncryptedData.nonceLength
let authenticationTagLength = Aes256GcmEncryptedData.authenticationTagLength
var remainingLength = concatenatedLength - nonceLength - authenticationTagLength
guard remainingLength >= 0 else {
throw OWSGenericError("ciphertext too short")
}
let nonceData = try readHandle._read(count: nonceLength)
let decryptor = try Aes256GcmDecryption(key: profileKey.keyData, nonce: nonceData, associatedData: [])
while remainingLength > 0 {
let kBatchLimit = 32768
var payloadData: Data = try readHandle._read(count: min(remainingLength, kBatchLimit))
try decryptor.decrypt(&payloadData)
try writeHandle._write(contentsOf: payloadData)
remainingLength -= payloadData.count
}
let authenticationTag = try readHandle._read(count: authenticationTagLength)
guard try decryptor.verifyTag(authenticationTag) else {
throw OWSGenericError("failed to decrypt")
}
}
}
private extension FileHandle {
func _seekToEnd() throws -> UInt64 {
if #available(iOS 13.4, *) {
return try seekToEnd()
} else {
return seekToEndOfFile()
}
}
func _seek(toOffset offset: UInt64) throws {
if #available(iOS 13.4, *) {
try seek(toOffset: offset)
} else {
seek(toFileOffset: offset)
}
}
func _read(count: Int) throws -> Data {
if #available(iOS 13.4, *) {
return try read(upToCount: count) ?? Data()
} else {
return readData(ofLength: count)
}
}
func _write(contentsOf data: Data) throws {
if #available(iOS 13.4, *) {
try write(contentsOf: data)
} else {
write(data)
}
}
}

View File

@ -51,7 +51,7 @@ public protocol ProfileManager: ProfileManagerProtocol {
func downloadAndDecryptAvatar(
avatarUrlPath: String,
profileKey: OWSAES256Key
) async throws -> Data
) async throws -> URL
func updateProfile(
address: SignalServiceAddress,

View File

@ -76,7 +76,6 @@ typedef NS_ENUM(NSUInteger, UserProfileWriter) {
- (nullable NSString *)profileAvatarURLPathForAddress:(SignalServiceAddress *)address
authedAccount:(AuthedAccount *)authedAccount
transaction:(SDSAnyReadTransaction *)transaction;
- (nullable NSURL *)writeAvatarDataToFile:(NSData *)avatarData NS_SWIFT_NAME(writeAvatarDataToFile(_:));
- (void)fillInProfileKeysForAllProfileKeys:(NSDictionary<SignalServiceAddress *, NSData *> *)allProfileKeys
authoritativeProfileKeys:(NSDictionary<SignalServiceAddress *, NSData *> *)authoritativeProfileKeys

View File

@ -49,7 +49,7 @@ extension OWSFakeProfileManager: ProfileManager {
return Promise(error: OWSGenericError("Not supported."))
}
public func downloadAndDecryptAvatar(avatarUrlPath: String, profileKey: OWSAES256Key) async throws -> Data {
public func downloadAndDecryptAvatar(avatarUrlPath: String, profileKey: OWSAES256Key) async throws -> URL {
throw OWSGenericError("Not supported.")
}
}

View File

@ -486,7 +486,7 @@ public final class OWSUserProfile: NSObject, NSCopying, SDSCodableModel, Decodab
if fileNameDidChange, let oldAvatarFileName, !oldAvatarFileName.isEmpty {
let oldAvatarFilePath = Self.profileAvatarFilePath(for: oldAvatarFileName)
DispatchQueue.global().async { OWSFileSystem.deleteFileIfExists(oldAvatarFilePath) }
OWSFileSystem.deleteFileIfExists(oldAvatarFilePath)
}
self.avatarUrlPath = newAvatarUrlPath
@ -494,6 +494,30 @@ public final class OWSUserProfile: NSObject, NSCopying, SDSCodableModel, Decodab
return true
}
/// Moves a temporary avatar file into its permanent location.
///
/// This method accepts a `tx` parameter to ensure that this move occurs
/// within a write transaction. Callers must ensure that it's the same write
/// transaction that assigns the result to an OWSUserProfile. In doing so,
/// callers ensure that the orphaned cleanup logic doesn't delete avatars
/// that are about to be referenced.
static func consumeTemporaryAvatarFileUrl(
_ avatarFileUrl: OptionalChange<URL?>,
tx: SDSAnyWriteTransaction
) throws -> OptionalChange<String?> {
switch avatarFileUrl {
case .noChange:
return .noChange
case .setTo(.none):
return .setTo(nil)
case .setTo(.some(let temporaryFileUrl)):
let filename = generateAvatarFilename()
let filePath = profileAvatarFilePath(for: filename)
try FileManager.default.moveItem(atPath: temporaryFileUrl.path, toPath: filePath)
return .setTo(filename)
}
}
@objc
public static var legacyProfileAvatarsDirPath: String {
return OWSFileSystem.appDocumentDirectoryPath().appendingPathComponent("ProfileAvatars")
@ -510,8 +534,12 @@ public final class OWSUserProfile: NSObject, NSCopying, SDSCodableModel, Decodab
return result
}()
static func generateAvatarFilename() -> String {
return UUID().uuidString + ".jpg"
}
@objc
public static func profileAvatarFilePath(for filename: String) -> String {
static func profileAvatarFilePath(for filename: String) -> String {
owsAssertDebug(!filename.isEmpty)
return Self.profileAvatarsDirPath.appendingPathComponent(filename)
}

View File

@ -270,9 +270,9 @@ public class ProfileFetcherJob: NSObject {
if didAlreadyDownloadAvatar {
return AvatarDownloadResult(remoteRelativePath: .noChange, localFileUrl: .noChange)
}
let avatarData: Data?
let temporaryAvatarUrl: URL?
do {
avatarData = try await profileManager.downloadAndDecryptAvatar(
temporaryAvatarUrl = try await profileManager.downloadAndDecryptAvatar(
avatarUrlPath: newAvatarUrlPath,
profileKey: profileKey
)
@ -294,11 +294,11 @@ public class ProfileFetcherJob: NSObject {
// afterward). This might be due to a race with an update that is in
// flight. We should eventually recover since profile updates are
// durable.
avatarData = nil
temporaryAvatarUrl = nil
}
return AvatarDownloadResult(
remoteRelativePath: .setTo(newAvatarUrlPath),
localFileUrl: .setTo(avatarData.flatMap { profileManager.writeAvatarDataToFile($0) })
localFileUrl: .setTo(temporaryAvatarUrl)
)
}
@ -336,11 +336,22 @@ public class ProfileFetcherJob: NSObject {
.map { $0.0 }
.filter { persistedBadgeIds.contains($0.badgeId) }
let avatarFilename: OptionalChange<String?>
do {
avatarFilename = try OWSUserProfile.consumeTemporaryAvatarFileUrl(
avatarDownloadResult.localFileUrl,
tx: transaction
)
} catch {
Logger.warn("Couldn't move downloaded avatar: \(error)")
avatarFilename = .noChange
}
self.profileManager.updateProfile(
address: SignalServiceAddress(serviceId),
decryptedProfile: fetchedProfile.decryptedProfile,
avatarUrlPath: avatarDownloadResult.remoteRelativePath,
avatarFileName: avatarDownloadResult.localFileUrl.map { $0?.lastPathComponent },
avatarFileName: avatarFilename,
profileBadges: profileBadgeMetadata,
lastFetchDate: Date(),
userProfileWriter: .profileFetch,