contexthub: Make client registration API safer

Modifies the registration API for new endpoint session capabilities to
return a new dedicated IEndpointCommunication interface tied to the
client registration. All endpoint/session APIs have been moved to the
new interface. This prevents security issues resulting from pid re-use.

Adds an unregister() API to IEndpointCommunication for a client to
signal that its associated hub will no longer be active. Subsequent
calls on IEndpointCommunication APIs will fail.

Bug: 384897139
Flag: android.chre.flags.offload_api
Test: Builds
Change-Id: Ibd8a28e96315ca146d62413b5337ece9912d14c4
diff --git a/contexthub/aidl/aidl_api/android.hardware.contexthub/current/android/hardware/contexthub/IContextHub.aidl b/contexthub/aidl/aidl_api/android.hardware.contexthub/current/android/hardware/contexthub/IContextHub.aidl
index 2646d15..2940745 100644
--- a/contexthub/aidl/aidl_api/android.hardware.contexthub/current/android/hardware/contexthub/IContextHub.aidl
+++ b/contexthub/aidl/aidl_api/android.hardware.contexthub/current/android/hardware/contexthub/IContextHub.aidl
@@ -51,14 +51,6 @@
   void sendMessageDeliveryStatusToHub(in int contextHubId, in android.hardware.contexthub.MessageDeliveryStatus messageDeliveryStatus);
   List<android.hardware.contexthub.HubInfo> getHubs();
   List<android.hardware.contexthub.EndpointInfo> getEndpoints();
-  void registerEndpoint(in android.hardware.contexthub.EndpointInfo endpoint);
-  void unregisterEndpoint(in android.hardware.contexthub.EndpointInfo endpoint);
-  void registerEndpointCallback(in android.hardware.contexthub.IEndpointCallback callback);
-  int[2] requestSessionIdRange(int size);
-  void openEndpointSession(int sessionId, in android.hardware.contexthub.EndpointId destination, in android.hardware.contexthub.EndpointId initiator, in @nullable String serviceDescriptor);
-  void sendMessageToEndpoint(int sessionId, in android.hardware.contexthub.Message msg);
-  void sendMessageDeliveryStatusToEndpoint(int sessionId, in android.hardware.contexthub.MessageDeliveryStatus msgStatus);
-  void closeEndpointSession(int sessionId, in android.hardware.contexthub.Reason reason);
-  void endpointSessionOpenComplete(int sessionId);
+  @PropagateAllowBlocking android.hardware.contexthub.IEndpointCommunication registerEndpointHub(in android.hardware.contexthub.IEndpointCallback callback, in android.hardware.contexthub.HubInfo hubInfo);
   const int EX_CONTEXT_HUB_UNSPECIFIED = (-1) /* -1 */;
 }
diff --git a/contexthub/aidl/aidl_api/android.hardware.contexthub/current/android/hardware/contexthub/IEndpointCommunication.aidl b/contexthub/aidl/aidl_api/android.hardware.contexthub/current/android/hardware/contexthub/IEndpointCommunication.aidl
new file mode 100644
index 0000000..8742415
--- /dev/null
+++ b/contexthub/aidl/aidl_api/android.hardware.contexthub/current/android/hardware/contexthub/IEndpointCommunication.aidl
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+///////////////////////////////////////////////////////////////////////////////
+// THIS FILE IS IMMUTABLE. DO NOT EDIT IN ANY CASE.                          //
+///////////////////////////////////////////////////////////////////////////////
+
+// This file is a snapshot of an AIDL file. Do not edit it manually. There are
+// two cases:
+// 1). this is a frozen version file - do not edit this in any case.
+// 2). this is a 'current' file. If you make a backwards compatible change to
+//     the interface (from the latest frozen version), the build system will
+//     prompt you to update this file with `m <name>-update-api`.
+//
+// You must not make a backward incompatible change to any AIDL file built
+// with the aidl_interface module type with versions property set. The module
+// type is used to build AIDL files in a way that they can be used across
+// independently updatable components of the system. If a device is shipped
+// with such a backward incompatible change, it has a high risk of breaking
+// later when a module using the interface is updated, e.g., Mainline modules.
+
+package android.hardware.contexthub;
+@VintfStability
+interface IEndpointCommunication {
+  void registerEndpoint(in android.hardware.contexthub.EndpointInfo endpoint);
+  void unregisterEndpoint(in android.hardware.contexthub.EndpointInfo endpoint);
+  int[2] requestSessionIdRange(int size);
+  void openEndpointSession(int sessionId, in android.hardware.contexthub.EndpointId destination, in android.hardware.contexthub.EndpointId initiator, in @nullable String serviceDescriptor);
+  void sendMessageToEndpoint(int sessionId, in android.hardware.contexthub.Message msg);
+  void sendMessageDeliveryStatusToEndpoint(int sessionId, in android.hardware.contexthub.MessageDeliveryStatus msgStatus);
+  void closeEndpointSession(int sessionId, in android.hardware.contexthub.Reason reason);
+  void endpointSessionOpenComplete(int sessionId);
+  void unregister();
+}
diff --git a/contexthub/aidl/android/hardware/contexthub/IContextHub.aidl b/contexthub/aidl/android/hardware/contexthub/IContextHub.aidl
index eb6d051..3fb452c 100644
--- a/contexthub/aidl/android/hardware/contexthub/IContextHub.aidl
+++ b/contexthub/aidl/android/hardware/contexthub/IContextHub.aidl
@@ -18,19 +18,16 @@
 
 import android.hardware.contexthub.ContextHubInfo;
 import android.hardware.contexthub.ContextHubMessage;
-import android.hardware.contexthub.EndpointId;
 import android.hardware.contexthub.EndpointInfo;
 import android.hardware.contexthub.HostEndpointInfo;
 import android.hardware.contexthub.HubInfo;
 import android.hardware.contexthub.IContextHubCallback;
 import android.hardware.contexthub.IEndpointCallback;
-import android.hardware.contexthub.Message;
+import android.hardware.contexthub.IEndpointCommunication;
 import android.hardware.contexthub.MessageDeliveryStatus;
 import android.hardware.contexthub.NanSessionStateUpdate;
 import android.hardware.contexthub.NanoappBinary;
 import android.hardware.contexthub.NanoappInfo;
-import android.hardware.contexthub.Reason;
-import android.hardware.contexthub.Service;
 import android.hardware.contexthub.Setting;
 
 @VintfStability
@@ -271,125 +268,22 @@
     List<EndpointInfo> getEndpoints();
 
     /**
-     * Publishes an endpoint from the calling side (e.g. Android). Endpoints must be registered
-     * prior to starting a session.
+     * Registers a new hub for endpoint communication which will receive events for its endpoints
+     * over the given callback. Returns an interface for the hub to register endpoints, start
+     * sessions, and send messages.
+     *
+     * It is valid for the same callback to be registered for multiple hubs, as the
+     * IEndpointCallback events provide sufficient information to determine which hub the event is
+     * intended for:
+     * * session ids are allocated to a specific hub and are unique
+     * * endpoints are identified by hub and endpoint id
+     *
+     * @param callback Interface to send endpoint events targeting the caller
+     * @param hubInfo Details of the hub being registered
+     * @return Interface for the hub to interact with other endpoint hubs
+     *
+     * @throws EX_ILLEGAL_STATE if hubInfo.hubId has already been registered
      */
-    void registerEndpoint(in EndpointInfo endpoint);
-
-    /**
-     * Teardown an endpoint from the calling side (e.g. Android). This endpoint must have already
-     * been published via registerEndpoint().
-     */
-    void unregisterEndpoint(in EndpointInfo endpoint);
-
-    /**
-     * Attaches a callback interface to receive events targeted at endpoints registered by the
-     * caller.
-     */
-    void registerEndpointCallback(in IEndpointCallback callback);
-
-    /**
-     * Request a range of session IDs for the caller to use when initiating sessions. This may be
-     * called more than once, but typical usage is to request a large enough range to accommodate
-     * the maximum expected number of concurrent sessions, but not overly large as to limit other
-     * clients.
-     *
-     * @param size The number of sessionId reserved for host-initiated sessions. This number should
-     *         be less than or equal to 1024.
-     *
-     * @return An array with two elements representing the smallest and largest possible session id
-     *         available for host.
-     *
-     * @throws EX_ILLEGAL_ARGUMENT if the size is invalid.
-     * @throws EX_SERVICE_SPECIFIC if the id range requested cannot be allocated.
-     */
-    int[2] requestSessionIdRange(int size);
-
-    /**
-     * Request to open a session for communication between an endpoint previously registered by the
-     * caller and a target endpoint found in getEndpoints(), optionally scoped to a service
-     * published by the target endpoint.
-     *
-     * Upon returning from this function, the session is in pending state, and the final result will
-     * be given by an asynchronous call to onEndpointSessionOpenComplete() on success, or
-     * onCloseEndpointSession() on failure.
-     *
-     * @param sessionId Caller-allocated session identifier, which must be unique across all active
-     *         sessions, and must fall in a range allocated via requestSessionIdRange().
-     * @param destination The EndpointId representing the destination side of the session.
-     * @param initiator The EndpointId representing the initiating side of the session, which
-     *         must've already been published through registerEndpoint().
-     * @param serviceDescriptor Descriptor for the service specification for scoping this session
-     *         (nullable). Null indicates a fully custom marshalling scheme. The value should match
-     *         a published descriptor for both destination and initiator.
-     *
-     * @throws EX_ILLEGAL_ARGUMENT if any of the arguments are invalid, or the combination of the
-     *         arguments is invalid.
-     * @throws EX_SERVICE_SPECIFIC on other errors
-     *         - EX_CONTEXT_HUB_UNSPECIFIED if the request failed for other reasons.
-     */
-    void openEndpointSession(int sessionId, in EndpointId destination, in EndpointId initiator,
-            in @nullable String serviceDescriptor);
-
-    /**
-     * Send a message from one endpoint to another on the (currently open) session.
-     *
-     * @param sessionId The integer representing the communication session, previously set in
-     *         openEndpointSession() or onEndpointSessionOpenRequest().
-     * @param msg The Message object representing a message to endpoint from the endpoint on host.
-     *
-     * @throws EX_ILLEGAL_ARGUMENT if any of the arguments are invalid, or the combination of the
-     *         arguments is invalid.
-     * @throws EX_SERVICE_SPECIFIC on other errors
-     *         - EX_CONTEXT_HUB_UNSPECIFIED if the request failed for other reasons.
-     */
-    void sendMessageToEndpoint(int sessionId, in Message msg);
-
-    /**
-     * Sends a message delivery status to the endpoint in response to receiving a Message with flag
-     * FLAG_REQUIRES_DELIVERY_STATUS. Each message with the flag should have a MessageDeliveryStatus
-     * response. This method sends the message delivery status back to the remote endpoint for a
-     * session.
-     *
-     * @param sessionId The integer representing the communication session, previously set in
-     *         openEndpointSession() or onEndpointSessionOpenRequest().
-     * @param msgStatus The MessageDeliveryStatus object representing the delivery status for a
-     *         specific message (identified by the sequenceNumber) within the session.
-     *
-     * @throws EX_UNSUPPORTED_OPERATION if ContextHubInfo.supportsReliableMessages is false for
-     *          the hub involved in this session.
-     */
-    void sendMessageDeliveryStatusToEndpoint(int sessionId, in MessageDeliveryStatus msgStatus);
-
-    /**
-     * Closes a session previously opened by openEndpointSession() or requested via
-     * onEndpointSessionOpenRequest(). Processing of session closure must be ordered/synchronized
-     * with message delivery, such that if this session was open, any messages previously passed to
-     * sendMessageToEndpoint() that are still in-flight must still be delivered before the session
-     * is closed. Any in-flight messages to the endpoint that requested to close the session will
-     * not be delivered.
-     *
-     * @param sessionId The integer representing the communication session, previously set in
-     *         openEndpointSession() or onEndpointSessionOpenRequest().
-     * @param reason The reason for this close endpoint session request.
-     *
-     * @throws EX_ILLEGAL_ARGUMENT if any of the arguments are invalid, or the combination of the
-     *         arguments is invalid.
-     * @throws EX_SERVICE_SPECIFIC on other errors
-     *         - EX_CONTEXT_HUB_UNSPECIFIED if the request failed for other reasons.
-     */
-    void closeEndpointSession(int sessionId, in Reason reason);
-
-    /**
-     * Notifies the HAL that the session requested by onEndpointSessionOpenRequest is ready to use.
-     *
-     * @param sessionId The integer representing the communication session, previously set in
-     *         onEndpointSessionOpenRequest(). This id is assigned by the HAL.
-     *
-     * @throws EX_ILLEGAL_ARGUMENT if any of the arguments are invalid, or the combination of the
-     *         arguments is invalid.
-     * @throws EX_SERVICE_SPECIFIC on other errors
-     *         - EX_CONTEXT_HUB_UNSPECIFIED if the request failed for other reasons.
-     */
-    void endpointSessionOpenComplete(int sessionId);
+    @PropagateAllowBlocking
+    IEndpointCommunication registerEndpointHub(in IEndpointCallback callback, in HubInfo hubInfo);
 }
diff --git a/contexthub/aidl/android/hardware/contexthub/IEndpointCommunication.aidl b/contexthub/aidl/android/hardware/contexthub/IEndpointCommunication.aidl
new file mode 100644
index 0000000..e5045ba
--- /dev/null
+++ b/contexthub/aidl/android/hardware/contexthub/IEndpointCommunication.aidl
@@ -0,0 +1,152 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.hardware.contexthub;
+
+import android.hardware.contexthub.EndpointId;
+import android.hardware.contexthub.EndpointInfo;
+import android.hardware.contexthub.IEndpointCallback;
+import android.hardware.contexthub.Message;
+import android.hardware.contexthub.MessageDeliveryStatus;
+import android.hardware.contexthub.Reason;
+import android.hardware.contexthub.Service;
+
+@VintfStability
+interface IEndpointCommunication {
+    /**
+     * Publishes an endpoint from the calling side (e.g. Android). Endpoints must be registered
+     * prior to starting a session.
+     */
+    void registerEndpoint(in EndpointInfo endpoint);
+
+    /**
+     * Teardown an endpoint from the calling side (e.g. Android). This endpoint must have already
+     * been published via registerEndpoint().
+     */
+    void unregisterEndpoint(in EndpointInfo endpoint);
+
+    /**
+     * Request a range of session IDs for the caller to use when initiating sessions. This may be
+     * called more than once, but typical usage is to request a large enough range to accommodate
+     * the maximum expected number of concurrent sessions, but not overly large as to limit other
+     * clients.
+     *
+     * @param size The number of sessionId reserved for host-initiated sessions. This number should
+     *         be less than or equal to 1024.
+     *
+     * @return An array with two elements representing the smallest and largest possible session id
+     *         available for host.
+     *
+     * @throws EX_ILLEGAL_ARGUMENT if the size is invalid.
+     * @throws EX_SERVICE_SPECIFIC if the id range requested cannot be allocated.
+     */
+    int[2] requestSessionIdRange(int size);
+
+    /**
+     * Request to open a session for communication between an endpoint previously registered by the
+     * caller and a target endpoint found in getEndpoints(), optionally scoped to a service
+     * published by the target endpoint.
+     *
+     * Upon returning from this function, the session is in pending state, and the final result will
+     * be given by an asynchronous call to onEndpointSessionOpenComplete() on success, or
+     * onCloseEndpointSession() on failure.
+     *
+     * @param sessionId Caller-allocated session identifier, which must be unique across all active
+     *         sessions, and must fall in a range allocated via requestSessionIdRange().
+     * @param destination The EndpointId representing the destination side of the session.
+     * @param initiator The EndpointId representing the initiating side of the session, which
+     *         must've already been published through registerEndpoint().
+     * @param serviceDescriptor Descriptor for the service specification for scoping this session
+     *         (nullable). Null indicates a fully custom marshalling scheme. The value should match
+     *         a published descriptor for both destination and initiator.
+     *
+     * @throws EX_ILLEGAL_ARGUMENT if any of the arguments are invalid, or the combination of the
+     *         arguments is invalid.
+     * @throws EX_SERVICE_SPECIFIC on other errors
+     *         - EX_CONTEXT_HUB_UNSPECIFIED if the request failed for other reasons.
+     */
+    void openEndpointSession(int sessionId, in EndpointId destination, in EndpointId initiator,
+            in @nullable String serviceDescriptor);
+
+    /**
+     * Send a message from one endpoint to another on the (currently open) session.
+     *
+     * @param sessionId The integer representing the communication session, previously set in
+     *         openEndpointSession() or onEndpointSessionOpenRequest().
+     * @param msg The Message object representing a message to endpoint from the endpoint on host.
+     *
+     * @throws EX_ILLEGAL_ARGUMENT if any of the arguments are invalid, or the combination of the
+     *         arguments is invalid.
+     * @throws EX_SERVICE_SPECIFIC on other errors
+     *         - EX_CONTEXT_HUB_UNSPECIFIED if the request failed for other reasons.
+     */
+    void sendMessageToEndpoint(int sessionId, in Message msg);
+
+    /**
+     * Sends a message delivery status to the endpoint in response to receiving a Message with flag
+     * FLAG_REQUIRES_DELIVERY_STATUS. Each message with the flag should have a MessageDeliveryStatus
+     * response. This method sends the message delivery status back to the remote endpoint for a
+     * session.
+     *
+     * @param sessionId The integer representing the communication session, previously set in
+     *         openEndpointSession() or onEndpointSessionOpenRequest().
+     * @param msgStatus The MessageDeliveryStatus object representing the delivery status for a
+     *         specific message (identified by the sequenceNumber) within the session.
+     *
+     * @throws EX_UNSUPPORTED_OPERATION if ContextHubInfo.supportsReliableMessages is false for
+     *          the hub involved in this session.
+     */
+    void sendMessageDeliveryStatusToEndpoint(int sessionId, in MessageDeliveryStatus msgStatus);
+
+    /**
+     * Closes a session previously opened by openEndpointSession() or requested via
+     * onEndpointSessionOpenRequest(). Processing of session closure must be ordered/synchronized
+     * with message delivery, such that if this session was open, any messages previously passed to
+     * sendMessageToEndpoint() that are still in-flight must still be delivered before the session
+     * is closed. Any in-flight messages to the endpoint that requested to close the session will
+     * not be delivered.
+     *
+     * @param sessionId The integer representing the communication session, previously set in
+     *         openEndpointSession() or onEndpointSessionOpenRequest().
+     * @param reason The reason for this close endpoint session request.
+     *
+     * @throws EX_ILLEGAL_ARGUMENT if any of the arguments are invalid, or the combination of the
+     *         arguments is invalid.
+     * @throws EX_SERVICE_SPECIFIC on other errors
+     *         - EX_CONTEXT_HUB_UNSPECIFIED if the request failed for other reasons.
+     */
+    void closeEndpointSession(int sessionId, in Reason reason);
+
+    /**
+     * Notifies the HAL that the session requested by onEndpointSessionOpenRequest is ready to use.
+     *
+     * @param sessionId The integer representing the communication session, previously set in
+     *         onEndpointSessionOpenRequest(). This id is assigned by the HAL.
+     *
+     * @throws EX_ILLEGAL_ARGUMENT if any of the arguments are invalid, or the combination of the
+     *         arguments is invalid.
+     * @throws EX_SERVICE_SPECIFIC on other errors
+     *         - EX_CONTEXT_HUB_UNSPECIFIED if the request failed for other reasons.
+     */
+    void endpointSessionOpenComplete(int sessionId);
+
+    /**
+     * Unregisters this hub. Subsequent calls on this interface will fail.
+     *
+     * @throws EX_ILLEGAL_STATE if this interface was already unregistered.
+     */
+    void unregister();
+}
diff --git a/contexthub/aidl/default/ContextHub.cpp b/contexthub/aidl/default/ContextHub.cpp
index 19d9639..433617e 100644
--- a/contexthub/aidl/default/ContextHub.cpp
+++ b/contexthub/aidl/default/ContextHub.cpp
@@ -150,10 +150,11 @@
 
 ScopedAStatus ContextHub::setTestMode(bool enable) {
     if (enable) {
-        std::unique_lock<std::mutex> lock(mEndpointMutex);
-        mEndpoints.clear();
-        mEndpointSessions.clear();
-        mEndpointCallback = nullptr;
+        std::lock_guard lock(mHostHubsLock);
+        for (auto& [id, hub] : mIdToHostHub) {
+            hub->mActive = false;
+        }
+        mIdToHostHub.clear();
     }
     return ScopedAStatus::ok();
 }
@@ -227,7 +228,23 @@
     return ScopedAStatus::ok();
 };
 
-ScopedAStatus ContextHub::registerEndpoint(const EndpointInfo& in_endpoint) {
+ScopedAStatus ContextHub::registerEndpointHub(
+        const std::shared_ptr<IEndpointCallback>& in_callback, const HubInfo& in_hubInfo,
+        std::shared_ptr<IEndpointCommunication>* _aidl_return) {
+    std::lock_guard lock(mHostHubsLock);
+    if (mIdToHostHub.count(in_hubInfo.hubId)) {
+        return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
+    }
+    auto hub = ndk::SharedRefBase::make<HubInterface>(*this, in_callback, in_hubInfo);
+    mIdToHostHub.insert({in_hubInfo.hubId, hub});
+    *_aidl_return = std::move(hub);
+    return ScopedAStatus::ok();
+}
+
+ScopedAStatus ContextHub::HubInterface::registerEndpoint(const EndpointInfo& in_endpoint) {
+    if (!mActive) {
+        return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
+    }
     std::unique_lock<std::mutex> lock(mEndpointMutex);
 
     for (const EndpointInfo& endpoint : mEndpoints) {
@@ -240,7 +257,10 @@
     return ScopedAStatus::ok();
 };
 
-ScopedAStatus ContextHub::unregisterEndpoint(const EndpointInfo& in_endpoint) {
+ScopedAStatus ContextHub::HubInterface::unregisterEndpoint(const EndpointInfo& in_endpoint) {
+    if (!mActive) {
+        return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
+    }
     std::unique_lock<std::mutex> lock(mEndpointMutex);
 
     for (auto it = mEndpoints.begin(); it != mEndpoints.end(); ++it) {
@@ -252,41 +272,47 @@
     return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
 };
 
-ScopedAStatus ContextHub::registerEndpointCallback(
-        const std::shared_ptr<IEndpointCallback>& in_callback) {
-    std::unique_lock<std::mutex> lock(mEndpointMutex);
-
-    mEndpointCallback = in_callback;
-    return ScopedAStatus::ok();
-};
-
-ScopedAStatus ContextHub::requestSessionIdRange(int32_t in_size,
-                                                std::array<int32_t, 2>* _aidl_return) {
+ScopedAStatus ContextHub::HubInterface::requestSessionIdRange(
+        int32_t in_size, std::array<int32_t, 2>* _aidl_return) {
+    if (!mActive) {
+        return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
+    }
     constexpr int32_t kMaxSize = 1024;
     if (in_size > kMaxSize || _aidl_return == nullptr) {
         return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
     }
 
+    uint16_t base = 0;
     {
-        std::lock_guard<std::mutex> lock(mEndpointMutex);
-        mMaxValidSessionId = in_size;
+        std::lock_guard lock(mHal.mHostHubsLock);
+        if (static_cast<int32_t>(USHRT_MAX) - mHal.mNextSessionIdBase + 1 < in_size) {
+            return ScopedAStatus::fromServiceSpecificError(EX_CONTEXT_HUB_UNSPECIFIED);
+        }
+        base = mHal.mNextSessionIdBase;
+        mHal.mNextSessionIdBase += in_size;
     }
 
-    (*_aidl_return)[0] = 0;
-    (*_aidl_return)[1] = in_size;
+    {
+        std::lock_guard<std::mutex> lock(mEndpointMutex);
+        (*_aidl_return)[0] = mBaseSessionId = base;
+        (*_aidl_return)[1] = mMaxSessionId = base + (in_size - 1);
+    }
     return ScopedAStatus::ok();
 };
 
-ScopedAStatus ContextHub::openEndpointSession(
+ScopedAStatus ContextHub::HubInterface::openEndpointSession(
         int32_t in_sessionId, const EndpointId& in_destination, const EndpointId& in_initiator,
         const std::optional<std::string>& in_serviceDescriptor) {
+    if (!mActive) {
+        return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
+    }
     // We are not calling onCloseEndpointSession on failure because the remote endpoints (our
     // mock endpoints) always accept the session.
 
     std::weak_ptr<IEndpointCallback> callback;
     {
         std::unique_lock<std::mutex> lock(mEndpointMutex);
-        if (in_sessionId > mMaxValidSessionId) {
+        if (in_sessionId < mBaseSessionId || in_sessionId > mMaxSessionId) {
             ALOGE("openEndpointSession: session ID %" PRId32 " is invalid", in_sessionId);
             return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
         }
@@ -346,7 +372,11 @@
     return ScopedAStatus::ok();
 };
 
-ScopedAStatus ContextHub::sendMessageToEndpoint(int32_t in_sessionId, const Message& in_msg) {
+ScopedAStatus ContextHub::HubInterface::sendMessageToEndpoint(int32_t in_sessionId,
+                                                              const Message& in_msg) {
+    if (!mActive) {
+        return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
+    }
     std::weak_ptr<IEndpointCallback> callback;
     {
         std::unique_lock<std::mutex> lock(mEndpointMutex);
@@ -393,12 +423,19 @@
     return ScopedAStatus::ok();
 };
 
-ScopedAStatus ContextHub::sendMessageDeliveryStatusToEndpoint(
+ScopedAStatus ContextHub::HubInterface::sendMessageDeliveryStatusToEndpoint(
         int32_t /* in_sessionId */, const MessageDeliveryStatus& /* in_msgStatus */) {
+    if (!mActive) {
+        return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
+    }
     return ScopedAStatus::ok();
 };
 
-ScopedAStatus ContextHub::closeEndpointSession(int32_t in_sessionId, Reason /* in_reason */) {
+ScopedAStatus ContextHub::HubInterface::closeEndpointSession(int32_t in_sessionId,
+                                                             Reason /* in_reason */) {
+    if (!mActive) {
+        return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
+    }
     std::unique_lock<std::mutex> lock(mEndpointMutex);
 
     for (auto it = mEndpointSessions.begin(); it != mEndpointSessions.end(); ++it) {
@@ -411,8 +448,20 @@
     return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_ARGUMENT);
 };
 
-ScopedAStatus ContextHub::endpointSessionOpenComplete(int32_t /* in_sessionId */) {
+ScopedAStatus ContextHub::HubInterface::endpointSessionOpenComplete(int32_t /* in_sessionId */) {
+    if (!mActive) {
+        return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
+    }
     return ScopedAStatus::ok();
 };
 
+ScopedAStatus ContextHub::HubInterface::unregister() {
+    if (!mActive.exchange(false)) {
+        return ScopedAStatus::fromExceptionCode(EX_ILLEGAL_STATE);
+    }
+    std::lock_guard lock(mHal.mHostHubsLock);
+    mHal.mIdToHostHub.erase(kInfo.hubId);
+    return ScopedAStatus::ok();
+}
+
 }  // namespace aidl::android::hardware::contexthub
diff --git a/contexthub/aidl/default/include/contexthub-impl/ContextHub.h b/contexthub/aidl/default/include/contexthub-impl/ContextHub.h
index 6da8bf2..65e84bb 100644
--- a/contexthub/aidl/default/include/contexthub-impl/ContextHub.h
+++ b/contexthub/aidl/default/include/contexthub-impl/ContextHub.h
@@ -17,7 +17,9 @@
 #pragma once
 
 #include <aidl/android/hardware/contexthub/BnContextHub.h>
+#include <aidl/android/hardware/contexthub/BnEndpointCommunication.h>
 
+#include <atomic>
 #include <mutex>
 #include <unordered_set>
 #include <vector>
@@ -56,54 +58,79 @@
 
     ::ndk::ScopedAStatus getHubs(std::vector<HubInfo>* _aidl_return) override;
     ::ndk::ScopedAStatus getEndpoints(std::vector<EndpointInfo>* _aidl_return) override;
-    ::ndk::ScopedAStatus registerEndpoint(const EndpointInfo& in_endpoint) override;
-    ::ndk::ScopedAStatus unregisterEndpoint(const EndpointInfo& in_endpoint) override;
-    ::ndk::ScopedAStatus registerEndpointCallback(
-            const std::shared_ptr<IEndpointCallback>& in_callback) override;
-    ::ndk::ScopedAStatus requestSessionIdRange(int32_t in_size,
-                                               std::array<int32_t, 2>* _aidl_return) override;
-    ::ndk::ScopedAStatus openEndpointSession(
-            int32_t in_sessionId, const EndpointId& in_destination, const EndpointId& in_initiator,
-            const std::optional<std::string>& in_serviceDescriptor) override;
-    ::ndk::ScopedAStatus sendMessageToEndpoint(int32_t in_sessionId,
-                                               const Message& in_msg) override;
-    ::ndk::ScopedAStatus sendMessageDeliveryStatusToEndpoint(
-            int32_t in_sessionId, const MessageDeliveryStatus& in_msgStatus) override;
-    ::ndk::ScopedAStatus closeEndpointSession(int32_t in_sessionId, Reason in_reason) override;
-    ::ndk::ScopedAStatus endpointSessionOpenComplete(int32_t in_sessionId) override;
+    ::ndk::ScopedAStatus registerEndpointHub(
+            const std::shared_ptr<IEndpointCallback>& in_callback, const HubInfo& in_hubInfo,
+            std::shared_ptr<IEndpointCommunication>* _aidl_return) override;
 
   private:
-    struct EndpointSession {
-        int32_t sessionId;
-        EndpointId initiator;
-        EndpointId peer;
-        std::optional<std::string> serviceDescriptor;
+    class HubInterface : public BnEndpointCommunication {
+      public:
+        HubInterface(ContextHub& hal, const std::shared_ptr<IEndpointCallback>& in_callback,
+                     const HubInfo& in_hubInfo)
+            : mHal(hal), mEndpointCallback(in_callback), kInfo(in_hubInfo) {}
+        ~HubInterface() = default;
+
+        ::ndk::ScopedAStatus registerEndpoint(const EndpointInfo& in_endpoint) override;
+        ::ndk::ScopedAStatus unregisterEndpoint(const EndpointInfo& in_endpoint) override;
+        ::ndk::ScopedAStatus requestSessionIdRange(int32_t in_size,
+                                                   std::array<int32_t, 2>* _aidl_return) override;
+        ::ndk::ScopedAStatus openEndpointSession(
+                int32_t in_sessionId, const EndpointId& in_destination,
+                const EndpointId& in_initiator,
+                const std::optional<std::string>& in_serviceDescriptor) override;
+        ::ndk::ScopedAStatus sendMessageToEndpoint(int32_t in_sessionId,
+                                                   const Message& in_msg) override;
+        ::ndk::ScopedAStatus sendMessageDeliveryStatusToEndpoint(
+                int32_t in_sessionId, const MessageDeliveryStatus& in_msgStatus) override;
+        ::ndk::ScopedAStatus closeEndpointSession(int32_t in_sessionId, Reason in_reason) override;
+        ::ndk::ScopedAStatus endpointSessionOpenComplete(int32_t in_sessionId) override;
+        ::ndk::ScopedAStatus unregister() override;
+
+      private:
+        friend class ContextHub;
+
+        struct EndpointSession {
+            int32_t sessionId;
+            EndpointId initiator;
+            EndpointId peer;
+            std::optional<std::string> serviceDescriptor;
+        };
+
+        //! Finds an endpoint in the range defined by the endpoints
+        //! @return whether the endpoint was found
+        template <typename Iter>
+        bool findEndpoint(const EndpointId& target, const Iter& begin, const Iter& end) {
+            for (auto iter = begin; iter != end; ++iter) {
+                if (iter->id.id == target.id && iter->id.hubId == target.hubId) {
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        //! Endpoint storage and information
+        ContextHub& mHal;
+        std::shared_ptr<IEndpointCallback> mEndpointCallback;
+        const HubInfo kInfo;
+
+        std::atomic<bool> mActive = true;
+
+        std::mutex mEndpointMutex;
+        std::vector<EndpointInfo> mEndpoints;
+        std::vector<EndpointSession> mEndpointSessions;
+        uint16_t mBaseSessionId;
+        uint16_t mMaxSessionId;
     };
 
     static constexpr uint32_t kMockHubId = 0;
 
-    //! Finds an endpoint in the range defined by the endpoints
-    //! @return whether the endpoint was found
-    template <typename Iter>
-    bool findEndpoint(const EndpointId& target, const Iter& begin, const Iter& end) {
-        for (auto iter = begin; iter != end; ++iter) {
-            if (iter->id.id == target.id && iter->id.hubId == target.hubId) {
-                return true;
-            }
-        }
-        return false;
-    }
-
     std::shared_ptr<IContextHubCallback> mCallback;
 
     std::unordered_set<char16_t> mConnectedHostEndpoints;
 
-    //! Endpoint storage and information
-    std::mutex mEndpointMutex;
-    std::vector<EndpointInfo> mEndpoints;
-    std::vector<EndpointSession> mEndpointSessions;
-    std::shared_ptr<IEndpointCallback> mEndpointCallback;
-    int32_t mMaxValidSessionId = 0;
+    std::mutex mHostHubsLock;
+    std::unordered_map<int64_t, std::shared_ptr<HubInterface>> mIdToHostHub;
+    int32_t mNextSessionIdBase = 0;
 };
 
 }  // namespace contexthub
diff --git a/contexthub/aidl/vts/VtsAidlHalContextHubTargetTest.cpp b/contexthub/aidl/vts/VtsAidlHalContextHubTargetTest.cpp
index 1df1a99..d7859d9 100644
--- a/contexthub/aidl/vts/VtsAidlHalContextHubTargetTest.cpp
+++ b/contexthub/aidl/vts/VtsAidlHalContextHubTargetTest.cpp
@@ -24,6 +24,7 @@
 #include <android/hardware/contexthub/IContextHub.h>
 #include <android/hardware/contexthub/IContextHubCallback.h>
 #include <android/hardware/contexthub/IEndpointCallback.h>
+#include <android/hardware/contexthub/IEndpointCommunication.h>
 #include <binder/IServiceManager.h>
 #include <binder/ProcessState.h>
 #include <log/log.h>
@@ -46,6 +47,7 @@
 using ::android::hardware::contexthub::HubInfo;
 using ::android::hardware::contexthub::IContextHub;
 using ::android::hardware::contexthub::IContextHubCallbackDefault;
+using ::android::hardware::contexthub::IEndpointCommunication;
 using ::android::hardware::contexthub::Message;
 using ::android::hardware::contexthub::MessageDeliveryStatus;
 using ::android::hardware::contexthub::NanoappBinary;
@@ -67,40 +69,66 @@
 
 const String16 kEchoServiceName{"android.hardware.contexthub.test.EchoService"};
 
+constexpr int64_t kDefaultHubId = 1;
+
+class TestEndpointCallback;
 class ContextHubAidl : public testing::TestWithParam<std::tuple<std::string, int32_t>> {
   public:
-    virtual void SetUp() override {
-        contextHub = android::waitForDeclaredService<IContextHub>(
+    void SetUp() override {
+        mContextHub = android::waitForDeclaredService<IContextHub>(
                 String16(std::get<0>(GetParam()).c_str()));
-        ASSERT_NE(contextHub, nullptr);
+        ASSERT_NE(mContextHub, nullptr);
+        mEndpointCb = sp<TestEndpointCallback>::make();
     }
 
     uint32_t getHubId() { return std::get<1>(GetParam()); }
 
-    sp<IContextHub> contextHub;
+    Status registerHub(int64_t id, sp<IEndpointCommunication>* hubInterface) {
+        HubInfo info;
+        info.hubId = id;
+        return mContextHub->registerEndpointHub(mEndpointCb, info, hubInterface);
+    }
+
+    bool registerDefaultHub() {
+        Status status = registerHub(kDefaultHubId, &mHubInterface);
+        if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
+            status.transactionError() == android::UNKNOWN_TRANSACTION) {
+            return false;
+        }
+        EXPECT_TRUE(status.isOk());
+        EXPECT_NE(mHubInterface, nullptr);
+        if (!mHubInterface) {
+            return false;
+        }
+        return true;
+    }
 
     void testSettingChanged(Setting setting);
+
+    sp<IContextHub> mContextHub;
+    sp<TestEndpointCallback> mEndpointCb;
+    sp<IEndpointCommunication> mHubInterface;
 };
 
 class ContextHubAidlWithTestMode : public ContextHubAidl {
   public:
-    virtual void SetUp() override {
+    void SetUp() override {
         ContextHubAidl::SetUp();
 
         // Best effort enable test mode - this may not be supported on older HALS, so we
         // ignore the return value.
-        contextHub->setTestMode(/* enable= */ true);
+        mContextHub->setTestMode(/* enable= */ true);
     }
 
-    virtual void TearDown() override {
-        contextHub->setTestMode(/* enable= */ false);
+    void TearDown() override {
+        mContextHub->setTestMode(/* enable= */ false);
         ContextHubAidl::TearDown();
     }
 };
 
 TEST_P(ContextHubAidl, TestGetHubs) {
     std::vector<ContextHubInfo> hubs;
-    ASSERT_TRUE(contextHub->getContextHubs(&hubs).isOk());
+    ASSERT_TRUE(mContextHub->getContextHubs(&hubs).isOk());
 
     ALOGD("System reports %zu hubs", hubs.size());
 
@@ -122,7 +150,7 @@
 }
 
 TEST_P(ContextHubAidl, TestEnableTestMode) {
-    Status status = contextHub->setTestMode(true);
+    Status status = mContextHub->setTestMode(true);
     if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
         status.transactionError() == android::UNKNOWN_TRANSACTION) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
@@ -132,7 +160,7 @@
 }
 
 TEST_P(ContextHubAidl, TestDisableTestMode) {
-    Status status = contextHub->setTestMode(false);
+    Status status = mContextHub->setTestMode(false);
     if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
         status.transactionError() == android::UNKNOWN_TRANSACTION) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
@@ -181,7 +209,7 @@
 
 TEST_P(ContextHubAidl, TestRegisterCallback) {
     sp<EmptyContextHubCallback> cb = sp<EmptyContextHubCallback>::make();
-    ASSERT_TRUE(contextHub->registerCallback(getHubId(), cb).isOk());
+    ASSERT_TRUE(mContextHub->registerCallback(getHubId(), cb).isOk());
 }
 
 // Helper callback that puts the async appInfo callback data into a promise
@@ -230,8 +258,8 @@
 // Calls queryApps() and checks the returned metadata
 TEST_P(ContextHubAidl, TestQueryApps) {
     sp<QueryAppsCallback> cb = sp<QueryAppsCallback>::make();
-    ASSERT_TRUE(contextHub->registerCallback(getHubId(), cb).isOk());
-    ASSERT_TRUE(contextHub->queryNanoapps(getHubId()).isOk());
+    ASSERT_TRUE(mContextHub->registerCallback(getHubId(), cb).isOk());
+    ASSERT_TRUE(mContextHub->queryNanoapps(getHubId()).isOk());
 
     std::vector<NanoappInfo> appInfoList;
     ASSERT_TRUE(waitForCallback(cb->promise.get_future(), &appInfoList));
@@ -252,7 +280,7 @@
 // Calls getPreloadedNanoappsIds() and verifies there are preloaded nanoapps
 TEST_P(ContextHubAidl, TestGetPreloadedNanoappIds) {
     std::vector<int64_t> preloadedNanoappIds;
-    Status status = contextHub->getPreloadedNanoappIds(getHubId(), &preloadedNanoappIds);
+    Status status = mContextHub->getPreloadedNanoappIds(getHubId(), &preloadedNanoappIds);
     if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
         status.transactionError() == android::UNKNOWN_TRANSACTION) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
@@ -315,7 +343,7 @@
   public:
     virtual void SetUp() override {
         ContextHubAidl::SetUp();
-        ASSERT_TRUE(contextHub->registerCallback(getHubId(), cb).isOk());
+        ASSERT_TRUE(mContextHub->registerCallback(getHubId(), cb).isOk());
     }
 
     sp<TransactionResultCallback> cb = sp<TransactionResultCallback>::make();
@@ -329,7 +357,7 @@
     std::fill(message.messageBody.begin(), message.messageBody.end(), 0);
 
     ALOGD("Sending message to non-existent nanoapp");
-    ASSERT_TRUE(contextHub->sendMessageToHub(getHubId(), message).isOk());
+    ASSERT_TRUE(mContextHub->sendMessageToHub(getHubId(), message).isOk());
 }
 
 TEST_P(ContextHubTransactionTest, TestLoadEmptyNanoapp) {
@@ -343,7 +371,7 @@
     emptyApp.targetChreApiMinorVersion = 0;
 
     ALOGD("Loading empty nanoapp");
-    bool success = contextHub->loadNanoapp(getHubId(), emptyApp, cb->expectedTransactionId).isOk();
+    bool success = mContextHub->loadNanoapp(getHubId(), emptyApp, cb->expectedTransactionId).isOk();
     if (success) {
         bool transactionSuccess;
         ASSERT_TRUE(waitForCallback(cb->promise.get_future(), &transactionSuccess));
@@ -356,7 +384,7 @@
 
     ALOGD("Unloading nonexistent nanoapp");
     bool success =
-            contextHub->unloadNanoapp(getHubId(), kNonExistentAppId, cb->expectedTransactionId)
+            mContextHub->unloadNanoapp(getHubId(), kNonExistentAppId, cb->expectedTransactionId)
                     .isOk();
     if (success) {
         bool transactionSuccess;
@@ -370,7 +398,7 @@
 
     ALOGD("Enabling nonexistent nanoapp");
     bool success =
-            contextHub->enableNanoapp(getHubId(), kNonExistentAppId, cb->expectedTransactionId)
+            mContextHub->enableNanoapp(getHubId(), kNonExistentAppId, cb->expectedTransactionId)
                     .isOk();
     if (success) {
         bool transactionSuccess;
@@ -384,7 +412,7 @@
 
     ALOGD("Disabling nonexistent nanoapp");
     bool success =
-            contextHub->disableNanoapp(getHubId(), kNonExistentAppId, cb->expectedTransactionId)
+            mContextHub->disableNanoapp(getHubId(), kNonExistentAppId, cb->expectedTransactionId)
                     .isOk();
     if (success) {
         bool transactionSuccess;
@@ -397,10 +425,10 @@
     // In VTS, we only test that sending the values doesn't cause things to blow up - GTS tests
     // verify the expected E2E behavior in CHRE
     sp<EmptyContextHubCallback> cb = sp<EmptyContextHubCallback>::make();
-    ASSERT_TRUE(contextHub->registerCallback(getHubId(), cb).isOk());
+    ASSERT_TRUE(mContextHub->registerCallback(getHubId(), cb).isOk());
 
-    ASSERT_TRUE(contextHub->onSettingChanged(setting, true /* enabled */).isOk());
-    ASSERT_TRUE(contextHub->onSettingChanged(setting, false /* enabled */).isOk());
+    ASSERT_TRUE(mContextHub->onSettingChanged(setting, true /* enabled */).isOk());
+    ASSERT_TRUE(mContextHub->onSettingChanged(setting, false /* enabled */).isOk());
 }
 
 TEST_P(ContextHubAidl, TestOnLocationSettingChanged) {
@@ -455,27 +483,27 @@
     hostEndpointInfo.type = HostEndpointInfo::Type::NATIVE;
     hostEndpointInfo.hostEndpointId = kHostEndpointId;
 
-    ASSERT_TRUE(contextHub->onHostEndpointConnected(hostEndpointInfo).isOk());
-    ASSERT_TRUE(contextHub->onHostEndpointDisconnected(kHostEndpointId).isOk());
+    ASSERT_TRUE(mContextHub->onHostEndpointConnected(hostEndpointInfo).isOk());
+    ASSERT_TRUE(mContextHub->onHostEndpointDisconnected(kHostEndpointId).isOk());
 }
 
 TEST_P(ContextHubTransactionTest, TestInvalidHostConnection) {
     constexpr char16_t kHostEndpointId = 1;
 
-    ASSERT_TRUE(contextHub->onHostEndpointDisconnected(kHostEndpointId).isOk());
+    ASSERT_TRUE(mContextHub->onHostEndpointDisconnected(kHostEndpointId).isOk());
 }
 
 TEST_P(ContextHubTransactionTest, TestNanSessionStateChange) {
     NanSessionStateUpdate update;
     update.state = true;
-    Status status = contextHub->onNanSessionStateChanged(update);
+    Status status = mContextHub->onNanSessionStateChanged(update);
     if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
         status.transactionError() == android::UNKNOWN_TRANSACTION) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
     } else {
         ASSERT_TRUE(status.isOk());
         update.state = false;
-        ASSERT_TRUE(contextHub->onNanSessionStateChanged(update).isOk());
+        ASSERT_TRUE(mContextHub->onNanSessionStateChanged(update).isOk());
     }
 }
 
@@ -484,7 +512,7 @@
     messageDeliveryStatus.messageSequenceNumber = 123;
     messageDeliveryStatus.errorCode = ErrorCode::OK;
 
-    Status status = contextHub->sendMessageDeliveryStatusToHub(getHubId(), messageDeliveryStatus);
+    Status status = mContextHub->sendMessageDeliveryStatusToHub(getHubId(), messageDeliveryStatus);
     if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
         status.transactionError() == android::UNKNOWN_TRANSACTION) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
@@ -557,7 +585,30 @@
     bool mWasOnEndpointSessionOpenCompleteCalled = false;
 };
 
+TEST_P(ContextHubAidlWithTestMode, RegisterHub) {
+    if (!registerDefaultHub()) {
+        GTEST_SKIP() << "Not supported -> old API; or not implemented";
+    }
+
+    sp<IEndpointCommunication> hub2;
+    Status status = registerHub(kDefaultHubId + 1, &hub2);
+    EXPECT_TRUE(status.isOk());
+
+    sp<IEndpointCommunication> hub3;
+    status = registerHub(kDefaultHubId + 1, &hub3);
+    ASSERT_FALSE(status.isOk());
+    EXPECT_EQ(status.exceptionCode(), Status::EX_ILLEGAL_STATE);
+
+    hub2->unregister();
+    status = registerHub(kDefaultHubId + 1, &hub3);
+    EXPECT_TRUE(status.isOk());
+}
+
 TEST_P(ContextHubAidlWithTestMode, RegisterEndpoint) {
+    if (!registerDefaultHub()) {
+        GTEST_SKIP() << "Not supported -> old API; or not implemented";
+    }
+
     EndpointInfo endpointInfo;
     endpointInfo.id.id = 1;
     endpointInfo.id.hubId = 0xCAFECAFECAFECAFE;
@@ -565,7 +616,7 @@
     endpointInfo.name = String16("Test host endpoint 1");
     endpointInfo.version = 42;
 
-    Status status = contextHub->registerEndpoint(endpointInfo);
+    Status status = mHubInterface->registerEndpoint(endpointInfo);
     if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
         status.transactionError() == android::UNKNOWN_TRANSACTION) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
@@ -575,6 +626,10 @@
 }
 
 TEST_P(ContextHubAidlWithTestMode, RegisterEndpointSameNameFailure) {
+    if (!registerDefaultHub()) {
+        GTEST_SKIP() << "Not supported -> old API; or not implemented";
+    }
+
     EndpointInfo endpointInfo;
     endpointInfo.id.id = 2;
     endpointInfo.id.hubId = 0xCAFECAFECAFECAFE;
@@ -589,7 +644,7 @@
     endpointInfo2.name = String16("Test host endpoint 2");
     endpointInfo2.version = 42;
 
-    Status status = contextHub->registerEndpoint(endpointInfo);
+    Status status = mHubInterface->registerEndpoint(endpointInfo);
     if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
         status.transactionError() == android::UNKNOWN_TRANSACTION) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
@@ -597,10 +652,14 @@
         EXPECT_TRUE(status.isOk());
     }
 
-    EXPECT_FALSE(contextHub->registerEndpoint(endpointInfo2).isOk());
+    EXPECT_FALSE(mHubInterface->registerEndpoint(endpointInfo2).isOk());
 }
 
 TEST_P(ContextHubAidlWithTestMode, RegisterEndpointSameIdFailure) {
+    if (!registerDefaultHub()) {
+        GTEST_SKIP() << "Not supported -> old API; or not implemented";
+    }
+
     EndpointInfo endpointInfo;
     endpointInfo.id.id = 4;
     endpointInfo.id.hubId = 0xCAFECAFECAFECAFE;
@@ -615,7 +674,7 @@
     endpointInfo2.name = String16("Test host endpoint - same ID test");
     endpointInfo2.version = 42;
 
-    Status status = contextHub->registerEndpoint(endpointInfo);
+    Status status = mHubInterface->registerEndpoint(endpointInfo);
     if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
         status.transactionError() == android::UNKNOWN_TRANSACTION) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
@@ -623,10 +682,14 @@
         EXPECT_TRUE(status.isOk());
     }
 
-    EXPECT_FALSE(contextHub->registerEndpoint(endpointInfo2).isOk());
+    EXPECT_FALSE(mHubInterface->registerEndpoint(endpointInfo2).isOk());
 }
 
 TEST_P(ContextHubAidlWithTestMode, UnregisterEndpoint) {
+    if (!registerDefaultHub()) {
+        GTEST_SKIP() << "Not supported -> old API; or not implemented";
+    }
+
     EndpointInfo endpointInfo;
     endpointInfo.id.id = 6;
     endpointInfo.id.hubId = 0xCAFECAFECAFECAFE;
@@ -634,7 +697,7 @@
     endpointInfo.name = String16("Test host endpoint 6");
     endpointInfo.version = 42;
 
-    Status status = contextHub->registerEndpoint(endpointInfo);
+    Status status = mHubInterface->registerEndpoint(endpointInfo);
     if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
         status.transactionError() == android::UNKNOWN_TRANSACTION) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
@@ -642,10 +705,14 @@
         EXPECT_TRUE(status.isOk());
     }
 
-    EXPECT_TRUE(contextHub->unregisterEndpoint(endpointInfo).isOk());
+    EXPECT_TRUE(mHubInterface->unregisterEndpoint(endpointInfo).isOk());
 }
 
 TEST_P(ContextHubAidlWithTestMode, UnregisterEndpointNonexistent) {
+    if (!registerDefaultHub()) {
+        GTEST_SKIP() << "Not supported -> old API; or not implemented";
+    }
+
     EndpointInfo endpointInfo;
     endpointInfo.id.id = 100;
     endpointInfo.id.hubId = 0xCAFECAFECAFECAFE;
@@ -653,7 +720,7 @@
     endpointInfo.name = String16("Test host endpoint 100");
     endpointInfo.version = 42;
 
-    Status status = contextHub->unregisterEndpoint(endpointInfo);
+    Status status = mHubInterface->unregisterEndpoint(endpointInfo);
     if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
         status.transactionError() == android::UNKNOWN_TRANSACTION) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
@@ -662,25 +729,9 @@
     }
 }
 
-TEST_P(ContextHubAidlWithTestMode, RegisterEndpointCallback) {
-    auto cb = sp<TestEndpointCallback>::make();
-    Status status = contextHub->registerEndpointCallback(cb);
-    if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
-        status.transactionError() == android::UNKNOWN_TRANSACTION) {
-        GTEST_SKIP() << "Not supported -> old API; or not implemented";
-    } else {
-        EXPECT_TRUE(status.isOk());
-    }
-}
-
 TEST_P(ContextHubAidlWithTestMode, OpenEndpointSessionInvalidRange) {
-    auto cb = sp<TestEndpointCallback>::make();
-    Status status = contextHub->registerEndpointCallback(cb);
-    if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
-        status.transactionError() == android::UNKNOWN_TRANSACTION) {
+    if (!registerDefaultHub()) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
-    } else {
-        EXPECT_TRUE(status.isOk());
     }
 
     // Register the endpoint
@@ -690,11 +741,11 @@
     initiatorEndpoint.type = EndpointInfo::EndpointType::NATIVE;
     initiatorEndpoint.name = String16("Test host endpoint 7");
     initiatorEndpoint.version = 42;
-    EXPECT_TRUE(contextHub->registerEndpoint(initiatorEndpoint).isOk());
+    EXPECT_TRUE(mHubInterface->registerEndpoint(initiatorEndpoint).isOk());
 
     // Find the destination, if it exists
     std::vector<EndpointInfo> endpoints;
-    EXPECT_TRUE(contextHub->getEndpoints(&endpoints).isOk());
+    EXPECT_TRUE(mContextHub->getEndpoints(&endpoints).isOk());
     const EndpointInfo* destinationEndpoint = nullptr;
     for (const EndpointInfo& endpoint : endpoints) {
         for (const Service& service : endpoint.services) {
@@ -711,13 +762,13 @@
     // Request the range
     constexpr int32_t requestedRange = 100;
     std::array<int32_t, 2> range;
-    ASSERT_TRUE(contextHub->requestSessionIdRange(requestedRange, &range).isOk());
+    ASSERT_TRUE(mHubInterface->requestSessionIdRange(requestedRange, &range).isOk());
     EXPECT_EQ(range.size(), 2);
     EXPECT_GE(range[1] - range[0] + 1, requestedRange);
 
     // Open the session
     int32_t sessionId = range[1] + 10;  // invalid
-    EXPECT_FALSE(contextHub
+    EXPECT_FALSE(mHubInterface
                          ->openEndpointSession(sessionId, destinationEndpoint->id,
                                                initiatorEndpoint.id,
                                                /* in_serviceDescriptor= */ kEchoServiceName)
@@ -725,16 +776,10 @@
 }
 
 TEST_P(ContextHubAidlWithTestMode, OpenEndpointSessionAndSendMessageEchoesBack) {
-    auto cb = sp<TestEndpointCallback>::make();
-    Status status = contextHub->registerEndpointCallback(cb);
-    if (status.exceptionCode() == Status::EX_UNSUPPORTED_OPERATION ||
-        status.transactionError() == android::UNKNOWN_TRANSACTION) {
+    if (!registerDefaultHub()) {
         GTEST_SKIP() << "Not supported -> old API; or not implemented";
-    } else {
-        EXPECT_TRUE(status.isOk());
     }
-
-    std::unique_lock<std::mutex> lock(cb->getMutex());
+    std::unique_lock<std::mutex> lock(mEndpointCb->getMutex());
 
     // Register the endpoint
     EndpointInfo initiatorEndpoint;
@@ -743,11 +788,11 @@
     initiatorEndpoint.type = EndpointInfo::EndpointType::NATIVE;
     initiatorEndpoint.name = String16("Test host endpoint 7");
     initiatorEndpoint.version = 42;
-    EXPECT_TRUE(contextHub->registerEndpoint(initiatorEndpoint).isOk());
+    EXPECT_TRUE(mHubInterface->registerEndpoint(initiatorEndpoint).isOk());
 
     // Find the destination, if it exists
     std::vector<EndpointInfo> endpoints;
-    EXPECT_TRUE(contextHub->getEndpoints(&endpoints).isOk());
+    EXPECT_TRUE(mContextHub->getEndpoints(&endpoints).isOk());
     const EndpointInfo* destinationEndpoint = nullptr;
     for (const EndpointInfo& endpoint : endpoints) {
         for (const Service& service : endpoint.services) {
@@ -764,32 +809,32 @@
     // Request the range
     constexpr int32_t requestedRange = 100;
     std::array<int32_t, 2> range;
-    ASSERT_TRUE(contextHub->requestSessionIdRange(requestedRange, &range).isOk());
+    ASSERT_TRUE(mHubInterface->requestSessionIdRange(requestedRange, &range).isOk());
     EXPECT_EQ(range.size(), 2);
     EXPECT_GE(range[1] - range[0] + 1, requestedRange);
 
     // Open the session
-    cb->resetWasOnEndpointSessionOpenCompleteCalled();
+    mEndpointCb->resetWasOnEndpointSessionOpenCompleteCalled();
     int32_t sessionId = range[0];
-    ASSERT_TRUE(contextHub
+    ASSERT_TRUE(mHubInterface
                         ->openEndpointSession(sessionId, destinationEndpoint->id,
                                               initiatorEndpoint.id,
                                               /* in_serviceDescriptor= */ kEchoServiceName)
                         .isOk());
-    cb->getCondVar().wait(lock);
-    EXPECT_TRUE(cb->wasOnEndpointSessionOpenCompleteCalled());
+    mEndpointCb->getCondVar().wait(lock);
+    EXPECT_TRUE(mEndpointCb->wasOnEndpointSessionOpenCompleteCalled());
 
     // Send the message
     Message message;
     message.flags = 0;
     message.sequenceNumber = 0;
     message.content.push_back(42);
-    ASSERT_TRUE(contextHub->sendMessageToEndpoint(sessionId, message).isOk());
+    ASSERT_TRUE(mHubInterface->sendMessageToEndpoint(sessionId, message).isOk());
 
     // Check for echo
-    cb->getCondVar().wait(lock);
-    EXPECT_FALSE(cb->getMessages().empty());
-    EXPECT_EQ(cb->getMessages().back().content.back(), 42);
+    mEndpointCb->getCondVar().wait(lock);
+    EXPECT_FALSE(mEndpointCb->getMessages().empty());
+    EXPECT_EQ(mEndpointCb->getMessages().back().content.back(), 42);
 }
 
 std::string PrintGeneratedTest(const testing::TestParamInfo<ContextHubAidl::ParamType>& info) {