From 941e65649dc0e532fa9c5a7bb1935f695ec788a5 Mon Sep 17 00:00:00 2001
From: Howie Kaye <howie@google.com>
Date: Thu, 2 May 2019 15:05:53 -0400
Subject: [PATCH] Attempt running ASR multiple times on failure (#66)

* Attempt running ASR multiple times (configurable in the profile) on failure.

* Document the ASRAttempts configuration option in the README file.

* fix nits.
---
 README.md                    | 10 +++++
 restord/ImageSessionServer.m | 79 ++++++++++++++++++++----------------
 2 files changed, 54 insertions(+), 35 deletions(-)

diff --git a/README.md b/README.md
index cb1e6ac..53b188c 100644
--- a/README.md
+++ b/README.md
@@ -127,6 +127,16 @@ sudo defaults write /Library/Preferences/com.google.corp.restor.plist DiskFilter
     "(bsdName != 'disk3s2')"
 ```
 
+### ASRAttempts
+
+__Optional__
+
+Set how many times Restor should attempt to run ASR until it runs successfully.
+This may help with issues where the Volume is not unmountable by ASR due to some
+external process (ie spotlight) holding onto it.  Defaults to 1.
+
+`sudo defaults write /Library/Preferences/com.google.corp.restor.plist ASRAttempts -int 5`
+
 ## 10.13 and APFS Note
 
 In order to restore an APFS 10.13 DMG to a machine, the host machine running
diff --git a/restord/ImageSessionServer.m b/restord/ImageSessionServer.m
index 780ec1c..0f04572 100644
--- a/restord/ImageSessionServer.m
+++ b/restord/ImageSessionServer.m
@@ -27,6 +27,8 @@
 // GUID Partition Table UUID Constants
 NSString * const kGPTAPFSUUID = @"7C3457EF-0000-11AA-AA11-00306543ECAC";
 NSString * const kGPTCoreStorageUUID = @"53746F72-6167-11AA-AA11-00306543ECAC";
+static NSString * const kPreferenceDomain = @"com.google.corp.restor";
+static NSString * const kASRAttempts = @"ASRAttempts";
 
 @interface ImageSessionServer ()
 
@@ -102,8 +104,6 @@ - (void)dealloc {
 }
 
 - (void)beginImaging {
-  [self unmountDisk:self.diskRef withOptions:kDADiskUnmountOptionWhole];
-
   // Remove any top level recovery partitions.
   [self removeRecovery];
 
@@ -260,39 +260,48 @@ - (NSString *)apfsOSDisk {
 - (int)applyImage {
   NSString *path = self.image.localURL.path;
 
-  self.asr = [[NSTask alloc] init];
-  self.asr.launchPath = @"/usr/sbin/asr";
-  self.asr.arguments = @[ @"restore",
-                          @"--buffersize",
-                          @"16m",
-                          @"--source",
-                          path,
-                          @"--target",
-                          self.destination.path,
-                          @"--erase",
-                          @"--noprompt",
-                          @"--noverify",
-                          @"--puppetstrings" ];
-
-  // Set task environment.
-  NSMutableDictionary *environment = [[[NSProcessInfo processInfo] environment] mutableCopy];
-  environment[@"NSUnbufferedIO"] = @"YES";
-  self.asr.environment = environment;
-
-  // Create output pipe & file handle.
-  self.asr.standardError = self.asr.standardOutput = [[NSPipe alloc] init];
-  NSFileHandle *outputFh = [self.asr.standardOutput fileHandleForReading];
-  outputFh.readabilityHandler = ^(NSFileHandle *h) {
-    NSData *availableData = [h availableData];
-    [self processOutput:(NSData *)availableData];
-  };
-
-  // Launch and wait for exit.
-  [self.asr launch];
-  [self.asr waitUntilExit];
-
-  // Clear readability handler or the file handle is never released.
-  outputFh.readabilityHandler = nil;
+  int retries = (int)[[[NSUserDefaults alloc] initWithSuiteName:kPreferenceDomain] integerForKey:kASRAttempts];
+  if (retries <= 0) retries = 1;
+  NSLog(@"Will attempt ASR up to %d times.", retries);
+  for (int i = 0; i < retries; ++i) {
+    [self unmountDisk:self.diskRef withOptions:kDADiskUnmountOptionWhole];
+
+    self.asr = [[NSTask alloc] init];
+    self.asr.launchPath = @"/usr/sbin/asr";
+    self.asr.arguments = @[ @"restore",
+                            @"--buffersize",
+                            @"16m",
+                            @"--source",
+                            path,
+                            @"--target",
+                            self.destination.path,
+                            @"--erase",
+                            @"--noprompt",
+                            @"--noverify",
+                            @"--puppetstrings" ];
+
+    // Set task environment.
+    NSMutableDictionary *environment = [[[NSProcessInfo processInfo] environment] mutableCopy];
+    environment[@"NSUnbufferedIO"] = @"YES";
+    self.asr.environment = environment;
+
+    // Create output pipe & file handle.
+    self.asr.standardError = self.asr.standardOutput = [[NSPipe alloc] init];
+    NSFileHandle *outputFh = [self.asr.standardOutput fileHandleForReading];
+    outputFh.readabilityHandler = ^(NSFileHandle *h) {
+      NSData *availableData = [h availableData];
+      [self processOutput:(NSData *)availableData];
+    };
+
+    // Launch and wait for exit.
+    [self.asr launch];
+    [self.asr waitUntilExit];
+
+    // Clear readability handler or the file handle is never released.
+    outputFh.readabilityHandler = nil;
+    if (self.asr && self.asr.terminationStatus == 0) break;
+    NSLog(@"%@ ASR attempt %d exit code: %d", self, i + 1, self.asr.terminationStatus);
+  }
   return self.asr ? self.asr.terminationStatus : -1;
 }