[AWARE] Create a locally maintained peer ID instead of firmware ID

The firmware provides a peer requestor ID + peer MAC to identify the
peer. The peer requestor ID alone is insufficient.

Original code used the peer requestor ID as the only identifier. This
failed once started testing with >2 devices.

Allocate a local (unique) peer ID which is mapped to the peer requestor
ID and peer MAC by the framework.

Bug: 63642354
Test: all unit tests + integration (sl4a) tests
Change-Id: Ic1b3a97b9675be8d8561f4500e8c777ddff42950

5 files changed, 215 insertions, 135 deletions

diff --git a/service/java/com/android/server/wifi/aware/WifiAwareDataPathStateManager.java b/service/java/com/android/server/wifi/aware/WifiAwareDataPathStateManager.java
index ca12d5d52..82a3264ee 100644
--- a/service/java/com/android/server/wifi/aware/WifiAwareDataPathStateManager.java
+++ b/service/java/com/android/server/wifi/aware/WifiAwareDataPathStateManager.java
@@ -673,7 +673,7 @@ public class WifiAwareDataPathStateManager {
                 }
 
                 nnri.interfaceName = selectInterfaceForRequest(nnri);
-                mMgr.initiateDataPathSetup(networkSpecifier, nnri.networkSpecifier.peerId,
+                mMgr.initiateDataPathSetup(networkSpecifier, nnri.peerInstanceId,
                         NanDataPathChannelCfg.CHANNEL_NOT_REQUESTED, selectChannelForRequest(nnri),
                         nnri.peerDiscoveryMac, nnri.interfaceName, nnri.networkSpecifier.pmk,
                         nnri.networkSpecifier.passphrase, nnri.networkSpecifier.isOutOfBand());
@@ -846,6 +846,7 @@ public class WifiAwareDataPathStateManager {
         public int uid;
         public String interfaceName;
         public int pubSubId = 0;
+        public int peerInstanceId = 0;
         public byte[] peerDiscoveryMac = null;
         public int ndpId;
         public byte[] peerDataMac;
@@ -857,6 +858,7 @@ public class WifiAwareDataPathStateManager {
         static AwareNetworkRequestInformation processNetworkSpecifier(WifiAwareNetworkSpecifier ns,
                 WifiAwareStateManager mgr, WifiPermissionsWrapper permissionWrapper) {
             int uid, pubSubId = 0;
+            int peerInstanceId = 0;
             byte[] peerMac = ns.peerMac;
 
             if (VDBG) {
@@ -919,15 +921,17 @@ public class WifiAwareDataPathStateManager {
 
                 if (ns.type == WifiAwareNetworkSpecifier.NETWORK_SPECIFIER_TYPE_IB) {
                     pubSubId = session.getPubSubId();
-                    String peerMacStr = session.getMac(ns.peerId, null);
-                    if (peerMacStr == null) {
+                    WifiAwareDiscoverySessionState.PeerInfo peerInfo = session.getPeerInfo(
+                            ns.peerId);
+                    if (peerInfo == null) {
                         Log.e(TAG, "processNetworkSpecifier: networkSpecifier=" + ns
-                                + " -- no MAC address associated with this peer id -- peerId="
+                                + " -- no peer info associated with this peer id -- peerId="
                                 + ns.peerId);
                         return null;
                     }
+                    peerInstanceId = peerInfo.mInstanceId;
                     try {
-                        peerMac = HexEncoding.decode(peerMacStr.toCharArray(), false);
+                        peerMac = peerInfo.mMac;
                         if (peerMac == null || peerMac.length != 6) {
                             Log.e(TAG, "processNetworkSpecifier: networkSpecifier="
                                     + ns + " -- invalid peer MAC address");
@@ -979,6 +983,7 @@ public class WifiAwareDataPathStateManager {
                     : AwareNetworkRequestInformation.STATE_RESPONDER_IDLE;
             nnri.uid = uid;
             nnri.pubSubId = pubSubId;
+            nnri.peerInstanceId = peerInstanceId;
             nnri.peerDiscoveryMac = peerMac;
             nnri.networkSpecifier = ns;
 
@@ -990,7 +995,8 @@ public class WifiAwareDataPathStateManager {
             StringBuilder sb = new StringBuilder("AwareNetworkRequestInformation: ");
             sb.append("state=").append(state).append(", ns=").append(networkSpecifier).append(
                     ", uid=").append(uid).append(", interfaceName=").append(interfaceName).append(
-                    ", pubSubId=").append(pubSubId).append(", peerDiscoveryMac=").append(
+                    ", pubSubId=").append(pubSubId).append(", peerInstanceId=").append(
+                    peerInstanceId).append(", peerDiscoveryMac=").append(
                     peerDiscoveryMac == null ? ""
                             : String.valueOf(HexEncoding.encode(peerDiscoveryMac))).append(
                     ", ndpId=").append(ndpId).append(", peerDataMac=").append(
diff --git a/service/java/com/android/server/wifi/aware/WifiAwareDiscoverySessionState.java b/service/java/com/android/server/wifi/aware/WifiAwareDiscoverySessionState.java
index ac3dee414..86f4e37c2 100644
--- a/service/java/com/android/server/wifi/aware/WifiAwareDiscoverySessionState.java
+++ b/service/java/com/android/server/wifi/aware/WifiAwareDiscoverySessionState.java
@@ -28,6 +28,7 @@ import libcore.util.HexEncoding;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
+import java.util.Arrays;
 
 /**
  * Manages the state of a single Aware discovery session (publish or subscribe).
@@ -40,6 +41,8 @@ public class WifiAwareDiscoverySessionState {
     private static final boolean DBG = false;
     private static final boolean VDBG = false; // STOPSHIP if true
 
+    private int mNextPeerIdToBeAllocated = 100; // used to create a unique peer ID
+
     private final WifiAwareNativeApi mWifiAwareNativeApi;
     private int mSessionId;
     private byte mPubSubId;
@@ -47,7 +50,24 @@ public class WifiAwareDiscoverySessionState {
     private boolean mIsPublishSession;
     private final long mCreationTime;
 
-    private final SparseArray<String> mMacByRequestorInstanceId = new SparseArray<>();
+    static class PeerInfo {
+        PeerInfo(int instanceId, byte[] mac) {
+            mInstanceId = instanceId;
+            mMac = mac;
+        }
+
+        int mInstanceId;
+        byte[] mMac;
+
+        @Override
+        public String toString() {
+            StringBuilder sb = new StringBuilder("instanceId [");
+            sb.append(mInstanceId).append(", mac=").append(HexEncoding.encode(mMac)).append("]");
+            return sb.toString();
+        }
+    }
+
+    private final SparseArray<PeerInfo> mPeerInfoByRequestorInstanceId = new SparseArray<>();
 
     public WifiAwareDiscoverySessionState(WifiAwareNativeApi wifiAwareNativeApi, int sessionId,
             byte pubSubId, IWifiAwareDiscoverySessionCallback callback, boolean isPublishSession,
@@ -81,16 +101,11 @@ public class WifiAwareDiscoverySessionState {
     }
 
     /**
-     * Return the MAC address (String) of the specified peer ID - or a null if no such address is
+     * Return the peer information of the specified peer ID - or a null if no such peer ID is
      * registered.
      */
-    public String getMac(int peerId, String sep) {
-        String mac = mMacByRequestorInstanceId.get(peerId);
-        if (mac != null && sep != null && !sep.isEmpty()) {
-            mac = new StringBuilder(mac).insert(10, sep).insert(8, sep).insert(6, sep)
-                    .insert(4, sep).insert(2, sep).toString();
-        }
-        return mac;
+    public PeerInfo getPeerInfo(int peerId) {
+        return mPeerInfoByRequestorInstanceId.get(peerId);
     }
 
     /**
@@ -190,8 +205,8 @@ public class WifiAwareDiscoverySessionState {
      *            callbacks related to the message (success/failure).
      */
     public boolean sendMessage(short transactionId, int peerId, byte[] message, int messageId) {
-        String peerMacStr = mMacByRequestorInstanceId.get(peerId);
-        if (peerMacStr == null) {
+        PeerInfo peerInfo = mPeerInfoByRequestorInstanceId.get(peerId);
+        if (peerInfo == null) {
             Log.e(TAG, "sendMessage: attempting to send a message to an address which didn't "
                     + "match/contact us");
             try {
@@ -201,10 +216,9 @@ public class WifiAwareDiscoverySessionState {
             }
             return false;
         }
-        byte[] peerMac = HexEncoding.decode(peerMacStr.toCharArray(), false);
 
-        boolean success = mWifiAwareNativeApi.sendMessage(transactionId, mPubSubId, peerId, peerMac,
-                message, messageId);
+        boolean success = mWifiAwareNativeApi.sendMessage(transactionId, mPubSubId,
+                peerInfo.mInstanceId, peerInfo.mMac, message, messageId);
         if (!success) {
             try {
                 mCallback.onMessageSendFail(messageId, NanStatusType.INTERNAL_FAILURE);
@@ -233,13 +247,10 @@ public class WifiAwareDiscoverySessionState {
      */
     public void onMatch(int requestorInstanceId, byte[] peerMac, byte[] serviceSpecificInfo,
             byte[] matchFilter) {
-        String prevMac = mMacByRequestorInstanceId.get(requestorInstanceId);
-        mMacByRequestorInstanceId.put(requestorInstanceId, new String(HexEncoding.encode(peerMac)));
-
-        if (DBG) Log.d(TAG, "onMatch: previous peer MAC replaced - " + prevMac);
+        int peerId = getPeerIdOrAddIfNew(requestorInstanceId, peerMac);
 
         try {
-            mCallback.onMatch(requestorInstanceId, serviceSpecificInfo, matchFilter);
+            mCallback.onMatch(peerId, serviceSpecificInfo, matchFilter);
         } catch (RemoteException e) {
             Log.w(TAG, "onMatch: RemoteException (FYI): " + e);
         }
@@ -256,20 +267,35 @@ public class WifiAwareDiscoverySessionState {
      * @param message The received message.
      */
     public void onMessageReceived(int requestorInstanceId, byte[] peerMac, byte[] message) {
-        String prevMac = mMacByRequestorInstanceId.get(requestorInstanceId);
-        mMacByRequestorInstanceId.put(requestorInstanceId, new String(HexEncoding.encode(peerMac)));
-
-        if (DBG) {
-            Log.d(TAG, "onMessageReceived: previous peer MAC replaced - " + prevMac);
-        }
+        int peerId = getPeerIdOrAddIfNew(requestorInstanceId, peerMac);
 
         try {
-            mCallback.onMessageReceived(requestorInstanceId, message);
+            mCallback.onMessageReceived(peerId, message);
         } catch (RemoteException e) {
             Log.w(TAG, "onMessageReceived: RemoteException (FYI): " + e);
         }
     }
 
+    private int getPeerIdOrAddIfNew(int requestorInstanceId, byte[] peerMac) {
+        for (int i = 0; i < mPeerInfoByRequestorInstanceId.size(); ++i) {
+            PeerInfo peerInfo = mPeerInfoByRequestorInstanceId.valueAt(i);
+            if (peerInfo.mInstanceId == requestorInstanceId && Arrays.equals(peerMac,
+                    peerInfo.mMac)) {
+                return mPeerInfoByRequestorInstanceId.keyAt(i);
+            }
+        }
+
+        int newPeerId = mNextPeerIdToBeAllocated++;
+        PeerInfo newPeerInfo = new PeerInfo(requestorInstanceId, peerMac);
+        mPeerInfoByRequestorInstanceId.put(newPeerId, newPeerInfo);
+
+        if (DBG) {
+            Log.d(TAG, "New peer info: peerId=" + newPeerId + ", peerInfo=" + newPeerInfo);
+        }
+
+        return newPeerId;
+    }
+
     /**
      * Dump the internal state of the class.
      */
@@ -278,6 +304,6 @@ public class WifiAwareDiscoverySessionState {
         pw.println("  mSessionId: " + mSessionId);
         pw.println("  mIsPublishSession: " + mIsPublishSession);
         pw.println("  mPubSubId: " + mPubSubId);
-        pw.println("  mMacByRequestorInstanceId: [" + mMacByRequestorInstanceId + "]");
+        pw.println("  mPeerInfoByRequestorInstanceId: [" + mPeerInfoByRequestorInstanceId + "]");
     }
 }
diff --git a/service/java/com/android/server/wifi/aware/WifiAwareStateManager.java b/service/java/com/android/server/wifi/aware/WifiAwareStateManager.java
index aa5a426ed..30aa42a26 100644
--- a/service/java/com/android/server/wifi/aware/WifiAwareStateManager.java
+++ b/service/java/com/android/server/wifi/aware/WifiAwareStateManager.java
@@ -48,6 +48,7 @@ import com.android.internal.util.MessageUtils;
 import com.android.internal.util.State;
 import com.android.internal.util.StateMachine;
 import com.android.internal.util.WakeupMessage;
+import com.android.server.wifi.util.NativeUtil;
 import com.android.server.wifi.util.WifiPermissionsWrapper;
 
 import libcore.util.HexEncoding;
@@ -2332,10 +2333,13 @@ public class WifiAwareStateManager implements WifiAwareShellCommand.DelegatedShe
         for (RttManager.RttParams param : params) {
             String peerIdStr = param.bssid;
             try {
-                param.bssid = session.getMac(Integer.parseInt(peerIdStr), ":");
-                if (param.bssid == null) {
+                WifiAwareDiscoverySessionState.PeerInfo peerInfo = session.getPeerInfo(
+                        Integer.parseInt(peerIdStr));
+                if (peerInfo == null || peerInfo.mMac == null) {
                     Log.d(TAG, "startRangingLocal: no MAC address for peer ID=" + peerIdStr);
                     param.bssid = "";
+                } else {
+                    param.bssid = NativeUtil.macAddressFromByteArray(peerInfo.mMac);
                 }
             } catch (NumberFormatException e) {
                 Log.e(TAG, "startRangingLocal: invalid peer ID specification (in bssid field): '"
diff --git a/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareDataPathStateManagerTest.java b/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareDataPathStateManagerTest.java
index 0ddf765d7..f3927357f 100644
--- a/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareDataPathStateManagerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareDataPathStateManagerTest.java
@@ -67,7 +67,6 @@ import android.os.PowerManager;
 import android.os.Process;
 import android.os.test.TestLooper;
 import android.test.suitebuilder.annotation.SmallTest;
-import android.util.Pair;
 
 import com.android.internal.util.AsyncChannel;
 import com.android.server.wifi.util.WifiPermissionsWrapper;
@@ -263,7 +262,7 @@ public class WifiAwareDataPathStateManagerTest {
         final int clientId = 123;
         final byte pubSubId = 55;
         final byte[] pmk = "01234567890123456789012345678901".getBytes();
-        final PeerHandle peerHandle = new PeerHandle(1341234);
+        final int requestorId = 1341234;
         final byte[] peerDiscoveryMac = HexEncoding.decode("000102030405".toCharArray(), false);
 
         InOrder inOrder = inOrder(mMockNative, mMockCm, mMockCallback, mMockSessionCallback);
@@ -273,18 +272,18 @@ public class WifiAwareDataPathStateManagerTest {
                 anyInt())).thenReturn(PackageManager.PERMISSION_DENIED);
 
         // (0) initialize
-        Pair<Integer, Messenger> res = initDataPathEndPoint(clientId, pubSubId, peerHandle,
+        DataPathEndPointInfo res = initDataPathEndPoint(clientId, pubSubId, requestorId,
                 peerDiscoveryMac, inOrder, inOrderM, false);
 
         // (1) request network
-        NetworkRequest nr = getSessionNetworkRequest(clientId, res.first, peerHandle, pmk, null,
-                false);
+        NetworkRequest nr = getSessionNetworkRequest(clientId, res.mSessionId, res.mPeerHandle, pmk,
+                null, false);
 
         Message reqNetworkMsg = Message.obtain();
         reqNetworkMsg.what = NetworkFactory.CMD_REQUEST_NETWORK;
         reqNetworkMsg.obj = nr;
         reqNetworkMsg.arg1 = 0;
-        res.second.send(reqNetworkMsg);
+        res.mMessenger.send(reqNetworkMsg);
         mMockLooper.dispatchAll();
 
         // failure: no interactions with connectivity manager or native manager
@@ -502,19 +501,19 @@ public class WifiAwareDataPathStateManagerTest {
         final byte pubSubId = 55;
         final int ndpId = 2;
         final byte[] pmk = "01234567890123456789012345678901".getBytes();
-        final PeerHandle peerHandle = new PeerHandle(1341234);
+        final int requestorId = 1341234;
         final byte[] peerDiscoveryMac = HexEncoding.decode("000102030405".toCharArray(), false);
 
         InOrder inOrder = inOrder(mMockNative, mMockCm, mMockCallback, mMockSessionCallback);
         InOrder inOrderM = inOrder(mAwareMetricsMock);
 
         // (0) initialize
-        Pair<Integer, Messenger> res = initDataPathEndPoint(clientId, pubSubId, peerHandle,
+        DataPathEndPointInfo res = initDataPathEndPoint(clientId, pubSubId, requestorId,
                 peerDiscoveryMac, inOrder, inOrderM, doPublish);
 
         // (1) request network
-        NetworkRequest nr = getSessionNetworkRequest(clientId, res.first, peerHandle, pmk, null,
-                doPublish);
+        NetworkRequest nr = getSessionNetworkRequest(clientId, res.mSessionId, res.mPeerHandle, pmk,
+                null, doPublish);
 
         // corrupt the network specifier: reverse the role (so it's mis-matched)
         WifiAwareNetworkSpecifier ns =
@@ -535,7 +534,7 @@ public class WifiAwareDataPathStateManagerTest {
         reqNetworkMsg.what = NetworkFactory.CMD_REQUEST_NETWORK;
         reqNetworkMsg.obj = nr;
         reqNetworkMsg.arg1 = 0;
-        res.second.send(reqNetworkMsg);
+        res.mMessenger.send(reqNetworkMsg);
         mMockLooper.dispatchAll();
 
         // consequences of failure:
@@ -558,19 +557,19 @@ public class WifiAwareDataPathStateManagerTest {
         final byte pubSubId = 56;
         final int ndpId = 2;
         final byte[] pmk = "01234567890123456789012345678901".getBytes();
-        final PeerHandle peerHandle = new PeerHandle(1341234);
+        final int requestorId = 1341234;
         final byte[] peerDiscoveryMac = HexEncoding.decode("000102030405".toCharArray(), false);
 
         InOrder inOrder = inOrder(mMockNative, mMockCm, mMockCallback, mMockSessionCallback);
         InOrder inOrderM = inOrder(mAwareMetricsMock);
 
         // (0) initialize
-        Pair<Integer, Messenger> res = initDataPathEndPoint(clientId, pubSubId, peerHandle,
+        DataPathEndPointInfo res = initDataPathEndPoint(clientId, pubSubId, requestorId,
                 peerDiscoveryMac, inOrder, inOrderM, doPublish);
 
         // (1) create network request
-        NetworkRequest nr = getSessionNetworkRequest(clientId, res.first, peerHandle, pmk, null,
-                doPublish);
+        NetworkRequest nr = getSessionNetworkRequest(clientId, res.mSessionId, res.mPeerHandle, pmk,
+                null, doPublish);
 
         // (2) corrupt request's UID
         WifiAwareNetworkSpecifier ns =
@@ -592,7 +591,7 @@ public class WifiAwareDataPathStateManagerTest {
         reqNetworkMsg.what = NetworkFactory.CMD_REQUEST_NETWORK;
         reqNetworkMsg.obj = nr;
         reqNetworkMsg.arg1 = 0;
-        res.second.send(reqNetworkMsg);
+        res.mMessenger.send(reqNetworkMsg);
         mMockLooper.dispatchAll();
 
         // consequences of failure:
@@ -614,7 +613,7 @@ public class WifiAwareDataPathStateManagerTest {
             boolean immediateHalFailure) throws Exception {
         final int clientId = 123;
         final byte pubSubId = 58;
-        final PeerHandle peerHandle = new PeerHandle(1341234);
+        final int requestorId = 1341234;
         final int ndpId = 2;
         final byte[] pmk = "01234567890123456789012345678901".getBytes();
         final String passphrase = "some passphrase";
@@ -640,7 +639,7 @@ public class WifiAwareDataPathStateManagerTest {
         }
 
         // (0) initialize
-        Pair<Integer, Messenger> res = initDataPathEndPoint(clientId, pubSubId, peerHandle,
+        DataPathEndPointInfo res = initDataPathEndPoint(clientId, pubSubId, requestorId,
                 peerDiscoveryMac, inOrder, inOrderM, false);
 
         // (1) request network
@@ -651,18 +650,19 @@ public class WifiAwareDataPathStateManagerTest {
                     provideMac ? peerDiscoveryMac : null, providePmk ? pmk : null,
                     providePassphrase ? passphrase : null);
         } else {
-            nr = getSessionNetworkRequest(clientId, res.first, provideMac ? peerHandle : null,
-                    providePmk ? pmk : null, providePassphrase ? passphrase : null, false);
+            nr = getSessionNetworkRequest(clientId, res.mSessionId,
+                    provideMac ? res.mPeerHandle : null, providePmk ? pmk : null,
+                    providePassphrase ? passphrase : null, false);
         }
 
         Message reqNetworkMsg = Message.obtain();
         reqNetworkMsg.what = NetworkFactory.CMD_REQUEST_NETWORK;
         reqNetworkMsg.obj = nr;
         reqNetworkMsg.arg1 = 0;
-        res.second.send(reqNetworkMsg);
+        res.mMessenger.send(reqNetworkMsg);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).initiateDataPath(transactionId.capture(),
-                eq(useDirect ? 0 : peerHandle.peerId),
+                eq(useDirect ? 0 : requestorId),
                 eq(CHANNEL_NOT_REQUESTED), anyInt(), eq(peerDiscoveryMac),
                 eq(sAwareInterfacePrefix + "0"), eq(providePmk ? pmk : null),
                 eq(providePassphrase ? passphrase : null), eq(useDirect), any());
@@ -703,7 +703,7 @@ public class WifiAwareDataPathStateManagerTest {
             Message endNetworkReqMsg = Message.obtain();
             endNetworkReqMsg.what = NetworkFactory.CMD_CANCEL_REQUEST;
             endNetworkReqMsg.obj = nr;
-            res.second.send(endNetworkReqMsg);
+            res.mMessenger.send(endNetworkReqMsg);
 
             Message endNetworkUsageMsg = Message.obtain();
             endNetworkUsageMsg.what = AsyncChannel.CMD_CHANNEL_DISCONNECTED;
@@ -725,7 +725,7 @@ public class WifiAwareDataPathStateManagerTest {
             throws Exception {
         final int clientId = 123;
         final byte pubSubId = 60;
-        final PeerHandle peerHandle = new PeerHandle(1341234);
+        final int requestorId = 1341234;
         final int ndpId = 2;
         final byte[] pmk = "01234567890123456789012345678901".getBytes();
         final String passphrase = "some passphrase";
@@ -745,7 +745,7 @@ public class WifiAwareDataPathStateManagerTest {
         }
 
         // (0) initialize
-        Pair<Integer, Messenger> res = initDataPathEndPoint(clientId, pubSubId, peerHandle,
+        DataPathEndPointInfo res = initDataPathEndPoint(clientId, pubSubId, requestorId,
                 peerDiscoveryMac, inOrder, inOrderM, true);
 
         // (1) request network
@@ -756,15 +756,16 @@ public class WifiAwareDataPathStateManagerTest {
                     provideMac ? peerDiscoveryMac : null, providePmk ? pmk : null,
                     providePassphrase ? passphrase : null);
         } else {
-            nr = getSessionNetworkRequest(clientId, res.first, provideMac ? peerHandle : null,
-                    providePmk ? pmk : null, providePassphrase ? passphrase : null, true);
+            nr = getSessionNetworkRequest(clientId, res.mSessionId,
+                    provideMac ? res.mPeerHandle : null, providePmk ? pmk : null,
+                    providePassphrase ? passphrase : null, true);
         }
 
         Message reqNetworkMsg = Message.obtain();
         reqNetworkMsg.what = NetworkFactory.CMD_REQUEST_NETWORK;
         reqNetworkMsg.obj = nr;
         reqNetworkMsg.arg1 = 0;
-        res.second.send(reqNetworkMsg);
+        res.mMessenger.send(reqNetworkMsg);
         mMockLooper.dispatchAll();
 
         // (2) get request & respond
@@ -802,7 +803,7 @@ public class WifiAwareDataPathStateManagerTest {
             Message endNetworkMsg = Message.obtain();
             endNetworkMsg.what = NetworkFactory.CMD_CANCEL_REQUEST;
             endNetworkMsg.obj = nr;
-            res.second.send(endNetworkMsg);
+            res.mMessenger.send(endNetworkMsg);
 
             Message endNetworkUsageMsg = Message.obtain();
             endNetworkUsageMsg.what = AsyncChannel.CMD_CHANNEL_DISCONNECTED;
@@ -944,8 +945,8 @@ public class WifiAwareDataPathStateManagerTest {
         return new NetworkRequest(nc, 0, 0, NetworkRequest.Type.REQUEST);
     }
 
-    private Pair<Integer, Messenger> initDataPathEndPoint(int clientId, byte pubSubId,
-            PeerHandle peerHandle, byte[] peerDiscoveryMac, InOrder inOrder, InOrder inOrderM,
+    private DataPathEndPointInfo initDataPathEndPoint(int clientId, byte pubSubId,
+            int requestorId, byte[] peerDiscoveryMac, InOrder inOrder, InOrder inOrderM,
             boolean doPublish)
             throws Exception {
         final int pid = 2000;
@@ -960,6 +961,7 @@ public class WifiAwareDataPathStateManagerTest {
 
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdCaptor = ArgumentCaptor.forClass(Integer.class);
         ArgumentCaptor<Messenger> messengerCaptor = ArgumentCaptor.forClass(Messenger.class);
         ArgumentCaptor<String> strCaptor = ArgumentCaptor.forClass(String.class);
 
@@ -1020,12 +1022,25 @@ public class WifiAwareDataPathStateManagerTest {
         inOrderM.verify(mAwareMetricsMock).recordDiscoveryStatus(Process.myUid(),
                 NanStatusType.SUCCESS, doPublish);
 
-        mDut.onMessageReceivedNotification(pubSubId, peerHandle.peerId, peerDiscoveryMac,
+        mDut.onMessageReceivedNotification(pubSubId, requestorId, peerDiscoveryMac,
                 someMsg.getBytes());
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockSessionCallback).onMessageReceived(peerHandle.peerId,
-                someMsg.getBytes());
+        inOrder.verify(mMockSessionCallback).onMessageReceived(peerIdCaptor.capture(),
+                eq(someMsg.getBytes()));
+
+        return new DataPathEndPointInfo(sessionId.getValue(), peerIdCaptor.getValue(),
+                messengerCaptor.getValue());
+    }
+
+    private static class DataPathEndPointInfo {
+        int mSessionId;
+        PeerHandle mPeerHandle;
+        Messenger mMessenger;
 
-        return new Pair<>(sessionId.getValue(), messengerCaptor.getValue());
+        DataPathEndPointInfo(int sessionId, int peerId, Messenger messenger) {
+            mSessionId = sessionId;
+            mPeerHandle = new PeerHandle(peerId);
+            mMessenger = messenger;
+        }
     }
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareStateManagerTest.java b/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareStateManagerTest.java
index 3652e7b52..faa87df2f 100644
--- a/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareStateManagerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareStateManagerTest.java
@@ -1211,6 +1211,7 @@ public class WifiAwareStateManagerTest {
                 IWifiAwareDiscoverySessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdCaptor = ArgumentCaptor.forClass(Integer.class);
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
         mDut.enableUsage();
@@ -1241,16 +1242,18 @@ public class WifiAwareStateManagerTest {
         mDut.onMatchNotification(subscribeId, requestorId, peerMac, peerSsi.getBytes(),
                 peerMatchFilter.getBytes());
         mMockLooper.dispatchAll();
-        inOrder.verify(mockSessionCallback).onMatch(requestorId, peerSsi.getBytes(),
-                peerMatchFilter.getBytes());
+        inOrder.verify(mockSessionCallback).onMatch(peerIdCaptor.capture(), eq(peerSsi.getBytes()),
+                eq(peerMatchFilter.getBytes()));
 
         // (3) message Rx
         mDut.onMessageReceivedNotification(subscribeId, requestorId, peerMac, peerMsg.getBytes());
         mMockLooper.dispatchAll();
-        inOrder.verify(mockSessionCallback).onMessageReceived(requestorId, peerMsg.getBytes());
+        inOrder.verify(mockSessionCallback).onMessageReceived(peerIdCaptor.getValue(),
+                peerMsg.getBytes());
 
         // (4) message Tx successful queuing
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(), messageId, 0);
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), ssi.getBytes(),
+                messageId, 0);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
                 eq(requestorId), eq(peerMac), eq(ssi.getBytes()), eq(messageId));
@@ -1259,8 +1262,8 @@ public class WifiAwareStateManagerTest {
         mMockLooper.dispatchAll();
 
         // (5) message Tx successful queuing
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(), messageId2,
-                0);
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), ssi.getBytes(),
+                messageId2, 0);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
                 eq(requestorId), eq(peerMac), eq(ssi.getBytes()), eq(messageId2));
@@ -1295,8 +1298,8 @@ public class WifiAwareStateManagerTest {
         final int masterPref = 0;
         final String serviceName = "some-service-name";
         final byte publishId = 88;
-        final int peerId1 = 568;
-        final int peerId2 = 873;
+        final int requestorId1 = 568;
+        final int requestorId2 = 873;
         final byte[] peerMac1 = HexEncoding.decode("000102030405".toCharArray(), false);
         final byte[] peerMac2 = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final String msgFromPeer1 = "hey from 000102...";
@@ -1315,6 +1318,7 @@ public class WifiAwareStateManagerTest {
 
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdCaptor = ArgumentCaptor.forClass(Integer.class);
         IWifiAwareEventCallback mockCallback = mock(IWifiAwareEventCallback.class);
         IWifiAwareDiscoverySessionCallback mockSessionCallback = mock(
                 IWifiAwareDiscoverySessionCallback.class);
@@ -1345,18 +1349,24 @@ public class WifiAwareStateManagerTest {
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
         // (3) message received from peers 1 & 2
-        mDut.onMessageReceivedNotification(publishId, peerId1, peerMac1, msgFromPeer1.getBytes());
-        mDut.onMessageReceivedNotification(publishId, peerId2, peerMac2, msgFromPeer2.getBytes());
+        mDut.onMessageReceivedNotification(publishId, requestorId1, peerMac1,
+                msgFromPeer1.getBytes());
+        mDut.onMessageReceivedNotification(publishId, requestorId2, peerMac2,
+                msgFromPeer2.getBytes());
         mMockLooper.dispatchAll();
-        inOrder.verify(mockSessionCallback).onMessageReceived(peerId1, msgFromPeer1.getBytes());
-        inOrder.verify(mockSessionCallback).onMessageReceived(peerId2, msgFromPeer2.getBytes());
+        inOrder.verify(mockSessionCallback).onMessageReceived(peerIdCaptor.capture(),
+                eq(msgFromPeer1.getBytes()));
+        int peerId1 = peerIdCaptor.getValue();
+        inOrder.verify(mockSessionCallback).onMessageReceived(peerIdCaptor.capture(),
+                eq(msgFromPeer2.getBytes()));
+        int peerId2 = peerIdCaptor.getValue();
 
         // (4) sending messages back to same peers: one Tx fails, other succeeds
         mDut.sendMessage(clientId, sessionId.getValue(), peerId2, msgToPeer2.getBytes(),
                 msgToPeerId2, 0);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId), eq(peerId2),
-                eq(peerMac2), eq(msgToPeer2.getBytes()), eq(msgToPeerId2));
+        inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId),
+                eq(requestorId2), eq(peerMac2), eq(msgToPeer2.getBytes()), eq(msgToPeerId2));
         short transactionIdVal = transactionId.getValue();
         mDut.onMessageSendQueuedSuccessResponse(transactionIdVal);
         mDut.onMessageSendSuccessNotification(transactionIdVal);
@@ -1365,8 +1375,8 @@ public class WifiAwareStateManagerTest {
                 msgToPeerId1, 0);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onMessageSendSuccess(msgToPeerId2);
-        inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId), eq(peerId1),
-                eq(peerMac1), eq(msgToPeer1.getBytes()), eq(msgToPeerId1));
+        inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId),
+                eq(requestorId1), eq(peerMac1), eq(msgToPeer1.getBytes()), eq(msgToPeerId1));
         transactionIdVal = transactionId.getValue();
         mDut.onMessageSendQueuedSuccessResponse(transactionIdVal);
         mDut.onMessageSendFailNotification(transactionIdVal, reason);
@@ -1393,7 +1403,7 @@ public class WifiAwareStateManagerTest {
         final int masterPref = 0;
         final String serviceName = "some-service-name";
         final byte publishId = 88;
-        final int peerId = 568;
+        final int requestorId = 568;
         final byte[] peerMacOrig = HexEncoding.decode("000102030405".toCharArray(), false);
         final byte[] peerMacLater = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final String msgFromPeer1 = "hey from 000102...";
@@ -1410,6 +1420,7 @@ public class WifiAwareStateManagerTest {
 
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerId = ArgumentCaptor.forClass(Integer.class);
         IWifiAwareEventCallback mockCallback = mock(IWifiAwareEventCallback.class);
         IWifiAwareDiscoverySessionCallback mockSessionCallback = mock(
                 IWifiAwareDiscoverySessionCallback.class);
@@ -1440,13 +1451,17 @@ public class WifiAwareStateManagerTest {
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
         // (3) message received & responded to
-        mDut.onMessageReceivedNotification(publishId, peerId, peerMacOrig, msgFromPeer1.getBytes());
-        mDut.sendMessage(clientId, sessionId.getValue(), peerId, msgToPeer1.getBytes(),
+        mDut.onMessageReceivedNotification(publishId, requestorId, peerMacOrig,
+                msgFromPeer1.getBytes());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onMessageReceived(peerId.capture(),
+                eq(msgFromPeer1.getBytes()));
+        mDut.sendMessage(clientId, sessionId.getValue(), peerId.getValue(), msgToPeer1.getBytes(),
                 msgToPeerId1, 0);
         mMockLooper.dispatchAll();
-        inOrder.verify(mockSessionCallback).onMessageReceived(peerId, msgFromPeer1.getBytes());
-        inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId), eq(peerId),
-                eq(peerMacOrig), eq(msgToPeer1.getBytes()), eq(msgToPeerId1));
+        inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId),
+                eq(requestorId), eq(peerMacOrig), eq(msgToPeer1.getBytes()),
+                eq(msgToPeerId1));
         mDut.onMessageSendQueuedSuccessResponse(transactionId.getValue());
         mDut.onMessageSendSuccessNotification(transactionId.getValue());
         mMockLooper.dispatchAll();
@@ -1454,14 +1469,17 @@ public class WifiAwareStateManagerTest {
         validateInternalSendMessageQueuesCleanedUp(msgToPeerId1);
 
         // (4) message received with same peer ID but different MAC
-        mDut.onMessageReceivedNotification(publishId, peerId, peerMacLater,
+        mDut.onMessageReceivedNotification(publishId, requestorId, peerMacLater,
                 msgFromPeer2.getBytes());
-        mDut.sendMessage(clientId, sessionId.getValue(), peerId, msgToPeer2.getBytes(),
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onMessageReceived(peerId.capture(),
+                eq(msgFromPeer2.getBytes()));
+        mDut.sendMessage(clientId, sessionId.getValue(), peerId.getValue(), msgToPeer2.getBytes(),
                 msgToPeerId2, 0);
         mMockLooper.dispatchAll();
-        inOrder.verify(mockSessionCallback).onMessageReceived(peerId, msgFromPeer2.getBytes());
-        inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId), eq(peerId),
-                eq(peerMacLater), eq(msgToPeer2.getBytes()), eq(msgToPeerId2));
+        inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId),
+                eq(requestorId), eq(peerMacLater), eq(msgToPeer2.getBytes()),
+                eq(msgToPeerId2));
         mDut.onMessageSendQueuedSuccessResponse(transactionId.getValue());
         mDut.onMessageSendSuccessNotification(transactionId.getValue());
         mMockLooper.dispatchAll();
@@ -1497,6 +1515,7 @@ public class WifiAwareStateManagerTest {
                 IWifiAwareDiscoverySessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdCaptor = ArgumentCaptor.forClass(Integer.class);
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
         mDut.enableUsage();
@@ -1524,12 +1543,12 @@ public class WifiAwareStateManagerTest {
                 peerMatchFilter.getBytes());
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
-        inOrder.verify(mockSessionCallback).onMatch(requestorId, peerSsi.getBytes(),
-                peerMatchFilter.getBytes());
+        inOrder.verify(mockSessionCallback).onMatch(peerIdCaptor.capture(), eq(peerSsi.getBytes()),
+                eq(peerMatchFilter.getBytes()));
 
         // (3) send message to invalid peer ID
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId + 5, ssi.getBytes(),
-                messageId, 0);
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue() + 5,
+                ssi.getBytes(), messageId, 0);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onMessageSendFail(messageId,
                 NanStatusType.INTERNAL_FAILURE);
@@ -1565,6 +1584,7 @@ public class WifiAwareStateManagerTest {
                 IWifiAwareDiscoverySessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdCaptor = ArgumentCaptor.forClass(Integer.class);
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
         mDut.enableUsage();
@@ -1592,11 +1612,11 @@ public class WifiAwareStateManagerTest {
                 peerMatchFilter.getBytes());
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
-        inOrder.verify(mockSessionCallback).onMatch(requestorId, peerSsi.getBytes(),
-                peerMatchFilter.getBytes());
+        inOrder.verify(mockSessionCallback).onMatch(peerIdCaptor.capture(), eq(peerSsi.getBytes()),
+                eq(peerMatchFilter.getBytes()));
 
         // (3) send 2 messages and enqueue successfully
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(),
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), ssi.getBytes(),
                 messageId, 0);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
@@ -1605,7 +1625,7 @@ public class WifiAwareStateManagerTest {
         mDut.onMessageSendQueuedSuccessResponse(transactionId1);
         mMockLooper.dispatchAll();
 
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(),
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), ssi.getBytes(),
                 messageId + 1, 0);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
@@ -1615,8 +1635,8 @@ public class WifiAwareStateManagerTest {
         mMockLooper.dispatchAll();
 
         // (4) send a message and get a queueing failure (not queue full)
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(), messageId + 2,
-                0);
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), ssi.getBytes(),
+                messageId + 2, 0);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
                 eq(requestorId), eq(peerMac), eq(ssi.getBytes()), eq(messageId + 2));
@@ -1631,8 +1651,8 @@ public class WifiAwareStateManagerTest {
         when(mMockNative.sendMessage(anyShort(), anyByte(), anyInt(), any(),
                 any(), anyInt())).thenReturn(false);
 
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(), messageId + 3,
-                0);
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), ssi.getBytes(),
+                messageId + 3, 0);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
                 eq(requestorId), eq(peerMac), eq(ssi.getBytes()), eq(messageId + 3));
@@ -1690,6 +1710,7 @@ public class WifiAwareStateManagerTest {
                 IWifiAwareDiscoverySessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdCaptor = ArgumentCaptor.forClass(Integer.class);
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
         mDut.enableUsage();
@@ -1717,11 +1738,11 @@ public class WifiAwareStateManagerTest {
                 peerMatchFilter.getBytes());
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
-        inOrder.verify(mockSessionCallback).onMatch(requestorId, peerSsi.getBytes(),
-                peerMatchFilter.getBytes());
+        inOrder.verify(mockSessionCallback).onMatch(peerIdCaptor.capture(), eq(peerSsi.getBytes()),
+                eq(peerMatchFilter.getBytes()));
 
         // (3) send message and enqueue successfully
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(),
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), ssi.getBytes(),
                 messageId, retryCount);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
@@ -1776,6 +1797,7 @@ public class WifiAwareStateManagerTest {
                 IWifiAwareDiscoverySessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdCaptor = ArgumentCaptor.forClass(Integer.class);
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
         mDut.enableUsage();
@@ -1803,12 +1825,12 @@ public class WifiAwareStateManagerTest {
                 peerMatchFilter.getBytes());
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
-        inOrder.verify(mockSessionCallback).onMatch(requestorId, peerSsi.getBytes(),
-                peerMatchFilter.getBytes());
+        inOrder.verify(mockSessionCallback).onMatch(peerIdCaptor.capture(), eq(peerSsi.getBytes()),
+                eq(peerMatchFilter.getBytes()));
 
         // (3) send message and enqueue successfully
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(), messageId,
-                retryCount);
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), ssi.getBytes(),
+                messageId, retryCount);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
                 eq(requestorId), eq(peerMac), eq(ssi.getBytes()), eq(messageId));
@@ -1863,6 +1885,7 @@ public class WifiAwareStateManagerTest {
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
         ArgumentCaptor<Integer> messageIdCaptor = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdCaptor = ArgumentCaptor.forClass(Integer.class);
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
         mDut.enableUsage();
@@ -1892,7 +1915,7 @@ public class WifiAwareStateManagerTest {
         // (2) match
         mDut.onMatchNotification(subscribeId, requestorId, peerMac, null, null);
         mMockLooper.dispatchAll();
-        inOrder.verify(mockSessionCallback).onMatch(requestorId, null, null);
+        inOrder.verify(mockSessionCallback).onMatch(peerIdCaptor.capture(), isNull(), isNull());
 
         // (3) transmit messages
         SendMessageQueueModelAnswer answerObj = new SendMessageQueueModelAnswer(queueDepth,
@@ -1902,8 +1925,8 @@ public class WifiAwareStateManagerTest {
 
         int remainingMessages = numberOfMessages;
         for (int i = 0; i < numberOfMessages; ++i) {
-            mDut.sendMessage(clientId, sessionId.getValue(), requestorId, null, messageIdBase + i,
-                    0);
+            mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), null,
+                    messageIdBase + i, 0);
             mMockLooper.dispatchAll();
             // at 1/2 interval have the system simulate transmitting a queued message over-the-air
             if (i % 2 == 1) {
@@ -1959,7 +1982,7 @@ public class WifiAwareStateManagerTest {
                 IWifiAwareDiscoverySessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
-        ArgumentCaptor<Integer> messageIdCaptor = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdCaptor = ArgumentCaptor.forClass(Integer.class);
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
         mDut.enableUsage();
@@ -1989,7 +2012,7 @@ public class WifiAwareStateManagerTest {
         // (2) match
         mDut.onMatchNotification(subscribeId, requestorId, peerMac, null, null);
         mMockLooper.dispatchAll();
-        inOrder.verify(mockSessionCallback).onMatch(requestorId, null, null);
+        inOrder.verify(mockSessionCallback).onMatch(peerIdCaptor.capture(), isNull(), isNull());
 
         // (3) transmit messages: configure a mix of failures/success
         Set<Integer> failQueueCommandImmediately = new HashSet<>();
@@ -2042,8 +2065,8 @@ public class WifiAwareStateManagerTest {
                 any(), anyInt())).thenAnswer(answerObj);
 
         for (int i = 0; i < numberOfMessages; ++i) {
-            mDut.sendMessage(clientId, sessionId.getValue(), requestorId, null, messageIdBase + i,
-                    retransmitCount);
+            mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), null,
+                    messageIdBase + i, retransmitCount);
             mMockLooper.dispatchAll();
         }
 
@@ -2090,6 +2113,7 @@ public class WifiAwareStateManagerTest {
                 IWifiAwareDiscoverySessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdCaptor = ArgumentCaptor.forClass(Integer.class);
         ArgumentCaptor<byte[]> byteArrayCaptor = ArgumentCaptor.forClass(byte[].class);
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
@@ -2121,11 +2145,12 @@ public class WifiAwareStateManagerTest {
         mDut.onMatchNotification(subscribeId, requestorId, peerMac, peerSsi.getBytes(),
                 peerMatchFilter.getBytes());
         mMockLooper.dispatchAll();
-        inOrder.verify(mockSessionCallback).onMatch(requestorId, peerSsi.getBytes(),
-                peerMatchFilter.getBytes());
+        inOrder.verify(mockSessionCallback).onMatch(peerIdCaptor.capture(), eq(peerSsi.getBytes()),
+                eq(peerMatchFilter.getBytes()));
 
         // (3) message null Tx successful queuing
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, null, messageId, 0);
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), null, messageId,
+                0);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
                 eq(requestorId), eq(peerMac), isNull(byte[].class), eq(messageId));
@@ -2140,7 +2165,8 @@ public class WifiAwareStateManagerTest {
         validateInternalSendMessageQueuesCleanedUp(messageId);
 
         // (5) message byte[0] Tx successful queuing
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, new byte[0], messageId, 0);
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), new byte[0],
+                messageId, 0);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
                 eq(requestorId), eq(peerMac), eq(new byte[0]), eq(messageId));
@@ -2155,7 +2181,8 @@ public class WifiAwareStateManagerTest {
         validateInternalSendMessageQueuesCleanedUp(messageId);
 
         // (7) message "" Tx successful queuing
-        mDut.sendMessage(clientId, sessionId.getValue(), requestorId, "".getBytes(), messageId, 0);
+        mDut.sendMessage(clientId, sessionId.getValue(), peerIdCaptor.getValue(), "".getBytes(),
+                messageId, 0);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
                 eq(requestorId), eq(peerMac), byteArrayCaptor.capture(), eq(messageId));
@@ -2285,10 +2312,6 @@ public class WifiAwareStateManagerTest {
         final String peerMatchFilter = "filter binary array represented as string";
         final int rangingId = 18423;
         final RttManager.RttParams[] params = new RttManager.RttParams[2];
-        params[0] = new RttManager.RttParams();
-        params[0].bssid = Integer.toString(requestorId);
-        params[1] = new RttManager.RttParams();
-        params[1].bssid = Integer.toString(requestorId + 5);
 
         ConfigRequest configRequest = new ConfigRequest.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
@@ -2299,6 +2322,7 @@ public class WifiAwareStateManagerTest {
 
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        ArgumentCaptor<Integer> peerIdIdCaptor = ArgumentCaptor.forClass(Integer.class);
         ArgumentCaptor<WifiAwareClientState> clientCaptor =
                 ArgumentCaptor.forClass(WifiAwareClientState.class);
         ArgumentCaptor<RttManager.RttParams[]> rttParamsCaptor =
@@ -2332,15 +2356,20 @@ public class WifiAwareStateManagerTest {
                 peerMatchFilter.getBytes());
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
-        inOrder.verify(mockSessionCallback).onMatch(requestorId, peerSsi.getBytes(),
-                peerMatchFilter.getBytes());
+        inOrder.verify(mockSessionCallback).onMatch(peerIdIdCaptor.capture(),
+                eq(peerSsi.getBytes()), eq(peerMatchFilter.getBytes()));
 
         // (3) start ranging: pass along a valid peer ID and an invalid one
+        params[0] = new RttManager.RttParams();
+        params[0].bssid = Integer.toString(peerIdIdCaptor.getValue());
+        params[1] = new RttManager.RttParams();
+        params[1].bssid = Integer.toString(peerIdIdCaptor.getValue() + 5);
+
         mDut.startRanging(clientId, sessionId.getValue(), params, rangingId);
         mMockLooper.dispatchAll();
         inOrder.verify(mMockAwareRttStateManager).startRanging(eq(rangingId),
                 clientCaptor.capture(), rttParamsCaptor.capture());
-        collector.checkThat("RttParams[0].bssid", "06:07:08:09:0A:0B",
+        collector.checkThat("RttParams[0].bssid", "06:07:08:09:0a:0b",
                 equalTo(rttParamsCaptor.getValue()[0].bssid));
         collector.checkThat("RttParams[1].bssid", "", equalTo(rttParamsCaptor.getValue()[1].bssid));