From 4f59bd921940964697619ab5dde166f4d5d25f48 Mon Sep 17 00:00:00 2001
From: Mikhail Burshteyn <michael.burshteyn@gmail.com>
Date: Mon, 29 Aug 2022 10:35:14 -0400
Subject: [PATCH] Update PR based on comments from bzbarsky.

---
 .../Framework/CHIP/MTRDeviceController.h      |  7 +++-
 .../Framework/CHIP/MTRDeviceController.mm     | 33 +++++++++----------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/src/darwin/Framework/CHIP/MTRDeviceController.h b/src/darwin/Framework/CHIP/MTRDeviceController.h
index f86443c77e7dab..f45af2407ddc95 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceController.h
+++ b/src/darwin/Framework/CHIP/MTRDeviceController.h
@@ -134,7 +134,12 @@ typedef void (^MTRDeviceConnectionCallback)(MTRBaseDevice * _Nullable device, NS
  */
 - (void)setNocChainIssuer:(id<MTRNOCChainIssuer>)nocChainIssuer queue:(dispatch_queue_t)queue;
 
-/** Return the attestation challenge for the secure session of the device being commissioned. */
+/** 
+ * Return the attestation challenge for the secure session of the device being commissioned. 
+ *
+ * Attempts to retreive the generated attestation challenge from a commissionee with the given Device ID.
+ * Returns nil if given Device ID does not match an active commissionee, or if a Secure Session is not availale.
+*/
 - (nullable NSData *)generateAttestationChallengeForDeviceId:(uint64_t)deviceId;
 
 /**
diff --git a/src/darwin/Framework/CHIP/MTRDeviceController.mm b/src/darwin/Framework/CHIP/MTRDeviceController.mm
index a908172ac65b7a..a70892d8d247b4 100644
--- a/src/darwin/Framework/CHIP/MTRDeviceController.mm
+++ b/src/darwin/Framework/CHIP/MTRDeviceController.mm
@@ -71,6 +71,8 @@
 static NSString * const kErrorGenerateNOC = @"Generating operational certificate failed";
 static NSString * const kErrorKeyAllocation = @"Generating new operational key failed";
 static NSString * const kErrorCSRValidation = @"Extracting public key from CSR failed";
+static NSString * const kErrorGetCommissionee = @"Failure obtaining device being commissioned"
+static NSString * const kErrorGetAttestationChallenge = @"Failure getting attestation challenge"
 
 @interface MTRDeviceController ()
 
@@ -705,30 +707,25 @@ - (nullable NSData *)computePaseVerifier:(uint32_t)setupPincode iterations:(uint
 
 - (nullable NSData *)generateAttestationChallengeForDeviceId:(uint64_t)deviceId
 {
-    __block CHIP_ERROR errorCode = CHIP_ERROR_INCORRECT_STATE;
-    if (![self isRunning]) {
-        [self checkForError:errorCode logMsg:kErrorNotRunning error:nil];
-        return nil;
-    }
+    VerifyOrReturn([self checkIsRunning], nil);
 
     __block NSData * attestationChallenge;
     dispatch_sync(_chipWorkQueue, ^{
-        if ([self isRunning]) {
-            chip::CommissioneeDeviceProxy * deviceProxy;
-            errorCode = self.cppCommissioner->GetDeviceBeingCommissioned(deviceId, &deviceProxy);
-            if (errorCode != CHIP_NO_ERROR) {
-                [self checkForError:errorCode logMsg:@"Invalid Attestation Challenge device ID." error:nil];
-                return;
-            }
+        VerifyOrReturn([self checkIsRunning]);
+
+        chip::CommissioneeDeviceProxy * deviceProxy;
+        errorCode = self.cppCommissioner->GetDeviceBeingCommissioned(deviceId, &deviceProxy);
+        auto success = ![self checkForError:errorCode logMsg:kErrorGetCommissionee error:nil];
+        VerifyOrReturn(success);
 
-            NSMutableData * challengeBuffer = [[NSMutableData alloc] initWithLength:chip::Crypto::kAES_CCM128_Key_Length];
-            chip::ByteSpan challenge((uint8_t *) [challengeBuffer mutableBytes], chip::Crypto::kAES_CCM128_Key_Length);
+        uint8_t challengeBuffer[chip::Crypto::kAES_CCM128_Key_Length];
+        chip::ByteSpan challenge(challengeBuffer);
 
-            errorCode = deviceProxy->GetAttestationChallenge(challenge);
-            MTR_LOG_ERROR("GetAttestationChallenge: %s", chip::ErrorStr(errorCode));
+        errorCode = deviceProxy->GetAttestationChallenge(challenge);
+        success = ![self checkForError:errorCode logMsg:kErrorGetAttestationChallenge error:nil];
+        VerifyOrReturn(success);
 
-            attestationChallenge = [NSData dataWithBytes:challenge.data() length:challenge.size()];
-        }
+        attestationChallenge = AsData(challenge);
     });
 
     return attestationChallenge;