Adds ContextHubClient to callback parameters

Bug: 67734082
Test: make
Change-Id: Id09322c5375ee87a7533e77696e4673854bae4be
diff --git a/core/java/android/hardware/location/ContextHubClient.java b/core/java/android/hardware/location/ContextHubClient.java
index fa17a55..4ba5e80 100644
--- a/core/java/android/hardware/location/ContextHubClient.java
+++ b/core/java/android/hardware/location/ContextHubClient.java
@@ -38,12 +38,7 @@
     /*
      * The proxy to the client interface at the service.
      */
-    private final IContextHubClient mClientProxy;
-
-    /*
-     * The callback interface associated with this client.
-     */
-    private final IContextHubClientCallback mCallbackInterface;
+    private IContextHubClient mClientProxy = null;
 
     /*
      * The Context Hub that this client is attached to.
@@ -54,16 +49,28 @@
 
     private final AtomicBoolean mIsClosed = new AtomicBoolean(false);
 
-    /* package */ ContextHubClient(
-            IContextHubClient clientProxy, IContextHubClientCallback callback,
-            ContextHubInfo hubInfo) {
-        mClientProxy = clientProxy;
-        mCallbackInterface = callback;
+    /* package */ ContextHubClient(ContextHubInfo hubInfo) {
         mAttachedHub = hubInfo;
         mCloseGuard.open("close");
     }
 
     /**
+     * Sets the proxy interface of the client at the service. This method should always be called
+     * by the ContextHubManager after the client is registered at the service, and should only be
+     * called once.
+     *
+     * @param clientProxy the proxy of the client at the service
+     */
+    /* package */ void setClientProxy(IContextHubClient clientProxy) {
+        Preconditions.checkNotNull(clientProxy, "IContextHubClient cannot be null");
+        if (mClientProxy != null) {
+            throw new IllegalStateException("Cannot change client proxy multiple times");
+        }
+
+        mClientProxy = clientProxy;
+    }
+
+    /**
      * Returns the hub that this client is attached to.
      *
      * @return the ContextHubInfo of the attached hub
diff --git a/core/java/android/hardware/location/ContextHubClientCallback.java b/core/java/android/hardware/location/ContextHubClientCallback.java
index ab19d54..6c48ee9c 100644
--- a/core/java/android/hardware/location/ContextHubClientCallback.java
+++ b/core/java/android/hardware/location/ContextHubClientCallback.java
@@ -38,48 +38,56 @@
      * The message contents of this callback may either be broadcasted or targeted to the
      * client receiving the invocation.
      *
+     * @param client the client that is associated with this callback
      * @param message the message sent by the nanoapp
      */
-    public void onMessageFromNanoApp(NanoAppMessage message) {}
+    public void onMessageFromNanoApp(ContextHubClient client, NanoAppMessage message) {}
 
     /**
      * Callback invoked when the attached Context Hub has reset.
+     *
+     * @param client the client that is associated with this callback
      */
-    public void onHubReset() {}
+    public void onHubReset(ContextHubClient client) {}
 
     /**
      * Callback invoked when a nanoapp aborts at the attached Context Hub.
      *
+     * @param client the client that is associated with this callback
      * @param nanoAppId the ID of the nanoapp that had aborted
      * @param abortCode the reason for nanoapp's abort, specific to each nanoapp
      */
-    public void onNanoAppAborted(long nanoAppId, int abortCode) {}
+    public void onNanoAppAborted(ContextHubClient client, long nanoAppId, int abortCode) {}
 
     /**
      * Callback invoked when a nanoapp is loaded at the attached Context Hub.
      *
+     * @param client the client that is associated with this callback
      * @param nanoAppId the ID of the nanoapp that had been loaded
      */
-    public void onNanoAppLoaded(long nanoAppId) {}
+    public void onNanoAppLoaded(ContextHubClient client, long nanoAppId) {}
 
     /**
      * Callback invoked when a nanoapp is unloaded from the attached Context Hub.
      *
+     * @param client the client that is associated with this callback
      * @param nanoAppId the ID of the nanoapp that had been unloaded
      */
-    public void onNanoAppUnloaded(long nanoAppId) {}
+    public void onNanoAppUnloaded(ContextHubClient client, long nanoAppId) {}
 
     /**
      * Callback invoked when a nanoapp is enabled at the attached Context Hub.
      *
+     * @param client the client that is associated with this callback
      * @param nanoAppId the ID of the nanoapp that had been enabled
      */
-    public void onNanoAppEnabled(long nanoAppId) {}
+    public void onNanoAppEnabled(ContextHubClient client, long nanoAppId) {}
 
     /**
      * Callback invoked when a nanoapp is disabled at the attached Context Hub.
      *
+     * @param client the client that is associated with this callback
      * @param nanoAppId the ID of the nanoapp that had been disabled
      */
-    public void onNanoAppDisabled(long nanoAppId) {}
+    public void onNanoAppDisabled(ContextHubClient client, long nanoAppId) {}
 }
diff --git a/core/java/android/hardware/location/ContextHubManager.java b/core/java/android/hardware/location/ContextHubManager.java
index 5d87aaf..237702b 100644
--- a/core/java/android/hardware/location/ContextHubManager.java
+++ b/core/java/android/hardware/location/ContextHubManager.java
@@ -544,47 +544,48 @@
     /**
      * Creates an interface to the ContextHubClient to send down to the service.
      *
+     * @param client the ContextHubClient object associated with this callback
      * @param callback the callback to invoke at the client process
      * @param executor the executor to invoke callbacks for this client
      *
      * @return the callback interface
      */
     private IContextHubClientCallback createClientCallback(
-            ContextHubClientCallback callback, Executor executor) {
+            ContextHubClient client, ContextHubClientCallback callback, Executor executor) {
         return new IContextHubClientCallback.Stub() {
             @Override
             public void onMessageFromNanoApp(NanoAppMessage message) {
-                executor.execute(() -> callback.onMessageFromNanoApp(message));
+                executor.execute(() -> callback.onMessageFromNanoApp(client, message));
             }
 
             @Override
             public void onHubReset() {
-                executor.execute(() -> callback.onHubReset());
+                executor.execute(() -> callback.onHubReset(client));
             }
 
             @Override
             public void onNanoAppAborted(long nanoAppId, int abortCode) {
-                executor.execute(() -> callback.onNanoAppAborted(nanoAppId, abortCode));
+                executor.execute(() -> callback.onNanoAppAborted(client, nanoAppId, abortCode));
             }
 
             @Override
             public void onNanoAppLoaded(long nanoAppId) {
-                executor.execute(() -> callback.onNanoAppLoaded(nanoAppId));
+                executor.execute(() -> callback.onNanoAppLoaded(client, nanoAppId));
             }
 
             @Override
             public void onNanoAppUnloaded(long nanoAppId) {
-                executor.execute(() -> callback.onNanoAppUnloaded(nanoAppId));
+                executor.execute(() -> callback.onNanoAppUnloaded(client, nanoAppId));
             }
 
             @Override
             public void onNanoAppEnabled(long nanoAppId) {
-                executor.execute(() -> callback.onNanoAppEnabled(nanoAppId));
+                executor.execute(() -> callback.onNanoAppEnabled(client, nanoAppId));
             }
 
             @Override
             public void onNanoAppDisabled(long nanoAppId) {
-                executor.execute(() -> callback.onNanoAppDisabled(nanoAppId));
+                executor.execute(() -> callback.onNanoAppDisabled(client, nanoAppId));
             }
         };
     }
@@ -615,16 +616,19 @@
         Preconditions.checkNotNull(hubInfo, "ContextHubInfo cannot be null");
         Preconditions.checkNotNull(executor, "Executor cannot be null");
 
-        IContextHubClientCallback clientInterface = createClientCallback(callback, executor);
+        ContextHubClient client = new ContextHubClient(hubInfo);
+        IContextHubClientCallback clientInterface = createClientCallback(
+                client, callback, executor);
 
-        IContextHubClient client;
+        IContextHubClient clientProxy;
         try {
-            client = mService.createClient(clientInterface, hubInfo.getId());
+            clientProxy = mService.createClient(clientInterface, hubInfo.getId());
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
 
-        return new ContextHubClient(client, clientInterface, hubInfo);
+        client.setClientProxy(clientProxy);
+        return client;
     }
 
     /**