From cb8ee3536aaea829fa57c92b8c1405a8bd6417ad Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Thu, 22 Mar 2018 14:05:12 -0400
Subject: [PATCH] Lazy restore attachments.

---
 Signal/src/util/OWSBackup.m          | 21 +++++++++++----------
 Signal/src/util/OWSBackupIO.h        |  2 ++
 Signal/src/util/OWSBackupIO.m        |  7 ++++++-
 Signal/src/util/OWSBackupImportJob.m |  3 ++-
 4 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/Signal/src/util/OWSBackup.m b/Signal/src/util/OWSBackup.m
index 94a76ff5a..1bc178db6 100644
--- a/Signal/src/util/OWSBackup.m
+++ b/Signal/src/util/OWSBackup.m
@@ -490,7 +490,7 @@ NS_ASSUME_NONNULL_BEGIN
 {
     NSMutableArray<NSString *> *recordNames = [NSMutableArray new];
     [self.dbConnection readWithBlock:^(YapDatabaseReadTransaction *transaction) {
-        id ext = [transaction ext:TSMessageDatabaseViewExtensionName];
+        id ext = [transaction ext:TSLazyRestoreAttachmentsDatabaseViewExtensionName];
         if (!ext) {
             OWSProdLogAndFail(@"%@ Could not load database view.", self.logTag);
             return;
@@ -524,7 +524,7 @@ NS_ASSUME_NONNULL_BEGIN
 {
     NSMutableArray<NSString *> *attachmentIds = [NSMutableArray new];
     [self.dbConnection readWithBlock:^(YapDatabaseReadTransaction *transaction) {
-        id ext = [transaction ext:TSMessageDatabaseViewExtensionName];
+        id ext = [transaction ext:TSLazyRestoreAttachmentsDatabaseViewExtensionName];
         if (!ext) {
             OWSProdLogAndFail(@"%@ Could not load database view.", self.logTag);
             return;
@@ -570,10 +570,7 @@ NS_ASSUME_NONNULL_BEGIN
     // will leverage successful file downloads from previous attempts.
     //
     // TODO: This will also require imports using a predictable jobTempDirPath.
-    NSString *_Nullable tempFilePath = [backupIO createTempFile];
-    if (!tempFilePath) {
-        return completion(NO);
-    }
+    NSString *tempFilePath = [backupIO generateTempFilePath];
 
     [OWSBackupAPI downloadFileFromCloudWithRecordName:lazyRestoreFragment.recordName
         toFileUrl:[NSURL fileURLWithPath:tempFilePath]
@@ -611,10 +608,7 @@ NS_ASSUME_NONNULL_BEGIN
         return completion(NO);
     }
 
-    NSString *_Nullable decryptedFilePath = [backupIO createTempFile];
-    if (!decryptedFilePath) {
-        return completion(NO);
-    }
+    NSString *decryptedFilePath = [backupIO generateTempFilePath];
 
     @autoreleasepool {
         if (![backupIO decryptFileAsFile:encryptedFilePath dstFilePath:decryptedFilePath encryptionKey:encryptionKey]) {
@@ -628,6 +622,13 @@ NS_ASSUME_NONNULL_BEGIN
         DDLogError(@"%@ Attachment has invalid file path.", self.logTag);
         return completion(NO);
     }
+
+    NSString *attachmentDirPath = [attachmentFilePath stringByDeletingLastPathComponent];
+    if (![OWSFileSystem ensureDirectoryExists:attachmentDirPath]) {
+        DDLogError(@"%@ Couldn't create directory for attachment file.", self.logTag);
+        return completion(NO);
+    }
+
     NSError *error;
     BOOL success =
         [NSFileManager.defaultManager moveItemAtPath:decryptedFilePath toPath:attachmentFilePath error:&error];
diff --git a/Signal/src/util/OWSBackupIO.h b/Signal/src/util/OWSBackupIO.h
index eb25f0277..ec9ef0575 100644
--- a/Signal/src/util/OWSBackupIO.h
+++ b/Signal/src/util/OWSBackupIO.h
@@ -20,6 +20,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (instancetype)initWithJobTempDirPath:(NSString *)jobTempDirPath;
 
+- (NSString *)generateTempFilePath;
+
 - (nullable NSString *)createTempFile;
 
 #pragma mark - Encrypt
diff --git a/Signal/src/util/OWSBackupIO.m b/Signal/src/util/OWSBackupIO.m
index 75c91bb5d..7044c284c 100644
--- a/Signal/src/util/OWSBackupIO.m
+++ b/Signal/src/util/OWSBackupIO.m
@@ -44,9 +44,14 @@ static const compression_algorithm SignalCompressionAlgorithm = COMPRESSION_LZMA
     return self;
 }
 
+- (NSString *)generateTempFilePath
+{
+    return [self.jobTempDirPath stringByAppendingPathComponent:[NSUUID UUID].UUIDString];
+}
+
 - (nullable NSString *)createTempFile
 {
-    NSString *filePath = [self.jobTempDirPath stringByAppendingPathComponent:[NSUUID UUID].UUIDString];
+    NSString *filePath = [self generateTempFilePath];
     if (![OWSFileSystem ensureFileExists:filePath]) {
         OWSProdLogAndFail(@"%@ could not create temp file.", self.logTag);
         return nil;
diff --git a/Signal/src/util/OWSBackupImportJob.m b/Signal/src/util/OWSBackupImportJob.m
index 7941b8790..7ba13dc9a 100644
--- a/Signal/src/util/OWSBackupImportJob.m
+++ b/Signal/src/util/OWSBackupImportJob.m
@@ -259,7 +259,8 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
     DDLogVerbose(@"%@ %s: %zd", self.logTag, __PRETTY_FUNCTION__, self.attachmentsItems.count);
 
     __block NSUInteger count = 0;
-    [self.primaryStorage.newDatabaseConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+    YapDatabaseConnection *dbConnection = self.primaryStorage.newDatabaseConnection;
+    [dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
         for (OWSBackupFragment *item in self.attachmentsItems) {
             if (self.isComplete) {
                 return;