Implement Discovery PLATFORM logic in Rust

Add RemoteAuthConnectionCache and RemoteAuthPlatform
with support to sendMessage

Design doc: go/remote-auth-manager-fishfood-design

Test: built successfully.
Bug: : 291333048
Change-Id: I17f73b4fb2e22924a484eeb3baa9b933ae980076

remoteauth/service/Android.bp

@@ -25,7 +25,7 @@ filegroup {
 java_library {
     name: "service-remoteauth-pre-jarjar",
     srcs: [":remoteauth-service-srcs"],
-    required: ["libremoteauth_jni_rust_defaults"],
+    required: ["libremoteauth_jni_rust"],
     defaults: [
         "enable-remoteauth-targets",
         "framework-system-server-module-defaults",

remoteauth/service/java/com/android/server/remoteauth/RemoteAuthConnectionCache.java

@@ -0,0 +1,141 @@
+/*
+ * Copyright (C) 2023 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 com.android.server.remoteauth;
+
+import android.annotation.NonNull;
+import android.util.Log;
+
+import com.android.internal.util.Preconditions;
+import com.android.server.remoteauth.connectivity.Connection;
+import com.android.server.remoteauth.connectivity.ConnectionException;
+import com.android.server.remoteauth.connectivity.ConnectionInfo;
+import com.android.server.remoteauth.connectivity.ConnectivityManager;
+import com.android.server.remoteauth.connectivity.EventListener;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Manages caching of remote devices {@link ConnectionInfo} and {@link Connection}.
+ *
+ * <p>Allows mapping between {@link ConnectionInfo#getConnectionId()} to {@link ConnectionInfo} and
+ * {@link Connection}
+ */
+public class RemoteAuthConnectionCache {
+    public static final String TAG = "RemoteAuthConCache";
+    private final Map<Integer, ConnectionInfo> mConnectionInfoMap = new ConcurrentHashMap<>();
+    private final Map<Integer, Connection> mConnectionMap = new ConcurrentHashMap<>();
+
+    private final ConnectivityManager mConnectivityManager;
+
+    public RemoteAuthConnectionCache(@NonNull ConnectivityManager connectivityManager) {
+        Preconditions.checkNotNull(connectivityManager);
+        this.mConnectivityManager = connectivityManager;
+    }
+
+    /** Returns the {@link ConnectivityManager}. */
+    ConnectivityManager getConnectivityManager() {
+        return mConnectivityManager;
+    }
+
+    /**
+     * Associates the connectionId with {@link ConnectionInfo}. Updates association with new value
+     * if already exists
+     *
+     * @param connectionInfo of the remote device
+     */
+    public void setConnectionInfo(@NonNull ConnectionInfo connectionInfo) {
+        Preconditions.checkNotNull(connectionInfo);
+        mConnectionInfoMap.put(connectionInfo.getConnectionId(), connectionInfo);
+    }
+
+    /** Returns {@link ConnectionInfo} associated with connectionId. */
+    public ConnectionInfo getConnectionInfo(int connectionId) {
+        return mConnectionInfoMap.get(connectionId);
+    }
+
+    /**
+     * Associates the connectionId with {@link Connection}. Updates association with new value if
+     * already exists
+     *
+     * @param connection to the remote device
+     */
+    public void setConnection(@NonNull Connection connection) {
+        Preconditions.checkNotNull(connection);
+        mConnectionMap.put(connection.getConnectionInfo().getConnectionId(), connection);
+    }
+
+    /**
+     * Returns {@link Connection} associated with connectionId. Uses {@link ConnectivityManager} to
+     * create and associate with new {@link Connection}, if mapping doesn't exist
+     *
+     * @param connectionId of the remote device
+     */
+    public Connection getConnection(int connectionId) {
+        return mConnectionMap.computeIfAbsent(
+                connectionId,
+                id -> {
+                    ConnectionInfo connectionInfo = getConnectionInfo(id);
+                    if (null == connectionInfo) {
+                        // TODO: Try accessing DB to fetch by connectionId
+                        Log.e(TAG, String.format("Unknown connectionId: %d", connectionId));
+                        return null;
+                    }
+                    try {
+                        Connection connection =
+                                mConnectivityManager.connect(
+                                        connectionInfo,
+                                        new EventListener() {
+                                            @Override
+                                            public void onDisconnect(
+                                                    @NonNull ConnectionInfo connectionInfo) {
+                                                removeConnection(connectionInfo.getConnectionId());
+                                                Log.i(
+                                                        TAG,
+                                                        String.format(
+                                                                "Disconnected from: %d",
+                                                                connectionInfo.getConnectionId()));
+                                            }
+                                        });
+                        if (null == connection) {
+                            Log.e(TAG, String.format("Failed to connect: %d", connectionId));
+                            return null;
+                        }
+                        return connection;
+                    } catch (ConnectionException e) {
+                        Log.e(
+                                TAG,
+                                String.format("Failed to create connection to %d.", connectionId),
+                                e);
+                        return null;
+                    }
+                });
+    }
+
+    /**
+     * Removes {@link Connection} from cache.
+     *
+     * @param connectionId of the remote device
+     */
+    public void removeConnection(int connectionId) {
+        if (null != mConnectionMap.remove(connectionId)) {
+            Log.i(
+                    TAG,
+                    String.format("Connection associated with id: %d was removed", connectionId));
+        }
+    }
+}

remoteauth/service/java/com/android/server/remoteauth/RemoteAuthPlatform.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2023 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 com.android.server.remoteauth;
+
+import android.util.Log;
+
+import com.android.server.remoteauth.connectivity.Connection;
+import com.android.server.remoteauth.jni.INativeRemoteAuthService;
+
+/** Implementation of the {@link INativeRemoteAuthService.IPlatform} interface. */
+public class RemoteAuthPlatform implements INativeRemoteAuthService.IPlatform {
+    public static final String TAG = "RemoteAuthPlatform";
+    private final RemoteAuthConnectionCache mConnectionCache;
+
+    public RemoteAuthPlatform(RemoteAuthConnectionCache connectionCache) {
+        mConnectionCache = connectionCache;
+    }
+
+    /**
+     * Sends message to the remote device via {@link Connection} created by
+     * {@link com.android.server.remoteauth.connectivity.ConnectivityManager}.
+     *
+     * @param connectionId connection ID of the {@link android.remoteauth.RemoteAuthenticator}
+     * @param request payload of the request
+     * @param callback to be used to pass the response result
+     * @return true if succeeded, false otherwise.
+     * @hide
+     */
+    @Override
+    public boolean sendRequest(int connectionId, byte[] request, ResponseCallback callback) {
+        Connection connection = mConnectionCache.getConnection(connectionId);
+        if (null == connection) {
+            Log.e(TAG, String.format("Failed to get a connection for: %d", connectionId));
+            return false;
+        }
+        connection.sendRequest(
+                request,
+                new Connection.MessageResponseCallback() {
+                    @Override
+                    public void onSuccess(byte[] buffer) {
+                        callback.onSuccess(buffer);
+                    }
+
+                    @Override
+                    public void onFailure(@Connection.ErrorCode int errorCode) {
+                        callback.onFailure(errorCode);
+                    }
+                });
+        return true;
+    }
+}

remoteauth/service/java/com/android/server/remoteauth/RemoteAuthService.java

@@ -65,7 +65,7 @@ public class RemoteAuthService extends IRemoteAuthService.Stub {
     }
 
     private static void checkPermission(Context context, String permission) {
-        context.enforceCallingOrSelfPermission(permission,
+        context.enforceCallingOrSelfPermission(
-                "Must have " + permission + " permission.");
+                permission, "Must have " + permission + " permission.");
     }
 }

remoteauth/service/java/com/android/server/remoteauth/connectivity/EventListener.java

@@ -21,7 +21,7 @@ import android.annotation.NonNull;
 /**
  * Listens to the events from underlying transport.
  */
-interface EventListener {
+public interface EventListener {
     /** Called when remote device is disconnected from the underlying transport. */
     void onDisconnect(@NonNull ConnectionInfo connectionInfo);
 }

remoteauth/service/java/com/android/server/remoteauth/jni/INativeRemoteAuthService.java

@@ -34,10 +34,10 @@ public interface INativeRemoteAuthService {
          * @param connectionId connection ID of the {@link android.remoteauth.RemoteAuthenticator}
          * @param request payload of the request
          * @param callback to be used to pass the response result
-         *
+         * @return true if succeeded, false otherwise.
          * @hide
          */
-        void sendRequest(int connectionId, byte[] request, ResponseCallback callback);
+        boolean sendRequest(int connectionId, byte[] request, ResponseCallback callback);
 
         /**
          * Interface for a callback to send a response back.
@@ -49,7 +49,6 @@ public interface INativeRemoteAuthService {
              * Invoked when message sending succeeds.
              *
              * @param response contains response
-             *
              * @hide
              */
             void onSuccess(byte[] response);
@@ -58,7 +57,6 @@ public interface INativeRemoteAuthService {
              * Invoked when message sending fails.
              *
              * @param errorCode indicating the error
-             *
              * @hide
              */
             void onFailure(int errorCode);

remoteauth/service/java/com/android/server/remoteauth/jni/NativeRemoteAuthService.java

@@ -16,6 +16,8 @@
 
 package com.android.server.remoteauth.jni;
 
+import android.util.Log;
+
 import com.android.internal.annotations.Keep;
 import com.android.server.remoteauth.jni.INativeRemoteAuthService.IPlatform;
 
@@ -53,12 +55,13 @@ public class NativeRemoteAuthService {
      *     platform
      * @param platformHandle a handle associated with the platform object, used to pass the response
      *     to the specific platform
-     *
      * @hide
      */
     @Keep
     public void sendRequest(
             int connectionId, byte[] request, long responseHandle, long platformHandle) {
+        Log.d(TAG, String.format("sendRequest with connectionId: %d, rh: %d, ph: %d",
+                connectionId, responseHandle, platformHandle));
         mPlatform.sendRequest(
                 connectionId,
                 request,

remoteauth/service/java/com/android/server/remoteauth/jni/PlatformBadHandleException.java

@@ -21,6 +21,7 @@
 package com.android.server.remoteauth.jni;
 
 import com.android.internal.annotations.Keep;
+
 /**
  * Exception thrown by native platform rust implementation of {@link
  * com.android.server.remoteauth.RemoteAuthService}.

remoteauth/service/jni/Android.bp

@@ -30,6 +30,12 @@ rust_defaults {
     host_supported: true,
 }
 
+rust_ffi_shared {
+    name: "libremoteauth_jni_rust",
+    defaults: ["libremoteauth_jni_rust_defaults"],
+    rustlibs: [],
+}
+
 rust_test {
     name: "libremoteauth_jni_rust_tests",
     defaults: ["libremoteauth_jni_rust_defaults"],

remoteauth/service/jni/src/lib.rs

@@ -21,5 +21,7 @@ mod jnames;
 mod unique_jvm;
 mod utils;
 
+/// Implementation of JNI platform functionality.
 pub mod remoteauth_jni_android_platform;
+/// Implementation of JNI protocol functionality.
 pub mod remoteauth_jni_android_protocol;

remoteauth/service/jni/src/remoteauth_jni_android_platform.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//! Implementation of JNI platform functionality.
 use crate::jnames::{SEND_REQUEST_MNAME, SEND_REQUEST_MSIG};
 use crate::unique_jvm;
 use anyhow::anyhow;
@@ -73,11 +74,15 @@ fn insert_platform_handle(handle: i64, item: Arc<Mutex<JavaPlatform>>) {
     HANDLE_MAPPING.lock().unwrap().insert(handle, Arc::clone(&item));
 }
 
+/// Reports a response from remote device.
 pub trait ResponseCallback {
+    /// Invoked upon successful response
     fn on_response(&mut self, response: Vec<u8>);
+    /// Invoked upon failure
     fn on_error(&mut self, error_code: i32);
 }
 
+/// Trait to platform functionality
 pub trait Platform {
     /// Send a binary message to the remote with the given connection id and return the response.
     fn send_request(
@@ -89,6 +94,7 @@ pub trait Platform {
 }
 //////////////////////////////////
 
+/// Implementation of Platform trait
 pub struct JavaPlatform {
     platform_handle: i64,
     vm: &'static Arc<JavaVM>,
@@ -99,7 +105,7 @@ pub struct JavaPlatform {
 }
 
 impl JavaPlatform {
-    // Method to create JavaPlatform
+    /// Creates JavaPlatform and associates with unique handle id
     pub fn create(
         java_platform_native: JObject<'_>,
     ) -> Result<Arc<Mutex<impl Platform>>, JNIError> {
@@ -219,6 +225,7 @@ impl JavaPlatform {
     }
 }
 
+/// Returns successful response from remote device
 #[no_mangle]
 pub extern "system" fn Java_com_android_server_remoteauth_jni_NativeRemoteAuthJavaPlatform_native_on_send_request_success(
     env: JNIEnv,
@@ -250,6 +257,7 @@ fn native_on_send_request_success(
     }
 }
 
+/// Notifies about failure to receive a response from remote device
 #[no_mangle]
 pub extern "system" fn Java_com_android_server_remoteauth_jni_NativeRemoteAuthJavaPlatform_native_on_send_request_error(
     env: JNIEnv,

remoteauth/service/jni/src/remoteauth_jni_android_protocol.rs

@@ -14,12 +14,14 @@
  * limitations under the License.
  */
 
+//! Implementation of JNI protocol functionality.
 use crate::unique_jvm;
 use crate::utils::get_boolean_result;
 use jni::objects::JObject;
 use jni::sys::jboolean;
 use jni::JNIEnv;
 
+/// Initialize native library. Captures Java VM:
 #[no_mangle]
 pub extern "system" fn Java_com_android_server_remoteauth_jni_NativeRemoteAuthJavaPlatform_native_init(
     env: JNIEnv,

remoteauth/tests/unit/src/com/android/server/remoteauth/RemoteAuthConnectionCacheTest.java

@@ -0,0 +1,123 @@
+/*
+ * Copyright (C) 2023 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 com.android.server.remoteauth;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.anyObject;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import com.android.server.remoteauth.connectivity.Connection;
+import com.android.server.remoteauth.connectivity.ConnectionException;
+import com.android.server.remoteauth.connectivity.ConnectionInfo;
+import com.android.server.remoteauth.connectivity.ConnectivityManager;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/** Unit test for {@link com.android.server.remoteauth.RemoteAuthConnectionCache} */
+@RunWith(AndroidJUnit4.class)
+public class RemoteAuthConnectionCacheTest {
+    @Mock private Connection mConnection;
+    @Mock private ConnectionInfo mConnectionInfo;
+    @Mock private ConnectivityManager mConnectivityManager;
+    private RemoteAuthConnectionCache mConnectionCache;
+
+    private static final int CONNECTION_ID = 1;
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+
+        doReturn(CONNECTION_ID).when(mConnectionInfo).getConnectionId();
+        doReturn(mConnectionInfo).when(mConnection).getConnectionInfo();
+        mConnectionCache = new RemoteAuthConnectionCache(mConnectivityManager);
+    }
+
+    @Test
+    public void testCreateCache_managerIsNull() {
+        assertThrows(NullPointerException.class, () -> new RemoteAuthConnectionCache(null));
+    }
+
+    @Test
+    public void testGetManager_managerExists() {
+        assertEquals(mConnectivityManager, mConnectionCache.getConnectivityManager());
+    }
+
+    @Test
+    public void testSetConnectionInfo_infoIsNull() {
+        assertThrows(NullPointerException.class, () -> mConnectionCache.setConnectionInfo(null));
+    }
+
+    @Test
+    public void testSetConnectionInfo_infoIsValid() {
+        mConnectionCache.setConnectionInfo(mConnectionInfo);
+
+        assertEquals(mConnectionInfo, mConnectionCache.getConnectionInfo(CONNECTION_ID));
+    }
+
+    @Test
+    public void testSetConnection_connectionIsNull() {
+        assertThrows(NullPointerException.class, () -> mConnectionCache.setConnection(null));
+    }
+
+    @Test
+    public void testGetConnection_connectionAlreadyExists() {
+        mConnectionCache.setConnection(mConnection);
+
+        assertEquals(mConnection, mConnectionCache.getConnection(CONNECTION_ID));
+    }
+
+    @Test
+    public void testGetConnection_connectionInfoDoesntExists() {
+        assertNull(mConnectionCache.getConnection(CONNECTION_ID));
+    }
+
+    @Test
+    public void testGetConnection_failedToConnect() {
+        mConnectionCache.setConnectionInfo(mConnectionInfo);
+        doReturn(null).when(mConnectivityManager).connect(eq(mConnectionInfo), anyObject());
+
+        assertNull(mConnectionCache.getConnection(CONNECTION_ID));
+    }
+
+    @Test
+    public void testGetConnection_failedToConnectException() {
+        mConnectionCache.setConnectionInfo(mConnectionInfo);
+        doThrow(ConnectionException.class)
+                .when(mConnectivityManager)
+                .connect(eq(mConnectionInfo), anyObject());
+
+        assertNull(mConnectionCache.getConnection(CONNECTION_ID));
+    }
+
+    @Test
+    public void testGetConnection_connectionSucceed() {
+        mConnectionCache.setConnectionInfo(mConnectionInfo);
+        doReturn(mConnection).when(mConnectivityManager).connect(eq(mConnectionInfo), anyObject());
+
+        assertEquals(mConnection, mConnectionCache.getConnection(CONNECTION_ID));
+    }
+}

remoteauth/tests/unit/src/com/android/server/remoteauth/RemoteAuthPlatformTest.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2023 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 com.android.server.remoteauth;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyObject;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import com.android.server.remoteauth.connectivity.Connection;
+import com.android.server.remoteauth.jni.INativeRemoteAuthService;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/** Unit test for {@link com.android.server.remoteauth.RemoteAuthPlatform} */
+@RunWith(AndroidJUnit4.class)
+public class RemoteAuthPlatformTest {
+    @Mock private Connection mConnection;
+    @Mock private RemoteAuthConnectionCache mConnectionCache;
+    private RemoteAuthPlatform mPlatform;
+    private static final int CONNECTION_ID = 1;
+    private static final byte[] REQUEST = new byte[] {(byte) 0x01};
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+
+        mPlatform = new RemoteAuthPlatform(mConnectionCache);
+    }
+
+    @Test
+    public void testSendRequest_connectionIsNull() {
+        doReturn(null).when(mConnectionCache).getConnection(anyInt());
+        assertFalse(
+                mPlatform.sendRequest(
+                        CONNECTION_ID,
+                        REQUEST,
+                        new INativeRemoteAuthService.IPlatform.ResponseCallback() {
+                            @Override
+                            public void onSuccess(byte[] response) {}
+
+                            @Override
+                            public void onFailure(int errorCode) {}
+                        }));
+    }
+
+    @Test
+    public void testSendRequest_connectionExists() {
+        doReturn(mConnection).when(mConnectionCache).getConnection(anyInt());
+        assertTrue(
+                mPlatform.sendRequest(
+                        CONNECTION_ID,
+                        REQUEST,
+                        new INativeRemoteAuthService.IPlatform.ResponseCallback() {
+                            @Override
+                            public void onSuccess(byte[] response) {}
+
+                            @Override
+                            public void onFailure(int errorCode) {}
+                        }));
+        verify(mConnection, times(1)).sendRequest(eq(REQUEST), anyObject());
+    }
+}