From e032073fc084a20e6d63ece71b78d7808eec9168 Mon Sep 17 00:00:00 2001
From: Kevin Coppock <47542933+g-coppock@users.noreply.github.com>
Date: Tue, 31 May 2022 08:54:33 -0500
Subject: [PATCH] Fix retrieval of attestation challenge in Java (#18867)

When not using the auto-commissioning flow, the mDeviceBeingCommissioned
proxy is no longer assigned, so the default behavior of retrieving the
attestation challenge fails.

Since at the time of use we already have a reference to the correct
DeviceProxy, added an overload to pass the proxy along to retrieve the
challenge.
---
 src/app/DeviceProxy.h                         | 19 +++++++++++++++++++
 src/controller/CHIPDeviceController.cpp       | 13 -------------
 src/controller/CHIPDeviceController.h         | 10 ----------
 .../java/CHIPDeviceController-JNI.cpp         |  3 +--
 4 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/src/app/DeviceProxy.h b/src/app/DeviceProxy.h
index 2b9bd0e67b24a7..843ebd9d364111 100644
--- a/src/app/DeviceProxy.h
+++ b/src/app/DeviceProxy.h
@@ -59,6 +59,25 @@ class DLL_EXPORT DeviceProxy
 
     const ReliableMessageProtocolConfig & GetRemoteMRPConfig() const { return mRemoteMRPConfig; }
 
+    /**
+     * @brief
+     *   This function returns the attestation challenge for the secure session.
+     *
+     * @param[out] attestationChallenge The output for the attestationChallenge
+     *
+     * @return CHIP_ERROR               CHIP_NO_ERROR on success, or CHIP_ERROR_INVALID_ARGUMENT if no secure session is active
+     */
+    virtual CHIP_ERROR GetAttestationChallenge(ByteSpan & attestationChallenge)
+    {
+        Optional<SessionHandle> secureSessionHandle;
+
+        secureSessionHandle = GetSecureSession();
+        VerifyOrReturnError(secureSessionHandle.HasValue(), CHIP_ERROR_INCORRECT_STATE);
+
+        attestationChallenge = secureSessionHandle.Value()->AsSecureSession()->GetCryptoContext().GetAttestationChallenge();
+        return CHIP_NO_ERROR;
+    }
+
 protected:
     virtual bool IsSecureConnected() const = 0;
 
diff --git a/src/controller/CHIPDeviceController.cpp b/src/controller/CHIPDeviceController.cpp
index 1e4ab7b6e320b4..f134987b0edfc4 100644
--- a/src/controller/CHIPDeviceController.cpp
+++ b/src/controller/CHIPDeviceController.cpp
@@ -736,19 +736,6 @@ DeviceCommissioner::ContinueCommissioningAfterDeviceAttestationFailure(DevicePro
     return CHIP_NO_ERROR;
 }
 
-CHIP_ERROR DeviceCommissioner::GetAttestationChallenge(ByteSpan & attestationChallenge)
-{
-    Optional<SessionHandle> secureSessionHandle;
-
-    VerifyOrReturnError(mDeviceBeingCommissioned != nullptr, CHIP_ERROR_INCORRECT_STATE);
-
-    secureSessionHandle = mDeviceBeingCommissioned->GetSecureSession();
-    VerifyOrReturnError(secureSessionHandle.HasValue(), CHIP_ERROR_INCORRECT_STATE);
-
-    attestationChallenge = secureSessionHandle.Value()->AsSecureSession()->GetCryptoContext().GetAttestationChallenge();
-    return CHIP_NO_ERROR;
-}
-
 CHIP_ERROR DeviceCommissioner::StopPairing(NodeId remoteDeviceId)
 {
     VerifyOrReturnError(mState == State::Initialized, CHIP_ERROR_INCORRECT_STATE);
diff --git a/src/controller/CHIPDeviceController.h b/src/controller/CHIPDeviceController.h
index 5b04057c4ff533..91838ed1a0e343 100644
--- a/src/controller/CHIPDeviceController.h
+++ b/src/controller/CHIPDeviceController.h
@@ -463,16 +463,6 @@ class DLL_EXPORT DeviceCommissioner : public DeviceController,
 
     CHIP_ERROR GetDeviceBeingCommissioned(NodeId deviceId, CommissioneeDeviceProxy ** device);
 
-    /**
-     * @brief
-     *   This function returns the attestation challenge for the secure session of the device being commissioned.
-     *
-     * @param[out] attestationChallenge The output for the attestationChallenge
-     *
-     * @return CHIP_ERROR               CHIP_NO_ERROR on success, or CHIP_ERROR_INVALID_ARGUMENT if no secure session is active
-     */
-    CHIP_ERROR GetAttestationChallenge(ByteSpan & attestationChallenge);
-
     /**
      * @brief
      *   This function stops a pairing process that's in progress. It does not delete the pairing of a previously
diff --git a/src/controller/java/CHIPDeviceController-JNI.cpp b/src/controller/java/CHIPDeviceController-JNI.cpp
index ef229d1626a158..c32103dccacccd 100644
--- a/src/controller/java/CHIPDeviceController-JNI.cpp
+++ b/src/controller/java/CHIPDeviceController-JNI.cpp
@@ -655,8 +655,7 @@ JNI_METHOD(jbyteArray, getAttestationChallenge)
         JniReferences::GetInstance().ThrowError(env, sChipDeviceControllerExceptionCls, CHIP_ERROR_INCORRECT_STATE);
     }
 
-    AndroidDeviceControllerWrapper * wrapper = AndroidDeviceControllerWrapper::FromJNIHandle(handle);
-    err                                      = wrapper->Controller()->GetAttestationChallenge(attestationChallenge);
+    err = chipDevice->GetAttestationChallenge(attestationChallenge);
     SuccessOrExit(err);
     VerifyOrExit(attestationChallenge.size() == 16, err = CHIP_ERROR_INVALID_ARGUMENT);