[CDM Transport] Exchange platform info to decide which Transport to

create

Bug: 253307662

Test: SystemDataTransferTest & SystemDataTransportTest
Change-Id: I26b03a8323d2870d0d772dbef34dd64ed8d4fe1b

8 files changed, 214 insertions, 116 deletions

diff --git a/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java b/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java
index 9f27f721ea83..3fffdbecd0de 100644
--- a/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java
+++ b/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java
@@ -23,6 +23,7 @@ import static android.companion.CompanionDeviceManager.COMPANION_DEVICE_DISCOVER
 import static android.content.ComponentName.createRelative;
 
 import static com.android.server.companion.Utils.prepareForIpc;
+import static com.android.server.companion.transport.Transport.MESSAGE_REQUEST_PERMISSION_RESTORE;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -91,7 +92,8 @@ public class SystemDataTransferProcessor {
         mAssociationStore = associationStore;
         mSystemDataTransferRequestStore = systemDataTransferRequestStore;
         mTransportManager = transportManager;
-        mTransportManager.setListener(this::onReceivePermissionRestore);
+        mTransportManager.addListener(MESSAGE_REQUEST_PERMISSION_RESTORE,
+                this::onReceivePermissionRestore);
         mPermissionControllerManager = mContext.getSystemService(PermissionControllerManager.class);
         mExecutor = Executors.newSingleThreadExecutor();
     }
diff --git a/services/companion/java/com/android/server/companion/securechannel/AttestationVerifier.java b/services/companion/java/com/android/server/companion/securechannel/AttestationVerifier.java
index adaee757b96a..1559a3f8fdf8 100644
--- a/services/companion/java/com/android/server/companion/securechannel/AttestationVerifier.java
+++ b/services/companion/java/com/android/server/companion/securechannel/AttestationVerifier.java
@@ -35,7 +35,7 @@ import java.util.function.BiConsumer;
 /**
  * Helper class to perform attestation verification synchronously.
  */
-class AttestationVerifier {
+public class AttestationVerifier {
     private static final long ATTESTATION_VERIFICATION_TIMEOUT_SECONDS = 10; // 10 seconds
     private static final String PARAM_OWNED_BY_SYSTEM = "android.key_owned_by_system";
 
diff --git a/services/companion/java/com/android/server/companion/securechannel/SecureChannel.java b/services/companion/java/com/android/server/companion/securechannel/SecureChannel.java
index 13dba84487e3..05b6022ce569 100644
--- a/services/companion/java/com/android/server/companion/securechannel/SecureChannel.java
+++ b/services/companion/java/com/android/server/companion/securechannel/SecureChannel.java
@@ -110,7 +110,7 @@ public class SecureChannel {
         this(in, out, callback, null, new AttestationVerifier(context));
     }
 
-    private SecureChannel(
+    public SecureChannel(
             final InputStream in,
             final OutputStream out,
             Callback callback,
@@ -381,9 +381,10 @@ public class SecureChannel {
 
     private void exchangeAuthentication()
             throws IOException, GeneralSecurityException, BadHandleException, CryptoException {
-        if (mVerifier == null) {
+        if (mPreSharedKey != null) {
             exchangePreSharedKey();
-        } else {
+        }
+        if (mVerifier != null) {
             exchangeAttestation();
         }
     }
diff --git a/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java b/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java
index 6a53adfeea9d..2abdcb172965 100644
--- a/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java
+++ b/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java
@@ -19,9 +19,9 @@ package com.android.server.companion.transport;
 import static android.Manifest.permission.DELIVER_COMPANION_MESSAGES;
 
 import static com.android.server.companion.transport.Transport.MESSAGE_REQUEST_PERMISSION_RESTORE;
+import static com.android.server.companion.transport.Transport.MESSAGE_REQUEST_PLATFORM_INFO;
 
 import android.annotation.NonNull;
-import android.annotation.Nullable;
 import android.annotation.SuppressLint;
 import android.app.ActivityManagerInternal;
 import android.content.Context;
@@ -30,12 +30,17 @@ import android.content.pm.PackageManager.NameNotFoundException;
 import android.os.Binder;
 import android.os.Build;
 import android.os.ParcelFileDescriptor;
+import android.util.Slog;
 import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.server.LocalServices;
+import com.android.server.companion.transport.Transport.Listener;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
 
@@ -44,6 +49,8 @@ public class CompanionTransportManager {
     private static final String TAG = "CDM_CompanionTransportManager";
     private static final boolean DEBUG = false;
 
+    private static final int NON_ANDROID = -1;
+
     private boolean mSecureTransportEnabled = true;
 
     private static boolean isRequest(int message) {
@@ -54,24 +61,29 @@ public class CompanionTransportManager {
         return (message & 0xFF000000) == 0x33000000;
     }
 
-    public interface Listener {
-        void onRequestPermissionRestore(byte[] data);
-    }
-
     private final Context mContext;
 
     @GuardedBy("mTransports")
     private final SparseArray<Transport> mTransports = new SparseArray<>();
 
-    @Nullable
-    private Listener mListener;
+    @NonNull
+    private final Map<Integer, Listener> mListeners = new HashMap<>();
+
+    private Transport mTempTransport;
 
     public CompanionTransportManager(Context context) {
         mContext = context;
     }
 
-    public void setListener(@NonNull Listener listener) {
-        mListener = listener;
+    /**
+     * Add a message listener when a message is received for the message type
+     */
+    @GuardedBy("mTransports")
+    public void addListener(int message, @NonNull Listener listener) {
+        mListeners.put(message, listener);
+        for (int i = 0; i < mTransports.size(); i++) {
+            mTransports.valueAt(i).addListener(message, listener);
+        }
     }
 
     /**
@@ -105,15 +117,7 @@ public class CompanionTransportManager {
                 detachSystemDataTransport(packageName, userId, associationId);
             }
 
-            final Transport transport;
-            if (isSecureTransportEnabled(associationId)) {
-                transport = new SecureTransport(associationId, fd, mContext, mListener);
-            } else {
-                transport = new RawTransport(associationId, fd, mContext, mListener);
-            }
-
-            transport.start();
-            mTransports.put(associationId, transport);
+            initializeTransport(associationId, fd);
         }
     }
 
@@ -128,13 +132,85 @@ public class CompanionTransportManager {
         }
     }
 
+    @GuardedBy("mTransports")
+    private void initializeTransport(int associationId, ParcelFileDescriptor fd) {
+        if (!isSecureTransportEnabled()) {
+            Transport transport = new RawTransport(associationId, fd, mContext);
+            for (Map.Entry<Integer, Listener> entry : mListeners.entrySet()) {
+                transport.addListener(entry.getKey(), entry.getValue());
+            }
+            transport.start();
+            mTransports.put(associationId, transport);
+            Slog.i(TAG, "RawTransport is created");
+            return;
+        }
+
+        // Exchange platform info to decide which transport should be created
+        mTempTransport = new RawTransport(associationId, fd, mContext);
+        for (Map.Entry<Integer, Listener> entry : mListeners.entrySet()) {
+            mTempTransport.addListener(entry.getKey(), entry.getValue());
+        }
+        mTempTransport.addListener(MESSAGE_REQUEST_PLATFORM_INFO, this::onPlatformInfoReceived);
+        mTempTransport.start();
+
+        int sdk = Build.VERSION.SDK_INT;
+        String release = Build.VERSION.RELEASE;
+        // data format: | SDK_INT (int) | release length (int) | release |
+        final ByteBuffer data = ByteBuffer.allocate(4 + 4 + release.getBytes().length)
+                .putInt(sdk)
+                .putInt(release.getBytes().length)
+                .put(release.getBytes());
+
+        // TODO: it should check if preSharedKey is given
+        mTempTransport.requestForResponse(MESSAGE_REQUEST_PLATFORM_INFO, data.array());
+    }
+
+    /**
+     * Depending on the remote platform info to decide which transport should be created
+     */
+    @GuardedBy("mTransports")
+    private void onPlatformInfoReceived(byte[] data) {
+        // TODO: it should check if preSharedKey is given
+
+        ByteBuffer buffer = ByteBuffer.wrap(data);
+        int remoteSdk = buffer.getInt();
+        byte[] remoteRelease = new byte[buffer.getInt()];
+        buffer.get(remoteRelease);
+
+        Slog.i(TAG, "Remote device SDK: " + remoteSdk + ", release:" + new String(remoteRelease));
+
+        Transport transport = mTempTransport;
+        mTempTransport = null;
+
+        int sdk = Build.VERSION.SDK_INT;
+        String release = Build.VERSION.RELEASE;
+        if (remoteSdk == NON_ANDROID) {
+            // TODO: pass in a real preSharedKey
+            transport = new SecureTransport(transport.getAssociationId(), transport.getFd(),
+                    mContext, null, null);
+        } else if (sdk < Build.VERSION_CODES.UPSIDE_DOWN_CAKE
+                || remoteSdk < Build.VERSION_CODES.UPSIDE_DOWN_CAKE) {
+            // TODO: depending on the release version, either
+            //       1) using a RawTransport for old T versions
+            //       2) or an Ukey2 handshaked transport for UKey2 backported T versions
+        } else {
+            Slog.i(TAG, "Creating a secure channel");
+            transport = new SecureTransport(transport.getAssociationId(), transport.getFd(),
+                    mContext);
+            for (Map.Entry<Integer, Listener> entry : mListeners.entrySet()) {
+                transport.addListener(entry.getKey(), entry.getValue());
+            }
+            transport.start();
+        }
+        mTransports.put(transport.getAssociationId(), transport);
+    }
+
     public Future<?> requestPermissionRestore(int associationId, byte[] data) {
         synchronized (mTransports) {
             final Transport transport = mTransports.get(associationId);
             if (transport == null) {
                 return CompletableFuture.failedFuture(new IOException("Missing transport"));
             }
-
             return transport.requestForResponse(MESSAGE_REQUEST_PERMISSION_RESTORE, data);
         }
     }
@@ -146,10 +222,9 @@ public class CompanionTransportManager {
         this.mSecureTransportEnabled = enabled;
     }
 
-    private boolean isSecureTransportEnabled(int associationId) {
+    private boolean isSecureTransportEnabled() {
         boolean enabled = !Build.IS_DEBUGGABLE || mSecureTransportEnabled;
 
-        // TODO: version comparison logic
         return enabled;
     }
 }
diff --git a/services/companion/java/com/android/server/companion/transport/CryptoManager.java b/services/companion/java/com/android/server/companion/transport/CryptoManager.java
index b08354afc8ad..a15939e52936 100644
--- a/services/companion/java/com/android/server/companion/transport/CryptoManager.java
+++ b/services/companion/java/com/android/server/companion/transport/CryptoManager.java
@@ -16,51 +16,51 @@
 
 package com.android.server.companion.transport;
 
-import android.security.keystore.KeyGenParameterSpec;
-import android.security.keystore.KeyProperties;
 import android.util.Slog;
 
-import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.security.InvalidAlgorithmParameterException;
 import java.security.InvalidKeyException;
-import java.security.KeyStore;
-import java.security.KeyStoreException;
 import java.security.NoSuchAlgorithmException;
-import java.security.UnrecoverableEntryException;
-import java.security.cert.CertificateException;
+import java.util.Arrays;
 
 import javax.crypto.BadPaddingException;
 import javax.crypto.Cipher;
 import javax.crypto.IllegalBlockSizeException;
-import javax.crypto.KeyGenerator;
 import javax.crypto.NoSuchPaddingException;
 import javax.crypto.SecretKey;
 import javax.crypto.spec.IvParameterSpec;
+import javax.crypto.spec.SecretKeySpec;
 
 /**
- * This class can be used to encrypt and decrypt bytes using Android Cryptography
+ * This class uses Java Cryptography to encrypt and decrypt messages
  */
 public class CryptoManager {
 
     private static final String TAG = "CDM_CryptoManager";
+    private static final int SECRET_KEY_LENGTH = 32;
+    private static final String ALGORITHM = "AES";
+    private static final String TRANSFORMATION = "AES/CBC/PKCS7Padding";
 
-    private static final String KEY_STORE_ALIAS = "cdm_secret";
-    private static final String ALGORITHM = KeyProperties.KEY_ALGORITHM_AES;
-    private static final String BLOCK_MODE = KeyProperties.BLOCK_MODE_CBC;
-    private static final String PADDING = KeyProperties.ENCRYPTION_PADDING_PKCS7;
-    private static final String TRANSFORMATION = ALGORITHM + "/" + BLOCK_MODE + "/" + PADDING;
+    private final byte[] mPreSharedKey;
+    private Cipher mEncryptCipher;
+    private Cipher mDecryptCipher;
 
-    private final KeyStore mKeyStore;
+    private SecretKey mSecretKey;
 
-    public CryptoManager() {
-        // Initialize KeyStore
+    public CryptoManager(byte[] preSharedKey) {
+        if (preSharedKey == null) {
+            mPreSharedKey = Arrays.copyOf(new byte[0], SECRET_KEY_LENGTH);
+        } else {
+            mPreSharedKey = Arrays.copyOf(preSharedKey, SECRET_KEY_LENGTH);
+        }
+        mSecretKey = new SecretKeySpec(mPreSharedKey, ALGORITHM);
         try {
-            mKeyStore = KeyStore.getInstance("AndroidKeyStore");
-            mKeyStore.load(null);
-        } catch (KeyStoreException | IOException | NoSuchAlgorithmException
-                 | CertificateException e) {
-            throw new RuntimeException(e);
+            mEncryptCipher = Cipher.getInstance(TRANSFORMATION);
+            mEncryptCipher.init(Cipher.ENCRYPT_MODE, mSecretKey);
+            mDecryptCipher = Cipher.getInstance(TRANSFORMATION);
+        } catch (NoSuchPaddingException | NoSuchAlgorithmException | InvalidKeyException e) {
+            Slog.e(TAG, e.getMessage());
         }
     }
 
@@ -69,21 +69,19 @@ public class CryptoManager {
      */
     public byte[] encrypt(byte[] input) {
         try {
-            // Encrypt using Cipher
-            Cipher encryptCipher = Cipher.getInstance(TRANSFORMATION);
-            encryptCipher.init(Cipher.ENCRYPT_MODE, getKey());
-            byte[] encryptedBytes = encryptCipher.doFinal(input);
+            if (mEncryptCipher == null) {
+                return null;
+            }
 
-            // Write to bytes
+            byte[] encryptedBytes = mEncryptCipher.doFinal(input);
             ByteBuffer buffer = ByteBuffer.allocate(
-                            4 + encryptCipher.getIV().length + 4 + encryptedBytes.length)
-                    .putInt(encryptCipher.getIV().length)
-                    .put(encryptCipher.getIV())
+                            4 + mEncryptCipher.getIV().length + 4 + encryptedBytes.length)
+                    .putInt(mEncryptCipher.getIV().length)
+                    .put(mEncryptCipher.getIV())
                     .putInt(encryptedBytes.length)
                     .put(encryptedBytes);
             return buffer.array();
-        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException
-                 | IllegalBlockSizeException | BadPaddingException e) {
+        } catch (IllegalBlockSizeException | BadPaddingException e) {
             Slog.e(TAG, e.getMessage());
             return null;
         }
@@ -99,45 +97,20 @@ public class CryptoManager {
         byte[] encryptedBytes = new byte[buffer.getInt()];
         buffer.get(encryptedBytes);
         try {
-            Cipher decryptCipher = Cipher.getInstance(TRANSFORMATION);
-            decryptCipher.init(Cipher.DECRYPT_MODE, getKey(), new IvParameterSpec(iv));
-            return decryptCipher.doFinal(encryptedBytes);
-        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException
-                 | InvalidAlgorithmParameterException | IllegalBlockSizeException
-                 | BadPaddingException e) {
+            mDecryptCipher.init(Cipher.DECRYPT_MODE, getKey(), new IvParameterSpec(iv));
+            return mDecryptCipher.doFinal(encryptedBytes);
+        } catch (InvalidKeyException | InvalidAlgorithmParameterException
+                 | IllegalBlockSizeException | BadPaddingException e) {
             Slog.e(TAG, e.getMessage());
             return null;
         }
     }
 
     private SecretKey getKey() {
-        try {
-            KeyStore.Entry keyEntry = mKeyStore.getEntry(KEY_STORE_ALIAS, null);
-            if (keyEntry instanceof KeyStore.SecretKeyEntry
-                    && ((KeyStore.SecretKeyEntry) keyEntry).getSecretKey() != null) {
-                return ((KeyStore.SecretKeyEntry) keyEntry).getSecretKey();
-            } else {
-                return createKey();
-            }
-        } catch (NoSuchAlgorithmException | UnrecoverableEntryException | KeyStoreException e) {
-            throw new RuntimeException(e);
-        }
-    }
-
-    private SecretKey createKey() {
-        try {
-            KeyGenerator keyGenerator = KeyGenerator.getInstance(ALGORITHM);
-            keyGenerator.init(
-                    new KeyGenParameterSpec.Builder(KEY_STORE_ALIAS,
-                            KeyProperties.PURPOSE_ENCRYPT | KeyProperties.PURPOSE_DECRYPT)
-                            .setBlockModes(BLOCK_MODE)
-                            .setEncryptionPaddings(PADDING)
-                            .setUserAuthenticationRequired(false)
-                            .setRandomizedEncryptionRequired(true)
-                            .build());
-            return keyGenerator.generateKey();
-        } catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException e) {
-            throw new RuntimeException(e);
+        if (mSecretKey != null) {
+            return mSecretKey;
         }
+        mSecretKey = new SecretKeySpec(mPreSharedKey, ALGORITHM);
+        return mSecretKey;
     }
 }
diff --git a/services/companion/java/com/android/server/companion/transport/RawTransport.java b/services/companion/java/com/android/server/companion/transport/RawTransport.java
index 7c0c7cf7ac68..4060f6efe0ca 100644
--- a/services/companion/java/com/android/server/companion/transport/RawTransport.java
+++ b/services/companion/java/com/android/server/companion/transport/RawTransport.java
@@ -21,8 +21,6 @@ import android.content.Context;
 import android.os.ParcelFileDescriptor;
 import android.util.Slog;
 
-import com.android.server.companion.transport.CompanionTransportManager.Listener;
-
 import libcore.io.IoUtils;
 import libcore.io.Streams;
 
@@ -32,8 +30,8 @@ import java.nio.ByteBuffer;
 class RawTransport extends Transport {
     private volatile boolean mStopped;
 
-    RawTransport(int associationId, ParcelFileDescriptor fd, Context context, Listener listener) {
-        super(associationId, fd, context, listener);
+    RawTransport(int associationId, ParcelFileDescriptor fd, Context context) {
+        super(associationId, fd, context);
     }
 
     @Override
@@ -64,7 +62,7 @@ class RawTransport extends Transport {
     protected void sendMessage(int message, int sequence, @NonNull byte[] data)
             throws IOException {
         if (DEBUG) {
-            Slog.d(TAG, "Sending message 0x" + Integer.toHexString(message)
+            Slog.e(TAG, "Sending message 0x" + Integer.toHexString(message)
                     + " sequence " + sequence + " length " + data.length
                     + " to association " + mAssociationId);
         }
diff --git a/services/companion/java/com/android/server/companion/transport/SecureTransport.java b/services/companion/java/com/android/server/companion/transport/SecureTransport.java
index 4194130f7e84..cca08435c0a5 100644
--- a/services/companion/java/com/android/server/companion/transport/SecureTransport.java
+++ b/services/companion/java/com/android/server/companion/transport/SecureTransport.java
@@ -21,8 +21,8 @@ import android.content.Context;
 import android.os.ParcelFileDescriptor;
 import android.util.Slog;
 
+import com.android.server.companion.securechannel.AttestationVerifier;
 import com.android.server.companion.securechannel.SecureChannel;
-import com.android.server.companion.transport.CompanionTransportManager.Listener;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -37,14 +37,17 @@ class SecureTransport extends Transport implements SecureChannel.Callback {
 
     private final BlockingQueue<byte[]> mRequestQueue = new ArrayBlockingQueue<>(100);
 
-    SecureTransport(int associationId,
-            ParcelFileDescriptor fd,
-            Context context,
-            Listener listener) {
-        super(associationId, fd, context, listener);
+    SecureTransport(int associationId, ParcelFileDescriptor fd, Context context) {
+        super(associationId, fd, context);
         mSecureChannel = new SecureChannel(mRemoteIn, mRemoteOut, this, context);
     }
 
+    SecureTransport(int associationId, ParcelFileDescriptor fd, Context context,
+            byte[] preSharedKey, AttestationVerifier verifier) {
+        super(associationId, fd, context);
+        mSecureChannel = new SecureChannel(mRemoteIn, mRemoteOut, this, preSharedKey, verifier);
+    }
+
     @Override
     public void start() {
         mSecureChannel.start();
diff --git a/services/companion/java/com/android/server/companion/transport/Transport.java b/services/companion/java/com/android/server/companion/transport/Transport.java
index 923d4243a34c..e984c637b44c 100644
--- a/services/companion/java/com/android/server/companion/transport/Transport.java
+++ b/services/companion/java/com/android/server/companion/transport/Transport.java
@@ -25,23 +25,28 @@ import android.util.Slog;
 import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
-import com.android.server.companion.transport.CompanionTransportManager.Listener;
 
 import libcore.util.EmptyArray;
 
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicInteger;
 
-abstract class Transport {
+/**
+ * This class represents the channel established between two devices.
+ */
+public abstract class Transport {
     protected static final String TAG = "CDM_CompanionTransport";
     protected static final boolean DEBUG = Build.IS_DEBUGGABLE;
 
     static final int MESSAGE_REQUEST_PING = 0x63807378; // ?PIN
-    static final int MESSAGE_REQUEST_PERMISSION_RESTORE = 0x63826983; // ?RES
+    public static final int MESSAGE_REQUEST_PLATFORM_INFO = 0x63807086; // ?PFV
+    public static final int MESSAGE_REQUEST_PERMISSION_RESTORE = 0x63826983; // ?RES
 
     static final int MESSAGE_RESPONSE_SUCCESS = 0x33838567; // !SUC
     static final int MESSAGE_RESPONSE_FAILURE = 0x33706573; // !FAI
@@ -49,11 +54,24 @@ abstract class Transport {
     protected static final int HEADER_LENGTH = 12;
 
     protected final int mAssociationId;
+    protected final ParcelFileDescriptor mFd;
     protected final InputStream mRemoteIn;
     protected final OutputStream mRemoteOut;
     protected final Context mContext;
 
-    private final Listener mListener;
+    /** Message type -> Listener */
+    private final Map<Integer, Listener> mListeners;
+
+    /**
+     * Message listener
+     */
+    public interface Listener {
+        /**
+         * Called when a message is received
+         * @param data data content in the message
+         */
+        void onDataReceived(byte[] data);
+    }
 
     private static boolean isRequest(int message) {
         return (message & 0xFF000000) == 0x63000000;
@@ -68,16 +86,36 @@ abstract class Transport {
             new SparseArray<>();
     protected final AtomicInteger mNextSequence = new AtomicInteger();
 
-    Transport(int associationId, ParcelFileDescriptor fd, Context context, Listener listener) {
-        this.mAssociationId = associationId;
-        this.mRemoteIn = new ParcelFileDescriptor.AutoCloseInputStream(fd);
-        this.mRemoteOut = new ParcelFileDescriptor.AutoCloseOutputStream(fd);
-        this.mContext = context;
-        this.mListener = listener;
+    Transport(int associationId, ParcelFileDescriptor fd, Context context) {
+        mAssociationId = associationId;
+        mFd = fd;
+        mRemoteIn = new ParcelFileDescriptor.AutoCloseInputStream(fd);
+        mRemoteOut = new ParcelFileDescriptor.AutoCloseOutputStream(fd);
+        mContext = context;
+        mListeners = new HashMap<>();
+    }
+
+    /**
+     * Add a listener when a message is received for the message type
+     * @param message Message type
+     * @param listener Execute when a message with the type is received
+     */
+    public void addListener(int message, Listener listener) {
+        mListeners.put(message, listener);
+    }
+
+    public int getAssociationId() {
+        return mAssociationId;
+    }
+
+    protected ParcelFileDescriptor getFd() {
+        return mFd;
     }
 
     public abstract void start();
     public abstract void stop();
+    protected abstract void sendMessage(int message, int sequence, @NonNull byte[] data)
+            throws IOException;
 
     public Future<byte[]> requestForResponse(int message, byte[] data) {
         if (DEBUG) Slog.d(TAG, "Requesting for response");
@@ -99,9 +137,6 @@ abstract class Transport {
         return pending;
     }
 
-    protected abstract void sendMessage(int message, int sequence, @NonNull byte[] data)
-            throws IOException;
-
     protected final void handleMessage(int message, int sequence, @NonNull byte[] data)
             throws IOException {
         if (DEBUG) {
@@ -130,6 +165,11 @@ abstract class Transport {
                 sendMessage(MESSAGE_RESPONSE_SUCCESS, sequence, data);
                 break;
             }
+            case MESSAGE_REQUEST_PLATFORM_INFO: {
+                callback(message, data);
+                sendMessage(MESSAGE_RESPONSE_SUCCESS, sequence, EmptyArray.BYTE);
+                break;
+            }
             case MESSAGE_REQUEST_PERMISSION_RESTORE: {
                 if (!mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_WATCH)
                         && !Build.isDebuggable()) {
@@ -138,7 +178,7 @@ abstract class Transport {
                     break;
                 }
                 try {
-                    mListener.onRequestPermissionRestore(data);
+                    callback(message, data);
                     sendMessage(MESSAGE_RESPONSE_SUCCESS, sequence, EmptyArray.BYTE);
                 } catch (Exception e) {
                     Slog.w(TAG, "Failed to restore permissions");
@@ -154,6 +194,12 @@ abstract class Transport {
         }
     }
 
+    private void callback(int message, byte[] data) {
+        if (mListeners.containsKey(message)) {
+            mListeners.get(message).onDataReceived(data);
+        }
+    }
+
     private void processResponse(int message, int sequence, byte[] data) {
         final CompletableFuture<byte[]> future;
         synchronized (mPendingRequests) {