[AWARE] Fix update publish/subscribe

When updating an existing publish/subscribe discovery session the
native API must take the session ID of the existing session. The code
provided 0 - which was an indication to create a new discovery session.

Bug: 38397363
Test: unit tests and integration (sl4a) tests
Change-Id: I8657493a39fe39d12a55af2bdee140b850351c1f

5 files changed, 128 insertions, 96 deletions

diff --git a/service/java/com/android/server/wifi/aware/WifiAwareDiscoverySessionState.java b/service/java/com/android/server/wifi/aware/WifiAwareDiscoverySessionState.java
index d006aa801..b4ccf2e8f 100644
--- a/service/java/com/android/server/wifi/aware/WifiAwareDiscoverySessionState.java
+++ b/service/java/com/android/server/wifi/aware/WifiAwareDiscoverySessionState.java
@@ -42,14 +42,14 @@ public class WifiAwareDiscoverySessionState {
 
     private final WifiAwareNativeApi mWifiAwareNativeApi;
     private int mSessionId;
-    private int mPubSubId;
+    private byte mPubSubId;
     private IWifiAwareDiscoverySessionCallback mCallback;
     private boolean mIsPublishSession;
 
     private final SparseArray<String> mMacByRequestorInstanceId = new SparseArray<>();
 
     public WifiAwareDiscoverySessionState(WifiAwareNativeApi wifiAwareNativeApi, int sessionId,
-            int pubSubId, IWifiAwareDiscoverySessionCallback callback, boolean isPublishSession) {
+            byte pubSubId, IWifiAwareDiscoverySessionCallback callback, boolean isPublishSession) {
         mWifiAwareNativeApi = wifiAwareNativeApi;
         mSessionId = sessionId;
         mPubSubId = pubSubId;
diff --git a/service/java/com/android/server/wifi/aware/WifiAwareNativeApi.java b/service/java/com/android/server/wifi/aware/WifiAwareNativeApi.java
index fb7a144e8..69f70a39f 100644
--- a/service/java/com/android/server/wifi/aware/WifiAwareNativeApi.java
+++ b/service/java/com/android/server/wifi/aware/WifiAwareNativeApi.java
@@ -365,9 +365,10 @@ public class WifiAwareNativeApi implements WifiAwareShellCommand.DelegatedShellC
      *            session.
      * @param publishConfig Configuration of the discovery session.
      */
-    public boolean publish(short transactionId, int publishId, PublishConfig publishConfig) {
+    public boolean publish(short transactionId, byte publishId, PublishConfig publishConfig) {
         if (VDBG) {
-            Log.d(TAG, "publish: transactionId=" + transactionId + ", config=" + publishConfig);
+            Log.d(TAG, "publish: transactionId=" + transactionId + ", publishId=" + publishId
+                    + ", config=" + publishConfig);
         }
 
         IWifiNanIface iface = mHal.getWifiNanIface();
@@ -377,7 +378,7 @@ public class WifiAwareNativeApi implements WifiAwareShellCommand.DelegatedShellC
         }
 
         NanPublishRequest req = new NanPublishRequest();
-        req.baseConfigs.sessionId = 0;
+        req.baseConfigs.sessionId = publishId;
         req.baseConfigs.ttlSec = (short) publishConfig.mTtlSec;
         req.baseConfigs.discoveryWindowPeriod = 1;
         req.baseConfigs.discoveryCount = 0;
@@ -426,10 +427,11 @@ public class WifiAwareNativeApi implements WifiAwareShellCommand.DelegatedShellC
      *            subscribe session.
      * @param subscribeConfig Configuration of the discovery session.
      */
-    public boolean subscribe(short transactionId, int subscribeId,
+    public boolean subscribe(short transactionId, byte subscribeId,
             SubscribeConfig subscribeConfig) {
         if (VDBG) {
-            Log.d(TAG, "subscribe: transactionId=" + transactionId + ", config=" + subscribeConfig);
+            Log.d(TAG, "subscribe: transactionId=" + transactionId + ", subscribeId=" + subscribeId
+                    + ", config=" + subscribeConfig);
         }
 
         IWifiNanIface iface = mHal.getWifiNanIface();
@@ -439,7 +441,7 @@ public class WifiAwareNativeApi implements WifiAwareShellCommand.DelegatedShellC
         }
 
         NanSubscribeRequest req = new NanSubscribeRequest();
-        req.baseConfigs.sessionId = 0;
+        req.baseConfigs.sessionId = subscribeId;
         req.baseConfigs.ttlSec = (short) subscribeConfig.mTtlSec;
         req.baseConfigs.discoveryWindowPeriod = 1;
         req.baseConfigs.discoveryCount = 0;
@@ -491,7 +493,7 @@ public class WifiAwareNativeApi implements WifiAwareShellCommand.DelegatedShellC
      * @param messageId Arbitary integer from host (not sent to HAL - useful for
      *                  testing/debugging at this level)
      */
-    public boolean sendMessage(short transactionId, int pubSubId, int requestorInstanceId,
+    public boolean sendMessage(short transactionId, byte pubSubId, int requestorInstanceId,
             byte[] dest, byte[] message, int messageId) {
         if (VDBG) {
             Log.d(TAG,
@@ -508,7 +510,7 @@ public class WifiAwareNativeApi implements WifiAwareShellCommand.DelegatedShellC
         }
 
         NanTransmitFollowupRequest req = new NanTransmitFollowupRequest();
-        req.discoverySessionId = (byte) pubSubId;
+        req.discoverySessionId = pubSubId;
         req.peerId = requestorInstanceId;
         copyArray(dest, req.addr);
         req.isHighPriority = false;
@@ -538,7 +540,7 @@ public class WifiAwareNativeApi implements WifiAwareShellCommand.DelegatedShellC
      * @param pubSubId ID of the publish/subscribe session - obtained when
      *            creating a session.
      */
-    public boolean stopPublish(short transactionId, int pubSubId) {
+    public boolean stopPublish(short transactionId, byte pubSubId) {
         if (VDBG) {
             Log.d(TAG, "stopPublish: transactionId=" + transactionId + ", pubSubId=" + pubSubId);
         }
@@ -550,7 +552,7 @@ public class WifiAwareNativeApi implements WifiAwareShellCommand.DelegatedShellC
         }
 
         try {
-            WifiStatus status = iface.stopPublishRequest(transactionId, (byte) pubSubId);
+            WifiStatus status = iface.stopPublishRequest(transactionId, pubSubId);
             if (status.code == WifiStatusCode.SUCCESS) {
                 return true;
             } else {
@@ -571,7 +573,7 @@ public class WifiAwareNativeApi implements WifiAwareShellCommand.DelegatedShellC
      * @param pubSubId ID of the publish/subscribe session - obtained when
      *            creating a session.
      */
-    public boolean stopSubscribe(short transactionId, int pubSubId) {
+    public boolean stopSubscribe(short transactionId, byte pubSubId) {
         if (VDBG) {
             Log.d(TAG, "stopSubscribe: transactionId=" + transactionId + ", pubSubId=" + pubSubId);
         }
@@ -583,7 +585,7 @@ public class WifiAwareNativeApi implements WifiAwareShellCommand.DelegatedShellC
         }
 
         try {
-            WifiStatus status = iface.stopSubscribeRequest(transactionId, (byte) pubSubId);
+            WifiStatus status = iface.stopSubscribeRequest(transactionId, pubSubId);
             if (status.code == WifiStatusCode.SUCCESS) {
                 return true;
             } else {
diff --git a/service/java/com/android/server/wifi/aware/WifiAwareStateManager.java b/service/java/com/android/server/wifi/aware/WifiAwareStateManager.java
index 9203b098b..1c4c07ffd 100644
--- a/service/java/com/android/server/wifi/aware/WifiAwareStateManager.java
+++ b/service/java/com/android/server/wifi/aware/WifiAwareStateManager.java
@@ -684,7 +684,7 @@ public class WifiAwareStateManager implements WifiAwareShellCommand.DelegatedShe
      * configuration (new or update) request succeeded.
      */
     public void onSessionConfigSuccessResponse(short transactionId, boolean isPublish,
-            int pubSubId) {
+            byte pubSubId) {
         Message msg = mSm.obtainMessage(MESSAGE_TYPE_RESPONSE);
         msg.arg1 = RESPONSE_TYPE_ON_SESSION_CONFIG_SUCCESS;
         msg.arg2 = transactionId;
@@ -794,7 +794,7 @@ public class WifiAwareStateManager implements WifiAwareShellCommand.DelegatedShe
 
     /**
      * Response from firmware to
-     * {@link #respondToDataPathRequest(boolean, int, String, byte[], String)}.
+     * {@link #respondToDataPathRequest(boolean, int, String, byte[], String, boolean)}
      */
     public void onRespondToDataPathSetupRequestResponse(short transactionId, boolean success,
             int reasonOnFailure) {
@@ -1582,7 +1582,7 @@ public class WifiAwareStateManager implements WifiAwareShellCommand.DelegatedShe
                     break;
                 }
                 case RESPONSE_TYPE_ON_SESSION_CONFIG_SUCCESS: {
-                    int pubSubId = (Integer) msg.obj;
+                    byte pubSubId = (Byte) msg.obj;
                     boolean isPublish = msg.getData().getBoolean(MESSAGE_BUNDLE_KEY_SESSION_TYPE);
 
                     onSessionConfigSuccessLocal(mCurrentCommand, pubSubId, isPublish);
@@ -2055,7 +2055,7 @@ public class WifiAwareStateManager implements WifiAwareShellCommand.DelegatedShe
             return false;
         }
 
-        boolean success = mWifiAwareNativeApi.publish(transactionId, 0, publishConfig);
+        boolean success = mWifiAwareNativeApi.publish(transactionId, (byte) 0, publishConfig);
         if (!success) {
             try {
                 callback.onSessionConfigFail(NanStatusType.INTERNAL_FAILURE);
@@ -2103,7 +2103,7 @@ public class WifiAwareStateManager implements WifiAwareShellCommand.DelegatedShe
             return false;
         }
 
-        boolean success = mWifiAwareNativeApi.subscribe(transactionId, 0, subscribeConfig);
+        boolean success = mWifiAwareNativeApi.subscribe(transactionId, (byte) 0, subscribeConfig);
         if (!success) {
             try {
                 callback.onSessionConfigFail(NanStatusType.INTERNAL_FAILURE);
@@ -2392,7 +2392,7 @@ public class WifiAwareStateManager implements WifiAwareShellCommand.DelegatedShe
         }
     }
 
-    private void onSessionConfigSuccessLocal(Message completedCommand, int pubSubId,
+    private void onSessionConfigSuccessLocal(Message completedCommand, byte pubSubId,
             boolean isPublish) {
         if (VDBG) {
             Log.v(TAG, "onSessionConfigSuccessLocal: completedCommand=" + completedCommand
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 f144da1e9..df54442ff 100644
--- a/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareDataPathStateManagerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareDataPathStateManagerTest.java
@@ -22,6 +22,7 @@ import static org.hamcrest.core.IsEqual.equalTo;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyByte;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyShort;
 import static org.mockito.ArgumentMatchers.eq;
@@ -136,13 +137,13 @@ public class WifiAwareDataPathStateManagerTest {
         when(mMockNative.enableAndConfigure(anyShort(), any(), anyBoolean(),
                 anyBoolean(), anyBoolean(), anyBoolean())).thenReturn(true);
         when(mMockNative.disable(anyShort())).thenReturn(true);
-        when(mMockNative.publish(anyShort(), anyInt(), any())).thenReturn(true);
-        when(mMockNative.subscribe(anyShort(), anyInt(), any()))
+        when(mMockNative.publish(anyShort(), anyByte(), any())).thenReturn(true);
+        when(mMockNative.subscribe(anyShort(), anyByte(), any()))
                 .thenReturn(true);
-        when(mMockNative.sendMessage(anyShort(), anyInt(), anyInt(), any(),
+        when(mMockNative.sendMessage(anyShort(), anyByte(), anyInt(), any(),
                 any(), anyInt())).thenReturn(true);
-        when(mMockNative.stopPublish(anyShort(), anyInt())).thenReturn(true);
-        when(mMockNative.stopSubscribe(anyShort(), anyInt())).thenReturn(true);
+        when(mMockNative.stopPublish(anyShort(), anyByte())).thenReturn(true);
+        when(mMockNative.stopSubscribe(anyShort(), anyByte())).thenReturn(true);
         when(mMockNative.createAwareNetworkInterface(anyShort(), any())).thenReturn(true);
         when(mMockNative.deleteAwareNetworkInterface(anyShort(), any())).thenReturn(true);
         when(mMockNative.initiateDataPath(anyShort(), anyInt(), anyInt(), anyInt(),
@@ -449,7 +450,7 @@ public class WifiAwareDataPathStateManagerTest {
 
     private void testDataPathInitiatorResponderMismatchUtility(boolean doPublish) throws Exception {
         final int clientId = 123;
-        final int pubSubId = 11234;
+        final byte pubSubId = 55;
         final int ndpId = 2;
         final byte[] pmk = "some bytes".getBytes();
         final PeerHandle peerHandle = new PeerHandle(1341234);
@@ -504,7 +505,7 @@ public class WifiAwareDataPathStateManagerTest {
     private void testDataPathInitiatorResponderInvalidUidUtility(boolean doPublish,
             boolean isUidSet) throws Exception {
         final int clientId = 123;
-        final int pubSubId = 11234;
+        final byte pubSubId = 56;
         final int ndpId = 2;
         final byte[] pmk = "some bytes".getBytes();
         final PeerHandle peerHandle = new PeerHandle(1341234);
@@ -561,7 +562,7 @@ public class WifiAwareDataPathStateManagerTest {
             boolean providePmk, boolean getConfirmation, boolean immediateHalFailure)
             throws Exception {
         final int clientId = 123;
-        final int pubSubId = 11234;
+        final byte pubSubId = 58;
         final PeerHandle peerHandle = new PeerHandle(1341234);
         final int ndpId = 2;
         final byte[] pmk = "some bytes".getBytes();
@@ -652,7 +653,7 @@ public class WifiAwareDataPathStateManagerTest {
     private void testDataPathResponderUtility(boolean useDirect, boolean provideMac,
             boolean providePmk, boolean getConfirmation) throws Exception {
         final int clientId = 123;
-        final int pubSubId = 11234;
+        final byte pubSubId = 60;
         final PeerHandle peerHandle = new PeerHandle(1341234);
         final int ndpId = 2;
         final byte[] pmk = "some bytes".getBytes();
@@ -851,7 +852,7 @@ public class WifiAwareDataPathStateManagerTest {
         return new NetworkRequest(nc, 0, 0, NetworkRequest.Type.REQUEST);
     }
 
-    private Pair<Integer, Messenger> initDataPathEndPoint(int clientId, int pubSubId,
+    private Pair<Integer, Messenger> initDataPathEndPoint(int clientId, byte pubSubId,
             PeerHandle peerHandle, byte[] peerDiscoveryMac, InOrder inOrder,
             boolean doPublish)
             throws Exception {
@@ -910,9 +911,10 @@ public class WifiAwareDataPathStateManagerTest {
         }
         mMockLooper.dispatchAll();
         if (doPublish) {
-            inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+            inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                    eq(publishConfig));
         } else {
-            inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0),
+            inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
                     eq(subscribeConfig));
         }
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), doPublish, pubSubId);
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 9bea93b0e..13d0e03b7 100644
--- a/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareStateManagerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/aware/WifiAwareStateManagerTest.java
@@ -23,6 +23,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyByte;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyShort;
 import static org.mockito.ArgumentMatchers.eq;
@@ -152,13 +153,13 @@ public class WifiAwareStateManagerTest {
         when(mMockNative.enableAndConfigure(anyShort(), any(), anyBoolean(),
                 anyBoolean(), anyBoolean(), anyBoolean())).thenReturn(true);
         when(mMockNative.disable(anyShort())).thenReturn(true);
-        when(mMockNative.publish(anyShort(), anyInt(), any())).thenReturn(true);
-        when(mMockNative.subscribe(anyShort(), anyInt(), any()))
+        when(mMockNative.publish(anyShort(), anyByte(), any())).thenReturn(true);
+        when(mMockNative.subscribe(anyShort(), anyByte(), any()))
                 .thenReturn(true);
-        when(mMockNative.sendMessage(anyShort(), anyInt(), anyInt(), any(),
+        when(mMockNative.sendMessage(anyShort(), anyByte(), anyInt(), any(),
                 any(), anyInt())).thenReturn(true);
-        when(mMockNative.stopPublish(anyShort(), anyInt())).thenReturn(true);
-        when(mMockNative.stopSubscribe(anyShort(), anyInt())).thenReturn(true);
+        when(mMockNative.stopPublish(anyShort(), anyByte())).thenReturn(true);
+        when(mMockNative.stopSubscribe(anyShort(), anyByte())).thenReturn(true);
         when(mMockNative.getCapabilities(anyShort())).thenReturn(true);
     }
 
@@ -517,7 +518,7 @@ public class WifiAwareStateManagerTest {
         // (2) publish + timeout
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(anyShort(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(anyShort(), eq((byte) 0), eq(publishConfig));
         assertTrue(mAlarmManager.dispatch(WifiAwareStateManager.HAL_COMMAND_TIMEOUT_TAG));
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionConfigFail(NanStatusType.INTERNAL_FAILURE);
@@ -526,8 +527,9 @@ public class WifiAwareStateManagerTest {
         // (3) publish + success
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
-        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, 9999);
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, (byte) 99);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(anyInt());
 
@@ -573,7 +575,8 @@ public class WifiAwareStateManagerTest {
         // (1) initial publish
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
 
         // (2) publish failure callback (i.e. firmware tried and failed)
         mDut.onSessionConfigFailResponse(transactionId.getValue(), true, reasonFail);
@@ -582,11 +585,12 @@ public class WifiAwareStateManagerTest {
         validateInternalNoSessions(clientId);
 
         // (3) publish and get immediate failure (i.e. HAL failed)
-        when(mMockNative.publish(anyShort(), anyInt(), any())).thenReturn(false);
+        when(mMockNative.publish(anyShort(), anyByte(), any())).thenReturn(false);
 
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
 
         inOrder.verify(mockSessionCallback).onSessionConfigFail(reasonFail);
         validateInternalNoSessions(clientId);
@@ -607,7 +611,7 @@ public class WifiAwareStateManagerTest {
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
         final int reasonTerminate = NanStatusType.SUCCESS;
-        final int publishId = 15;
+        final byte publishId = 15;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().build();
         PublishConfig publishConfig = new PublishConfig.Builder().build();
@@ -637,7 +641,8 @@ public class WifiAwareStateManagerTest {
         // (1) initial publish
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
 
         // (2) publish success
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
@@ -680,7 +685,7 @@ public class WifiAwareStateManagerTest {
         final int uid = 1000;
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
-        final int publishId = 15;
+        final byte publishId = 15;
         final int reasonFail = NanStatusType.INTERNAL_FAILURE;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().build();
@@ -711,7 +716,8 @@ public class WifiAwareStateManagerTest {
         // (1) initial publish
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
 
         // (2) publish success
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
@@ -741,7 +747,7 @@ public class WifiAwareStateManagerTest {
         inOrder.verify(mockSessionCallback).onSessionConfigSuccess();
 
         // (7) another update + immediate failure
-        when(mMockNative.publish(anyShort(), anyInt(), any())).thenReturn(false);
+        when(mMockNative.publish(anyShort(), anyByte(), any())).thenReturn(false);
 
         mDut.updatePublish(clientId, sessionId.getValue(), publishConfig);
         mMockLooper.dispatchAll();
@@ -764,7 +770,7 @@ public class WifiAwareStateManagerTest {
         final int uid = 1000;
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
-        final int publishId = 15;
+        final byte publishId = 15;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().build();
         PublishConfig publishConfig = new PublishConfig.Builder().build();
@@ -793,7 +799,8 @@ public class WifiAwareStateManagerTest {
         // (1) initial publish
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
 
         // (2) disconnect (but doesn't get executed until get response for
         // publish command)
@@ -851,7 +858,8 @@ public class WifiAwareStateManagerTest {
         // (1) initial subscribe
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
 
         // (2) subscribe failure
         mDut.onSessionConfigFailResponse(transactionId.getValue(), false, reasonFail);
@@ -860,12 +868,13 @@ public class WifiAwareStateManagerTest {
         validateInternalNoSessions(clientId);
 
         // (3) subscribe and get immediate failure (i.e. HAL failed)
-        when(mMockNative.subscribe(anyShort(), anyInt(), any()))
+        when(mMockNative.subscribe(anyShort(), anyByte(), any()))
                 .thenReturn(false);
 
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
 
         inOrder.verify(mockSessionCallback).onSessionConfigFail(reasonFail);
         validateInternalNoSessions(clientId);
@@ -886,7 +895,7 @@ public class WifiAwareStateManagerTest {
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
         final int reasonTerminate = NanStatusType.SUCCESS;
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
@@ -916,7 +925,8 @@ public class WifiAwareStateManagerTest {
         // (1) initial subscribe
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
 
         // (2) subscribe success
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
@@ -958,7 +968,7 @@ public class WifiAwareStateManagerTest {
         final int uid = 1000;
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
         final int reasonFail = NanStatusType.INTERNAL_FAILURE;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().build();
@@ -989,7 +999,8 @@ public class WifiAwareStateManagerTest {
         // (1) initial subscribe
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
 
         // (2) subscribe success
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
@@ -1019,7 +1030,7 @@ public class WifiAwareStateManagerTest {
         inOrder.verify(mockSessionCallback).onSessionConfigSuccess();
 
         // (7) another update + immediate failure
-        when(mMockNative.subscribe(anyShort(), anyInt(), any()))
+        when(mMockNative.subscribe(anyShort(), anyByte(), any()))
                 .thenReturn(false);
 
         mDut.updateSubscribe(clientId, sessionId.getValue(), subscribeConfig);
@@ -1043,7 +1054,7 @@ public class WifiAwareStateManagerTest {
         final int uid = 1000;
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
@@ -1072,7 +1083,8 @@ public class WifiAwareStateManagerTest {
         // (1) initial subscribe
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
 
         // (2) disconnect (but doesn't get executed until get response for
         // subscribe command)
@@ -1104,7 +1116,7 @@ public class WifiAwareStateManagerTest {
         final String serviceName = "some-service-name";
         final String ssi = "some much longer and more arbitrary data";
         final int reasonFail = NanStatusType.INTERNAL_FAILURE;
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
         final int requestorId = 22;
         final byte[] peerMac = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final String peerSsi = "some peer ssi data";
@@ -1144,7 +1156,8 @@ public class WifiAwareStateManagerTest {
         // (1) subscribe
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
@@ -1206,7 +1219,7 @@ public class WifiAwareStateManagerTest {
         final int clusterHigh = 7;
         final int masterPref = 0;
         final String serviceName = "some-service-name";
-        final int publishId = 88;
+        final byte publishId = 88;
         final int peerId1 = 568;
         final int peerId2 = 873;
         final byte[] peerMac1 = HexEncoding.decode("000102030405".toCharArray(), false);
@@ -1250,7 +1263,8 @@ public class WifiAwareStateManagerTest {
         // (2) publish
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
@@ -1303,7 +1317,7 @@ public class WifiAwareStateManagerTest {
         final int clusterHigh = 7;
         final int masterPref = 0;
         final String serviceName = "some-service-name";
-        final int publishId = 88;
+        final byte publishId = 88;
         final int peerId = 568;
         final byte[] peerMacOrig = HexEncoding.decode("000102030405".toCharArray(), false);
         final byte[] peerMacLater = HexEncoding.decode("060708090A0B".toCharArray(), false);
@@ -1344,7 +1358,8 @@ public class WifiAwareStateManagerTest {
         // (2) publish
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
@@ -1392,7 +1407,7 @@ public class WifiAwareStateManagerTest {
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
         final String ssi = "some much longer and more arbitrary data";
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
         final int requestorId = 22;
         final byte[] peerMac = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final String peerSsi = "some peer ssi data";
@@ -1427,7 +1442,8 @@ public class WifiAwareStateManagerTest {
         // (2) subscribe & match
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mDut.onMatchNotification(subscribeId, requestorId, peerMac, peerSsi.getBytes(),
                 peerMatchFilter.getBytes());
@@ -1459,7 +1475,7 @@ public class WifiAwareStateManagerTest {
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
         final String ssi = "some much longer and more arbitrary data";
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
         final int requestorId = 22;
         final byte[] peerMac = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final String peerSsi = "some peer ssi data";
@@ -1494,7 +1510,8 @@ public class WifiAwareStateManagerTest {
         // (2) subscribe & match
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mDut.onMatchNotification(subscribeId, requestorId, peerMac, peerSsi.getBytes(),
                 peerMatchFilter.getBytes());
@@ -1536,7 +1553,7 @@ public class WifiAwareStateManagerTest {
         validateInternalSendMessageQueuesCleanedUp(messageId + 2);
 
         // (5) send a message and get an immediate failure (configure first)
-        when(mMockNative.sendMessage(anyShort(), anyInt(), anyInt(), any(),
+        when(mMockNative.sendMessage(anyShort(), anyByte(), anyInt(), any(),
                 any(), anyInt())).thenReturn(false);
 
         mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(), messageId + 3,
@@ -1582,7 +1599,7 @@ public class WifiAwareStateManagerTest {
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
         final String ssi = "some much longer and more arbitrary data";
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
         final int requestorId = 22;
         final byte[] peerMac = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final String peerSsi = "some peer ssi data";
@@ -1618,7 +1635,8 @@ public class WifiAwareStateManagerTest {
         // (2) subscribe & match
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mDut.onMatchNotification(subscribeId, requestorId, peerMac, peerSsi.getBytes(),
                 peerMatchFilter.getBytes());
@@ -1667,7 +1685,7 @@ public class WifiAwareStateManagerTest {
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
         final String ssi = "some much longer and more arbitrary data";
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
         final int requestorId = 22;
         final byte[] peerMac = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final String peerSsi = "some peer ssi data";
@@ -1703,7 +1721,8 @@ public class WifiAwareStateManagerTest {
         // (2) subscribe & match
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mDut.onMatchNotification(subscribeId, requestorId, peerMac, peerSsi.getBytes(),
                 peerMatchFilter.getBytes());
@@ -1752,7 +1771,7 @@ public class WifiAwareStateManagerTest {
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
         final String serviceName = "some-service-name";
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
         final int requestorId = 22;
         final byte[] peerMac = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final int messageIdBase = 6948;
@@ -1789,7 +1808,8 @@ public class WifiAwareStateManagerTest {
         // (1) subscribe
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
@@ -1802,7 +1822,7 @@ public class WifiAwareStateManagerTest {
         // (3) transmit messages
         SendMessageQueueModelAnswer answerObj = new SendMessageQueueModelAnswer(queueDepth,
                 null, null, null);
-        when(mMockNative.sendMessage(anyShort(), anyInt(), anyInt(), any(),
+        when(mMockNative.sendMessage(anyShort(), anyByte(), anyInt(), any(),
                 any(), anyInt())).thenAnswer(answerObj);
 
         int remainingMessages = numberOfMessages;
@@ -1847,7 +1867,7 @@ public class WifiAwareStateManagerTest {
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
         final String serviceName = "some-service-name";
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
         final int requestorId = 22;
         final byte[] peerMac = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final int messageIdBase = 6948;
@@ -1885,7 +1905,8 @@ public class WifiAwareStateManagerTest {
         // (1) subscribe
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
@@ -1942,7 +1963,7 @@ public class WifiAwareStateManagerTest {
 
         SendMessageQueueModelAnswer answerObj = new SendMessageQueueModelAnswer(queueDepth,
                 failQueueCommandImmediately, failQueueCommandLater, numberOfRetries);
-        when(mMockNative.sendMessage(anyShort(), anyInt(), anyInt(), any(),
+        when(mMockNative.sendMessage(anyShort(), anyByte(), anyInt(), any(),
                 any(), anyInt())).thenAnswer(answerObj);
 
         for (int i = 0; i < numberOfMessages; ++i) {
@@ -1976,7 +1997,7 @@ public class WifiAwareStateManagerTest {
         final String callingPackage = "com.google.somePackage";
         final String serviceName = "some-service-name";
         final String ssi = "some much longer and more arbitrary data";
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
         final int requestorId = 22;
         final byte[] peerMac = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final String peerSsi = "some peer ssi data";
@@ -2015,7 +2036,8 @@ public class WifiAwareStateManagerTest {
         // (1) subscribe
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
@@ -2107,7 +2129,7 @@ public class WifiAwareStateManagerTest {
             }
         }
 
-        public boolean answer(short transactionId, int pubSubId, int requestorInstanceId,
+        public boolean answer(short transactionId, byte pubSubId, int requestorInstanceId,
                 byte[] dest, byte[] message, int messageId) throws Exception {
             if (mFailQueueCommandImmediately != null && mFailQueueCommandImmediately.contains(
                     messageId)) {
@@ -2181,7 +2203,7 @@ public class WifiAwareStateManagerTest {
         final int uid = 1000;
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
-        final int subscribeId = 15;
+        final byte subscribeId = 15;
         final int requestorId = 22;
         final byte[] peerMac = HexEncoding.decode("060708090A0B".toCharArray(), false);
         final String peerSsi = "some peer ssi data";
@@ -2228,7 +2250,8 @@ public class WifiAwareStateManagerTest {
         // (2) subscribe & match
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mDut.onMatchNotification(subscribeId, requestorId, peerMac, peerSsi.getBytes(),
                 peerMatchFilter.getBytes());
@@ -2378,7 +2401,7 @@ public class WifiAwareStateManagerTest {
         final int masterPref = 111;
         final String serviceName = "some-service-name";
         final String ssi = "some much longer and more arbitrary data";
-        final int publishId = 22;
+        final byte publishId = 22;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().setClusterLow(clusterLow)
                 .setClusterHigh(clusterHigh).setMasterPreference(masterPref).build();
@@ -2411,7 +2434,8 @@ public class WifiAwareStateManagerTest {
         // (2) publish (no response yet)
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
 
         // (3) disconnect (but doesn't get executed until get a RESPONSE to the
         // previous publish)
@@ -2485,7 +2509,8 @@ public class WifiAwareStateManagerTest {
         // (2) publish - no response
         mDut.publish(clientId, publishConfig, mockPublishSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
 
         verifyNoMoreInteractions(mMockNative, mockCallback, mockPublishSessionCallback);
     }
@@ -2532,7 +2557,7 @@ public class WifiAwareStateManagerTest {
      */
     @Test
     public void testInvalidCallbackIdParameters() throws Exception {
-        final int pubSubId = 1235;
+        final byte pubSubId = 125;
         final int clientId = 132;
         final int uid = 1000;
         final int pid = 2000;
@@ -2587,7 +2612,7 @@ public class WifiAwareStateManagerTest {
         final int uid = 1000;
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
-        final int publishId = 25;
+        final byte publishId = 25;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().build();
         PublishConfig publishConfig = new PublishConfig.Builder().build();
@@ -2618,7 +2643,8 @@ public class WifiAwareStateManagerTest {
         // (2) publish
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                eq(publishConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
@@ -2641,7 +2667,7 @@ public class WifiAwareStateManagerTest {
         final int uid = 1000;
         final int pid = 2000;
         final String callingPackage = "com.google.somePackage";
-        final int subscribeId = 25;
+        final byte subscribeId = 25;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().build();
         PublishConfig publishConfig = new PublishConfig.Builder().build();
@@ -2672,7 +2698,8 @@ public class WifiAwareStateManagerTest {
         // (2) subscribe
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq((byte) 0),
+                eq(subscribeConfig));
         mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mMockLooper.dispatchAll();
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
@@ -2726,10 +2753,11 @@ public class WifiAwareStateManagerTest {
             // (2) publish
             mDut.publish(clientId, publishConfig, mockSessionCallback);
             mMockLooper.dispatchAll();
-            inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+            inOrder.verify(mMockNative).publish(transactionId.capture(), eq((byte) 0),
+                    eq(publishConfig));
 
             // (3) publish-success
-            mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, i + 1);
+            mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, (byte) (i + 1));
             mMockLooper.dispatchAll();
             inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());