Fix race conditions in HubEndpoint

The current logic is error-prone because it's possible that mServiceToken becomes null between null checks and calling the interface. Instead, this CL modfiies the logic such that mServiceToken is never nullified, and we rely on the service implementation's error handling (IllegalStateException, as noted in the javadocs). For callbacks, it's always possible to hit this race condition if we were still handling a previously invoked callback before the endpoint was unregistered, in which case we log a warning.

Bug: 378974199
Flag: android.chre.flags.offload_implementation
Test: Tests pass

Change-Id: Ic2c9eeb1684caddacbcdfb4bdc431afa5b663c36

1 files changed, 32 insertions, 67 deletions

diff --git a/core/java/android/hardware/contexthub/HubEndpoint.java b/core/java/android/hardware/contexthub/HubEndpoint.java
index 70d63efed2fb..14911de6d032 100644
--- a/core/java/android/hardware/contexthub/HubEndpoint.java
+++ b/core/java/android/hardware/contexthub/HubEndpoint.java
@@ -107,6 +107,13 @@ public class HubEndpoint {
     @GuardedBy("mLock")
     private final SparseArray<HubEndpointSession> mActiveSessions = new SparseArray<>();
 
+    /*
+     * Internal interface used to invoke IContextHubEndpoint calls.
+     */
+    interface EndpointConsumer {
+        void accept(IContextHubEndpoint endpoint) throws RemoteException;
+    }
+
     private final IContextHubEndpointCallback mServiceCallback =
             new IContextHubEndpointCallback.Stub() {
                 @Override
@@ -229,12 +236,12 @@ public class HubEndpoint {
                 private void sendMessageDeliveryStatus(
                         int sessionId, HubMessage message, byte errorCode) {
                     if (message.isResponseRequired()) {
-                        try {
-                            mServiceToken.sendMessageDeliveryStatus(
-                                    sessionId, message.getMessageSequenceNumber(), errorCode);
-                        } catch (RemoteException e) {
-                            e.rethrowFromSystemServer();
-                        }
+                        invokeCallback(
+                                (callback) ->
+                                        callback.sendMessageDeliveryStatus(
+                                                sessionId,
+                                                message.getMessageSequenceNumber(),
+                                                errorCode));
                     }
                     invokeCallbackFinished();
                 }
@@ -270,11 +277,6 @@ public class HubEndpoint {
                         int sessionId,
                         HubEndpointInfo initiator,
                         @Nullable String serviceDescriptor) {
-                    if (mServiceToken == null || mAssignedHubEndpointInfo == null) {
-                        // No longer registered?
-                        return;
-                    }
-
                     // Retrieve the active session
                     HubEndpointSession activeSession;
                     synchronized (mLock) {
@@ -295,14 +297,9 @@ public class HubEndpoint {
                                         mAssignedHubEndpointInfo,
                                         initiator,
                                         serviceDescriptor);
-                        try {
-                            // oneway call to notify system service that the request is completed
-                            mServiceToken.openSessionRequestComplete(sessionId);
-                        } catch (RemoteException e) {
-                            Log.e(TAG, "onSessionOpenRequestResult: ", e);
-                            return;
-                        }
 
+                        invokeCallback(
+                                (callback) -> callback.openSessionRequestComplete(sessionId));
                         mActiveSessions.put(sessionId, activeSession);
                     }
 
@@ -316,22 +313,25 @@ public class HubEndpoint {
                 }
 
                 private void rejectSession(int sessionId) {
-                    if (mServiceToken == null || mAssignedHubEndpointInfo == null) {
-                        // No longer registered?
-                        return;
-                    }
-
-                    try {
-                        mServiceToken.closeSession(
-                                sessionId, REASON_OPEN_ENDPOINT_SESSION_REQUEST_REJECTED);
-                    } catch (RemoteException e) {
-                        e.rethrowFromSystemServer();
-                    }
+                    invokeCallback(
+                            (callback) ->
+                                    callback.closeSession(
+                                            sessionId,
+                                            REASON_OPEN_ENDPOINT_SESSION_REQUEST_REJECTED));
                 }
 
                 private void invokeCallbackFinished() {
+                    invokeCallback((callback) -> callback.onCallbackFinished());
+                }
+
+                private void invokeCallback(EndpointConsumer consumer) {
                     try {
-                        mServiceToken.onCallbackFinished();
+                        consumer.accept(mServiceToken);
+                    } catch (IllegalStateException e) {
+                        // It's possible to hit this exception if the endpoint was unregistered
+                        // while processing the callback. It's not a fatal error so we just log
+                        // a warning.
+                        Log.w(TAG, "IllegalStateException while calling callback", e);
                     } catch (RemoteException e) {
                         e.rethrowFromSystemServer();
                     }
@@ -359,11 +359,6 @@ public class HubEndpoint {
 
     /** @hide */
     public void register(IContextHubService service) {
-        // TODO(b/378974199): Consider refactor these assertions
-        if (mServiceToken != null) {
-            // Already registered
-            return;
-        }
         try {
             IContextHubEndpoint serviceToken =
                     service.registerEndpoint(
@@ -381,13 +376,6 @@ public class HubEndpoint {
 
     /** @hide */
     public void unregister() {
-        IContextHubEndpoint serviceToken = mServiceToken;
-        // TODO(b/378974199): Consider refactor these assertions
-        if (serviceToken == null) {
-            // Not yet registered
-            return;
-        }
-
         try {
             synchronized (mLock) {
                 // Don't call HubEndpointSession.close() here.
@@ -400,20 +388,11 @@ public class HubEndpoint {
         } catch (RemoteException e) {
             Log.e(TAG, "unregisterEndpoint: failed to unregister endpoint", e);
             e.rethrowFromSystemServer();
-        } finally {
-            mServiceToken = null;
-            mAssignedHubEndpointInfo = null;
         }
     }
 
     /** @hide */
     public void openSession(HubEndpointInfo destinationInfo, @Nullable String serviceDescriptor) {
-        // TODO(b/378974199): Consider refactor these assertions
-        if (mServiceToken == null || mAssignedHubEndpointInfo == null) {
-            // No longer registered?
-            return;
-        }
-
         HubEndpointSession newSession;
         try {
             synchronized (mLock) {
@@ -439,13 +418,6 @@ public class HubEndpoint {
 
     /** @hide */
     public void closeSession(HubEndpointSession session) {
-        IContextHubEndpoint serviceToken = mServiceToken;
-        // TODO(b/378974199): Consider refactor these assertions
-        if (serviceToken == null || mAssignedHubEndpointInfo == null) {
-            // Not registered
-            return;
-        }
-
         synchronized (mLock) {
             if (!mActiveSessions.contains(session.getId())) {
                 // Already closed?
@@ -456,8 +428,7 @@ public class HubEndpoint {
         }
 
         try {
-            // Oneway notification to system service
-            serviceToken.closeSession(session.getId(), REASON_CLOSE_ENDPOINT_SESSION_REQUESTED);
+            mServiceToken.closeSession(session.getId(), REASON_CLOSE_ENDPOINT_SESSION_REQUESTED);
         } catch (RemoteException e) {
             Log.e(TAG, "closeSession: failed to close session " + session, e);
             e.rethrowFromSystemServer();
@@ -468,14 +439,8 @@ public class HubEndpoint {
             HubEndpointSession session,
             HubMessage message,
             @Nullable IContextHubTransactionCallback transactionCallback) {
-        IContextHubEndpoint serviceToken = mServiceToken;
-        if (serviceToken == null) {
-            // Not registered
-            return;
-        }
-
         try {
-            serviceToken.sendMessage(session.getId(), message, transactionCallback);
+            mServiceToken.sendMessage(session.getId(), message, transactionCallback);
         } catch (RemoteException e) {
             Log.e(TAG, "sendMessage: failed to send message session=" + session, e);
             e.rethrowFromSystemServer();