[NAN] Re-factor connect/config flow
am: 9864521

* commit '9864521b92325bad1d20510d99d4e967b6f3d4eb':
  [NAN] Re-factor connect/config flow

Change-Id: I492fee202b76e033579be8805bf34fdbb0273c07
diff --git a/service/java/com/android/server/wifi/nan/WifiNanClientState.java b/service/java/com/android/server/wifi/nan/WifiNanClientState.java
index c2a5ed2..9e70ef8 100644
--- a/service/java/com/android/server/wifi/nan/WifiNanClientState.java
+++ b/service/java/com/android/server/wifi/nan/WifiNanClientState.java
@@ -18,7 +18,6 @@
 
 import android.net.wifi.nan.ConfigRequest;
 import android.net.wifi.nan.IWifiNanEventCallback;
-import android.net.wifi.nan.IWifiNanSessionCallback;
 import android.os.RemoteException;
 import android.util.Log;
 import android.util.SparseArray;
@@ -49,9 +48,11 @@
     private int mClientId;
     private ConfigRequest mConfigRequest;
 
-    public WifiNanClientState(int clientId, IWifiNanEventCallback callback) {
+    public WifiNanClientState(int clientId, IWifiNanEventCallback callback,
+            ConfigRequest configRequest) {
         mClientId = clientId;
         mCallback = callback;
+        mConfigRequest = configRequest;
     }
 
     /**
@@ -59,7 +60,6 @@
      * the client. Destroys all discovery sessions belonging to this client.
      */
     public void destroy() {
-        mCallback = null;
         for (int i = 0; i < mSessions.size(); ++i) {
             mSessions.valueAt(i).terminate();
         }
@@ -67,10 +67,6 @@
         mConfigRequest = null;
     }
 
-    public void setConfigRequest(ConfigRequest configRequest) {
-        mConfigRequest = configRequest;
-    }
-
     public ConfigRequest getConfigRequest() {
         return mConfigRequest;
     }
@@ -79,6 +75,10 @@
         return mClientId;
     }
 
+    public IWifiNanEventCallback getCallback() {
+        return mCallback;
+    }
+
     /**
      * Searches the discovery sessions of this client and returns the one
      * corresponding to the publish/subscribe ID. Used on callbacks from HAL to
@@ -99,19 +99,33 @@
     }
 
     /**
-     * Create a new discovery session.
+     * Add the session to the client database.
      *
-     * @param sessionId Session ID of the new discovery session.
-     * @param callback Singleton session callback.
+     * @param session Session to be added.
      */
-    public void createSession(int sessionId, IWifiNanSessionCallback callback,
-            boolean isPublishSession) {
-        WifiNanSessionState session = mSessions.get(sessionId);
-        if (session != null) {
-            Log.e(TAG, "createSession: sessionId already exists (replaced) - " + sessionId);
+    public void addSession(WifiNanSessionState session) {
+        int sessionId = session.getSessionId();
+        if (mSessions.get(sessionId) != null) {
+            Log.w(TAG, "createSession: sessionId already exists (replaced) - " + sessionId);
         }
 
-        mSessions.put(sessionId, new WifiNanSessionState(sessionId, callback, isPublishSession));
+        mSessions.put(sessionId, session);
+    }
+
+    /**
+     * Remove the specified session from the client database - without doing a
+     * terminate on the session. The assumption is that it is already
+     * terminated.
+     *
+     * @param sessionId The session ID of the session to be removed.
+     */
+    public void removeSession(int sessionId) {
+        if (mSessions.get(sessionId) == null) {
+            Log.e(TAG, "removeSession: sessionId doesn't exist - " + sessionId);
+            return;
+        }
+
+        mSessions.delete(sessionId);
     }
 
     /**
@@ -127,8 +141,8 @@
             return;
         }
 
-        mSessions.delete(sessionId);
         session.terminate();
+        mSessions.delete(sessionId);
     }
 
     /**
@@ -143,71 +157,16 @@
     }
 
     /**
-     * Called to dispatch the configuration completed event to the client.
-     * Dispatched if the client registered for this event.
-     *
-     * @param completedConfig The configuration which was completed.
-     */
-    public void onConfigCompleted(ConfigRequest completedConfig) {
-        if (mCallback != null) {
-            try {
-                mCallback.onConfigCompleted(completedConfig);
-            } catch (RemoteException e) {
-                Log.w(TAG, "onConfigCompleted: RemoteException - ignored: " + e);
-            }
-        }
-    }
-
-    /**
-     * Called to dispatch the configuration failed event to the client.
-     * Dispatched if the client registered for this event.
-     *
-     * @param failedConfig The configuration which failed.
-     * @param reason The failure reason.
-     */
-    public void onConfigFailed(ConfigRequest failedConfig, int reason) {
-        if (mCallback != null) {
-            try {
-                mCallback.onConfigFailed(failedConfig, reason);
-            } catch (RemoteException e) {
-                Log.w(TAG, "onConfigFailed: RemoteException - ignored: " + e);
-            }
-        }
-    }
-
-    /**
-     * Called to dispatch the NAN down event to the client. Dispatched if the
-     * client registered for this event.
-     *
-     * @param reason The reason code for NAN going down.
-     * @return A 1 if registered to listen for event, 0 otherwise.
-     */
-    public int onNanDown(int reason) {
-        if (mCallback != null) {
-            try {
-                mCallback.onNanDown(reason);
-            } catch (RemoteException e) {
-                Log.w(TAG, "onNanDown: RemoteException - ignored: " + e);
-            }
-
-            return 1;
-        }
-
-        return 0;
-    }
-
-    /**
      * Called to dispatch the NAN interface address change to the client - as an
      * identity change (interface address information not propagated to client -
-     * privacy concerns). Dispatched if the client registered for the identity
-     * changed event.
+     * privacy concerns).
      *
      * @param mac The new MAC address of the discovery interface - not
      *            propagated to client!
      * @return A 1 if registered to listen for event, 0 otherwise.
      */
     public int onInterfaceAddressChange(byte[] mac) {
-        if (mCallback != null && mConfigRequest.mEnableIdentityChangeCallback) {
+        if (mConfigRequest.mEnableIdentityChangeCallback) {
             try {
                 mCallback.onIdentityChanged();
             } catch (RemoteException e) {
@@ -232,7 +191,7 @@
      * @return A 1 if registered to listen for event, 0 otherwise.
      */
     public int onClusterChange(int flag, byte[] mac) {
-        if (mCallback != null && mConfigRequest.mEnableIdentityChangeCallback) {
+        if (mConfigRequest.mEnableIdentityChangeCallback) {
             try {
                 mCallback.onIdentityChanged();
             } catch (RemoteException e) {
diff --git a/service/java/com/android/server/wifi/nan/WifiNanNative.java b/service/java/com/android/server/wifi/nan/WifiNanNative.java
index f166e6c..1c775d0 100644
--- a/service/java/com/android/server/wifi/nan/WifiNanNative.java
+++ b/service/java/com/android/server/wifi/nan/WifiNanNative.java
@@ -19,6 +19,7 @@
 import android.net.wifi.nan.ConfigRequest;
 import android.net.wifi.nan.PublishConfig;
 import android.net.wifi.nan.SubscribeConfig;
+import android.net.wifi.nan.WifiNanEventCallback;
 import android.net.wifi.nan.WifiNanSessionCallback;
 import android.util.Log;
 
@@ -99,9 +100,10 @@
         }
     }
 
-    /* package */ static native int initNanHandlersNative(Object cls, int iface);
+    /* package */ static native int initNanHandlersNative(Class<WifiNative> cls, int iface);
 
-    private static native int getCapabilitiesNative(short transactionId, Object cls, int iface);
+    private static native int getCapabilitiesNative(short transactionId, Class<WifiNative> cls,
+            int iface);
 
     private boolean isNanInit(boolean tryToInit) {
         if (!tryToInit || sNanNativeInit) {
@@ -120,9 +122,7 @@
                 sNanNativeInit = ret == WIFI_SUCCESS;
 
                 if (sNanNativeInit) {
-                    ret = getCapabilitiesNative(
-                            WifiNanStateManager.getInstance().createNextTransactionId(),
-                            WifiNative.class,
+                    ret = getCapabilitiesNative((short) 0, WifiNative.class,
                             WifiNative.sWlan0Index);
                     if (DBG) Log.d(TAG, "getCapabilitiesNative: res=" + ret);
                 }
@@ -139,8 +139,11 @@
         // do nothing
     }
 
-    private static native int enableAndConfigureNative(short transactionId, Object cls, int iface,
-            ConfigRequest configRequest);
+    private static native int enableAndConfigureNative(short transactionId, Class<WifiNative> cls,
+            int iface, ConfigRequest configRequest);
+
+    private static native int updateConfigurationNative(short transactionId, Class<WifiNative> cls,
+            int iface, ConfigRequest configRequest);
 
     /**
      * Enable and configure NAN.
@@ -148,18 +151,29 @@
      * @param transactionId Transaction ID for the transaction - used in the
      *            async callback to match with the original request.
      * @param configRequest Requested NAN configuration.
+     * @param initialConfiguration Specifies whether initial configuration
+     *            (true) or an update (false) to the configuration.
      */
-    public void enableAndConfigure(short transactionId, ConfigRequest configRequest) {
+    public void enableAndConfigure(short transactionId, ConfigRequest configRequest,
+            boolean initialConfiguration) {
         boolean success;
 
         if (VDBG) Log.d(TAG, "enableAndConfigure: configRequest=" + configRequest);
         if (isNanInit(true)) {
             int ret;
-            synchronized (WifiNative.sLock) {
-                ret = enableAndConfigureNative(transactionId, WifiNative.class,
-                        WifiNative.sWlan0Index, configRequest);
+            if (initialConfiguration) {
+                synchronized (WifiNative.sLock) {
+                    ret = enableAndConfigureNative(transactionId, WifiNative.class,
+                            WifiNative.sWlan0Index, configRequest);
+                }
+                if (DBG) Log.d(TAG, "enableAndConfigureNative: ret=" + ret);
+            } else {
+                synchronized (WifiNative.sLock) {
+                    ret = updateConfigurationNative(transactionId, WifiNative.class,
+                            WifiNative.sWlan0Index, configRequest);
+                }
+                if (DBG) Log.d(TAG, "updateConfigurationNative: ret=" + ret);
             }
-            if (DBG) Log.d(TAG, "enableAndConfigureNative: ret=" + ret);
             success = ret == WIFI_SUCCESS;
         } else {
             Log.w(TAG, "enableAndConfigure: NanInit fails");
@@ -170,7 +184,7 @@
         // TODO: do something on !success - send failure message back
     }
 
-    private static native int disableNative(short transactionId, Object cls, int iface);
+    private static native int disableNative(short transactionId, Class<WifiNative> cls, int iface);
 
     /**
      * Disable NAN.
@@ -197,8 +211,8 @@
         // TODO: do something on !success - send failure message back
     }
 
-    private static native int publishNative(short transactionId, int publishId, Object cls,
-            int iface, PublishConfig publishConfig);
+    private static native int publishNative(short transactionId, int publishId,
+            Class<WifiNative> cls, int iface, PublishConfig publishConfig);
 
     /**
      * Start or modify a service publish session.
@@ -232,8 +246,8 @@
         // TODO: do something on !success - send failure message back
     }
 
-    private static native int subscribeNative(short transactionId, int subscribeId, Object cls,
-            int iface, SubscribeConfig subscribeConfig);
+    private static native int subscribeNative(short transactionId, int subscribeId,
+            Class<WifiNative> cls, int iface, SubscribeConfig subscribeConfig);
 
     /**
      * Start or modify a service subscription session.
@@ -267,8 +281,9 @@
         // TODO: do something on !success - send failure message back
     }
 
-    private static native int sendMessageNative(short transactionId, Object cls, int iface,
-            int pubSubId, int requestorInstanceId, byte[] dest, byte[] message, int messageLength);
+    private static native int sendMessageNative(short transactionId, Class<WifiNative> cls,
+            int iface, int pubSubId, int requestorInstanceId, byte[] dest, byte[] message,
+            int messageLength);
 
     /**
      * Send a message through an existing discovery session.
@@ -311,8 +326,8 @@
         // TODO: do something on !success - send failure message back
     }
 
-    private static native int stopPublishNative(short transactionId, Object cls, int iface,
-            int pubSubId);
+    private static native int stopPublishNative(short transactionId, Class<WifiNative> cls,
+            int iface, int pubSubId);
 
     /**
      * Terminate a publish discovery session.
@@ -345,8 +360,8 @@
         // TODO: do something on !success - send failure message back
     }
 
-    private static native int stopSubscribeNative(short transactionId, Object cls, int iface,
-            int pubSubId);
+    private static native int stopSubscribeNative(short transactionId, Class<WifiNative> cls,
+            int iface, int pubSubId);
 
     /**
      * Terminate a subscribe discovery session.
@@ -388,9 +403,11 @@
     public static final int NAN_RESPONSE_TRANSMIT_FOLLOWUP = 4;
     public static final int NAN_RESPONSE_SUBSCRIBE = 5;
     public static final int NAN_RESPONSE_SUBSCRIBE_CANCEL = 6;
+    public static final int NAN_RESPONSE_CONFIG = 8;
     public static final int NAN_RESPONSE_GET_CAPABILITIES = 12;
 
     // direct copy from wifi_nan.h: need to keep in sync
+    /* NAN Protocol Response Codes */
     public static final int NAN_STATUS_SUCCESS = 0;
     public static final int NAN_STATUS_TIMEOUT = 1;
     public static final int NAN_STATUS_DE_FAILURE = 2;
@@ -411,8 +428,13 @@
     public static final int NAN_STATUS_INVALID_TLV_VALUE = 17;
     public static final int NAN_STATUS_INVALID_TX_PRIORITY = 18;
     public static final int NAN_STATUS_INVALID_CONNECTION_MAP = 19;
+    public static final int NAN_STATUS_INVALID_TCA_ID = 20;
+    public static final int NAN_STATUS_INVALID_STATS_ID = 21;
+    public static final int NAN_STATUS_NAN_NOT_ALLOWED = 22;
+    public static final int NAN_STATUS_NO_OTA_ACK = 23;
+    public static final int NAN_STATUS_TX_FAIL = 24;
 
-    // NAN Configuration Response codes
+    /* NAN Configuration Response codes */
     public static final int NAN_STATUS_INVALID_RSSI_CLOSE_VALUE = 4096;
     public static final int NAN_STATUS_INVALID_RSSI_MIDDLE_VALUE = 4097;
     public static final int NAN_STATUS_INVALID_HOP_COUNT_LIMIT = 4098;
@@ -435,8 +457,12 @@
     public static final int NAN_STATUS_INVALID_POST_NAN_DISCOVERY_BITMAP_VALUE = 4115;
     public static final int NAN_STATUS_MISSING_FUTHER_AVAILABILITY_MAP = 4116;
     public static final int NAN_STATUS_INVALID_BAND_CONFIG_FLAGS = 4117;
+    public static final int NAN_STATUS_INVALID_RANDOM_FACTOR_UPDATE_TIME_VALUE = 4118;
+    public static final int NAN_STATUS_INVALID_ONGOING_SCAN_PERIOD = 4119;
+    public static final int NAN_STATUS_INVALID_DW_INTERVAL_VALUE = 4120;
+    public static final int NAN_STATUS_INVALID_DB_INTERVAL_VALUE = 4121;
 
-    // publish/subscribe termination reasons
+    /* publish/subscribe termination reasons */
     public static final int NAN_TERMINATED_REASON_INVALID = 8192;
     public static final int NAN_TERMINATED_REASON_TIMEOUT = 8193;
     public static final int NAN_TERMINATED_REASON_USER_REQUEST = 8194;
@@ -448,31 +474,10 @@
     public static final int NAN_TERMINATED_REASON_POST_DISC_LEN_EXCEEDED = 8200;
     public static final int NAN_TERMINATED_REASON_FURTHER_AVAIL_MAP_EMPTY = 8201;
 
-    private static int translateHalStatusToPublicStatus(int halStatus) {
+    private static int translateHalStatusToNanEventCallbackReason(int halStatus) {
         switch (halStatus) {
-            case NAN_STATUS_NO_SPACE_AVAILABLE:
-                return WifiNanSessionCallback.FAIL_REASON_NO_RESOURCES;
-
-            case NAN_STATUS_TIMEOUT:
-            case NAN_STATUS_DE_FAILURE:
-            case NAN_STATUS_DISABLE_IN_PROGRESS:
-                return WifiNanSessionCallback.FAIL_REASON_OTHER;
-
-            case NAN_STATUS_INVALID_MSG_VERSION:
-            case NAN_STATUS_INVALID_MSG_LEN:
-            case NAN_STATUS_INVALID_MSG_ID:
-            case NAN_STATUS_INVALID_HANDLE:
-            case NAN_STATUS_INVALID_PUBLISH_TYPE:
-            case NAN_STATUS_INVALID_TX_TYPE:
-            case NAN_STATUS_INVALID_MATCH_ALGORITHM:
-            case NAN_STATUS_INVALID_TLV_LEN:
-            case NAN_STATUS_INVALID_TLV_TYPE:
-            case NAN_STATUS_MISSING_TLV_TYPE:
-            case NAN_STATUS_INVALID_TOTAL_TLVS_LEN:
-            case NAN_STATUS_INVALID_MATCH_HANDLE:
-            case NAN_STATUS_INVALID_TLV_VALUE:
-            case NAN_STATUS_INVALID_TX_PRIORITY:
-            case NAN_STATUS_INVALID_CONNECTION_MAP:
+            case NAN_STATUS_SUCCESS:
+                return WifiNanEventCallback.REASON_REQUESTED;
             case NAN_STATUS_INVALID_RSSI_CLOSE_VALUE:
             case NAN_STATUS_INVALID_RSSI_MIDDLE_VALUE:
             case NAN_STATUS_INVALID_HOP_COUNT_LIMIT:
@@ -495,9 +500,18 @@
             case NAN_STATUS_INVALID_POST_NAN_DISCOVERY_BITMAP_VALUE:
             case NAN_STATUS_MISSING_FUTHER_AVAILABILITY_MAP:
             case NAN_STATUS_INVALID_BAND_CONFIG_FLAGS:
-                return WifiNanSessionCallback.FAIL_REASON_INVALID_ARGS;
+            case NAN_STATUS_INVALID_RANDOM_FACTOR_UPDATE_TIME_VALUE:
+            case NAN_STATUS_INVALID_ONGOING_SCAN_PERIOD:
+            case NAN_STATUS_INVALID_DW_INTERVAL_VALUE:
+            case NAN_STATUS_INVALID_DB_INTERVAL_VALUE:
+                return WifiNanEventCallback.REASON_INVALID_ARGS;
+        }
 
-                // publish/subscribe termination reasons
+        return WifiNanEventCallback.REASON_OTHER;
+    }
+
+    private static int translateHalStatusToNanSessionCallbackTerminate(int halStatus) {
+        switch (halStatus) {
             case NAN_TERMINATED_REASON_TIMEOUT:
             case NAN_TERMINATED_REASON_USER_REQUEST:
             case NAN_TERMINATED_REASON_COUNT_REACHED:
@@ -513,7 +527,43 @@
                 return WifiNanSessionCallback.TERMINATE_REASON_FAIL;
         }
 
-        return WifiNanSessionCallback.FAIL_REASON_OTHER;
+        return WifiNanSessionCallback.TERMINATE_REASON_FAIL;
+    }
+
+    private static int translateHalStatusToNanSessionCallbackReason(int halStatus) {
+        switch (halStatus) {
+            case NAN_STATUS_TIMEOUT:
+            case NAN_STATUS_DE_FAILURE:
+            case NAN_STATUS_INVALID_MSG_VERSION:
+            case NAN_STATUS_INVALID_MSG_LEN:
+            case NAN_STATUS_INVALID_MSG_ID:
+            case NAN_STATUS_INVALID_HANDLE:
+                return WifiNanSessionCallback.REASON_OTHER;
+            case NAN_STATUS_NO_SPACE_AVAILABLE:
+                return WifiNanSessionCallback.REASON_NO_RESOURCES;
+            case NAN_STATUS_INVALID_PUBLISH_TYPE:
+            case NAN_STATUS_INVALID_TX_TYPE:
+            case NAN_STATUS_INVALID_MATCH_ALGORITHM:
+                return WifiNanSessionCallback.REASON_INVALID_ARGS;
+            case NAN_STATUS_DISABLE_IN_PROGRESS:
+            case NAN_STATUS_INVALID_TLV_LEN:
+            case NAN_STATUS_INVALID_TLV_TYPE:
+            case NAN_STATUS_MISSING_TLV_TYPE:
+            case NAN_STATUS_INVALID_TOTAL_TLVS_LEN:
+            case NAN_STATUS_INVALID_MATCH_HANDLE:
+            case NAN_STATUS_INVALID_TLV_VALUE:
+            case NAN_STATUS_INVALID_TX_PRIORITY:
+            case NAN_STATUS_INVALID_CONNECTION_MAP:
+            case NAN_STATUS_INVALID_TCA_ID:
+            case NAN_STATUS_INVALID_STATS_ID:
+            case NAN_STATUS_NAN_NOT_ALLOWED:
+                return WifiNanSessionCallback.REASON_OTHER;
+            case NAN_STATUS_NO_OTA_ACK:
+            case NAN_STATUS_TX_FAIL:
+                return WifiNanSessionCallback.REASON_TX_FAIL;
+        }
+
+        return WifiNanSessionCallback.REASON_OTHER;
     }
 
     // callback from native
@@ -528,11 +578,13 @@
 
         switch (responseType) {
             case NAN_RESPONSE_ENABLED:
+                /* fall through */
+            case NAN_RESPONSE_CONFIG:
                 if (status == NAN_STATUS_SUCCESS) {
-                    stateMgr.onConfigCompleted(transactionId);
+                    stateMgr.onConfigSuccessResponse(transactionId);
                 } else {
-                    stateMgr.onConfigFailed(transactionId,
-                            translateHalStatusToPublicStatus(status));
+                    stateMgr.onConfigFailedResponse(transactionId,
+                            translateHalStatusToNanEventCallbackReason(status));
                 }
                 break;
             case NAN_RESPONSE_PUBLISH_CANCEL:
@@ -540,14 +592,13 @@
                     Log.e(TAG, "onNanNotifyResponse: NAN_RESPONSE_PUBLISH_CANCEL error - status="
                             + status + ", value=" + value);
                 }
-                stateMgr.onNoOpTransaction(transactionId);
                 break;
             case NAN_RESPONSE_TRANSMIT_FOLLOWUP:
                 if (status == NAN_STATUS_SUCCESS) {
-                    stateMgr.onMessageSendSuccess(transactionId);
+                    stateMgr.onMessageSendSuccessResponse(transactionId);
                 } else {
-                    stateMgr.onMessageSendFail(transactionId,
-                            translateHalStatusToPublicStatus(status));
+                    stateMgr.onMessageSendFailResponse(transactionId,
+                            translateHalStatusToNanSessionCallbackReason(status));
                 }
                 break;
             case NAN_RESPONSE_SUBSCRIBE_CANCEL:
@@ -555,10 +606,9 @@
                     Log.e(TAG, "onNanNotifyResponse: NAN_RESPONSE_PUBLISH_CANCEL error - status="
                             + status + ", value=" + value);
                 }
-                stateMgr.onNoOpTransaction(transactionId);
                 break;
             default:
-                stateMgr.onUnknownTransaction(responseType, transactionId, status);
+                Log.e(TAG, "onNanNotifyResponse: unclassified responseType=" + responseType);
                 break;
         }
     }
@@ -575,23 +625,25 @@
         switch (responseType) {
             case NAN_RESPONSE_PUBLISH:
                 if (status == NAN_STATUS_SUCCESS) {
-                    WifiNanStateManager.getInstance().onPublishSuccess(transactionId, pubSubId);
+                    WifiNanStateManager.getInstance().onSessionConfigSuccessResponse(transactionId,
+                            true, pubSubId);
                 } else {
-                    WifiNanStateManager.getInstance().onPublishFail(transactionId,
-                            translateHalStatusToPublicStatus(status));
+                    WifiNanStateManager.getInstance().onSessionConfigFailResponse(transactionId,
+                            true, translateHalStatusToNanSessionCallbackReason(status));
                 }
                 break;
             case NAN_RESPONSE_SUBSCRIBE:
                 if (status == NAN_STATUS_SUCCESS) {
-                    WifiNanStateManager.getInstance().onSubscribeSuccess(transactionId, pubSubId);
+                    WifiNanStateManager.getInstance().onSessionConfigSuccessResponse(transactionId,
+                            false, pubSubId);
                 } else {
-                    WifiNanStateManager.getInstance().onSubscribeFail(transactionId,
-                            translateHalStatusToPublicStatus(status));
+                    WifiNanStateManager.getInstance().onSessionConfigFailResponse(transactionId,
+                            false, translateHalStatusToNanSessionCallbackReason(status));
                 }
                 break;
             default:
-                WifiNanStateManager.getInstance().onUnknownTransaction(responseType, transactionId,
-                        status);
+                Log.wtf(TAG, "onNanNotifyResponsePublishSubscribe: unclassified responseType="
+                        + responseType);
                 break;
         }
     }
@@ -604,7 +656,7 @@
         }
 
         if (status == NAN_STATUS_SUCCESS) {
-            WifiNanStateManager.getInstance().onCapabilitiesUpdate(transactionId, capabilities);
+            WifiNanStateManager.getInstance().onCapabilitiesUpdateNotification(capabilities);
         } else {
             Log.e(TAG,
                     "onNanNotifyResponseCapabilities: error status=" + status + ", value=" + value);
@@ -623,13 +675,13 @@
         }
 
         if (eventType == NAN_EVENT_ID_DISC_MAC_ADDR) {
-            WifiNanStateManager.getInstance().onInterfaceAddressChange(mac);
+            WifiNanStateManager.getInstance().onInterfaceAddressChangeNotification(mac);
         } else if (eventType == NAN_EVENT_ID_STARTED_CLUSTER) {
-            WifiNanStateManager.getInstance()
-                    .onClusterChange(WifiNanClientState.CLUSTER_CHANGE_EVENT_STARTED, mac);
+            WifiNanStateManager.getInstance().onClusterChangeNotification(
+                    WifiNanClientState.CLUSTER_CHANGE_EVENT_STARTED, mac);
         } else if (eventType == NAN_EVENT_ID_JOINED_CLUSTER) {
-            WifiNanStateManager.getInstance()
-                    .onClusterChange(WifiNanClientState.CLUSTER_CHANGE_EVENT_JOINED, mac);
+            WifiNanStateManager.getInstance().onClusterChangeNotification(
+                    WifiNanClientState.CLUSTER_CHANGE_EVENT_JOINED, mac);
         } else {
             Log.w(TAG, "onDiscoveryEngineEvent: invalid eventType=" + eventType);
         }
@@ -646,7 +698,7 @@
                     + matchFilterLength + ", matchFilter=" + matchFilter);
         }
 
-        WifiNanStateManager.getInstance().onMatch(pubSubId, requestorInstanceId, mac,
+        WifiNanStateManager.getInstance().onMatchNotification(pubSubId, requestorInstanceId, mac,
                 serviceSpecificInfo, serviceSpecificInfoLength, matchFilter, matchFilterLength);
     }
 
@@ -654,8 +706,8 @@
     private static void onPublishTerminated(int publishId, int status) {
         if (VDBG) Log.v(TAG, "onPublishTerminated: publishId=" + publishId + ", status=" + status);
 
-        WifiNanStateManager.getInstance().onPublishTerminated(publishId,
-                translateHalStatusToPublicStatus(status));
+        WifiNanStateManager.getInstance().onSessionTerminatedNotification(publishId,
+                translateHalStatusToNanSessionCallbackTerminate(status), true);
     }
 
     // callback from native
@@ -664,8 +716,8 @@
             Log.v(TAG, "onSubscribeTerminated: subscribeId=" + subscribeId + ", status=" + status);
         }
 
-        WifiNanStateManager.getInstance().onSubscribeTerminated(subscribeId,
-                translateHalStatusToPublicStatus(status));
+        WifiNanStateManager.getInstance().onSessionTerminatedNotification(subscribeId,
+                translateHalStatusToNanSessionCallbackTerminate(status), false);
     }
 
     // callback from native
@@ -677,14 +729,15 @@
                     + ", messageLength=" + messageLength);
         }
 
-        WifiNanStateManager.getInstance().onMessageReceived(pubSubId, requestorInstanceId, mac,
-                message, messageLength);
+        WifiNanStateManager.getInstance().onMessageReceivedNotification(pubSubId,
+                requestorInstanceId, mac, message, messageLength);
     }
 
     // callback from native
     private static void onDisabledEvent(int status) {
         if (VDBG) Log.v(TAG, "onDisabledEvent: status=" + status);
 
-        WifiNanStateManager.getInstance().onNanDown(translateHalStatusToPublicStatus(status));
+        WifiNanStateManager.getInstance()
+                .onNanDownNotification(translateHalStatusToNanEventCallbackReason(status));
     }
 }
diff --git a/service/java/com/android/server/wifi/nan/WifiNanServiceImpl.java b/service/java/com/android/server/wifi/nan/WifiNanServiceImpl.java
index 93c1465..a0fe8fe 100644
--- a/service/java/com/android/server/wifi/nan/WifiNanServiceImpl.java
+++ b/service/java/com/android/server/wifi/nan/WifiNanServiceImpl.java
@@ -24,6 +24,7 @@
 import android.net.wifi.nan.IWifiNanSessionCallback;
 import android.net.wifi.nan.PublishConfig;
 import android.net.wifi.nan.SubscribeConfig;
+import android.net.wifi.nan.WifiNanEventCallback;
 import android.os.Binder;
 import android.os.HandlerThread;
 import android.os.IBinder;
@@ -47,7 +48,6 @@
 
     private Context mContext;
     private WifiNanStateManager mStateManager;
-    private final boolean mNanSupported;
 
     private final Object mLock = new Object();
     private final SparseArray<IBinder.DeathRecipient> mDeathRecipientsByClientId =
@@ -57,11 +57,6 @@
 
     public WifiNanServiceImpl(Context context) {
         mContext = context.getApplicationContext();
-
-        mNanSupported = mContext.getPackageManager()
-                .hasSystemFeature(PackageManager.FEATURE_WIFI_NAN);
-        if (DBG) Log.w(TAG, "WifiNanServiceImpl: mNanSupported=" + mNanSupported);
-
         mStateManager = WifiNanStateManager.getInstance();
     }
 
@@ -80,7 +75,7 @@
     public void start() {
         Log.i(TAG, "Starting Wi-Fi NAN service");
 
-        // TODO: share worker thread with other Wi-Fi handlers
+        // TODO: share worker thread with other Wi-Fi handlers (b/27924886)
         HandlerThread wifiNanThread = new HandlerThread("wifiNanService");
         wifiNanThread.start();
 
@@ -88,22 +83,38 @@
     }
 
     @Override
-    public int connect(final IBinder binder, IWifiNanEventCallback callback) {
+    public int connect(final IBinder binder, IWifiNanEventCallback callback,
+            ConfigRequest configRequest) {
         enforceAccessPermission();
         enforceChangePermission();
         if (callback == null) {
             throw new IllegalArgumentException("Callback must not be null");
         }
+        if (binder == null) {
+            throw new IllegalArgumentException("Binder must not be null");
+        }
+
+        if (configRequest != null) {
+            /*
+             * TODO: enforce additional permissions if configuration is
+             * non-standard (i.e. the system API). (b/27696149)
+             */
+        } else {
+            configRequest = new ConfigRequest.Builder().build();
+        }
+        configRequest.validate();
 
         final int uid = getMockableCallingUid();
 
         final int clientId;
         synchronized (mLock) {
             clientId = mNextClientId++;
-            mUidByClientId.put(clientId, uid);
         }
 
-        if (VDBG) Log.v(TAG, "connect: uid=" + uid + ", clientId=" + clientId);
+        if (VDBG) {
+            Log.v(TAG, "connect: uid=" + uid + ", clientId=" + clientId + ", configRequest"
+                    + configRequest);
+        }
 
         IBinder.DeathRecipient dr = new IBinder.DeathRecipient() {
             @Override
@@ -119,16 +130,25 @@
                 mStateManager.disconnect(clientId);
             }
         };
-        synchronized (mLock) {
-            mDeathRecipientsByClientId.put(clientId, dr);
-        }
+
         try {
             binder.linkToDeath(dr, 0);
         } catch (RemoteException e) {
-            Log.w(TAG, "Error on linkToDeath - " + e);
+            Log.e(TAG, "Error on linkToDeath - " + e);
+            try {
+                callback.onConnectFail(WifiNanEventCallback.REASON_OTHER);
+            } catch (RemoteException e1) {
+                Log.e(TAG, "Error on onConnectFail()");
+            }
+            return 0;
         }
 
-        mStateManager.connect(clientId, callback);
+        synchronized (mLock) {
+            mDeathRecipientsByClientId.put(clientId, dr);
+            mUidByClientId.put(clientId, uid);
+        }
+
+        mStateManager.connect(clientId, callback, configRequest);
 
         return clientId;
     }
@@ -142,6 +162,10 @@
         enforceClientValidity(uid, clientId);
         if (VDBG) Log.v(TAG, "disconnect: uid=" + uid + ", clientId=" + clientId);
 
+        if (binder == null) {
+            throw new IllegalArgumentException("Binder must not be null");
+        }
+
         synchronized (mLock) {
             IBinder.DeathRecipient dr = mDeathRecipientsByClientId.get(clientId);
             if (dr != null) {
@@ -155,21 +179,6 @@
     }
 
     @Override
-    public void requestConfig(int clientId, ConfigRequest configRequest) {
-        enforceAccessPermission();
-        enforceChangePermission();
-
-        int uid = getMockableCallingUid();
-        enforceClientValidity(uid, clientId);
-        if (VDBG) {
-            Log.v(TAG, "requestConfig: uid=" + uid + "clientId=" + clientId + ", configRequest="
-                    + configRequest);
-        }
-
-        mStateManager.requestConfig(clientId, configRequest);
-    }
-
-    @Override
     public void terminateSession(int clientId, int sessionId) {
         enforceAccessPermission();
         enforceChangePermission();
@@ -189,9 +198,14 @@
             IWifiNanSessionCallback callback) {
         enforceAccessPermission();
         enforceChangePermission();
+
         if (callback == null) {
             throw new IllegalArgumentException("Callback must not be null");
         }
+        if (publishConfig == null) {
+            throw new IllegalArgumentException("PublishConfig must not be null");
+        }
+        publishConfig.validate();
 
         int uid = getMockableCallingUid();
         enforceClientValidity(uid, clientId);
@@ -208,6 +222,11 @@
         enforceAccessPermission();
         enforceChangePermission();
 
+        if (publishConfig == null) {
+            throw new IllegalArgumentException("PublishConfig must not be null");
+        }
+        publishConfig.validate();
+
         int uid = getMockableCallingUid();
         enforceClientValidity(uid, clientId);
         if (VDBG) {
@@ -223,9 +242,14 @@
             IWifiNanSessionCallback callback) {
         enforceAccessPermission();
         enforceChangePermission();
+
         if (callback == null) {
             throw new IllegalArgumentException("Callback must not be null");
         }
+        if (subscribeConfig == null) {
+            throw new IllegalArgumentException("SubscribeConfig must not be null");
+        }
+        subscribeConfig.validate();
 
         int uid = getMockableCallingUid();
         enforceClientValidity(uid, clientId);
@@ -242,6 +266,11 @@
         enforceAccessPermission();
         enforceChangePermission();
 
+        if (subscribeConfig == null) {
+            throw new IllegalArgumentException("SubscribeConfig must not be null");
+        }
+        subscribeConfig.validate();
+
         int uid = getMockableCallingUid();
         enforceClientValidity(uid, clientId);
         if (VDBG) {
@@ -258,6 +287,11 @@
         enforceAccessPermission();
         enforceChangePermission();
 
+        if (messageLength != 0 && (message == null || message.length < messageLength)) {
+            throw new IllegalArgumentException(
+                    "Non-matching combination of message and messageLength");
+        }
+
         int uid = getMockableCallingUid();
         enforceClientValidity(uid, clientId);
         if (VDBG) {
@@ -279,10 +313,11 @@
             return;
         }
         pw.println("Wi-Fi NAN Service");
-        pw.println("  mNanSupported: " + mNanSupported);
-        pw.println("  mNextClientId: " + mNextClientId);
-        pw.println("  mDeathRecipientsByClientId: " + mDeathRecipientsByClientId);
-        pw.println("  mUidByClientId: " + mUidByClientId);
+        synchronized (mLock) {
+            pw.println("  mNextClientId: " + mNextClientId);
+            pw.println("  mDeathRecipientsByClientId: " + mDeathRecipientsByClientId);
+            pw.println("  mUidByClientId: " + mUidByClientId);
+        }
         mStateManager.dump(fd, pw, args);
     }
 
diff --git a/service/java/com/android/server/wifi/nan/WifiNanSessionState.java b/service/java/com/android/server/wifi/nan/WifiNanSessionState.java
index 91e48c6..9531cfb 100644
--- a/service/java/com/android/server/wifi/nan/WifiNanSessionState.java
+++ b/service/java/com/android/server/wifi/nan/WifiNanSessionState.java
@@ -40,47 +40,41 @@
     private static final boolean DBG = false;
     private static final boolean VDBG = false; // STOPSHIP if true
 
-    private final SparseArray<String> mMacByRequestorInstanceId = new SparseArray<>();
-
     private int mSessionId;
+    private int mPubSubId;
     private IWifiNanSessionCallback mCallback;
     private boolean mIsPublishSession;
 
-    private boolean mSessionValid = false;
-    private int mPubSubId;
+    private final SparseArray<String> mMacByRequestorInstanceId = new SparseArray<>();
 
-    public WifiNanSessionState(int sessionId, IWifiNanSessionCallback callback,
+    public WifiNanSessionState(int sessionId, int pubSubId, IWifiNanSessionCallback callback,
             boolean isPublishSession) {
         mSessionId = sessionId;
+        mPubSubId = pubSubId;
         mCallback = callback;
         mIsPublishSession = isPublishSession;
     }
 
+    public int getSessionId() {
+        return mSessionId;
+    }
+
+    public IWifiNanSessionCallback getCallback() {
+        return mCallback;
+    }
+
     /**
      * Destroy the current discovery session - stops publishing or subscribing
      * if currently active.
      */
     public void terminate() {
-        if (!mSessionValid) {
-            if (DBG) {
-                Log.d(TAG, "terminate: attempting to terminate an already terminated session");
-            }
-            return;
-        }
-
-        short transactionId = WifiNanStateManager.getInstance().createNextTransactionId();
-        if (mIsPublishSession) {
-            WifiNanNative.getInstance().stopPublish(transactionId, mPubSubId);
-        } else {
-            WifiNanNative.getInstance().stopSubscribe(transactionId, mPubSubId);
-        }
-
-        mSessionValid = false;
         mCallback = null;
-    }
 
-    public int getSessionId() {
-        return mSessionId;
+        if (mIsPublishSession) {
+            WifiNanNative.getInstance().stopPublish((short) 0, mPubSubId);
+        } else {
+            WifiNanNative.getInstance().stopSubscribe((short) 0, mPubSubId);
+        }
     }
 
     /**
@@ -91,18 +85,7 @@
      * @return true if corresponds to this session, false otherwise.
      */
     public boolean isPubSubIdSession(int pubSubId) {
-        return mSessionValid && mPubSubId == pubSubId;
-    }
-
-    /**
-     * Start a publish discovery session.
-     *
-     * @param transactionId Transaction ID for the transaction - used in the
-     *            async callback to match with the original request.
-     * @param config Configuration of the publish session.
-     */
-    public void publish(short transactionId, PublishConfig config) {
-        WifiNanNative.getInstance().publish(transactionId, 0, config);
+        return mPubSubId == pubSubId;
     }
 
     /**
@@ -112,31 +95,19 @@
      *            async callback to match with the original request.
      * @param config Configuration of the publish session.
      */
-    public void updatePublish(short transactionId, PublishConfig config) {
+    public boolean updatePublish(short transactionId, PublishConfig config) {
         if (!mIsPublishSession) {
             Log.e(TAG, "A SUBSCRIBE session is being used to publish");
-            WifiNanStateManager.getInstance().onPublishFail(transactionId,
-                    WifiNanSessionCallback.FAIL_REASON_OTHER);
-            return;
-        }
-        if (!mSessionValid) {
-            Log.e(TAG, "Attempting a re-publish on a terminated session");
-            onPublishFail(WifiNanSessionCallback.FAIL_REASON_SESSION_TERMINATED);
-            return;
+            try {
+                mCallback.onSessionConfigFail(WifiNanSessionCallback.REASON_OTHER);
+            } catch (RemoteException e) {
+                Log.e(TAG, "updatePublish: RemoteException=" + e);
+            }
+            return false;
         }
 
         WifiNanNative.getInstance().publish(transactionId, mPubSubId, config);
-    }
-
-    /**
-     * Start a subscribe discovery session.
-     *
-     * @param transactionId Transaction ID for the transaction - used in the
-     *            async callback to match with the original request.
-     * @param config Configuration of the subscribe session.
-     */
-    public void subscribe(short transactionId, SubscribeConfig config) {
-        WifiNanNative.getInstance().subscribe(transactionId, 0, config);
+        return true;
     }
 
     /**
@@ -146,20 +117,19 @@
      *            async callback to match with the original request.
      * @param config Configuration of the subscribe session.
      */
-    public void updateSubscribe(short transactionId, SubscribeConfig config) {
+    public boolean updateSubscribe(short transactionId, SubscribeConfig config) {
         if (mIsPublishSession) {
             Log.e(TAG, "A PUBLISH session is being used to subscribe");
-            WifiNanStateManager.getInstance().onSubscribeFail(transactionId,
-                    WifiNanSessionCallback.FAIL_REASON_OTHER);
-            return;
-        }
-        if (!mSessionValid) {
-            Log.e(TAG, "Attempting a re-subscribe on a terminated session");
-            onSubscribeFail(WifiNanSessionCallback.FAIL_REASON_SESSION_TERMINATED);
-            return;
+            try {
+                mCallback.onSessionConfigFail(WifiNanSessionCallback.REASON_OTHER);
+            } catch (RemoteException e) {
+                Log.e(TAG, "updateSubscribe: RemoteException=" + e);
+            }
+            return false;
         }
 
         WifiNanNative.getInstance().subscribe(transactionId, mPubSubId, config);
+        return true;
     }
 
     /**
@@ -174,173 +144,25 @@
      * @param messageId A message ID provided by caller to be used in any
      *            callbacks related to the message (success/failure).
      */
-    public void sendMessage(short transactionId, int peerId, byte[] message, int messageLength,
+    public boolean sendMessage(short transactionId, int peerId, byte[] message, int messageLength,
             int messageId) {
-        if (!mSessionValid) {
-            Log.e(TAG, "sendMessage: attempting to send a message on a terminated session "
-                    + "(no successful publish or subscribe");
-            onMessageSendFail(messageId, WifiNanSessionCallback.FAIL_REASON_SESSION_TERMINATED);
-            return;
-        }
-
         String peerMacStr = mMacByRequestorInstanceId.get(peerId);
         if (peerMacStr == null) {
             Log.e(TAG, "sendMessage: attempting to send a message to an address which didn't "
                     + "match/contact us");
-            onMessageSendFail(messageId, WifiNanSessionCallback.FAIL_REASON_NO_MATCH_SESSION);
-            return;
+            try {
+                mCallback.onMessageSendFail(messageId,
+                        WifiNanSessionCallback.REASON_NO_MATCH_SESSION);
+            } catch (RemoteException e) {
+                Log.e(TAG, "sendMessage: RemoteException=" + e);
+            }
+            return false;
         }
         byte[] peerMac = HexEncoding.decode(peerMacStr.toCharArray(), false);
 
         WifiNanNative.getInstance().sendMessage(transactionId, mPubSubId, peerId, peerMac, message,
                 messageLength);
-    }
-
-    /**
-     * Callback from HAL updating session in case of publish session creation
-     * success (i.e. publish session configured correctly and is active).
-     *
-     * @param publishId The HAL id of the (now active) publish session.
-     */
-    public void onPublishSuccess(int publishId) {
-        if (mSessionValid) {
-            // indicates an update-publish: no need to inform client
-            return;
-        }
-        mPubSubId = publishId;
-        mSessionValid = true;
-        try {
-            if (mCallback != null) {
-                mCallback.onSessionStarted(mSessionId);
-            }
-        } catch (RemoteException e) {
-            Log.w(TAG, "onPublishSuccess: RemoteException (FYI): " + e);
-        }
-    }
-
-    /**
-     * Callback from HAL updating session in case of publish session creation
-     * fail. Propagates call to client if registered.
-     *
-     * @param status Reason code for failure.
-     */
-    public void onPublishFail(int status) {
-        try {
-            if (mCallback != null) {
-                mCallback.onSessionConfigFail(status);
-            }
-        } catch (RemoteException e) {
-            Log.w(TAG, "onPublishFail: RemoteException (FYI): " + e);
-        }
-    }
-
-    /**
-     * Callback from HAL updating session when the publish session has
-     * terminated (per plan or due to failure). Propagates call to client if
-     * registered.
-     *
-     * @param status Reason code for session termination.
-     */
-    public void onPublishTerminated(int status) {
-        mSessionValid = false;
-        try {
-            if (mCallback != null) {
-                mCallback.onSessionTerminated(status);
-            }
-        } catch (RemoteException e) {
-            Log.w(TAG, "onPublishTerminated: RemoteException (FYI): " + e);
-        }
-    }
-
-    /**
-     * Callback from HAL updating session in case of subscribe session creation
-     * success (i.e. subscribe session configured correctly and is active).
-     *
-     * @param subscribeId The HAL id of the (now active) subscribe session.
-     */
-    public void onSubscribeSuccess(int subscribeId) {
-        if (mSessionValid) {
-            // indicates an update-publish: no need to inform client
-            return;
-        }
-
-        mPubSubId = subscribeId;
-        mSessionValid = true;
-        try {
-            if (mCallback != null) {
-                mCallback.onSessionStarted(mSessionId);
-            }
-        } catch (RemoteException e) {
-            Log.w(TAG, "onSubscribeSuccess: RemoteException (FYI): " + e);
-        }
-    }
-
-    /**
-     * Callback from HAL updating session in case of subscribe session creation
-     * fail. Propagates call to client if registered.
-     *
-     * @param status Reason code for failure.
-     */
-    public void onSubscribeFail(int status) {
-        try {
-            if (mCallback != null) {
-                mCallback.onSessionConfigFail(status);
-            }
-        } catch (RemoteException e) {
-            Log.w(TAG, "onSubscribeFail: RemoteException (FYI): " + e);
-        }
-    }
-
-    /**
-     * Callback from HAL updating session when the subscribe session has
-     * terminated (per plan or due to failure). Propagates call to client if
-     * registered.
-     *
-     * @param status Reason code for session termination.
-     */
-    public void onSubscribeTerminated(int status) {
-        mSessionValid = false;
-        try {
-            if (mCallback != null) {
-                mCallback.onSessionTerminated(status);
-            }
-        } catch (RemoteException e) {
-            Log.w(TAG, "onSubscribeTerminated: RemoteException (FYI): " + e);
-        }
-    }
-
-    /**
-     * Callback from HAL when message is sent successfully (i.e. an ACK was
-     * received). Propagates call to client if registered.
-     *
-     * @param messageId ID provided by caller with the message.
-     */
-    public void onMessageSendSuccess(int messageId) {
-        try {
-            if (mCallback != null) {
-                mCallback.onMessageSendSuccess(messageId);
-            }
-        } catch (RemoteException e) {
-            Log.w(TAG, "onMessageSendSuccess: RemoteException (FYI): " + e);
-        }
-    }
-
-    /**
-     * Callback from HAL when message fails to be transmitted - including when
-     * transmitted but no ACK received from intended receiver. Propagates call
-     * to client if registered.
-     *
-     * @param messageId ID provided by caller with the message.
-     * @param status Reason code for transmit failure.
-     */
-    public void onMessageSendFail(int messageId, int status) {
-        try {
-            if (mCallback != null) {
-                mCallback.onMessageSendFail(messageId, status);
-            }
-        } catch (RemoteException e) {
-            Log.w(TAG, "onMessageSendFail: RemoteException (FYI): " + e);
-        }
+        return true;
     }
 
     /**
@@ -367,10 +189,8 @@
         if (DBG) Log.d(TAG, "onMatch: previous peer MAC replaced - " + prevMac);
 
         try {
-            if (mCallback != null) {
-                mCallback.onMatch(requestorInstanceId, serviceSpecificInfo,
-                        serviceSpecificInfoLength, matchFilter, matchFilterLength);
-            }
+            mCallback.onMatch(requestorInstanceId, serviceSpecificInfo, serviceSpecificInfoLength,
+                    matchFilter, matchFilterLength);
         } catch (RemoteException e) {
             Log.w(TAG, "onMatch: RemoteException (FYI): " + e);
         }
@@ -397,9 +217,7 @@
         }
 
         try {
-            if (mCallback != null) {
-                mCallback.onMessageReceived(requestorInstanceId, message, messageLength);
-            }
+            mCallback.onMessageReceived(requestorInstanceId, message, messageLength);
         } catch (RemoteException e) {
             Log.w(TAG, "onMessageReceived: RemoteException (FYI): " + e);
         }
@@ -412,7 +230,7 @@
         pw.println("NanSessionState:");
         pw.println("  mSessionId: " + mSessionId);
         pw.println("  mIsPublishSession: " + mIsPublishSession);
-        pw.println("  mPubSubId: " + (mSessionValid ? Integer.toString(mPubSubId) : "not valid"));
+        pw.println("  mPubSubId: " + mPubSubId);
         pw.println("  mMacByRequestorInstanceId: [" + mMacByRequestorInstanceId + "]");
     }
 }
diff --git a/service/java/com/android/server/wifi/nan/WifiNanStateManager.java b/service/java/com/android/server/wifi/nan/WifiNanStateManager.java
index 3ec237f..d69e848 100644
--- a/service/java/com/android/server/wifi/nan/WifiNanStateManager.java
+++ b/service/java/com/android/server/wifi/nan/WifiNanStateManager.java
@@ -21,19 +21,22 @@
 import android.net.wifi.nan.IWifiNanSessionCallback;
 import android.net.wifi.nan.PublishConfig;
 import android.net.wifi.nan.SubscribeConfig;
+import android.net.wifi.nan.WifiNanEventCallback;
 import android.os.Bundle;
-import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
+import android.os.RemoteException;
 import android.util.Log;
+import android.util.Pair;
 import android.util.SparseArray;
 
+import com.android.internal.util.State;
+import com.android.internal.util.StateMachine;
+
 import libcore.util.HexEncoding;
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
-import java.util.ArrayList;
-import java.util.List;
 
 /**
  * Manages the state of the Wi-Fi NAN system service.
@@ -45,34 +48,48 @@
 
     private static WifiNanStateManager sNanStateManagerSingleton;
 
-    private static final int MESSAGE_CONNECT = 0;
-    private static final int MESSAGE_DISCONNECT = 1;
-    private static final int MESSAGE_REQUEST_CONFIG = 2;
-    private static final int MESSAGE_TERMINATE_SESSION = 3;
-    private static final int MESSAGE_PUBLISH = 4;
-    private static final int MESSAGE_UPDATE_PUBLISH = 5;
-    private static final int MESSAGE_SUBSCRIBE = 6;
-    private static final int MESSAGE_UPDATE_SUBSCRIBE = 7;
-    private static final int MESSAGE_SEND_MESSAGE = 8;
-    private static final int MESSAGE_ON_CONFIG_COMPLETED = 9;
-    private static final int MESSAGE_ON_CONFIG_FAILED = 10;
-    private static final int MESSAGE_ON_NAN_DOWN = 11;
-    private static final int MESSAGE_ON_INTERFACE_CHANGE = 12;
-    private static final int MESSAGE_ON_CLUSTER_CHANGE = 13;
-    private static final int MESSAGE_ON_PUBLISH_SUCCESS = 14;
-    private static final int MESSAGE_ON_PUBLISH_FAIL = 15;
-    private static final int MESSAGE_ON_PUBLISH_TERMINATED = 16;
-    private static final int MESSAGE_ON_SUBSCRIBE_SUCCESS = 17;
-    private static final int MESSAGE_ON_SUBSCRIBE_FAIL = 18;
-    private static final int MESSAGE_ON_SUBSCRIBE_TERMINATED = 19;
-    private static final int MESSAGE_ON_MESSAGE_SEND_SUCCESS = 20;
-    private static final int MESSAGE_ON_MESSAGE_SEND_FAIL = 21;
-    private static final int MESSAGE_ON_UNKNOWN_TRANSACTION = 22;
-    private static final int MESSAGE_ON_MATCH = 23;
-    private static final int MESSAGE_ON_MESSAGE_RECEIVED = 24;
-    private static final int MESSAGE_ON_CAPABILITIES_UPDATED = 25;
-    private static final int MESSAGE_ON_NO_OP_TRANSACTION = 26;
+    /*
+     * State machine message types. There are sub-types for the messages (except
+     * for TIMEOUT). Format: - Message.arg1: contains message sub-type -
+     * Message.arg2: contains transaction ID for RESPONSES & TIMEOUT
+     */
+    private static final int MESSAGE_TYPE_NOTIFICATION = 1;
+    private static final int MESSAGE_TYPE_COMMAND = 2;
+    private static final int MESSAGE_TYPE_RESPONSE = 3;
+    private static final int MESSAGE_TYPE_TIMEOUT = 4;
 
+    /*
+     * Message sub-types:
+     */
+    private static final int COMMAND_TYPE_CONNECT = 0;
+    private static final int COMMAND_TYPE_DISCONNECT = 1;
+    private static final int COMMAND_TYPE_TERMINATE_SESSION = 2;
+    private static final int COMMAND_TYPE_PUBLISH = 3;
+    private static final int COMMAND_TYPE_UPDATE_PUBLISH = 4;
+    private static final int COMMAND_TYPE_SUBSCRIBE = 5;
+    private static final int COMMAND_TYPE_UPDATE_SUBSCRIBE = 6;
+    private static final int COMMAND_TYPE_SEND_MESSAGE = 7;
+
+    private static final int RESPONSE_TYPE_ON_CONFIG_SUCCESS = 0;
+    private static final int RESPONSE_TYPE_ON_CONFIG_FAIL = 1;
+    private static final int RESPONSE_TYPE_ON_SESSION_CONFIG_SUCCESS = 2;
+    private static final int RESPONSE_TYPE_ON_SESSION_CONFIG_FAIL = 3;
+    private static final int RESPONSE_TYPE_ON_MESSAGE_SEND_SUCCESS = 4;
+    private static final int RESPONSE_TYPE_ON_MESSAGE_SEND_FAIL = 5;
+
+    private static final int NOTIFICATION_TYPE_CAPABILITIES_UPDATED = 0;
+    private static final int NOTIFICATION_TYPE_INTERFACE_CHANGE = 1;
+    private static final int NOTIFICATION_TYPE_CLUSTER_CHANGE = 2;
+    private static final int NOTIFICATION_TYPE_MATCH = 3;
+    private static final int NOTIFICATION_TYPE_SESSION_TERMINATED = 4;
+    private static final int NOTIFICATION_TYPE_MESSAGE_RECEIVED = 5;
+    private static final int NOTIFICATION_TYPE_NAN_DOWN = 6;
+
+    /*
+     * Keys used when passing (some) arguments to the Handler thread (too many
+     * arguments to pass in the short-cut Message members).
+     */
+    private static final String MESSAGE_BUNDLE_KEY_SESSION_TYPE = "session_type";
     private static final String MESSAGE_BUNDLE_KEY_SESSION_ID = "session_id";
     private static final String MESSAGE_BUNDLE_KEY_CONFIG = "config";
     private static final String MESSAGE_BUNDLE_KEY_MESSAGE = "message";
@@ -86,16 +103,17 @@
     private static final String MESSAGE_BUNDLE_KEY_MAC_ADDRESS = "mac_address";
     private static final String MESSAGE_BUNDLE_KEY_MESSAGE_DATA = "message_data";
     private static final String MESSAGE_BUNDLE_KEY_MESSAGE_LENGTH = "message_length";
+    private static final String MESSAGE_BUNDLE_KEY_REQ_INSTANCE_ID = "req_instance_id";
 
+    /*
+     * All state is only accessed through the state machine handler thread: no
+     * need to synchronize.
+     */
     private WifiNanNative.Capabilities mCapabilities;
+    private WifiNanStateMachine mSm;
 
-    private WifiNanStateHandler mHandler;
-
-    // no synchronization necessary: only access through Handler
     private final SparseArray<WifiNanClientState> mClients = new SparseArray<>();
-    private final SparseArray<TransactionInfoBase> mPendingResponses = new SparseArray<>();
-    private short mNextTransactionId = 1;
-    private int mNextSessionId = 1;
+    private ConfigRequest mCurrentNanConfiguration = null;
 
     private WifiNanStateManager() {
         // EMPTY: singleton pattern
@@ -124,470 +142,448 @@
     public void start(Looper looper) {
         Log.i(TAG, "start()");
 
-        mHandler = new WifiNanStateHandler(looper);
+        mSm = new WifiNanStateMachine(TAG, looper);
+        mSm.setDbg(DBG);
+        mSm.start();
     }
 
-    /**
-     * Place a request for a new client connection on the handler queue.
+    /*
+     * COMMANDS
      */
-    public void connect(int clientId, IWifiNanEventCallback callback) {
-        Message msg = mHandler.obtainMessage(MESSAGE_CONNECT);
-        msg.arg1 = clientId;
+
+    /**
+     * Place a request for a new client connection on the state machine queue.
+     */
+    public void connect(int clientId, IWifiNanEventCallback callback, ConfigRequest configRequest) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_COMMAND);
+        msg.arg1 = COMMAND_TYPE_CONNECT;
+        msg.arg2 = clientId;
         msg.obj = callback;
-        mHandler.sendMessage(msg);
+        msg.getData().putParcelable(MESSAGE_BUNDLE_KEY_CONFIG, configRequest);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a request to disconnect (destroy) an existing client on the handler
-     * queue.
+     * Place a request to disconnect (destroy) an existing client on the state
+     * machine queue.
      */
     public void disconnect(int clientId) {
-        Message msg = mHandler.obtainMessage(MESSAGE_DISCONNECT);
-        msg.arg1 = clientId;
-        mHandler.sendMessage(msg);
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_COMMAND);
+        msg.arg1 = COMMAND_TYPE_DISCONNECT;
+        msg.arg2 = clientId;
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a request to enable and configure NAN on the handler queue.
-     */
-    public void requestConfig(int clientId, ConfigRequest configRequest) {
-        Message msg = mHandler.obtainMessage(MESSAGE_REQUEST_CONFIG);
-        msg.arg1 = clientId;
-        msg.obj = configRequest;
-        mHandler.sendMessage(msg);
-    }
-
-    /**
-     * Place a request to stop a discovery session on the handler queue.
+     * Place a request to stop a discovery session on the state machine queue.
      */
     public void terminateSession(int clientId, int sessionId) {
-        Message msg = mHandler.obtainMessage(MESSAGE_TERMINATE_SESSION);
-        msg.arg1 = clientId;
-        msg.arg2 = sessionId;
-        mHandler.sendMessage(msg);
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_COMMAND);
+        msg.arg1 = COMMAND_TYPE_TERMINATE_SESSION;
+        msg.arg2 = clientId;
+        msg.obj = new Integer(sessionId);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a request to start a new publish discovery session on the handler
-     * queue.
+     * Place a request to start a new publish discovery session on the state
+     * machine queue.
      */
     public void publish(int clientId, PublishConfig publishConfig,
             IWifiNanSessionCallback callback) {
-        Bundle data = new Bundle();
-        data.putParcelable(MESSAGE_BUNDLE_KEY_CONFIG, publishConfig);
-
-        Message msg = mHandler.obtainMessage(MESSAGE_PUBLISH);
-        msg.arg1 = clientId;
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_COMMAND);
+        msg.arg1 = COMMAND_TYPE_PUBLISH;
+        msg.arg2 = clientId;
         msg.obj = callback;
-        msg.setData(data);
-        mHandler.sendMessage(msg);
+        msg.getData().putParcelable(MESSAGE_BUNDLE_KEY_CONFIG, publishConfig);
+        mSm.sendMessage(msg);
     }
 
     /**
      * Place a request to modify an existing publish discovery session on the
-     * handler queue.
+     * state machine queue.
      */
     public void updatePublish(int clientId, int sessionId, PublishConfig publishConfig) {
-        Message msg = mHandler.obtainMessage(MESSAGE_UPDATE_PUBLISH);
-        msg.arg1 = clientId;
-        msg.arg2 = sessionId;
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_COMMAND);
+        msg.arg1 = COMMAND_TYPE_UPDATE_PUBLISH;
+        msg.arg2 = clientId;
         msg.obj = publishConfig;
-        mHandler.sendMessage(msg);
+        msg.getData().putInt(MESSAGE_BUNDLE_KEY_SESSION_ID, sessionId);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a request to start a new subscribe discovery session on the handler
-     * queue.
+     * Place a request to start a new subscribe discovery session on the state
+     * machine queue.
      */
     public void subscribe(int clientId, SubscribeConfig subscribeConfig,
             IWifiNanSessionCallback callback) {
-        Bundle data = new Bundle();
-        data.putParcelable(MESSAGE_BUNDLE_KEY_CONFIG, subscribeConfig);
-
-        Message msg = mHandler.obtainMessage(MESSAGE_SUBSCRIBE);
-        msg.arg1 = clientId;
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_COMMAND);
+        msg.arg1 = COMMAND_TYPE_SUBSCRIBE;
+        msg.arg2 = clientId;
         msg.obj = callback;
-        msg.setData(data);
-        mHandler.sendMessage(msg);
+        msg.getData().putParcelable(MESSAGE_BUNDLE_KEY_CONFIG, subscribeConfig);
+        mSm.sendMessage(msg);
     }
 
     /**
      * Place a request to modify an existing subscribe discovery session on the
-     * handler queue.
+     * state machine queue.
      */
     public void updateSubscribe(int clientId, int sessionId, SubscribeConfig subscribeConfig) {
-        Message msg = mHandler.obtainMessage(MESSAGE_UPDATE_SUBSCRIBE);
-        msg.arg1 = clientId;
-        msg.arg2 = sessionId;
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_COMMAND);
+        msg.arg1 = COMMAND_TYPE_UPDATE_SUBSCRIBE;
+        msg.arg2 = clientId;
         msg.obj = subscribeConfig;
-        mHandler.sendMessage(msg);
+        msg.getData().putInt(MESSAGE_BUNDLE_KEY_SESSION_ID, sessionId);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a request to send a message on a discovery session on the handler
-     * queue.
+     * Place a request to send a message on a discovery session on the state
+     * machine queue.
      */
     public void sendMessage(int clientId, int sessionId, int peerId, byte[] message,
-            int messageLength,
-            int messageId) {
-        Bundle data = new Bundle();
-        data.putInt(MESSAGE_BUNDLE_KEY_SESSION_ID, sessionId);
-        data.putInt(MESSAGE_BUNDLE_KEY_MESSAGE_PEER_ID, peerId);
-        data.putByteArray(MESSAGE_BUNDLE_KEY_MESSAGE, message);
-        data.putInt(MESSAGE_BUNDLE_KEY_MESSAGE_ID, messageId);
-
-        Message msg = mHandler.obtainMessage(MESSAGE_SEND_MESSAGE);
-        msg.arg1 = clientId;
-        msg.arg2 = messageLength;
-        msg.setData(data);
-        mHandler.sendMessage(msg);
+            int messageLength, int messageId) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_COMMAND);
+        msg.arg1 = COMMAND_TYPE_SEND_MESSAGE;
+        msg.arg2 = clientId;
+        msg.getData().putInt(MESSAGE_BUNDLE_KEY_SESSION_ID, sessionId);
+        msg.getData().putInt(MESSAGE_BUNDLE_KEY_MESSAGE_PEER_ID, peerId);
+        msg.getData().putByteArray(MESSAGE_BUNDLE_KEY_MESSAGE, message);
+        msg.getData().putInt(MESSAGE_BUNDLE_KEY_MESSAGE_ID, messageId);
+        msg.getData().putInt(MESSAGE_BUNDLE_KEY_MESSAGE_LENGTH, messageLength);
+        mSm.sendMessage(msg);
     }
 
-    /**
-     * Place a callback request on the handler queue: update vendor capabilities
-     * of the NAN stack.
+    /*
+     * RESPONSES
      */
-    public void onCapabilitiesUpdate(short transactionId, WifiNanNative.Capabilities capabilities) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_CAPABILITIES_UPDATED);
-        msg.arg1 = transactionId;
-        msg.obj = capabilities;
-        mHandler.sendMessage(msg);
-    }
 
     /**
-     * Place a callback request on the handler queue: configuration request
-     * completed (successfully).
+     * Place a callback request on the state machine queue: configuration
+     * request completed (successfully).
      */
-    public void onConfigCompleted(short transactionId) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_CONFIG_COMPLETED);
-        msg.arg1 = transactionId;
-        mHandler.sendMessage(msg);
+    public void onConfigSuccessResponse(short transactionId) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_RESPONSE);
+        msg.arg1 = RESPONSE_TYPE_ON_CONFIG_SUCCESS;
+        msg.arg2 = transactionId;
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: configuration request
-     * failed.
-     */
-    public void onConfigFailed(short transactionId, int reason) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_CONFIG_FAILED);
-        msg.arg1 = transactionId;
-        msg.arg2 = reason;
-        mHandler.sendMessage(msg);
-    }
-
-    /**
-     * Place a callback request on the handler queue: publish session creation
-     * request succeeded (i.e. session is now active).
-     */
-    public void onPublishSuccess(short transactionId, int publishId) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_PUBLISH_SUCCESS);
-        msg.arg1 = transactionId;
-        msg.arg2 = publishId;
-        mHandler.sendMessage(msg);
-    }
-
-    /**
-     * Place a callback request on the handler queue: publish session creation
+     * Place a callback request on the state machine queue: configuration
      * request failed.
      */
-    public void onPublishFail(short transactionId, int status) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_PUBLISH_FAIL);
-        msg.arg1 = transactionId;
-        msg.arg2 = status;
-        mHandler.sendMessage(msg);
+    public void onConfigFailedResponse(short transactionId, int reason) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_RESPONSE);
+        msg.arg1 = RESPONSE_TYPE_ON_CONFIG_FAIL;
+        msg.arg2 = transactionId;
+        msg.obj = new Integer(reason);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: message has been sent
-     * successfully (i.e. an ACK was received from the targeted peer).
+     * Place a callback request on the state machine queue: session
+     * configuration (new or update) request succeeded.
      */
-    public void onMessageSendSuccess(short transactionId) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_MESSAGE_SEND_SUCCESS);
-        msg.arg1 = transactionId;
-        mHandler.sendMessage(msg);
+    public void onSessionConfigSuccessResponse(short transactionId, boolean isPublish,
+            int pubSubId) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_RESPONSE);
+        msg.arg1 = RESPONSE_TYPE_ON_SESSION_CONFIG_SUCCESS;
+        msg.arg2 = transactionId;
+        msg.obj = new Integer(pubSubId);
+        msg.getData().putBoolean(MESSAGE_BUNDLE_KEY_SESSION_TYPE, isPublish);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: attempt to send a message
-     * has failed.
+     * Place a callback request on the state machine queue: session
+     * configuration (new or update) request failed.
      */
-    public void onMessageSendFail(short transactionId, int status) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_MESSAGE_SEND_FAIL);
-        msg.arg1 = transactionId;
-        msg.arg2 = status;
-        mHandler.sendMessage(msg);
+    public void onSessionConfigFailResponse(short transactionId, boolean isPublish, int reason) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_RESPONSE);
+        msg.arg1 = RESPONSE_TYPE_ON_SESSION_CONFIG_FAIL;
+        msg.arg2 = transactionId;
+        msg.obj = new Integer(reason);
+        msg.getData().putBoolean(MESSAGE_BUNDLE_KEY_SESSION_TYPE, isPublish);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: subscribe session creation
-     * has succeeded (i.e. session is now active).
+     * Place a callback request on the state machine queue: message has been
+     * sent successfully (i.e. an ACK was received from the targeted peer).
      */
-    public void onSubscribeSuccess(short transactionId, int subscribeId) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_SUBSCRIBE_SUCCESS);
-        msg.arg1 = transactionId;
-        msg.arg2 = subscribeId;
-        mHandler.sendMessage(msg);
+    public void onMessageSendSuccessResponse(short transactionId) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_RESPONSE);
+        msg.arg1 = RESPONSE_TYPE_ON_MESSAGE_SEND_SUCCESS;
+        msg.arg2 = transactionId;
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: subscribe session creation
-     * has failed.
+     * Place a callback request on the state machine queue: attempt to send a
+     * message has failed.
      */
-    public void onSubscribeFail(short transactionId, int status) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_SUBSCRIBE_FAIL);
-        msg.arg1 = transactionId;
-        msg.arg2 = status;
-        mHandler.sendMessage(msg);
+    public void onMessageSendFailResponse(short transactionId, int reason) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_RESPONSE);
+        msg.arg1 = RESPONSE_TYPE_ON_MESSAGE_SEND_FAIL;
+        msg.arg2 = transactionId;
+        msg.obj = new Integer(reason);
+        mSm.sendMessage(msg);
+    }
+
+    /*
+     * NOTIFICATIONS
+     */
+
+    /**
+     * Place a callback request on the state machine queue: update vendor
+     * capabilities of the NAN stack. This is actually a RESPONSE from the HAL -
+     * but treated as a NOTIFICATION.
+     */
+    public void onCapabilitiesUpdateNotification(WifiNanNative.Capabilities capabilities) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_NOTIFICATION);
+        msg.arg1 = NOTIFICATION_TYPE_CAPABILITIES_UPDATED;
+        msg.obj = capabilities;
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: HAL callback with a NOP
-     * operation - the only purpose is to clean-up the state of the pending
-     * transaction ID.
-     *
-     * @param transactionId Transaction ID of the operation to be cleaned-up.
+     * Place a callback request on the state machine queue: the discovery
+     * interface has changed.
      */
-    public void onNoOpTransaction(short transactionId) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_NO_OP_TRANSACTION);
-        msg.arg1 = transactionId;
-        mHandler.sendMessage(msg);
-    }
-
-    /**
-     * Place a callback request on the handler queue: HAL callback with an
-     * unknown transaction type.
-     */
-    public void onUnknownTransaction(int responseType, short transactionId, int status) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_UNKNOWN_TRANSACTION);
-        Bundle data = new Bundle();
-        data.putInt(MESSAGE_BUNDLE_KEY_RESPONSE_TYPE, responseType);
-        msg.setData(data);
-        msg.arg1 = transactionId;
-        msg.arg2 = status;
-        mHandler.sendMessage(msg);
-    }
-
-    /**
-     * Place a callback request on the handler queue: the discovery interface
-     * has changed.
-     */
-    public void onInterfaceAddressChange(byte[] mac) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_INTERFACE_CHANGE);
+    public void onInterfaceAddressChangeNotification(byte[] mac) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_NOTIFICATION);
+        msg.arg1 = NOTIFICATION_TYPE_INTERFACE_CHANGE;
         msg.obj = mac;
-        mHandler.sendMessage(msg);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: the cluster membership has
-     * changed (e.g. due to starting a new cluster or joining another cluster).
+     * Place a callback request on the state machine queue: the cluster
+     * membership has changed (e.g. due to starting a new cluster or joining
+     * another cluster).
      */
-    public void onClusterChange(int flag, byte[] clusterId) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_CLUSTER_CHANGE);
-        msg.arg1 = flag;
+    public void onClusterChangeNotification(int flag, byte[] clusterId) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_NOTIFICATION);
+        msg.arg1 = NOTIFICATION_TYPE_CLUSTER_CHANGE;
+        msg.arg2 = flag;
         msg.obj = clusterId;
-        mHandler.sendMessage(msg);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: a discovery match has
-     * occurred - e.g. our subscription discovered someone else publishing a
+     * Place a callback request on the state machine queue: a discovery match
+     * has occurred - e.g. our subscription discovered someone else publishing a
      * matching service (to the one we were looking for).
      */
-    public void onMatch(int pubSubId, int requestorInstanceId, byte[] peerMac,
+    public void onMatchNotification(int pubSubId, int requestorInstanceId, byte[] peerMac,
             byte[] serviceSpecificInfo, int serviceSpecificInfoLength, byte[] matchFilter,
             int matchFilterLength) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_MATCH);
-        msg.arg1 = pubSubId;
-        msg.arg2 = requestorInstanceId;
-        Bundle data = new Bundle();
-        data.putByteArray(MESSAGE_BUNDLE_KEY_MAC_ADDRESS, peerMac);
-        data.putByteArray(MESSAGE_BUNDLE_KEY_SSI_DATA, serviceSpecificInfo);
-        data.putInt(MESSAGE_BUNDLE_KEY_SSI_LENGTH, serviceSpecificInfoLength);
-        data.putByteArray(MESSAGE_BUNDLE_KEY_FILTER_DATA, matchFilter);
-        data.putInt(MESSAGE_BUNDLE_KEY_FILTER_LENGTH, matchFilterLength);
-        msg.setData(data);
-        mHandler.sendMessage(msg);
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_NOTIFICATION);
+        msg.arg1 = NOTIFICATION_TYPE_MATCH;
+        msg.arg2 = pubSubId;
+        msg.getData().putInt(MESSAGE_BUNDLE_KEY_REQ_INSTANCE_ID, requestorInstanceId);
+        msg.getData().putByteArray(MESSAGE_BUNDLE_KEY_MAC_ADDRESS, peerMac);
+        msg.getData().putByteArray(MESSAGE_BUNDLE_KEY_SSI_DATA, serviceSpecificInfo);
+        msg.getData().putInt(MESSAGE_BUNDLE_KEY_SSI_LENGTH, serviceSpecificInfoLength);
+        msg.getData().putByteArray(MESSAGE_BUNDLE_KEY_FILTER_DATA, matchFilter);
+        msg.getData().putInt(MESSAGE_BUNDLE_KEY_FILTER_LENGTH, matchFilterLength);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: a publish session has
-     * terminated (per plan or due to an error).
+     * Place a callback request on the state machine queue: a session (publish
+     * or subscribe) has terminated (per plan or due to an error).
      */
-    public void onPublishTerminated(int publishId, int status) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_PUBLISH_TERMINATED);
-        msg.arg1 = publishId;
-        msg.arg2 = status;
-        mHandler.sendMessage(msg);
+    public void onSessionTerminatedNotification(int pubSubId, int reason, boolean isPublish) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_NOTIFICATION);
+        msg.arg1 = NOTIFICATION_TYPE_SESSION_TERMINATED;
+        msg.arg2 = pubSubId;
+        msg.obj = new Integer(reason);
+        msg.getData().putBoolean(MESSAGE_BUNDLE_KEY_SESSION_TYPE, isPublish);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: a subscribe session has
-     * terminated (per plan or due to an error).
-     */
-    public void onSubscribeTerminated(int subscribeId, int status) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_SUBSCRIBE_TERMINATED);
-        msg.arg1 = subscribeId;
-        msg.arg2 = status;
-        mHandler.sendMessage(msg);
-    }
-
-    /**
-     * Place a callback request on the handler queue: a message has been
+     * Place a callback request on the state machine queue: a message has been
      * received as part of a discovery session.
      */
-    public void onMessageReceived(int pubSubId, int requestorInstanceId, byte[] peerMac,
+    public void onMessageReceivedNotification(int pubSubId, int requestorInstanceId, byte[] peerMac,
             byte[] message, int messageLength) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_MESSAGE_RECEIVED);
-        msg.arg1 = pubSubId;
-        msg.arg2 = requestorInstanceId;
-        Bundle data = new Bundle();
-        data.putByteArray(MESSAGE_BUNDLE_KEY_MAC_ADDRESS, peerMac);
-        data.putByteArray(MESSAGE_BUNDLE_KEY_MESSAGE_DATA, message);
-        data.putInt(MESSAGE_BUNDLE_KEY_MESSAGE_LENGTH, messageLength);
-        msg.setData(data);
-        mHandler.sendMessage(msg);
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_NOTIFICATION);
+        msg.arg1 = NOTIFICATION_TYPE_MESSAGE_RECEIVED;
+        msg.arg2 = pubSubId;
+        msg.obj = new Integer(requestorInstanceId);
+        msg.getData().putByteArray(MESSAGE_BUNDLE_KEY_MAC_ADDRESS, peerMac);
+        msg.getData().putByteArray(MESSAGE_BUNDLE_KEY_MESSAGE_DATA, message);
+        msg.getData().putInt(MESSAGE_BUNDLE_KEY_MESSAGE_LENGTH, messageLength);
+        mSm.sendMessage(msg);
     }
 
     /**
-     * Place a callback request on the handler queue: NAN is going down.
+     * Place a callback request on the state machine queue: NAN is going down.
      */
-    public void onNanDown(int reason) {
-        Message msg = mHandler.obtainMessage(MESSAGE_ON_NAN_DOWN);
-        msg.arg1 = reason;
-        mHandler.sendMessage(msg);
+    public void onNanDownNotification(int reason) {
+        Message msg = mSm.obtainMessage(MESSAGE_TYPE_NOTIFICATION);
+        msg.arg1 = NOTIFICATION_TYPE_NAN_DOWN;
+        msg.arg2 = reason;
+        mSm.sendMessage(msg);
     }
 
-    private class WifiNanStateHandler extends Handler {
-        WifiNanStateHandler(android.os.Looper looper) {
-            super(looper);
+    /**
+     * State machine.
+     */
+    private class WifiNanStateMachine extends StateMachine {
+        private static final int TRANSACTION_ID_IGNORE = 0;
+
+        private DefaultState mDefaultState = new DefaultState();
+        private WaitState mWaitState = new WaitState();
+        private WaitForResponseState mWaitForResponseState = new WaitForResponseState();
+
+        private short mNextTransactionId = 1;
+        public int mNextSessionId = 1;
+
+        private Message mCurrentCommand;
+        private short mCurrentTransactionId = TRANSACTION_ID_IGNORE;
+
+        WifiNanStateMachine(String name, Looper looper) {
+            super(name, looper);
+
+            addState(mDefaultState);
+            /* --> */ addState(mWaitState, mDefaultState);
+            /* --> */ addState(mWaitForResponseState, mDefaultState);
+
+            setInitialState(mWaitState);
         }
 
-        @Override
-        public void handleMessage(Message msg) {
-            if (DBG) {
-                Log.d(TAG, "Message: " + msg.what);
+        private class DefaultState extends State {
+            @Override
+            public boolean processMessage(Message msg) {
+                if (VDBG) {
+                    Log.v(TAG, getName() + msg.toString());
+                }
+
+                switch (msg.what) {
+                    case MESSAGE_TYPE_NOTIFICATION:
+                        processNotification(msg);
+                        return HANDLED;
+                    default:
+                        /* fall-through */
+                }
+
+                Log.wtf(TAG,
+                        "DefaultState: should not get non-NOTIFICATION in this state: msg=" + msg);
+                return NOT_HANDLED;
             }
-            switch (msg.what) {
-                case MESSAGE_CONNECT: {
-                    if (VDBG) {
-                        Log.d(TAG, "NAN connection request received");
-                    }
-                    connectLocal(msg.arg1, (IWifiNanEventCallback) msg.obj);
-                    break;
-                }
-                case MESSAGE_DISCONNECT: {
-                    if (VDBG) {
-                        Log.d(TAG, "NAN disconnection request received");
-                    }
-                    disconnectLocal(msg.arg1);
-                    break;
-                }
-                case MESSAGE_REQUEST_CONFIG: {
-                    if (VDBG) {
-                        Log.d(TAG, "NAN configuration request received");
-                    }
-                    requestConfigLocal(msg.arg1, (ConfigRequest) msg.obj);
-                    break;
-                }
-                case MESSAGE_TERMINATE_SESSION: {
-                    if (VDBG) {
-                        Log.d(TAG, "Terminate session");
-                    }
-                    terminateSessionLocal(msg.arg1, msg.arg2);
-                    break;
-                }
-                case MESSAGE_PUBLISH: {
-                    Bundle data = msg.getData();
-                    publishLocal(msg.arg1,
-                            (PublishConfig) data.getParcelable(MESSAGE_BUNDLE_KEY_CONFIG),
-                            (IWifiNanSessionCallback) msg.obj);
-                    break;
-                }
-                case MESSAGE_UPDATE_PUBLISH: {
-                    updatePublishLocal(msg.arg1, msg.arg2, (PublishConfig) msg.obj);
-                    break;
-                }
-                case MESSAGE_SUBSCRIBE: {
-                    Bundle data = msg.getData();
-                    subscribeLocal(msg.arg1,
-                            (SubscribeConfig) data.getParcelable(MESSAGE_BUNDLE_KEY_CONFIG),
-                            (IWifiNanSessionCallback) msg.obj);
-                    break;
-                }
-                case MESSAGE_UPDATE_SUBSCRIBE: {
-                    updateSubscribeLocal(msg.arg1, msg.arg2, (SubscribeConfig) msg.obj);
-                    break;
-                }
-                case MESSAGE_SEND_MESSAGE: {
-                    Bundle data = msg.getData();
-                    int sessionId = msg.getData().getInt(MESSAGE_BUNDLE_KEY_SESSION_ID);
-                    int peerId = data.getInt(MESSAGE_BUNDLE_KEY_MESSAGE_PEER_ID);
-                    byte[] message = data.getByteArray(MESSAGE_BUNDLE_KEY_MESSAGE);
-                    int messageId = data.getInt(MESSAGE_BUNDLE_KEY_MESSAGE_ID);
+        }
 
-                    if (VDBG) {
-                        Log.d(TAG, "Send Message: message='" + message + "' (ID=" + messageId
-                                + ") to peerId=" + peerId);
-                    }
+        private class WaitState extends State {
+            @Override
+            public boolean processMessage(Message msg) {
+                if (VDBG) {
+                    Log.v(TAG, getName() + msg.toString());
+                }
 
-                    sendFollowonMessageLocal(msg.arg1, sessionId, peerId, message, msg.arg2,
-                            messageId);
+                switch (msg.what) {
+                    case MESSAGE_TYPE_COMMAND:
+                        if (processCommand(msg)) {
+                            transitionTo(mWaitForResponseState);
+                        }
+                        return HANDLED;
+                    case MESSAGE_TYPE_RESPONSE:
+                        /* fall-through */
+                    case MESSAGE_TYPE_TIMEOUT:
+                        /*
+                         * remnants/delayed/out-of-sync messages - but let
+                         * WaitForResponseState deal with them (identified as
+                         * out-of-date by transaction ID).
+                         */
+                        deferMessage(msg);
+                        return HANDLED;
+                    default:
+                        /* fall-through */
+                }
+
+                return NOT_HANDLED;
+            }
+        }
+
+        private class WaitForResponseState extends State {
+            @Override
+            public boolean processMessage(Message msg) {
+                if (VDBG) {
+                    Log.v(TAG, getName() + msg.toString());
+                }
+
+                switch (msg.what) {
+                    case MESSAGE_TYPE_COMMAND:
+                        /*
+                         * don't want COMMANDs in this state - defer until back
+                         * in WaitState
+                         */
+                        deferMessage(msg);
+                        return HANDLED;
+                    case MESSAGE_TYPE_RESPONSE:
+                        if (msg.arg2 == mCurrentTransactionId) {
+                            processResponse(msg);
+                            transitionTo(mWaitState);
+                        } else {
+                            Log.w(TAG,
+                                    "WaitForResponseState: processMessage: non-matching "
+                                            + "transaction ID on RESPONSE (a very late "
+                                            + "response) -- msg=" + msg);
+                            /* no transition */
+                        }
+                        return HANDLED;
+                    case MESSAGE_TYPE_TIMEOUT:
+                        if (msg.arg2 == mCurrentTransactionId) {
+                            processTimeout(msg);
+                            transitionTo(mWaitState);
+                        } else {
+                            Log.w(TAG, "WaitForResponseState: processMessage: non-matching "
+                                    + "transaction ID on TIMEOUT (either a non-cancelled "
+                                    + "timeout or a race condition with cancel) -- msg=" + msg);
+                            /* no transition */
+                        }
+                        return HANDLED;
+                    default:
+                        /* fall-through */
+                }
+
+                return NOT_HANDLED;
+            }
+        }
+
+        private void processNotification(Message msg) {
+            if (VDBG) {
+                Log.v(TAG, "processNotification: msg=" + msg);
+            }
+
+            switch (msg.arg1) {
+                case NOTIFICATION_TYPE_CAPABILITIES_UPDATED: {
+                    WifiNanNative.Capabilities capabilities = (WifiNanNative.Capabilities) msg.obj;
+                    onCapabilitiesUpdatedLocal(capabilities);
                     break;
                 }
-                case MESSAGE_ON_CAPABILITIES_UPDATED:
-                    onCapabilitiesUpdatedLocal((short) msg.arg1,
-                            (WifiNanNative.Capabilities) msg.obj);
+                case NOTIFICATION_TYPE_INTERFACE_CHANGE: {
+                    byte[] mac = (byte[]) msg.obj;
+
+                    onInterfaceAddressChangeLocal(mac);
                     break;
-                case MESSAGE_ON_CONFIG_COMPLETED:
-                    onConfigCompletedLocal((short) msg.arg1);
+                }
+                case NOTIFICATION_TYPE_CLUSTER_CHANGE: {
+                    int flag = msg.arg2;
+                    byte[] clusterId = (byte[]) msg.obj;
+
+                    onClusterChangeLocal(flag, clusterId);
                     break;
-                case MESSAGE_ON_CONFIG_FAILED:
-                    onConfigFailedLocal((short) msg.arg1, msg.arg2);
-                    break;
-                case MESSAGE_ON_NAN_DOWN:
-                    onNanDownLocal(msg.arg1);
-                    break;
-                case MESSAGE_ON_INTERFACE_CHANGE:
-                    onInterfaceAddressChangeLocal((byte[]) msg.obj);
-                    break;
-                case MESSAGE_ON_CLUSTER_CHANGE:
-                    onClusterChangeLocal(msg.arg1, (byte[]) msg.obj);
-                    break;
-                case MESSAGE_ON_PUBLISH_SUCCESS:
-                    onPublishSuccessLocal((short) msg.arg1, msg.arg2);
-                    break;
-                case MESSAGE_ON_PUBLISH_FAIL:
-                    onPublishFailLocal((short) msg.arg1, msg.arg2);
-                    break;
-                case MESSAGE_ON_PUBLISH_TERMINATED:
-                    onPublishTerminatedLocal(msg.arg1, msg.arg2);
-                    break;
-                case MESSAGE_ON_SUBSCRIBE_SUCCESS:
-                    onSubscribeSuccessLocal((short) msg.arg1, msg.arg2);
-                    break;
-                case MESSAGE_ON_SUBSCRIBE_FAIL:
-                    onSubscribeFailLocal((short) msg.arg1, msg.arg2);
-                    break;
-                case MESSAGE_ON_SUBSCRIBE_TERMINATED:
-                    onSubscribeTerminatedLocal(msg.arg1, msg.arg2);
-                    break;
-                case MESSAGE_ON_MESSAGE_SEND_SUCCESS:
-                    onMessageSendSuccessLocal((short) msg.arg1);
-                    break;
-                case MESSAGE_ON_MESSAGE_SEND_FAIL:
-                    onMessageSendFailLocal((short) msg.arg1, msg.arg2);
-                    break;
-                case MESSAGE_ON_UNKNOWN_TRANSACTION:
-                    onUnknownTransactionLocal(
-                            msg.getData().getInt(MESSAGE_BUNDLE_KEY_RESPONSE_TYPE),
-                            (short) msg.arg1, msg.arg2);
-                    break;
-                case MESSAGE_ON_NO_OP_TRANSACTION:
-                    onNoOpNotificationLocal((short) msg.arg1);
-                    break;
-                case MESSAGE_ON_MATCH: {
-                    int pubSubId = msg.arg1;
-                    int requestorInstanceId = msg.arg2;
+                }
+                case NOTIFICATION_TYPE_MATCH: {
+                    int pubSubId = msg.arg2;
+                    int requestorInstanceId = msg.getData()
+                            .getInt(MESSAGE_BUNDLE_KEY_REQ_INSTANCE_ID);
                     byte[] peerMac = msg.getData().getByteArray(MESSAGE_BUNDLE_KEY_MAC_ADDRESS);
                     byte[] serviceSpecificInfo = msg.getData()
                             .getByteArray(MESSAGE_BUNDLE_KEY_SSI_DATA);
@@ -595,207 +591,310 @@
                             .getInt(MESSAGE_BUNDLE_KEY_SSI_LENGTH);
                     byte[] matchFilter = msg.getData().getByteArray(MESSAGE_BUNDLE_KEY_FILTER_DATA);
                     int matchFilterLength = msg.getData().getInt(MESSAGE_BUNDLE_KEY_FILTER_LENGTH);
+
                     onMatchLocal(pubSubId, requestorInstanceId, peerMac, serviceSpecificInfo,
                             serviceSpecificInfoLength, matchFilter, matchFilterLength);
                     break;
                 }
-                case MESSAGE_ON_MESSAGE_RECEIVED: {
-                    int pubSubId = msg.arg1;
-                    int requestorInstanceId = msg.arg2;
+                case NOTIFICATION_TYPE_SESSION_TERMINATED: {
+                    int pubSubId = msg.arg2;
+                    int reason = (Integer) msg.obj;
+                    boolean isPublish = msg.getData().getBoolean(MESSAGE_BUNDLE_KEY_SESSION_TYPE);
+
+                    onSessionTerminatedLocal(pubSubId, isPublish, reason);
+                    break;
+                }
+                case NOTIFICATION_TYPE_MESSAGE_RECEIVED: {
+                    int pubSubId = msg.arg2;
+                    int requestorInstanceId = (Integer) msg.obj;
                     byte[] peerMac = msg.getData().getByteArray(MESSAGE_BUNDLE_KEY_MAC_ADDRESS);
                     byte[] message = msg.getData().getByteArray(MESSAGE_BUNDLE_KEY_MESSAGE_DATA);
                     int messageLength = msg.getData().getInt(MESSAGE_BUNDLE_KEY_MESSAGE_LENGTH);
+
                     onMessageReceivedLocal(pubSubId, requestorInstanceId, peerMac, message,
                             messageLength);
                     break;
                 }
+                case NOTIFICATION_TYPE_NAN_DOWN: {
+                    int reason = msg.arg2;
+
+                    onNanDownLocal(reason);
+
+                    /*
+                     * Possibly also clear out all remaining Messages in queue
+                     * but they should also function correctly after this
+                     * (primarily a NOP).
+                     */
+                    mCurrentCommand = null;
+                    mCurrentTransactionId = TRANSACTION_ID_IGNORE;
+                    transitionTo(mWaitState);
+                    break;
+                }
                 default:
-                    Log.e(TAG, "Unknown message code: " + msg.what);
+                    Log.wtf(TAG, "processNotification: this isn't a NOTIFICATION -- msg=" + msg);
+                    return;
             }
         }
+
+        /**
+         * Execute the command specified by the input Message. Returns a true if
+         * need to wait for a RESPONSE, otherwise a false. We may not have to
+         * wait for a RESPONSE if there was an error in the state (so no command
+         * is sent to HAL) OR if we choose not to wait for response - e.g. for
+         * disconnected/terminate commands failure is not possible.
+         */
+        private boolean processCommand(Message msg) {
+            if (VDBG) {
+                Log.v(TAG, "processCommand: msg=" + msg);
+            }
+
+            if (mCurrentCommand != null) {
+                Log.wtf(TAG,
+                        "processCommand: receiving a command (msg=" + msg
+                                + ") but current (previous) command isn't null (prev_msg="
+                                + mCurrentCommand + ")");
+                mCurrentCommand = null;
+            }
+
+            mCurrentTransactionId = mNextTransactionId++;
+
+            boolean waitForResponse = true;
+
+            switch (msg.arg1) {
+                case COMMAND_TYPE_CONNECT: {
+                    int clientId = msg.arg2;
+                    IWifiNanEventCallback callback = (IWifiNanEventCallback) msg.obj;
+                    ConfigRequest configRequest = (ConfigRequest) msg.getData()
+                            .getParcelable(MESSAGE_BUNDLE_KEY_CONFIG);
+
+                    waitForResponse = connectLocal(mCurrentTransactionId, clientId, callback,
+                            configRequest);
+                    break;
+                }
+                case COMMAND_TYPE_DISCONNECT: {
+                    int clientId = msg.arg2;
+
+                    waitForResponse = disconnectLocal(mCurrentTransactionId, clientId);
+                    break;
+                }
+                case COMMAND_TYPE_TERMINATE_SESSION: {
+                    int clientId = msg.arg2;
+                    int sessionId = (Integer) msg.obj;
+
+                    terminateSessionLocal(clientId, sessionId);
+                    waitForResponse = false;
+                    break;
+                }
+                case COMMAND_TYPE_PUBLISH: {
+                    int clientId = msg.arg2;
+                    IWifiNanSessionCallback callback = (IWifiNanSessionCallback) msg.obj;
+                    PublishConfig publishConfig = (PublishConfig) msg.getData()
+                            .getParcelable(MESSAGE_BUNDLE_KEY_CONFIG);
+
+                    waitForResponse = publishLocal(mCurrentTransactionId, clientId, publishConfig,
+                            callback);
+                    break;
+                }
+                case COMMAND_TYPE_UPDATE_PUBLISH: {
+                    int clientId = msg.arg2;
+                    int sessionId = msg.getData().getInt(MESSAGE_BUNDLE_KEY_SESSION_ID);
+                    PublishConfig publishConfig = (PublishConfig) msg.obj;
+
+                    waitForResponse = updatePublishLocal(mCurrentTransactionId, clientId, sessionId,
+                            publishConfig);
+                    break;
+                }
+                case COMMAND_TYPE_SUBSCRIBE: {
+                    int clientId = msg.arg2;
+                    IWifiNanSessionCallback callback = (IWifiNanSessionCallback) msg.obj;
+                    SubscribeConfig subscribeConfig = (SubscribeConfig) msg.getData()
+                            .getParcelable(MESSAGE_BUNDLE_KEY_CONFIG);
+
+                    waitForResponse = subscribeLocal(mCurrentTransactionId, clientId,
+                            subscribeConfig, callback);
+                    break;
+                }
+                case COMMAND_TYPE_UPDATE_SUBSCRIBE: {
+                    int clientId = msg.arg2;
+                    int sessionId = msg.getData().getInt(MESSAGE_BUNDLE_KEY_SESSION_ID);
+                    SubscribeConfig subscribeConfig = (SubscribeConfig) msg.obj;
+
+                    waitForResponse = updateSubscribeLocal(mCurrentTransactionId, clientId,
+                            sessionId, subscribeConfig);
+                    break;
+                }
+                case COMMAND_TYPE_SEND_MESSAGE: {
+                    Bundle data = msg.getData();
+
+                    int clientId = msg.arg2;
+                    int sessionId = msg.getData().getInt(MESSAGE_BUNDLE_KEY_SESSION_ID);
+                    int peerId = data.getInt(MESSAGE_BUNDLE_KEY_MESSAGE_PEER_ID);
+                    byte[] message = data.getByteArray(MESSAGE_BUNDLE_KEY_MESSAGE);
+                    int messageId = data.getInt(MESSAGE_BUNDLE_KEY_MESSAGE_ID);
+                    int messageLength = data.getInt(MESSAGE_BUNDLE_KEY_MESSAGE_LENGTH);
+
+                    waitForResponse = sendFollowonMessageLocal(mCurrentTransactionId, clientId,
+                            sessionId, peerId, message, messageLength, messageId);
+                    break;
+                }
+                default:
+                    waitForResponse = false;
+                    Log.wtf(TAG, "processCommand: this isn't a COMMAND -- msg=" + msg);
+                    /* fall-through */
+            }
+
+            if (!waitForResponse) {
+                mCurrentTransactionId = TRANSACTION_ID_IGNORE;
+            } else {
+                mCurrentCommand = obtainMessage(msg.what);
+                mCurrentCommand.copyFrom(msg);
+                // TODO: create a TIMEOUT signal b/28021222
+            }
+
+            return waitForResponse;
+        }
+
+        private void processResponse(Message msg) {
+            if (VDBG) {
+                Log.v(TAG, "processResponse: msg=" + msg);
+            }
+
+            if (mCurrentCommand == null) {
+                Log.wtf(TAG, "processResponse: no existing command stored!? msg=" + msg);
+                mCurrentTransactionId = TRANSACTION_ID_IGNORE;
+                return;
+            }
+
+            switch (msg.arg1) {
+                case RESPONSE_TYPE_ON_CONFIG_SUCCESS:
+                    onConfigCompletedLocal(mCurrentCommand);
+                    break;
+                case RESPONSE_TYPE_ON_CONFIG_FAIL: {
+                    int reason = (Integer) msg.obj;
+
+                    onConfigFailedLocal(mCurrentCommand, reason);
+                    break;
+                }
+                case RESPONSE_TYPE_ON_SESSION_CONFIG_SUCCESS: {
+                    int pubSubId = (Integer) msg.obj;
+                    boolean isPublish = msg.getData().getBoolean(MESSAGE_BUNDLE_KEY_SESSION_TYPE);
+
+                    onSessionConfigSuccessLocal(mCurrentCommand, pubSubId, isPublish);
+                    break;
+                }
+                case RESPONSE_TYPE_ON_SESSION_CONFIG_FAIL: {
+                    int reason = (Integer) msg.obj;
+                    boolean isPublish = msg.getData().getBoolean(MESSAGE_BUNDLE_KEY_SESSION_TYPE);
+
+                    onSessionConfigFailLocal(mCurrentCommand, isPublish, reason);
+                    break;
+                }
+                case RESPONSE_TYPE_ON_MESSAGE_SEND_SUCCESS:
+                    onMessageSendSuccessLocal(mCurrentCommand);
+                    break;
+                case RESPONSE_TYPE_ON_MESSAGE_SEND_FAIL: {
+                    int reason = (Integer) msg.obj;
+
+                    onMessageSendFailLocal(mCurrentCommand, reason);
+                    break;
+                }
+                default:
+                    Log.wtf(TAG, "processResponse: this isn't a RESPONSE -- msg=" + msg);
+                    mCurrentCommand = null;
+                    mCurrentTransactionId = TRANSACTION_ID_IGNORE;
+                    return;
+            }
+
+            mCurrentCommand = null;
+            mCurrentTransactionId = TRANSACTION_ID_IGNORE;
+
+            // TODO: disable TIMEOUT b/28021222
+        }
+
+        private void processTimeout(Message msg) {
+            // TODO: b/28021222
+        }
+
+        @Override
+        public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
+            pw.println("WifiNanStateMachine:");
+            pw.println("  mNextTransactionId: " + mNextTransactionId);
+            pw.println("  mNextSessionId: " + mNextSessionId);
+            pw.println("  mCurrentCommand: " + mCurrentCommand);
+            pw.println("  mCurrentTransaction: " + mCurrentTransactionId);
+            super.dump(fd, pw, args);
+        }
     }
 
     /*
-     * Transaction management classes & operations
+     * COMMANDS
      */
-
-    // non-synchronized (should be ok as long as only used from NanStateManager,
-    // NanClientState, and NanSessionState)
-    /* package */ short createNextTransactionId() {
-        return mNextTransactionId++;
-    }
-
-    private static class TransactionInfoBase {
-        short mTransactionId;
-    }
-
-    private static class TransactionInfoSession extends TransactionInfoBase {
-        public WifiNanClientState mClient;
-        public WifiNanSessionState mSession;
-    }
-
-    private static class TransactionInfoMessage extends TransactionInfoSession {
-        public int mMessageId;
-    }
-
-    private static class TransactionInfoConfig extends TransactionInfoBase {
-        public ConfigRequest mConfig;
-    }
-
-    private void allocateAndRegisterTransactionId(TransactionInfoBase info) {
-        info.mTransactionId = createNextTransactionId();
-
-        mPendingResponses.put(info.mTransactionId, info);
-    }
-
-    private boolean fillInTransactionInfoSession(TransactionInfoSession info, int clientId,
-            int sessionId) {
-        WifiNanClientState client = mClients.get(clientId);
-        if (client == null) {
-            Log.w(TAG, "getAndRegisterTransactionId: no client exists for clientId=" + clientId);
-            return false;
-        }
-        info.mClient = client;
-
-        WifiNanSessionState session = info.mClient.getSession(sessionId);
-        if (session == null) {
-            Log.w(TAG, "getAndRegisterSessionTransactionId: no session exists for clientId="
-                    + clientId + ", sessionId=" + sessionId);
-            return false;
-        }
-        info.mSession = session;
-        return true;
-    }
-
-    private TransactionInfoBase createTransactionInfo() {
-        TransactionInfoBase info = new TransactionInfoBase();
-        allocateAndRegisterTransactionId(info);
-        return info;
-    }
-
-    private TransactionInfoSession createTransactionInfoSession(int clientId, int sessionId) {
-        TransactionInfoSession info = new TransactionInfoSession();
-        if (!fillInTransactionInfoSession(info, clientId, sessionId)) {
-            return null;
-        }
-        allocateAndRegisterTransactionId(info);
-        return info;
-    }
-
-    private TransactionInfoMessage createTransactionInfoMessage(int clientId, int sessionId,
-            int messageId) {
-        TransactionInfoMessage info = new TransactionInfoMessage();
-        if (!fillInTransactionInfoSession(info, clientId, sessionId)) {
-            return null;
-        }
-        info.mMessageId = messageId;
-        allocateAndRegisterTransactionId(info);
-        return info;
-    }
-
-    private TransactionInfoConfig createTransactionInfoConfig(ConfigRequest configRequest) {
-        TransactionInfoConfig info = new TransactionInfoConfig();
-        info.mConfig = configRequest;
-        allocateAndRegisterTransactionId(info);
-        return info;
-    }
-
-    private TransactionInfoBase getAndRemovePendingResponseTransactionInfo(short transactionId) {
-        TransactionInfoBase transInfo = mPendingResponses.get(transactionId);
-        if (transInfo != null) {
-            mPendingResponses.remove(transactionId);
-        }
-
-        return transInfo;
-    }
-
-    private WifiNanSessionState getNanSessionStateForPubSubId(int pubSubId) {
-        for (int i = 0; i < mClients.size(); ++i) {
-            WifiNanSessionState session = mClients.valueAt(i)
-                    .getNanSessionStateForPubSubId(pubSubId);
-            if (session != null) {
-                return session;
-            }
-        }
-
-        return null;
-    }
-
-    /*
-     * Actions (calls from API to service)
-     */
-    private void connectLocal(int clientId, IWifiNanEventCallback callback) {
+    private boolean connectLocal(short transactionId, int clientId, IWifiNanEventCallback callback,
+            ConfigRequest configRequest) {
         if (VDBG) {
-            Log.v(TAG, "connect(): clientId=" + clientId + ", callback=" + callback);
+            Log.v(TAG, "connectLocal(): transactionId=" + transactionId + ", clientId=" + clientId
+                    + ", callback=" + callback + ", configRequest=" + configRequest);
         }
 
         if (mClients.get(clientId) != null) {
-            Log.e(TAG, "connect: entry already exists for clientId=" + clientId);
-            return;
+            Log.e(TAG, "connectLocal: entry already exists for clientId=" + clientId);
         }
 
-        WifiNanClientState client = new WifiNanClientState(clientId, callback);
-        mClients.put(clientId, client);
+        if (mCurrentNanConfiguration != null
+                && !mCurrentNanConfiguration.equalsOnTheAir(configRequest)) {
+            try {
+                callback.onConnectFail(
+                        WifiNanEventCallback.REASON_ALREADY_CONNECTED_INCOMPAT_CONFIG);
+            } catch (RemoteException e) {
+                Log.w(TAG, "connectLocal onConnectFail(): RemoteException (FYI): " + e);
+            }
+            return false;
+        }
+
+        ConfigRequest merged = mergeConfigRequests(configRequest);
+        if (mCurrentNanConfiguration != null && mCurrentNanConfiguration.equals(merged)) {
+            try {
+                callback.onConnectSuccess();
+                mClients.append(clientId,
+                        new WifiNanClientState(clientId, callback, configRequest));
+            } catch (RemoteException e) {
+                Log.w(TAG, "connectLocal onConnectSuccess(): RemoteException (FYI): " + e);
+            }
+            return false;
+        }
+
+        WifiNanNative.getInstance().enableAndConfigure(transactionId, merged,
+                mCurrentNanConfiguration == null);
+        return true;
     }
 
-    private void disconnectLocal(int clientId) {
+    private boolean disconnectLocal(short transactionId, int clientId) {
         if (VDBG) {
-            Log.v(TAG, "disconnect(): clientId=" + clientId);
+            Log.v(TAG,
+                    "disconnectLocal(): transactionId=" + transactionId + ", clientId=" + clientId);
         }
 
         WifiNanClientState client = mClients.get(clientId);
-        mClients.delete(clientId);
-
         if (client == null) {
-            Log.e(TAG, "disconnect: no entry for clientId=" + clientId);
-            return;
+            Log.e(TAG, "disconnectLocal: no entry for clientId=" + clientId);
+            return false;
         }
-
-        List<Integer> toRemove = new ArrayList<>();
-        for (int i = 0; i < mPendingResponses.size(); ++i) {
-            TransactionInfoBase info = mPendingResponses.valueAt(i);
-            if (!(info instanceof TransactionInfoSession)) {
-                continue;
-            }
-            if (((TransactionInfoSession) info).mClient.getClientId() == clientId) {
-                toRemove.add(i);
-            }
-        }
-        for (Integer id : toRemove) {
-            mPendingResponses.removeAt(id);
-        }
-
+        mClients.delete(clientId);
         client.destroy();
 
         if (mClients.size() == 0) {
-            WifiNanNative.getInstance().disable(createTransactionInfo().mTransactionId);
-            return;
+            mCurrentNanConfiguration = null;
+            WifiNanNative.getInstance().disable((short) 0);
+            return false;
         }
 
-        ConfigRequest merged = mergeConfigRequests();
-
-        WifiNanNative.getInstance()
-                .enableAndConfigure(createTransactionInfoConfig(merged).mTransactionId, merged);
-    }
-
-    private void requestConfigLocal(int clientId, ConfigRequest configRequest) {
-        if (VDBG) {
-            Log.v(TAG,
-                    "requestConfig(): clientId=" + clientId + ", configRequest=" + configRequest);
+        ConfigRequest merged = mergeConfigRequests(null);
+        if (merged.equals(mCurrentNanConfiguration)) {
+            return false;
         }
 
-        WifiNanClientState client = mClients.get(clientId);
-        if (client == null) {
-            Log.e(TAG, "requestConfig: no client exists for clientId=" + clientId);
-            return;
-        }
-
-        client.setConfigRequest(configRequest);
-
-        ConfigRequest merged = mergeConfigRequests();
-
-        WifiNanNative.getInstance()
-                .enableAndConfigure(createTransactionInfoConfig(merged).mTransactionId, merged);
+        WifiNanNative.getInstance().enableAndConfigure(transactionId, merged, false);
+        return true;
     }
 
     private void terminateSessionLocal(int clientId, int sessionId) {
@@ -810,202 +909,353 @@
             return;
         }
 
-        List<Integer> toRemove = new ArrayList<>();
-        for (int i = 0; i < mPendingResponses.size(); ++i) {
-            TransactionInfoBase info = mPendingResponses.valueAt(i);
-            if (!(info instanceof TransactionInfoSession)) {
-                continue;
-            }
-            TransactionInfoSession infoSession = (TransactionInfoSession) info;
-            if (infoSession.mClient.getClientId() == clientId
-                    && infoSession.mSession.getSessionId() == sessionId) {
-                toRemove.add(i);
-            }
-        }
-        for (Integer id : toRemove) {
-            mPendingResponses.removeAt(id);
-        }
-
         client.terminateSession(sessionId);
     }
 
-    private void publishLocal(int clientId, PublishConfig publishConfig,
+    private boolean publishLocal(short transactionId, int clientId, PublishConfig publishConfig,
             IWifiNanSessionCallback callback) {
         if (VDBG) {
-            Log.v(TAG, "publishLocal(): clientId=" + clientId + ", config=" + publishConfig
-                    + ", callback=" + callback);
+            Log.v(TAG, "publishLocal(): transactionId=" + transactionId + ", clientId=" + clientId
+                    + ", publishConfig=" + publishConfig + ", callback=" + callback);
         }
 
         WifiNanClientState client = mClients.get(clientId);
         if (client == null) {
             Log.e(TAG, "publishLocal: no client exists for clientId=" + clientId);
-            return;
-        }
-        int sessionId = mNextSessionId++;
-        client.createSession(sessionId, callback, true);
-
-        TransactionInfoSession info = createTransactionInfoSession(clientId, sessionId);
-        if (info == null) {
-            return;
+            return false;
         }
 
-        info.mSession.publish(info.mTransactionId, publishConfig);
+        WifiNanNative.getInstance().publish(transactionId, 0, publishConfig);
+        return true;
     }
 
-    private void updatePublishLocal(int clientId, int sessionId, PublishConfig publishConfig) {
+    private boolean updatePublishLocal(short transactionId, int clientId, int sessionId,
+            PublishConfig publishConfig) {
         if (VDBG) {
-            Log.v(TAG, "updatePublishLocal(): clientId=" + clientId + ", sessionId=" + sessionId
-                    + ", config=" + publishConfig);
+            Log.v(TAG, "updatePublishLocal(): transactionId=" + transactionId + ", clientId="
+                    + clientId + ", sessionId=" + sessionId + ", publishConfig=" + publishConfig);
         }
 
-        TransactionInfoSession info = createTransactionInfoSession(clientId, sessionId);
-        if (info == null) {
-            return;
+        WifiNanClientState client = mClients.get(clientId);
+        if (client == null) {
+            Log.e(TAG, "updatePublishLocal: no client exists for clientId=" + clientId);
+            return false;
         }
 
-        info.mSession.updatePublish(info.mTransactionId, publishConfig);
+        WifiNanSessionState session = client.getSession(sessionId);
+        if (session == null) {
+            Log.e(TAG, "updatePublishLocal: no session exists for clientId=" + clientId
+                    + ", sessionId=" + sessionId);
+            return false;
+        }
+
+        return session.updatePublish(transactionId, publishConfig);
     }
 
-    private void subscribeLocal(int clientId, SubscribeConfig subscribeConfig,
-            IWifiNanSessionCallback callback) {
+    private boolean subscribeLocal(short transactionId, int clientId,
+            SubscribeConfig subscribeConfig, IWifiNanSessionCallback callback) {
         if (VDBG) {
-            Log.v(TAG, "subscribe(): clientId=" + clientId + ", config=" + subscribeConfig
-                    + ", callback=" + callback);
+            Log.v(TAG, "subscribeLocal(): transactionId=" + transactionId + ", clientId=" + clientId
+                    + ", subscribeConfig=" + subscribeConfig + ", callback=" + callback);
         }
 
         WifiNanClientState client = mClients.get(clientId);
         if (client == null) {
             Log.e(TAG, "subscribeLocal: no client exists for clientId=" + clientId);
-            return;
-        }
-        int sessionId = mNextSessionId++;
-        client.createSession(sessionId, callback, false);
-
-        TransactionInfoSession info = createTransactionInfoSession(clientId, sessionId);
-        if (info == null) {
-            return;
+            return false;
         }
 
-        info.mSession.subscribe(info.mTransactionId, subscribeConfig);
+        WifiNanNative.getInstance().subscribe(transactionId, 0, subscribeConfig);
+        return true;
     }
 
-    private void updateSubscribeLocal(int clientId, int sessionId,
+    private boolean updateSubscribeLocal(short transactionId, int clientId, int sessionId,
             SubscribeConfig subscribeConfig) {
         if (VDBG) {
-            Log.v(TAG, "updateSubscribeLocal(): clientId=" + clientId + ", sessionId=" + sessionId
-                    + ", config=" + subscribeConfig);
+            Log.v(TAG,
+                    "updateSubscribeLocal(): transactionId=" + transactionId + ", clientId="
+                            + clientId + ", sessionId=" + sessionId + ", subscribeConfig="
+                            + subscribeConfig);
         }
 
-        TransactionInfoSession info = createTransactionInfoSession(clientId, sessionId);
-        if (info == null) {
-            return;
+        WifiNanClientState client = mClients.get(clientId);
+        if (client == null) {
+            Log.e(TAG, "updateSubscribeLocal: no client exists for clientId=" + clientId);
+            return false;
         }
 
-        info.mSession.updateSubscribe(info.mTransactionId, subscribeConfig);
+        WifiNanSessionState session = client.getSession(sessionId);
+        if (session == null) {
+            Log.e(TAG, "updateSubscribeLocal: no session exists for clientId=" + clientId
+                    + ", sessionId=" + sessionId);
+            return false;
+        }
+
+        return session.updateSubscribe(transactionId, subscribeConfig);
     }
 
-    private void sendFollowonMessageLocal(int clientId, int sessionId, int peerId, byte[] message,
-            int messageLength, int messageId) {
+    private boolean sendFollowonMessageLocal(short transactionId, int clientId, int sessionId,
+            int peerId, byte[] message, int messageLength, int messageId) {
         if (VDBG) {
             Log.v(TAG,
-                    "sendMessage(): clientId=" + clientId + ", sessionId=" + sessionId + ", peerId="
-                            + peerId + ", messageLength=" + messageLength + ", messageId="
-                            + messageId);
+                    "sendFollowonMessageLocal(): transactionId=" + transactionId + ", clientId="
+                            + clientId + ", sessionId=" + sessionId + ", peerId=" + peerId
+                            + ", messageLength=" + messageLength + ", messageId=" + messageId);
         }
 
-        TransactionInfoMessage info = createTransactionInfoMessage(clientId, sessionId, messageId);
-        if (info == null) {
-            return;
+        WifiNanClientState client = mClients.get(clientId);
+        if (client == null) {
+            Log.e(TAG, "sendFollowonMessageLocal: no client exists for clientId=" + clientId);
+            return false;
         }
 
-        info.mSession.sendMessage(info.mTransactionId, peerId, message, messageLength, messageId);
+        WifiNanSessionState session = client.getSession(sessionId);
+        if (session == null) {
+            Log.e(TAG, "sendFollowonMessageLocal: no session exists for clientId=" + clientId
+                    + ", sessionId=" + sessionId);
+            return false;
+        }
+
+        return session.sendMessage(transactionId, peerId, message, messageLength, messageId);
     }
 
     /*
-     * Callbacks (calls from HAL/Native to service)
+     * RESPONSES
      */
 
-    private void onCapabilitiesUpdatedLocal(short transactionId,
-            WifiNanNative.Capabilities capabilities) {
+    private void onConfigCompletedLocal(Message completedCommand) {
         if (VDBG) {
-            Log.v(TAG, "onCapabilitiesUpdatedLocal: transactionId=" + transactionId
-                    + ", capabilites=" + capabilities);
+            Log.v(TAG, "onConfigCompleted: completedCommand=" + completedCommand);
+        }
+
+        if (completedCommand.arg1 == COMMAND_TYPE_CONNECT) {
+            Bundle data = completedCommand.getData();
+
+            int clientId = completedCommand.arg2;
+            IWifiNanEventCallback callback = (IWifiNanEventCallback) completedCommand.obj;
+            ConfigRequest configRequest = (ConfigRequest) data
+                    .getParcelable(MESSAGE_BUNDLE_KEY_CONFIG);
+
+            mClients.put(clientId, new WifiNanClientState(clientId, callback, configRequest));
+            try {
+                callback.onConnectSuccess();
+            } catch (RemoteException e) {
+                Log.w(TAG,
+                        "onConfigCompletedLocal onConnectSuccess(): RemoteException (FYI): " + e);
+            }
+        } else if (completedCommand.arg1 == COMMAND_TYPE_DISCONNECT) {
+            /*
+             * NOP (i.e. updated configuration after disconnecting a client)
+             */
+        } else {
+            Log.wtf(TAG, "onConfigCompletedLocal: unexpected completedCommand=" + completedCommand);
+            return;
+        }
+
+        mCurrentNanConfiguration = mergeConfigRequests(null);
+    }
+
+    private void onConfigFailedLocal(Message failedCommand, int reason) {
+        if (VDBG) {
+            Log.v(TAG,
+                    "onConfigFailedLocal: failedCommand=" + failedCommand + ", reason=" + reason);
+        }
+
+        if (failedCommand.arg1 == COMMAND_TYPE_CONNECT) {
+            IWifiNanEventCallback callback = (IWifiNanEventCallback) failedCommand.obj;
+
+            try {
+                callback.onConnectFail(reason);
+            } catch (RemoteException e) {
+                Log.w(TAG, "onConfigFailedLocal onConnectFail(): RemoteException (FYI): " + e);
+            }
+        } else if (failedCommand.arg1 == COMMAND_TYPE_DISCONNECT) {
+            /*
+             * NOP (tried updating configuration after disconnecting a client -
+             * shouldn't fail but there's nothing to do - the old configuration
+             * is still up-and-running).
+             */
+        } else {
+            Log.wtf(TAG, "onConfigFailedLocal: unexpected failedCommand=" + failedCommand);
+            return;
+        }
+
+    }
+
+    private void onSessionConfigSuccessLocal(Message completedCommand, int pubSubId,
+            boolean isPublish) {
+        if (VDBG) {
+            Log.v(TAG, "onSessionConfigSuccessLocal: completedCommand=" + completedCommand
+                    + ", pubSubId=" + pubSubId + ", isPublish=" + isPublish);
+        }
+
+        if (completedCommand.arg1 == COMMAND_TYPE_PUBLISH
+                || completedCommand.arg1 == COMMAND_TYPE_SUBSCRIBE) {
+            int clientId = completedCommand.arg2;
+            IWifiNanSessionCallback callback = (IWifiNanSessionCallback) completedCommand.obj;
+
+            WifiNanClientState client = mClients.get(clientId);
+            if (client == null) {
+                Log.e(TAG,
+                        "onSessionConfigSuccessLocal: no client exists for clientId=" + clientId);
+                return;
+            }
+
+            int sessionId = mSm.mNextSessionId++;
+            try {
+                callback.onSessionStarted(sessionId);
+            } catch (RemoteException e) {
+                Log.e(TAG, "onSessionConfigSuccessLocal: onSessionStarted() RemoteException=" + e);
+                return;
+            }
+
+            WifiNanSessionState session = new WifiNanSessionState(sessionId, pubSubId, callback,
+                    isPublish);
+            client.addSession(session);
+        } else if (completedCommand.arg1 == COMMAND_TYPE_UPDATE_PUBLISH
+                || completedCommand.arg1 == COMMAND_TYPE_UPDATE_SUBSCRIBE) {
+            int clientId = completedCommand.arg2;
+            int sessionId = completedCommand.getData().getInt(MESSAGE_BUNDLE_KEY_SESSION_ID);
+
+            WifiNanClientState client = mClients.get(clientId);
+            if (client == null) {
+                Log.e(TAG,
+                        "onSessionConfigSuccessLocal: no client exists for clientId=" + clientId);
+                return;
+            }
+
+            WifiNanSessionState session = client.getSession(sessionId);
+            if (session == null) {
+                Log.e(TAG, "onSessionConfigSuccessLocal: no session exists for clientId=" + clientId
+                        + ", sessionId=" + sessionId);
+                return;
+            }
+
+            try {
+                session.getCallback().onSessionConfigSuccess();
+            } catch (RemoteException e) {
+                Log.e(TAG, "onSessionConfigSuccessLocal: onSessionConfigSuccess() RemoteException="
+                        + e);
+            }
+        } else {
+            Log.wtf(TAG,
+                    "onSessionConfigSuccessLocal: unexpected completedCommand=" + completedCommand);
+        }
+    }
+
+    private void onSessionConfigFailLocal(Message failedCommand, boolean isPublish, int reason) {
+        if (VDBG) {
+            Log.v(TAG, "onSessionConfigFailLocal: failedCommand=" + failedCommand + ", isPublish="
+                    + isPublish + ", reason=" + reason);
+        }
+
+        if (failedCommand.arg1 == COMMAND_TYPE_PUBLISH
+                || failedCommand.arg1 == COMMAND_TYPE_SUBSCRIBE) {
+            IWifiNanSessionCallback callback = (IWifiNanSessionCallback) failedCommand.obj;
+            try {
+                callback.onSessionConfigFail(reason);
+            } catch (RemoteException e) {
+                Log.w(TAG, "onSessionConfigFailLocal onSessionConfigFail(): RemoteException (FYI): "
+                        + e);
+            }
+        } else if (failedCommand.arg1 == COMMAND_TYPE_UPDATE_PUBLISH
+                || failedCommand.arg1 == COMMAND_TYPE_UPDATE_SUBSCRIBE) {
+            int clientId = failedCommand.arg2;
+            int sessionId = failedCommand.getData().getInt(MESSAGE_BUNDLE_KEY_SESSION_ID);
+
+            WifiNanClientState client = mClients.get(clientId);
+            if (client == null) {
+                Log.e(TAG, "onSessionConfigFailLocal: no client exists for clientId=" + clientId);
+                return;
+            }
+
+            WifiNanSessionState session = client.getSession(sessionId);
+            if (session == null) {
+                Log.e(TAG, "onSessionConfigFailLocal: no session exists for clientId=" + clientId
+                        + ", sessionId=" + sessionId);
+                return;
+            }
+
+            try {
+                session.getCallback().onSessionConfigFail(reason);
+            } catch (RemoteException e) {
+                Log.e(TAG, "onSessionConfigFailLocal: onSessionConfigFail() RemoteException=" + e);
+            }
+        } else {
+            Log.wtf(TAG, "onSessionConfigFailLocal: unexpected failedCommand=" + failedCommand);
+        }
+    }
+
+    private void onMessageSendSuccessLocal(Message completedCommand) {
+        if (VDBG) {
+            Log.v(TAG, "onMessageSendSuccess: completedCommand=" + completedCommand);
+        }
+
+        int clientId = completedCommand.arg2;
+        int sessionId = completedCommand.getData().getInt(MESSAGE_BUNDLE_KEY_SESSION_ID);
+        int messageId = completedCommand.getData().getInt(MESSAGE_BUNDLE_KEY_MESSAGE_ID);
+
+        WifiNanClientState client = mClients.get(clientId);
+        if (client == null) {
+            Log.e(TAG, "onMessageSendSuccessLocal: no client exists for clientId=" + clientId);
+            return;
+        }
+
+        WifiNanSessionState session = client.getSession(sessionId);
+        if (session == null) {
+            Log.e(TAG, "onMessageSendSuccessLocal: no session exists for clientId=" + clientId
+                    + ", sessionId=" + sessionId);
+            return;
+        }
+
+        try {
+            session.getCallback().onMessageSendSuccess(messageId);
+        } catch (RemoteException e) {
+            Log.w(TAG, "onMessageSendSuccessLocal: RemoteException (FYI): " + e);
+        }
+    }
+
+    private void onMessageSendFailLocal(Message failedCommand, int reason) {
+        if (VDBG) {
+            Log.v(TAG, "onMessageSendFail: failedCommand=" + failedCommand + ", reason=" + reason);
+        }
+
+        int clientId = failedCommand.arg2;
+        int sessionId = failedCommand.getData().getInt(MESSAGE_BUNDLE_KEY_SESSION_ID);
+        int messageId = failedCommand.getData().getInt(MESSAGE_BUNDLE_KEY_MESSAGE_ID);
+
+        WifiNanClientState client = mClients.get(clientId);
+        if (client == null) {
+            Log.e(TAG, "onMessageSendFailLocal: no client exists for clientId=" + clientId);
+            return;
+        }
+
+        WifiNanSessionState session = client.getSession(sessionId);
+        if (session == null) {
+            Log.e(TAG, "onMessageSendFailLocal: no session exists for clientId=" + clientId
+                    + ", sessionId=" + sessionId);
+            return;
+        }
+
+        try {
+            session.getCallback().onMessageSendFail(messageId, reason);
+        } catch (RemoteException e) {
+            Log.e(TAG, "onMessageSendFailLocal: onMessageSendFail RemoteException=" + e);
+        }
+    }
+
+    /*
+     * NOTIFICATIONS
+     */
+
+    private void onCapabilitiesUpdatedLocal(WifiNanNative.Capabilities capabilities) {
+        if (VDBG) {
+            Log.v(TAG, "onCapabilitiesUpdatedLocal: capabilites=" + capabilities);
         }
 
         mCapabilities = capabilities;
     }
 
-    private void onConfigCompletedLocal(short transactionId) {
-        if (VDBG) {
-            Log.v(TAG, "onConfigCompleted: transactionId=" + transactionId);
-        }
-
-        TransactionInfoBase info = getAndRemovePendingResponseTransactionInfo(transactionId);
-        if (info == null) {
-            Log.e(TAG, "onConfigCompleted: no transaction info for transactionId=" + transactionId);
-            return;
-        }
-        if (!(info instanceof TransactionInfoConfig)) {
-            Log.e(TAG, "onConfigCompleted: invalid info structure stored for transactionId="
-                    + transactionId);
-            return;
-        }
-        TransactionInfoConfig infoConfig = (TransactionInfoConfig) info;
-
-        if (DBG) {
-            Log.d(TAG, "onConfigCompleted: request=" + infoConfig.mConfig);
-        }
-
-        for (int i = 0; i < mClients.size(); ++i) {
-            WifiNanClientState client = mClients.valueAt(i);
-            client.onConfigCompleted(infoConfig.mConfig);
-        }
-    }
-
-    private void onConfigFailedLocal(short transactionId, int reason) {
-        if (VDBG) {
-            Log.v(TAG, "onEnableFailed: transactionId=" + transactionId + ", reason=" + reason);
-        }
-
-        TransactionInfoBase info = getAndRemovePendingResponseTransactionInfo(transactionId);
-        if (info == null) {
-            Log.e(TAG, "onConfigFailed: no transaction info for transactionId=" + transactionId);
-            return;
-        }
-        if (!(info instanceof TransactionInfoConfig)) {
-            Log.e(TAG, "onConfigCompleted: invalid info structure stored for transactionId="
-                    + transactionId);
-            return;
-        }
-        TransactionInfoConfig infoConfig = (TransactionInfoConfig) info;
-
-        if (DBG) {
-            Log.d(TAG, "onConfigFailed: request=" + infoConfig.mConfig);
-        }
-
-        for (int i = 0; i < mClients.size(); ++i) {
-            WifiNanClientState client = mClients.valueAt(i);
-            client.onConfigFailed(infoConfig.mConfig, reason);
-        }
-    }
-
-    private void onNanDownLocal(int reason) {
-        if (VDBG) {
-            Log.v(TAG, "onNanDown: reason=" + reason);
-        }
-
-        int interested = 0;
-        for (int i = 0; i < mClients.size(); ++i) {
-            WifiNanClientState client = mClients.valueAt(i);
-            interested += client.onNanDown(reason);
-        }
-
-        if (interested == 0) {
-            Log.e(TAG, "onNanDown: event received but no callbacks registered for this event "
-                    + "- should be disabled from fw!");
-        }
-    }
-
     private void onInterfaceAddressChangeLocal(byte[] mac) {
         if (VDBG) {
             Log.v(TAG, "onInterfaceAddressChange: mac=" + String.valueOf(HexEncoding.encode(mac)));
@@ -1041,185 +1291,6 @@
         }
     }
 
-    private void onPublishSuccessLocal(short transactionId, int publishId) {
-        if (VDBG) {
-            Log.v(TAG, "onPublishSuccess: transactionId=" + transactionId + ", publishId="
-                    + publishId);
-        }
-
-        TransactionInfoBase info = getAndRemovePendingResponseTransactionInfo(transactionId);
-        if (info == null) {
-            Log.w(TAG, "onPublishSuccess(): no info registered for transactionId=" + transactionId);
-            WifiNanNative.getInstance().stopPublish(createNextTransactionId(), publishId);
-            return;
-        }
-        if (!(info instanceof TransactionInfoSession)) {
-            Log.e(TAG, "onPublishSuccess: invalid info structure stored for transactionId="
-                    + transactionId);
-            return;
-        }
-
-        TransactionInfoSession infoSession = (TransactionInfoSession) info;
-
-        infoSession.mSession.onPublishSuccess(publishId);
-    }
-
-    private void onPublishFailLocal(short transactionId, int status) {
-        if (VDBG) {
-            Log.v(TAG, "onPublishFail: transactionId=" + transactionId + ", status=" + status);
-        }
-
-        TransactionInfoBase info = getAndRemovePendingResponseTransactionInfo(transactionId);
-        if (info == null) {
-            Log.e(TAG, "onPublishFail(): no info registered for transactionId=" + transactionId);
-            return;
-        }
-        if (!(info instanceof TransactionInfoSession)) {
-            Log.e(TAG, "onPublishFail: invalid info structure stored for transactionId="
-                    + transactionId);
-            return;
-        }
-        TransactionInfoSession infoSession = (TransactionInfoSession) info;
-
-        infoSession.mSession.onPublishFail(status);
-    }
-
-    private void onPublishTerminatedLocal(int publishId, int status) {
-        if (VDBG) {
-            Log.v(TAG, "onPublishTerminated: publishId=" + publishId + ", status=" + status);
-        }
-
-        WifiNanSessionState session = getNanSessionStateForPubSubId(publishId);
-        if (session == null) {
-            Log.e(TAG, "onPublishTerminated: no session found for publishId=" + publishId);
-            return;
-        }
-
-        session.onPublishTerminated(status);
-    }
-
-    private void onSubscribeSuccessLocal(short transactionId, int subscribeId) {
-        if (VDBG) {
-            Log.v(TAG, "onSubscribeSuccess: transactionId=" + transactionId + ", subscribeId="
-                    + subscribeId);
-        }
-
-        TransactionInfoBase info = getAndRemovePendingResponseTransactionInfo(transactionId);
-        if (info == null) {
-            Log.w(TAG,
-                    "onSubscribeSuccess(): no info registered for transactionId=" + transactionId);
-            WifiNanNative.getInstance().stopSubscribe(createNextTransactionId(), subscribeId);
-            return;
-        }
-        if (!(info instanceof TransactionInfoSession)) {
-            Log.e(TAG, "onSubscribeSuccess: invalid info structure stored for transactionId="
-                    + transactionId);
-            return;
-        }
-        TransactionInfoSession infoSession = (TransactionInfoSession) info;
-
-        infoSession.mSession.onSubscribeSuccess(subscribeId);
-    }
-
-    private void onSubscribeFailLocal(short transactionId, int status) {
-        if (VDBG) {
-            Log.v(TAG, "onSubscribeFail: transactionId=" + transactionId + ", status=" + status);
-        }
-
-        TransactionInfoBase info = getAndRemovePendingResponseTransactionInfo(transactionId);
-        if (info == null) {
-            Log.e(TAG, "onSubscribeFail(): no info registered for transactionId=" + transactionId);
-            return;
-        }
-        if (!(info instanceof TransactionInfoSession)) {
-            Log.e(TAG, "onSubscribeFail: invalid info structure stored for transactionId="
-                    + transactionId);
-            return;
-        }
-        TransactionInfoSession infoSession = (TransactionInfoSession) info;
-
-        infoSession.mSession.onSubscribeFail(status);
-    }
-
-    private void onSubscribeTerminatedLocal(int subscribeId, int status) {
-        if (VDBG) {
-            Log.v(TAG, "onPublishTerminated: subscribeId=" + subscribeId + ", status=" + status);
-        }
-
-        WifiNanSessionState session = getNanSessionStateForPubSubId(subscribeId);
-        if (session == null) {
-            Log.e(TAG, "onSubscribeTerminated: no session found for subscribeId=" + subscribeId);
-            return;
-        }
-
-        session.onSubscribeTerminated(status);
-    }
-
-    private void onMessageSendSuccessLocal(short transactionId) {
-        if (VDBG) {
-            Log.v(TAG, "onMessageSendSuccess: transactionId=" + transactionId);
-        }
-
-        TransactionInfoBase info = getAndRemovePendingResponseTransactionInfo(transactionId);
-        if (info == null) {
-            Log.e(TAG, "onMessageSendSuccess(): no info registered for transactionId="
-                    + transactionId);
-            return;
-        }
-        if (!(info instanceof TransactionInfoMessage)) {
-            Log.e(TAG, "onMessageSendSuccess: invalid info structure stored for transactionId="
-                    + transactionId);
-            return;
-        }
-        TransactionInfoMessage infoMessage = (TransactionInfoMessage) info;
-
-        infoMessage.mSession.onMessageSendSuccess(infoMessage.mMessageId);
-    }
-
-    private void onMessageSendFailLocal(short transactionId, int status) {
-        if (VDBG) {
-            Log.v(TAG, "onMessageSendFail: transactionId=" + transactionId + ", status=" + status);
-        }
-
-        TransactionInfoBase info = getAndRemovePendingResponseTransactionInfo(transactionId);
-        if (info == null) {
-            Log.e(TAG,
-                    "onMessageSendFail(): no info registered for transactionId=" + transactionId);
-            return;
-        }
-        if (!(info instanceof TransactionInfoMessage)) {
-            Log.e(TAG, "onMessageSendFail: invalid info structure stored for transactionId="
-                    + transactionId);
-            return;
-        }
-        TransactionInfoMessage infoMessage = (TransactionInfoMessage) info;
-
-        infoMessage.mSession.onMessageSendFail(infoMessage.mMessageId, status);
-    }
-
-    private void onUnknownTransactionLocal(int responseType, short transactionId, int status) {
-        Log.e(TAG, "onUnknownTransaction: responseType=" + responseType + ", transactionId="
-                + transactionId + ", status=" + status);
-
-        TransactionInfoBase info = getAndRemovePendingResponseTransactionInfo(transactionId);
-        if (info == null) {
-            Log.e(TAG, "onUnknownTransaction(): no info registered for transactionId="
-                    + transactionId);
-        }
-    }
-
-    private void onNoOpNotificationLocal(short transactionId) {
-        if (VDBG) {
-            Log.v(TAG, "onNoOpNotificationLocal: transactionId=" + transactionId);
-        }
-
-        TransactionInfoBase info = getAndRemovePendingResponseTransactionInfo(transactionId);
-        if (info == null) {
-            Log.e(TAG, "onNoOpNotificationLocal(): no info registered for transactionId="
-                    + transactionId);
-        }
-    }
-
     private void onMatchLocal(int pubSubId, int requestorInstanceId, byte[] peerMac,
             byte[] serviceSpecificInfo, int serviceSpecificInfoLength, byte[] matchFilter,
             int matchFilterLength) {
@@ -1231,59 +1302,121 @@
                     + ", matchFilterLength=" + matchFilterLength + ", matchFilter=" + matchFilter);
         }
 
-        WifiNanSessionState session = getNanSessionStateForPubSubId(pubSubId);
-        if (session == null) {
+        Pair<WifiNanClientState, WifiNanSessionState> data = getClientSessionForPubSubId(pubSubId);
+        if (data == null) {
             Log.e(TAG, "onMatch: no session found for pubSubId=" + pubSubId);
             return;
         }
 
-        session.onMatch(requestorInstanceId, peerMac, serviceSpecificInfo,
+        data.second.onMatch(requestorInstanceId, peerMac, serviceSpecificInfo,
                 serviceSpecificInfoLength, matchFilter, matchFilterLength);
     }
 
+    private void onSessionTerminatedLocal(int pubSubId, boolean isPublish, int reason) {
+        if (VDBG) {
+            Log.v(TAG, "onSessionTerminatedLocal: pubSubId=" + pubSubId + ", isPublish=" + isPublish
+                    + ", reason=" + reason);
+        }
+
+        Pair<WifiNanClientState, WifiNanSessionState> data = getClientSessionForPubSubId(pubSubId);
+        if (data == null) {
+            Log.e(TAG, "onSessionTerminatedLocal: no session found for pubSubId=" + pubSubId);
+            return;
+        }
+
+        try {
+            data.second.getCallback().onSessionTerminated(reason);
+        } catch (RemoteException e) {
+            Log.w(TAG,
+                    "onSessionTerminatedLocal onSessionTerminated(): RemoteException (FYI): " + e);
+        }
+        data.first.removeSession(data.second.getSessionId());
+    }
+
     private void onMessageReceivedLocal(int pubSubId, int requestorInstanceId, byte[] peerMac,
             byte[] message, int messageLength) {
         if (VDBG) {
             Log.v(TAG,
-                    "onMessageReceived: pubSubId=" + pubSubId + ", requestorInstanceId="
+                    "onMessageReceivedLocal: pubSubId=" + pubSubId + ", requestorInstanceId="
                             + requestorInstanceId + ", peerMac="
                             + String.valueOf(HexEncoding.encode(peerMac)) + ", messageLength="
                             + messageLength);
         }
 
-        WifiNanSessionState session = getNanSessionStateForPubSubId(pubSubId);
-        if (session == null) {
-            Log.e(TAG, "onMessageReceived: no session found for pubSubId=" + pubSubId);
+        Pair<WifiNanClientState, WifiNanSessionState> data = getClientSessionForPubSubId(pubSubId);
+        if (data == null) {
+            Log.e(TAG, "onMessageReceivedLocal: no session found for pubSubId=" + pubSubId);
             return;
         }
 
-        session.onMessageReceived(requestorInstanceId, peerMac, message, messageLength);
+        data.second.onMessageReceived(requestorInstanceId, peerMac, message, messageLength);
     }
 
-    private ConfigRequest mergeConfigRequests() {
+    private void onNanDownLocal(int reason) {
         if (VDBG) {
-            Log.v(TAG, "mergeConfigRequests(): mClients=[" + mClients + "]");
+            Log.v(TAG, "onNanDown: reason=" + reason);
         }
 
-        if (mClients.size() == 0) {
+        for (int i = 0; i < mClients.size(); ++i) {
+            WifiNanClientState client = mClients.valueAt(i);
+            try {
+                client.getCallback().onNanDown(reason);
+            } catch (RemoteException e) {
+                Log.w(TAG, "onNanDownLocal onNanDown(): RemoteException (FYI): " + e);
+            }
+        }
+        mClients.clear();
+        mCurrentNanConfiguration = null;
+    }
+
+    /*
+     * Utilities
+     */
+
+    private Pair<WifiNanClientState, WifiNanSessionState> getClientSessionForPubSubId(
+            int pubSubId) {
+        for (int i = 0; i < mClients.size(); ++i) {
+            WifiNanClientState client = mClients.valueAt(i);
+            WifiNanSessionState session = client.getNanSessionStateForPubSubId(pubSubId);
+            if (session != null) {
+                return new Pair<>(client, session);
+            }
+        }
+
+        return null;
+    }
+
+    private ConfigRequest mergeConfigRequests(ConfigRequest configRequest) {
+        if (VDBG) {
+            Log.v(TAG, "mergeConfigRequests(): mClients=[" + mClients + "], configRequest="
+                    + configRequest);
+        }
+
+        if (mClients.size() == 0 && configRequest == null) {
             Log.e(TAG, "mergeConfigRequests: invalid state - called with 0 clients registered!");
             return null;
         }
 
-        if (mClients.size() == 1) {
-            return mClients.valueAt(0).getConfigRequest();
-        }
-
         // TODO: continue working on merge algorithm:
         // - if any request 5g: enable
         // - maximal master preference
         // - cluster range covering all requests: assume that [0,max] is a
         // non-request
+        // - if any request identity change: enable
         boolean support5gBand = false;
         int masterPreference = 0;
         boolean clusterIdValid = false;
         int clusterLow = 0;
         int clusterHigh = ConfigRequest.CLUSTER_ID_MAX;
+        boolean identityChange = false;
+        if (configRequest != null) {
+            support5gBand = configRequest.mSupport5gBand;
+            masterPreference = configRequest.mMasterPreference;
+            clusterIdValid = true;
+            clusterLow = configRequest.mClusterLow;
+            clusterHigh = configRequest.mClusterHigh;
+            identityChange = configRequest.mEnableIdentityChangeCallback;
+        }
         for (int i = 0; i < mClients.size(); ++i) {
             ConfigRequest cr = mClients.valueAt(i).getConfigRequest();
 
@@ -1303,12 +1436,15 @@
                 }
                 clusterIdValid = true;
             }
-        }
-        ConfigRequest.Builder builder = new ConfigRequest.Builder();
-        builder.setSupport5gBand(support5gBand).setMasterPreference(masterPreference)
-                .setClusterLow(clusterLow).setClusterHigh(clusterHigh);
 
-        return builder.build();
+            if (cr.mEnableIdentityChangeCallback) {
+                identityChange = true;
+            }
+        }
+        return new ConfigRequest.Builder().setSupport5gBand(support5gBand)
+                .setMasterPreference(masterPreference).setClusterLow(clusterLow)
+                .setClusterHigh(clusterHigh).setEnableIdentityChangeCallback(identityChange)
+                .build();
     }
 
     /**
@@ -1317,12 +1453,11 @@
     public void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
         pw.println("NanStateManager:");
         pw.println("  mClients: [" + mClients + "]");
-        pw.println("  mPendingResponses: [" + mPendingResponses + "]");
         pw.println("  mCapabilities: [" + mCapabilities + "]");
-        pw.println("  mNextTransactionId: " + mNextTransactionId);
-        pw.println("  mNextSessionId: " + mNextSessionId);
+        pw.println("  mCurrentNanConfiguration: " + mCurrentNanConfiguration);
         for (int i = 0; i < mClients.size(); ++i) {
             mClients.valueAt(i).dump(fd, pw, args);
         }
+        mSm.dump(fd, pw, args);
     }
 }
diff --git a/service/jni/com_android_server_wifi_nan_WifiNanNative.cpp b/service/jni/com_android_server_wifi_nan_WifiNanNative.cpp
index caa889e..7b7f3e1 100644
--- a/service/jni/com_android_server_wifi_nan_WifiNanNative.cpp
+++ b/service/jni/com_android_server_wifi_nan_WifiNanNative.cpp
@@ -260,6 +260,27 @@
     return hal_fn.wifi_nan_enable_request(transaction_id, handle, &msg);
 }
 
+static jint android_net_wifi_nan_config_request(JNIEnv *env, jclass cls,
+                                                jshort transaction_id,
+                                                jclass wifi_native_cls,
+                                                jint iface,
+                                                jobject config_request) {
+    JNIHelper helper(env);
+    wifi_interface_handle handle = getIfaceHandle(helper, wifi_native_cls, iface);
+
+    ALOGD("android_net_wifi_nan_config_request handle=%p, id=%d",
+          handle, transaction_id);
+
+    NanConfigRequest msg;
+    memset(&msg, 0, sizeof(NanConfigRequest));
+
+    /* configurable settings */
+    msg.config_master_pref = 1;
+    msg.master_pref = helper.getIntField(config_request, "mMasterPreference");
+
+    return hal_fn.wifi_nan_config_request(transaction_id, handle, &msg);
+}
+
 static jint android_net_wifi_nan_get_capabilities(JNIEnv *env, jclass cls,
                                                   jshort transaction_id,
                                                   jclass wifi_native_cls,
@@ -510,16 +531,16 @@
 
 static JNINativeMethod gWifiNanMethods[] = {
     /* name, signature, funcPtr */
-
-    {"initNanHandlersNative", "(Ljava/lang/Object;I)I", (void*)android_net_wifi_nan_register_handler },
-    {"getCapabilitiesNative", "(SLjava/lang/Object;I)I", (void*)android_net_wifi_nan_get_capabilities },
-    {"enableAndConfigureNative", "(SLjava/lang/Object;ILandroid/net/wifi/nan/ConfigRequest;)I", (void*)android_net_wifi_nan_enable_request },
-    {"disableNative", "(SLjava/lang/Object;I)I", (void*)android_net_wifi_nan_disable_request },
-    {"publishNative", "(SILjava/lang/Object;ILandroid/net/wifi/nan/PublishConfig;)I", (void*)android_net_wifi_nan_publish },
-    {"subscribeNative", "(SILjava/lang/Object;ILandroid/net/wifi/nan/SubscribeConfig;)I", (void*)android_net_wifi_nan_subscribe },
-    {"sendMessageNative", "(SLjava/lang/Object;III[B[BI)I", (void*)android_net_wifi_nan_send_message },
-    {"stopPublishNative", "(SLjava/lang/Object;II)I", (void*)android_net_wifi_nan_stop_publish },
-    {"stopSubscribeNative", "(SLjava/lang/Object;II)I", (void*)android_net_wifi_nan_stop_subscribe },
+    {"initNanHandlersNative", "(Ljava/lang/Class;I)I", (void*)android_net_wifi_nan_register_handler },
+    {"getCapabilitiesNative", "(SLjava/lang/Class;I)I", (void*)android_net_wifi_nan_get_capabilities },
+    {"enableAndConfigureNative", "(SLjava/lang/Class;ILandroid/net/wifi/nan/ConfigRequest;)I", (void*)android_net_wifi_nan_enable_request },
+    {"updateConfigurationNative", "(SLjava/lang/Class;ILandroid/net/wifi/nan/ConfigRequest;)I", (void*)android_net_wifi_nan_config_request },
+    {"disableNative", "(SLjava/lang/Class;I)I", (void*)android_net_wifi_nan_disable_request },
+    {"publishNative", "(SILjava/lang/Class;ILandroid/net/wifi/nan/PublishConfig;)I", (void*)android_net_wifi_nan_publish },
+    {"subscribeNative", "(SILjava/lang/Class;ILandroid/net/wifi/nan/SubscribeConfig;)I", (void*)android_net_wifi_nan_subscribe },
+    {"sendMessageNative", "(SLjava/lang/Class;III[B[BI)I", (void*)android_net_wifi_nan_send_message },
+    {"stopPublishNative", "(SLjava/lang/Class;II)I", (void*)android_net_wifi_nan_stop_publish },
+    {"stopSubscribeNative", "(SLjava/lang/Class;II)I", (void*)android_net_wifi_nan_stop_subscribe },
 };
 
 /* User to register native functions */
diff --git a/tests/wifitests/jni/wifi_nan_hal_mock.cpp b/tests/wifitests/jni/wifi_nan_hal_mock.cpp
index 5c0cdc2..052a5be 100644
--- a/tests/wifitests/jni/wifi_nan_hal_mock.cpp
+++ b/tests/wifitests/jni/wifi_nan_hal_mock.cpp
@@ -288,8 +288,41 @@
 wifi_error wifi_nan_config_request_mock(transaction_id id,
                                         wifi_interface_handle iface,
                                         NanConfigRequest* msg) {
+  JNIHelper helper(mock_mVM);
+
   ALOGD("wifi_nan_config_request_mock");
-  return WIFI_ERROR_UNINITIALIZED;
+  HalMockJsonWriter jsonW;
+  jsonW.put_int("config_sid_beacon", msg->config_sid_beacon);
+  jsonW.put_int("sid_beacon", msg->sid_beacon);
+  jsonW.put_int("config_rssi_proximity", msg->config_rssi_proximity);
+  jsonW.put_int("rssi_proximity", msg->rssi_proximity);
+  jsonW.put_int("config_master_pref", msg->config_master_pref);
+  jsonW.put_int("master_pref", msg->master_pref);
+  jsonW.put_int("config_5g_rssi_close_proximity", msg->config_5g_rssi_close_proximity);
+  jsonW.put_int("rssi_close_proximity_5g_val", msg->rssi_close_proximity_5g_val);
+  jsonW.put_int("config_rssi_window_size", msg->config_rssi_window_size);
+  jsonW.put_int("rssi_window_size_val", msg->rssi_window_size_val);
+  jsonW.put_int("config_cluster_attribute_val", msg->config_cluster_attribute_val);
+  jsonW.put_int("config_scan_params", msg->config_scan_params);
+  // TODO: NanSocialChannelScanParams scan_params_val
+  jsonW.put_int("config_random_factor_force", msg->config_random_factor_force);
+  jsonW.put_int("random_factor_force_val", msg->random_factor_force_val);
+  jsonW.put_int("config_hop_count_force", msg->config_hop_count_force);
+  jsonW.put_int("hop_count_force_val", msg->hop_count_force_val);
+  jsonW.put_int("config_conn_capability", msg->config_conn_capability);
+  // TODO: NanTransmitPostConnectivityCapability conn_capability_val
+  jsonW.put_int("num_config_discovery_attr", msg->num_config_discovery_attr);
+  // TODO: NanTransmitPostDiscovery discovery_attr_val[NAN_MAX_POSTDISCOVERY_LEN]
+  jsonW.put_int("config_fam", msg->config_fam);
+  // TODO: NanFurtherAvailabilityMap fam_val
+  std::string str = jsonW.to_string();
+
+  JNIObject < jstring > json_write_string = helper.newStringUTF(str.c_str());
+
+  helper.callMethod(mock_mObj, "configHalMockNative", "(SLjava/lang/String;)V",
+                    (short) id, json_write_string.get());
+
+  return WIFI_SUCCESS;
 }
 
 wifi_error wifi_nan_tca_request_mock(transaction_id id,
diff --git a/tests/wifitests/src/com/android/server/wifi/nan/WifiNanHalMock.java b/tests/wifitests/src/com/android/server/wifi/nan/WifiNanHalMock.java
index 4f3ba4d..5cbeb27 100644
--- a/tests/wifitests/src/com/android/server/wifi/nan/WifiNanHalMock.java
+++ b/tests/wifitests/src/com/android/server/wifi/nan/WifiNanHalMock.java
@@ -39,6 +39,10 @@
         throw new IllegalStateException("Please mock this class!");
     }
 
+    public void configHalMockNative(short transactionId, String jsonArgs) {
+        throw new IllegalStateException("Please mock this class!");
+    }
+
     public void disableHalMockNative(short transactionId) {
         throw new IllegalStateException("Please mock this class!");
     }
diff --git a/tests/wifitests/src/com/android/server/wifi/nan/WifiNanHalTest.java b/tests/wifitests/src/com/android/server/wifi/nan/WifiNanHalTest.java
index 00ff340..d982887 100644
--- a/tests/wifitests/src/com/android/server/wifi/nan/WifiNanHalTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/nan/WifiNanHalTest.java
@@ -19,11 +19,13 @@
 import static org.hamcrest.core.IsEqual.equalTo;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import android.net.wifi.nan.ConfigRequest;
 import android.net.wifi.nan.PublishConfig;
 import android.net.wifi.nan.SubscribeConfig;
 import android.net.wifi.nan.TlvBufferUtils;
+import android.net.wifi.nan.WifiNanEventCallback;
 import android.net.wifi.nan.WifiNanSessionCallback;
 import android.os.Bundle;
 import android.test.suitebuilder.annotation.SmallTest;
@@ -94,6 +96,14 @@
     }
 
     @Test
+    public void testConfigCall() throws JSONException {
+        final short transactionId = 31235;
+        final short masterPref = 157;
+
+        testConfig(transactionId, masterPref);
+    }
+
+    @Test
     public void testDisable() {
         final short transactionId = 5478;
 
@@ -292,8 +302,7 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onCapabilitiesUpdate(eq(transactionId),
-                capabilitiesCapture.capture());
+        verify(mNanStateManager).onCapabilitiesUpdateNotification(capabilitiesCapture.capture());
         WifiNanNative.Capabilities capabilities = capabilitiesCapture.getValue();
         collector.checkThat("max_concurrent_nan_clusters", capabilities.maxConcurrentNanClusters,
                 equalTo(max_concurrent_nan_clusters));
@@ -317,6 +326,7 @@
                 equalTo(max_ndp_sessions));
         collector.checkThat("max_app_info_len", capabilities.maxAppInfoLen,
                 equalTo(max_app_info_len));
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -331,7 +341,8 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onConfigCompleted(transactionId);
+        verify(mNanStateManager).onConfigSuccessResponse(transactionId);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -346,8 +357,9 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onConfigFailed(transactionId,
-                WifiNanSessionCallback.FAIL_REASON_INVALID_ARGS);
+        verify(mNanStateManager).onConfigFailedResponse(transactionId,
+                WifiNanEventCallback.REASON_INVALID_ARGS);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -364,7 +376,8 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onPublishSuccess(transactionId, publishId);
+        verify(mNanStateManager).onSessionConfigSuccessResponse(transactionId, true, publishId);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -381,8 +394,9 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onPublishFail(transactionId,
-                WifiNanSessionCallback.FAIL_REASON_NO_RESOURCES);
+        verify(mNanStateManager).onSessionConfigFailResponse(transactionId, true,
+                WifiNanSessionCallback.REASON_NO_RESOURCES);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -397,7 +411,7 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onNoOpTransaction(transactionId);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -414,7 +428,8 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onSubscribeSuccess(transactionId, subscribeId);
+        verify(mNanStateManager).onSessionConfigSuccessResponse(transactionId, false, subscribeId);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -431,8 +446,9 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onSubscribeFail(transactionId,
-                WifiNanSessionCallback.FAIL_REASON_OTHER);
+        verify(mNanStateManager).onSessionConfigFailResponse(transactionId, false,
+                WifiNanSessionCallback.REASON_OTHER);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -447,7 +463,7 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onNoOpTransaction(transactionId);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -462,7 +478,8 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onMessageSendSuccess(transactionId);
+        verify(mNanStateManager).onMessageSendSuccessResponse(transactionId);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -477,8 +494,9 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onMessageSendFail(transactionId,
-                WifiNanSessionCallback.FAIL_REASON_OTHER);
+        verify(mNanStateManager).onMessageSendFailResponse(transactionId,
+                WifiNanSessionCallback.REASON_OTHER);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -494,8 +512,7 @@
         WifiNanHalMock.callNotifyResponse(transactionId,
                 HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onUnknownTransaction(invalidTransactionId, transactionId,
-                WifiNanNative.NAN_STATUS_SUCCESS);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -508,8 +525,9 @@
 
         WifiNanHalMock.callPublishTerminated(HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onPublishTerminated(publishId,
-                WifiNanSessionCallback.TERMINATE_REASON_DONE);
+        verify(mNanStateManager).onSessionTerminatedNotification(publishId,
+                WifiNanSessionCallback.TERMINATE_REASON_DONE, true);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -522,8 +540,9 @@
 
         WifiNanHalMock.callSubscribeTerminated(HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onSubscribeTerminated(subscribeId,
-                WifiNanSessionCallback.TERMINATE_REASON_FAIL);
+        verify(mNanStateManager).onSessionTerminatedNotification(subscribeId,
+                WifiNanSessionCallback.TERMINATE_REASON_FAIL, false);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -543,8 +562,9 @@
 
         WifiNanHalMock.callFollowup(HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onMessageReceived(pubSubId, reqInstanceId, peer,
+        verify(mNanStateManager).onMessageReceivedNotification(pubSubId, reqInstanceId, peer,
                 message.getBytes(), message.length());
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -566,8 +586,9 @@
 
         WifiNanHalMock.callMatch(HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onMatch(pubSubId, reqInstanceId, peer, ssi.getBytes(),
+        verify(mNanStateManager).onMatchNotification(pubSubId, reqInstanceId, peer, ssi.getBytes(),
                 ssi.length(), filter.getBytes(), filter.length());
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -580,7 +601,8 @@
 
         WifiNanHalMock.callDiscEngEvent(HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onInterfaceAddressChange(mac);
+        verify(mNanStateManager).onInterfaceAddressChangeNotification(mac);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -593,8 +615,9 @@
 
         WifiNanHalMock.callDiscEngEvent(HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onClusterChange(WifiNanClientState.CLUSTER_CHANGE_EVENT_JOINED,
-                mac);
+        verify(mNanStateManager)
+                .onClusterChangeNotification(WifiNanClientState.CLUSTER_CHANGE_EVENT_JOINED, mac);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -607,8 +630,9 @@
 
         WifiNanHalMock.callDiscEngEvent(HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onClusterChange(WifiNanClientState.CLUSTER_CHANGE_EVENT_STARTED,
-                mac);
+        verify(mNanStateManager)
+                .onClusterChangeNotification(WifiNanClientState.CLUSTER_CHANGE_EVENT_STARTED, mac);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     @Test
@@ -618,7 +642,8 @@
 
         WifiNanHalMock.callDisabled(HalMockUtils.convertBundleToJson(args).toString());
 
-        verify(mNanStateManager).onNanDown(WifiNanSessionCallback.FAIL_REASON_OTHER);
+        verify(mNanStateManager).onNanDownNotification(WifiNanEventCallback.REASON_OTHER);
+        verifyNoMoreInteractions(mNanStateManager);
     }
 
     /*
@@ -631,7 +656,7 @@
                 .setClusterHigh(clusterHigh).setMasterPreference(masterPref)
                 .setSupport5gBand(enable5g).build();
 
-        mDut.enableAndConfigure(transactionId, configRequest);
+        mDut.enableAndConfigure(transactionId, configRequest, true);
 
         verify(mNanHalMock).enableHalMockNative(eq(transactionId), mArgs.capture());
 
@@ -680,6 +705,52 @@
                 equalTo(0));
     }
 
+    private void testConfig(short transactionId, int masterPref) throws JSONException {
+        ConfigRequest configRequest = new ConfigRequest.Builder().setMasterPreference(masterPref)
+                .build();
+
+        mDut.enableAndConfigure(transactionId, configRequest, false);
+
+        verify(mNanHalMock).configHalMockNative(eq(transactionId), mArgs.capture());
+
+        Bundle argsData = HalMockUtils.convertJsonToBundle(mArgs.getValue());
+
+        collector.checkThat("config_master_pref", argsData.getInt("config_master_pref"),
+                equalTo(1));
+        collector.checkThat("master_pref", argsData.getInt("master_pref"), equalTo(masterPref));
+
+        collector.checkThat("config_sid_beacon", argsData.getInt("config_sid_beacon"), equalTo(0));
+        collector.checkThat("sid_beacon", argsData.getInt("sid_beacon"), equalTo(0));
+        collector.checkThat("config_rssi_proximity", argsData.getInt("config_rssi_proximity"),
+                equalTo(0));
+        collector.checkThat("rssi_proximity", argsData.getInt("rssi_proximity"), equalTo(0));
+        collector.checkThat("config_5g_rssi_close_proximity",
+                argsData.getInt("config_5g_rssi_close_proximity"), equalTo(0));
+        collector.checkThat("rssi_close_proximity_5g_val",
+                argsData.getInt("rssi_close_proximity_5g_val"), equalTo(0));
+        collector.checkThat("config_rssi_window_size", argsData.getInt("config_rssi_window_size"),
+                equalTo(0));
+        collector.checkThat("rssi_window_size_val", argsData.getInt("rssi_window_size_val"),
+                equalTo(0));
+        collector.checkThat("config_cluster_attribute_val",
+                argsData.getInt("config_cluster_attribute_val"), equalTo(0));
+        collector.checkThat("config_scan_params", argsData.getInt("config_scan_params"),
+                equalTo(0));
+        collector.checkThat("config_random_factor_force",
+                argsData.getInt("config_random_factor_force"), equalTo(0));
+        collector.checkThat("random_factor_force_val", argsData.getInt("random_factor_force_val"),
+                equalTo(0));
+        collector.checkThat("config_hop_count_force", argsData.getInt("config_hop_count_force"),
+                equalTo(0));
+        collector.checkThat("hop_count_force_val", argsData.getInt("hop_count_force_val"),
+                equalTo(0));
+        collector.checkThat("config_conn_capability", argsData.getInt("config_conn_capability"),
+                equalTo(0));
+        collector.checkThat("num_config_discovery_attr",
+                argsData.getInt("num_config_discovery_attr"), equalTo(0));
+        collector.checkThat("config_fam", argsData.getInt("config_fam"), equalTo(0));
+    }
+
     private void testPublish(short transactionId, int publishId, int publishType,
             String serviceName, String ssi, TlvBufferUtils.TlvConstructor tlvTx,
             TlvBufferUtils.TlvConstructor tlvRx, int publishCount, int publishTtl,
diff --git a/tests/wifitests/src/com/android/server/wifi/nan/WifiNanManagerTest.java b/tests/wifitests/src/com/android/server/wifi/nan/WifiNanManagerTest.java
index 2330a61..9f5356f 100644
--- a/tests/wifitests/src/com/android/server/wifi/nan/WifiNanManagerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/nan/WifiNanManagerTest.java
@@ -20,6 +20,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.eq;
+import static org.mockito.Matchers.isNull;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.when;
 
@@ -84,13 +85,195 @@
     }
 
     /*
+     * WifiNanEventCallbackProxy Tests
+     */
+
+    /**
+     * Validate the successful connect flow: (1) try subscribing (2) connect +
+     * success (3) publish, (4) disconnect (5) try publishing (6) connect again
+     */
+    @Test
+    public void testConnectFlow() throws Exception {
+        final int clientId = 4565;
+
+        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class),
+                any(ConfigRequest.class))).thenReturn(clientId);
+
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mockNanService);
+        ArgumentCaptor<IWifiNanEventCallback> clientProxyCallback = ArgumentCaptor
+                .forClass(IWifiNanEventCallback.class);
+        ArgumentCaptor<IBinder> binder = ArgumentCaptor.forClass(IBinder.class);
+
+        // (1) try subscribing on an unconnected manager: fails silently
+        mDut.subscribe(new SubscribeConfig.Builder().build(), mockSessionCallback);
+
+        // (2) connect + success
+        mDut.connect(mMockLooper.getLooper(), mockCallback);
+        inOrder.verify(mockNanService).connect(binder.capture(),
+                clientProxyCallback.capture(), (ConfigRequest) isNull());
+        clientProxyCallback.getValue().onConnectSuccess();
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        // (3) publish - should succeed
+        PublishConfig publishConfig = new PublishConfig.Builder().build();
+        mDut.publish(publishConfig, mockSessionCallback);
+        inOrder.verify(mockNanService).publish(eq(clientId), eq(publishConfig),
+                any(IWifiNanSessionCallback.class));
+
+        // (4) disconnect
+        mDut.disconnect();
+        inOrder.verify(mockNanService).disconnect(eq(clientId), eq(binder.getValue()));
+
+        // (5) try publishing again - fails silently
+        mDut.publish(new PublishConfig.Builder().build(), mockSessionCallback);
+
+        // (6) connect
+        mDut.connect(mMockLooper.getLooper(), mockCallback);
+        inOrder.verify(mockNanService).connect(binder.capture(), any(IWifiNanEventCallback.class),
+                (ConfigRequest) isNull());
+
+        inOrder.verifyNoMoreInteractions();
+    }
+
+    /**
+     * Validate the failed connect flow: (1) connect + failure, (2) try
+     * publishing (3) connect + success (4) subscribe
+     */
+    @Test
+    public void testConnectFailure() throws Exception {
+        final int clientId = 4565;
+        final int reason = WifiNanEventCallback.REASON_OTHER;
+
+        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class),
+                any(ConfigRequest.class))).thenReturn(clientId);
+
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mockNanService);
+        ArgumentCaptor<IWifiNanEventCallback> clientProxyCallback = ArgumentCaptor
+                .forClass(IWifiNanEventCallback.class);
+
+        // (1) connect + failure
+        mDut.connect(mMockLooper.getLooper(), mockCallback);
+        inOrder.verify(mockNanService).connect(any(IBinder.class), clientProxyCallback.capture(),
+                (ConfigRequest) isNull());
+        clientProxyCallback.getValue().onConnectFail(reason);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectFail(reason);
+
+        // (2) try publishing - silent failure (since already know that no
+        // connection)
+        mDut.publish(new PublishConfig.Builder().build(), mockSessionCallback);
+
+        // (3) connect + success
+        mDut.connect(mMockLooper.getLooper(), mockCallback);
+        inOrder.verify(mockNanService).connect(any(IBinder.class), clientProxyCallback.capture(),
+                (ConfigRequest) isNull());
+        clientProxyCallback.getValue().onConnectSuccess();
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        // (4) subscribe: should succeed
+        SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
+        mDut.subscribe(subscribeConfig, mockSessionCallback);
+        inOrder.verify(mockNanService).subscribe(eq(clientId), eq(subscribeConfig),
+                any(IWifiNanSessionCallback.class));
+
+        inOrder.verifyNoMoreInteractions();
+    }
+
+    /**
+     * Validate connect flow when NAN is shutting down: (1) connect + success
+     * (2) publish (3) NAN down (4) try updating the publish (5) connect again
+     */
+    @Test
+    public void testConnectNanDownFlow() throws Exception {
+        final int clientId = 4565;
+        final int sessionId = 123;
+        final int reason = WifiNanEventCallback.REASON_REQUESTED;
+
+        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class),
+                any(ConfigRequest.class))).thenReturn(clientId);
+
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mockNanService);
+        ArgumentCaptor<IWifiNanEventCallback> clientProxyCallback = ArgumentCaptor
+                .forClass(IWifiNanEventCallback.class);
+        ArgumentCaptor<IWifiNanSessionCallback> sessionProxyCallback = ArgumentCaptor
+                .forClass(IWifiNanSessionCallback.class);
+        ArgumentCaptor<WifiNanPublishSession> publishSession = ArgumentCaptor
+                .forClass(WifiNanPublishSession.class);
+
+        // (1) connect + success
+        mDut.connect(mMockLooper.getLooper(), mockCallback);
+        inOrder.verify(mockNanService).connect(any(IBinder.class), clientProxyCallback.capture(),
+                (ConfigRequest) isNull());
+        clientProxyCallback.getValue().onConnectSuccess();
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        // (2) publish - should succeed
+        PublishConfig publishConfig = new PublishConfig.Builder().build();
+        mDut.publish(publishConfig, mockSessionCallback);
+        inOrder.verify(mockNanService).publish(eq(clientId), eq(publishConfig),
+                sessionProxyCallback.capture());
+        sessionProxyCallback.getValue().onSessionStarted(sessionId);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onPublishStarted(publishSession.capture());
+
+        // (3) NAN down
+        clientProxyCallback.getValue().onNanDown(reason);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onNanDown(reason);
+
+        // (4) try updating the publish - silent failure (already informed that
+        // NAN down)
+        publishSession.getValue().updatePublish(publishConfig);
+
+        // (5) connect - should try the service (i.e. succeed)
+        mDut.connect(mMockLooper.getLooper(), mockCallback);
+        inOrder.verify(mockNanService).connect(any(IBinder.class), any(IWifiNanEventCallback.class),
+                (ConfigRequest) isNull());
+
+        inOrder.verifyNoMoreInteractions();
+    }
+
+    /**
+     * Validate that cannot call connect on an existing connection: (1) connect
+     * + success, (2) try connect again
+     */
+    @Test
+    public void testInvalidConnectSequence() throws Exception {
+        final int clientId = 4565;
+
+        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class),
+                any(ConfigRequest.class))).thenReturn(clientId);
+
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mockNanService);
+        ArgumentCaptor<IWifiNanEventCallback> clientProxyCallback = ArgumentCaptor
+                .forClass(IWifiNanEventCallback.class);
+
+        // (1) connect + success
+        mDut.connect(mMockLooper.getLooper(), mockCallback);
+        inOrder.verify(mockNanService).connect(any(IBinder.class), clientProxyCallback.capture(),
+                (ConfigRequest) isNull());
+        clientProxyCallback.getValue().onConnectSuccess();
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        // (2) connect - fails silently
+        mDut.connect(mMockLooper.getLooper(), mockCallback);
+
+        inOrder.verifyNoMoreInteractions();
+    }
+
+    /*
      * WifiNanSessionCallbackProxy Tests
      */
 
     /**
-     * Validate the publish flow: (1) publish, (2) success creates session, (3)
-     * pass through everything, (4) update publish through session, (5)
-     * terminate locally, (6) try another command for local failure feedback.
+     * Validate the publish flow: (0) connect + success, (1) publish, (2)
+     * success creates session, (3) pass through everything, (4) update publish
+     * through session, (5) terminate locally, (6) try another command -
+     * ignored.
      */
     @Test
     public void testPublishFlow() throws Exception {
@@ -102,67 +285,81 @@
         final String string1 = "hey from here...";
         final String string2 = "some other arbitrary string...";
         final int messageId = 2123;
-        final int reason = WifiNanSessionCallback.FAIL_REASON_OTHER;
+        final int reason = WifiNanSessionCallback.REASON_OTHER;
 
-        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class)))
-                .thenReturn(clientId);
+        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class),
+                eq(configRequest))).thenReturn(clientId);
 
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mockNanService,
                 mockPublishSession);
+        ArgumentCaptor<IWifiNanEventCallback> clientProxyCallback = ArgumentCaptor
+                .forClass(IWifiNanEventCallback.class);
         ArgumentCaptor<IWifiNanSessionCallback> sessionProxyCallback = ArgumentCaptor
                 .forClass(IWifiNanSessionCallback.class);
         ArgumentCaptor<WifiNanPublishSession> publishSession = ArgumentCaptor
                 .forClass(WifiNanPublishSession.class);
 
-        mDut.connect(mMockLooper.getLooper(), mockCallback);
-        mDut.requestConfig(configRequest);
-        mDut.publish(publishConfig, mockSessionCallback);
-
+        // (0) connect + success
+        mDut.connect(mMockLooper.getLooper(), mockCallback, configRequest);
         inOrder.verify(mockNanService).connect(any(IBinder.class),
-                any(IWifiNanEventCallback.class));
-        inOrder.verify(mockNanService).requestConfig(eq(clientId), eq(configRequest));
+                clientProxyCallback.capture(), eq(configRequest));
+        clientProxyCallback.getValue().onConnectSuccess();
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        // (1) publish
+        mDut.publish(publishConfig, mockSessionCallback);
         inOrder.verify(mockNanService).publish(eq(clientId), eq(publishConfig),
                 sessionProxyCallback.capture());
 
+        // (2) publish session created
         sessionProxyCallback.getValue().onSessionStarted(sessionId);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mockSessionCallback).onPublishStarted(publishSession.capture());
 
+        // (3) ...
         publishSession.getValue().sendMessage(peerId, string1.getBytes(), string1.length(),
                 messageId);
-        publishSession.getValue().updatePublish(publishConfig);
         sessionProxyCallback.getValue().onMatch(peerId, string1.getBytes(),
                 string1.length(), string2.getBytes(), string2.length());
         sessionProxyCallback.getValue().onMessageReceived(peerId, string1.getBytes(),
                 string1.length());
         sessionProxyCallback.getValue().onMessageSendFail(messageId, reason);
         sessionProxyCallback.getValue().onMessageSendSuccess(messageId);
-        sessionProxyCallback.getValue().onSessionConfigFail(reason);
-        sessionProxyCallback.getValue().onSessionTerminated(reason);
         mMockLooper.dispatchAll();
 
         inOrder.verify(mockNanService).sendMessage(eq(clientId), eq(sessionId), eq(peerId),
                 eq(string1.getBytes()), eq(string1.length()), eq(messageId));
-        inOrder.verify(mockNanService).updatePublish(eq(clientId), eq(sessionId),
-                eq(publishConfig));
         inOrder.verify(mockSessionCallback).onMatch(eq(peerId), eq(string1.getBytes()),
                 eq(string1.length()), eq(string2.getBytes()), eq(string2.length()));
         inOrder.verify(mockSessionCallback).onMessageReceived(eq(peerId), eq(string1.getBytes()),
                 eq(string1.length()));
         inOrder.verify(mockSessionCallback).onMessageSendFail(eq(messageId), eq(reason));
         inOrder.verify(mockSessionCallback).onMessageSendSuccess(eq(messageId));
+
+        // (4) update publish
+        publishSession.getValue().updatePublish(publishConfig);
+        sessionProxyCallback.getValue().onSessionConfigFail(reason);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockNanService).updatePublish(eq(clientId), eq(sessionId),
+                eq(publishConfig));
         inOrder.verify(mockSessionCallback).onSessionConfigFail(eq(reason));
-        inOrder.verify(mockSessionCallback).onSessionTerminated(eq(reason));
-        inOrder.verify(mockNanService).terminateSession(eq(clientId), eq(sessionId));
+
+        // (5) terminate
+        publishSession.getValue().terminate();
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockNanService).terminateSession(clientId, sessionId);
+
+        // (6) try an update (nothing)
+        publishSession.getValue().updatePublish(publishConfig);
+        mMockLooper.dispatchAll();
 
         inOrder.verifyNoMoreInteractions();
     }
 
     /**
-     * Validate that if an active publish receives a termination (error or done)
-     * then it feeds back to the service an actual termination to free service
-     * data.
+     * Validate race condition of session terminate and session action: (1)
+     * connect, (2) publish success + terminate, (3) update.
      */
     @Test
     public void testPublishRemoteTerminate() throws Exception {
@@ -172,45 +369,47 @@
         final PublishConfig publishConfig = new PublishConfig.Builder().build();
         final int reason = WifiNanSessionCallback.TERMINATE_REASON_DONE;
 
-        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class)))
-                .thenReturn(clientId);
+        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class),
+                eq(configRequest))).thenReturn(clientId);
 
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mockNanService,
                 mockPublishSession);
+        ArgumentCaptor<IWifiNanEventCallback> clientProxyCallback = ArgumentCaptor
+                .forClass(IWifiNanEventCallback.class);
         ArgumentCaptor<IWifiNanSessionCallback> sessionProxyCallback = ArgumentCaptor
                 .forClass(IWifiNanSessionCallback.class);
         ArgumentCaptor<WifiNanPublishSession> publishSession = ArgumentCaptor
                 .forClass(WifiNanPublishSession.class);
 
-        mDut.connect(mMockLooper.getLooper(), mockCallback);
-        mDut.requestConfig(configRequest);
-        mDut.publish(publishConfig, mockSessionCallback);
+        // (1) connect successfully
+        mDut.connect(mMockLooper.getLooper(), mockCallback, configRequest);
+        inOrder.verify(mockNanService).connect(any(IBinder.class), clientProxyCallback.capture(),
+                eq(configRequest));
+        clientProxyCallback.getValue().onConnectSuccess();
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
-        inOrder.verify(mockNanService).connect(any(IBinder.class),
-                any(IWifiNanEventCallback.class));
-        inOrder.verify(mockNanService).requestConfig(eq(clientId), eq(configRequest));
+        // (2) publish: successfully - then terminated
+        mDut.publish(publishConfig, mockSessionCallback);
         inOrder.verify(mockNanService).publish(eq(clientId), eq(publishConfig),
                 sessionProxyCallback.capture());
-
         sessionProxyCallback.getValue().onSessionStarted(sessionId);
         sessionProxyCallback.getValue().onSessionTerminated(reason);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mockSessionCallback).onPublishStarted(publishSession.capture());
         inOrder.verify(mockSessionCallback).onSessionTerminated(reason);
-        inOrder.verify(mockNanService).terminateSession(clientId, sessionId);
 
+        // (3) failure when trying to update: NOP
         publishSession.getValue().updatePublish(publishConfig);
-        inOrder.verify(mockSessionCallback)
-                .onSessionConfigFail(WifiNanSessionCallback.FAIL_REASON_SESSION_TERMINATED);
 
         inOrder.verifyNoMoreInteractions();
     }
 
     /**
-     * Validate the subscribe flow: (1) subscribe, (2) success creates session,
-     * (3) pass through everything, (4) update subscribe through session, (5)
-     * terminate locally, (6) try another command for local failure feedback.
+     * Validate the subscribe flow: (0) connect + success, (1) subscribe, (2)
+     * success creates session, (3) pass through everything, (4) update
+     * subscribe through session, (5) terminate locally, (6) try another command
+     * - ignored.
      */
     @Test
     public void testSubscribeFlow() throws Exception {
@@ -222,67 +421,81 @@
         final String string1 = "hey from here...";
         final String string2 = "some other arbitrary string...";
         final int messageId = 2123;
-        final int reason = WifiNanSessionCallback.FAIL_REASON_OTHER;
+        final int reason = WifiNanSessionCallback.REASON_OTHER;
 
-        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class)))
-                .thenReturn(clientId);
+        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class),
+                eq(configRequest))).thenReturn(clientId);
 
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mockNanService,
                 mockSubscribeSession);
+        ArgumentCaptor<IWifiNanEventCallback> clientProxyCallback = ArgumentCaptor
+                .forClass(IWifiNanEventCallback.class);
         ArgumentCaptor<IWifiNanSessionCallback> sessionProxyCallback = ArgumentCaptor
                 .forClass(IWifiNanSessionCallback.class);
         ArgumentCaptor<WifiNanSubscribeSession> subscribeSession = ArgumentCaptor
                 .forClass(WifiNanSubscribeSession.class);
 
-        mDut.connect(mMockLooper.getLooper(), mockCallback);
-        mDut.requestConfig(configRequest);
-        mDut.subscribe(subscribeConfig, mockSessionCallback);
+        // (0) connect + success
+        mDut.connect(mMockLooper.getLooper(), mockCallback, configRequest);
+        inOrder.verify(mockNanService).connect(any(IBinder.class), clientProxyCallback.capture(),
+                eq(configRequest));
+        clientProxyCallback.getValue().onConnectSuccess();
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
-        inOrder.verify(mockNanService).connect(any(IBinder.class),
-                any(IWifiNanEventCallback.class));
-        inOrder.verify(mockNanService).requestConfig(eq(clientId), eq(configRequest));
+        // (1) subscribe
+        mDut.subscribe(subscribeConfig, mockSessionCallback);
         inOrder.verify(mockNanService).subscribe(eq(clientId), eq(subscribeConfig),
                 sessionProxyCallback.capture());
 
+        // (2) subscribe session created
         sessionProxyCallback.getValue().onSessionStarted(sessionId);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mockSessionCallback).onSubscribeStarted(subscribeSession.capture());
 
+        // (3) ...
         subscribeSession.getValue().sendMessage(peerId, string1.getBytes(), string1.length(),
                 messageId);
-        subscribeSession.getValue().updateSubscribe(subscribeConfig);
         sessionProxyCallback.getValue().onMatch(peerId, string1.getBytes(), string1.length(),
                 string2.getBytes(), string2.length());
         sessionProxyCallback.getValue().onMessageReceived(peerId, string1.getBytes(),
                 string1.length());
         sessionProxyCallback.getValue().onMessageSendFail(messageId, reason);
         sessionProxyCallback.getValue().onMessageSendSuccess(messageId);
-        sessionProxyCallback.getValue().onSessionConfigFail(reason);
-        sessionProxyCallback.getValue().onSessionTerminated(reason);
         mMockLooper.dispatchAll();
 
         inOrder.verify(mockNanService).sendMessage(eq(clientId), eq(sessionId), eq(peerId),
                 eq(string1.getBytes()), eq(string1.length()), eq(messageId));
-        inOrder.verify(mockNanService).updateSubscribe(eq(clientId), eq(sessionId),
-                eq(subscribeConfig));
         inOrder.verify(mockSessionCallback).onMatch(eq(peerId), eq(string1.getBytes()),
                 eq(string1.length()), eq(string2.getBytes()), eq(string2.length()));
         inOrder.verify(mockSessionCallback).onMessageReceived(eq(peerId), eq(string1.getBytes()),
                 eq(string1.length()));
         inOrder.verify(mockSessionCallback).onMessageSendFail(eq(messageId), eq(reason));
         inOrder.verify(mockSessionCallback).onMessageSendSuccess(eq(messageId));
+
+        // (4) update subscribe
+        subscribeSession.getValue().updateSubscribe(subscribeConfig);
+        sessionProxyCallback.getValue().onSessionConfigFail(reason);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockNanService).updateSubscribe(eq(clientId), eq(sessionId),
+                eq(subscribeConfig));
         inOrder.verify(mockSessionCallback).onSessionConfigFail(eq(reason));
-        inOrder.verify(mockSessionCallback).onSessionTerminated(eq(reason));
-        inOrder.verify(mockNanService).terminateSession(eq(clientId), eq(sessionId));
+
+        // (5) terminate
+        subscribeSession.getValue().terminate();
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockNanService).terminateSession(clientId, sessionId);
+
+        // (6) try an update (nothing)
+        subscribeSession.getValue().updateSubscribe(subscribeConfig);
+        mMockLooper.dispatchAll();
 
         inOrder.verifyNoMoreInteractions();
     }
 
     /**
-     * Validate that if an active subscribe receives a termination (error or
-     * done) then it feeds back to the service an actual termination to free
-     * service data.
+     * Validate race condition of session terminate and session action: (1)
+     * connect, (2) subscribe success + terminate, (3) update.
      */
     @Test
     public void testSubscribeRemoteTerminate() throws Exception {
@@ -292,37 +505,38 @@
         final SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
         final int reason = WifiNanSessionCallback.TERMINATE_REASON_DONE;
 
-        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class)))
-                .thenReturn(clientId);
+        when(mockNanService.connect(any(IBinder.class), any(IWifiNanEventCallback.class),
+                eq(configRequest))).thenReturn(clientId);
 
         InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mockNanService,
                 mockSubscribeSession);
+        ArgumentCaptor<IWifiNanEventCallback> clientProxyCallback = ArgumentCaptor
+                .forClass(IWifiNanEventCallback.class);
         ArgumentCaptor<IWifiNanSessionCallback> sessionProxyCallback = ArgumentCaptor
                 .forClass(IWifiNanSessionCallback.class);
         ArgumentCaptor<WifiNanSubscribeSession> subscribeSession = ArgumentCaptor
                 .forClass(WifiNanSubscribeSession.class);
 
-        mDut.connect(mMockLooper.getLooper(), mockCallback);
-        mDut.requestConfig(configRequest);
-        mDut.subscribe(subscribeConfig, mockSessionCallback);
+        // (1) connect successfully
+        mDut.connect(mMockLooper.getLooper(), mockCallback, configRequest);
+        inOrder.verify(mockNanService).connect(any(IBinder.class), clientProxyCallback.capture(),
+                eq(configRequest));
+        clientProxyCallback.getValue().onConnectSuccess();
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
-        inOrder.verify(mockNanService).connect(any(IBinder.class),
-                any(IWifiNanEventCallback.class));
-        inOrder.verify(mockNanService).requestConfig(eq(clientId), eq(configRequest));
+        // (2) subscribe: successfully - then terminated
+        mDut.subscribe(subscribeConfig, mockSessionCallback);
         inOrder.verify(mockNanService).subscribe(eq(clientId), eq(subscribeConfig),
                 sessionProxyCallback.capture());
-
         sessionProxyCallback.getValue().onSessionStarted(sessionId);
         sessionProxyCallback.getValue().onSessionTerminated(reason);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mockSessionCallback).onSubscribeStarted(subscribeSession.capture());
         inOrder.verify(mockSessionCallback).onSessionTerminated(reason);
-        inOrder.verify(mockNanService).terminateSession(clientId, sessionId);
 
+        // (3) failure when trying to update: NOP
         subscribeSession.getValue().updateSubscribe(subscribeConfig);
-        inOrder.verify(mockSessionCallback)
-                .onSessionConfigFail(WifiNanSessionCallback.FAIL_REASON_SESSION_TERMINATED);
 
         inOrder.verifyNoMoreInteractions();
     }
diff --git a/tests/wifitests/src/com/android/server/wifi/nan/WifiNanServiceImplTest.java b/tests/wifitests/src/com/android/server/wifi/nan/WifiNanServiceImplTest.java
index 31b99b1..428e862 100644
--- a/tests/wifitests/src/com/android/server/wifi/nan/WifiNanServiceImplTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/nan/WifiNanServiceImplTest.java
@@ -61,7 +61,9 @@
     @Mock
     private IBinder mBinderMock;
     @Mock
-    IWifiNanEventCallback mCallbackMock;
+    private IWifiNanEventCallback mCallbackMock;
+    @Mock
+    private IWifiNanSessionCallback mSessionCallbackMock;
 
     /**
      * Using instead of spy to avoid native crash failures - possibly due to
@@ -121,6 +123,21 @@
     }
 
     /**
+     * Validate connect() when a non-null config is passed.
+     */
+    @Test
+    public void testConnectWithConfig() {
+        ConfigRequest configRequest = new ConfigRequest.Builder().setMasterPreference(55).build();
+
+        int returnedClientId = mDut.connect(mBinderMock, mCallbackMock, configRequest);
+
+        ArgumentCaptor<Integer> clientId = ArgumentCaptor.forClass(Integer.class);
+        verify(mNanStateManagerMock).connect(clientId.capture(), eq(mCallbackMock),
+                eq(configRequest));
+        assertEquals(returnedClientId, (int) clientId.getValue());
+    }
+
+    /**
      * Validate disconnect() - correct pass-through args.
      *
      * @throws Exception
@@ -145,8 +162,8 @@
     }
 
     /**
-     * Validate that security exception thrown when attempting operation using
-     * an a client ID which was already cleared-up.
+     * Validate that security exception thrown when attempting operation using a
+     * client ID which was already cleared-up.
      */
     @Test(expected = SecurityException.class)
     public void testFailOnClearedUpClientId() throws Exception {
@@ -185,10 +202,10 @@
 
         mDut.fakeUid = mDefaultUid;
 
-        ConfigRequest configRequest = new ConfigRequest.Builder().build();
-        mDut.requestConfig(clientId, configRequest);
+        PublishConfig publishConfig = new PublishConfig.Builder().build();
+        mDut.publish(clientId, publishConfig, mSessionCallbackMock);
 
-        verify(mNanStateManagerMock).requestConfig(clientId, configRequest);
+        verify(mNanStateManagerMock).publish(clientId, publishConfig, mSessionCallbackMock);
         assertTrue("SecurityException for invalid access from wrong UID thrown", failsAsExpected);
     }
 
@@ -217,7 +234,7 @@
 
         int prevId = 0;
         for (int i = 0; i < loopCount; ++i) {
-            int id = mDut.connect(mBinderMock, mCallbackMock);
+            int id = mDut.connect(mBinderMock, mCallbackMock, null);
             if (i != 0) {
                 assertTrue("Client ID incrementing", id > prevId);
             }
@@ -226,23 +243,10 @@
     }
 
     /**
-     * Validate requestConfig() - correct pass-through args.
-     */
-    @Test
-    public void testRequestConfig() {
-        int clientId = doConnect();
-        ConfigRequest configRequest = new ConfigRequest.Builder().build();
-
-        mDut.requestConfig(clientId, configRequest);
-
-        verify(mNanStateManagerMock).requestConfig(clientId, configRequest);
-    }
-
-    /**
      * Validate terminateSession() - correct pass-through args.
      */
     @Test
-    public void testStopSession() {
+    public void testTerminateSession() {
         int sessionId = 1024;
         int clientId = doConnect();
 
@@ -266,6 +270,20 @@
     }
 
     /**
+     * Validate updatePublish() - correct pass-through args.
+     */
+    @Test
+    public void testUpdatePublish() {
+        int sessionId = 1232;
+        PublishConfig publishConfig = new PublishConfig.Builder().build();
+        int clientId = doConnect();
+
+        mDut.updatePublish(clientId, sessionId, publishConfig);
+
+        verify(mNanStateManagerMock).updatePublish(clientId, sessionId, publishConfig);
+    }
+
+    /**
      * Validate subscribe() - correct pass-through args.
      */
     @Test
@@ -280,6 +298,20 @@
     }
 
     /**
+     * Validate updateSubscribe() - correct pass-through args.
+     */
+    @Test
+    public void testUpdateSubscribe() {
+        int sessionId = 1232;
+        SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
+        int clientId = doConnect();
+
+        mDut.updateSubscribe(clientId, sessionId, subscribeConfig);
+
+        verify(mNanStateManagerMock).updateSubscribe(clientId, sessionId, subscribeConfig);
+    }
+
+    /**
      * Validate sendMessage() - correct pass-through args.
      */
     @Test
@@ -316,10 +348,11 @@
      */
 
     private int doConnect() {
-        int returnedClientId = mDut.connect(mBinderMock, mCallbackMock);
+        int returnedClientId = mDut.connect(mBinderMock, mCallbackMock, null);
 
         ArgumentCaptor<Integer> clientId = ArgumentCaptor.forClass(Integer.class);
-        verify(mNanStateManagerMock).connect(clientId.capture(), eq(mCallbackMock));
+        verify(mNanStateManagerMock).connect(clientId.capture(), eq(mCallbackMock),
+                eq(new ConfigRequest.Builder().build()));
         assertEquals(returnedClientId, (int) clientId.getValue());
 
         return returnedClientId;
diff --git a/tests/wifitests/src/com/android/server/wifi/nan/WifiNanStateManagerTest.java b/tests/wifitests/src/com/android/server/wifi/nan/WifiNanStateManagerTest.java
index 6bdde3a..43549fa 100644
--- a/tests/wifitests/src/com/android/server/wifi/nan/WifiNanStateManagerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/nan/WifiNanStateManagerTest.java
@@ -17,22 +17,21 @@
 package com.android.server.wifi.nan;
 
 import static org.hamcrest.core.IsEqual.equalTo;
-import static org.hamcrest.core.IsNot.not;
 import static org.hamcrest.core.IsNull.notNullValue;
 import static org.hamcrest.core.IsNull.nullValue;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.Matchers.anyShort;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.mockito.Mockito.verifyZeroInteractions;
 
 import android.net.wifi.nan.ConfigRequest;
 import android.net.wifi.nan.IWifiNanEventCallback;
 import android.net.wifi.nan.IWifiNanSessionCallback;
 import android.net.wifi.nan.PublishConfig;
 import android.net.wifi.nan.SubscribeConfig;
+import android.net.wifi.nan.WifiNanEventCallback;
 import android.net.wifi.nan.WifiNanSessionCallback;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.util.SparseArray;
@@ -65,6 +64,9 @@
     @Rule
     public ErrorCollector collector = new ErrorCollector();
 
+    /**
+     * Pre-test configuration. Initialize and install mocks.
+     */
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
@@ -78,156 +80,112 @@
     }
 
     /**
-     * Validates that all events are delivered with correct arguments.
+     * Validates that all events are delivered with correct arguments. Validates
+     * that IdentityChanged not delivered if configuration disables delivery.
      */
     @Test
-    public void testNanEventsDelivered() throws Exception {
-        final int clientId = 1005;
-        final int clusterLow1 = 5;
-        final int clusterHigh1 = 100;
-        final int masterPref1 = 111;
-        final int clusterLow2 = 7;
-        final int clusterHigh2 = 155;
-        final int masterPref2 = 0;
-        final int reason = WifiNanSessionCallback.FAIL_REASON_NO_RESOURCES;
-        final byte[] someMac = HexEncoding.decode("000102030405".toCharArray(), false);
-
-        ConfigRequest configRequest1 = new ConfigRequest.Builder().setClusterLow(clusterLow1)
-                .setClusterHigh(clusterHigh1).setMasterPreference(masterPref1).build();
-
-        ConfigRequest configRequest2 = new ConfigRequest.Builder().setClusterLow(clusterLow2)
-                .setClusterHigh(clusterHigh2).setMasterPreference(masterPref2)
-                .setEnableIdentityChangeCallback(true).build();
-
-        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
-        ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
-
-        mDut.connect(clientId, mockCallback);
-        mDut.requestConfig(clientId, configRequest1);
-        mMockLooper.dispatchAll();
-
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest1));
-        short transactionId1 = transactionId.getValue();
-
-        mDut.requestConfig(clientId, configRequest2);
-        mMockLooper.dispatchAll();
-
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest2));
-        short transactionId2 = transactionId.getValue();
-
-        mDut.onClusterChange(WifiNanClientState.CLUSTER_CHANGE_EVENT_STARTED, someMac);
-        mDut.onConfigCompleted(transactionId1);
-        mDut.onConfigFailed(transactionId2, reason);
-        mDut.onInterfaceAddressChange(someMac);
-        mDut.onNanDown(reason);
-        mMockLooper.dispatchAll();
-
-        inOrder.verify(mockCallback).onIdentityChanged();
-        inOrder.verify(mockCallback).onConfigCompleted(configRequest1);
-        inOrder.verify(mockCallback).onConfigFailed(configRequest2, reason);
-        inOrder.verify(mockCallback).onIdentityChanged();
-        inOrder.verify(mockCallback).onNanDown(reason);
-        verifyNoMoreInteractions(mockCallback);
-
-        validateInternalTransactionInfoCleanedUp(transactionId1);
-        validateInternalTransactionInfoCleanedUp(transactionId2);
-    }
-
-    /**
-     * Test that the configuration disabling Identity Change notification works:
-     * trigger changes and validate that aren't delivered.
-     */
-    @Test
-    public void testNanOnIdentityEventsNotDelivered() throws Exception {
-        final int clientId = 1005;
+    public void testNanEventsDelivery() throws Exception {
+        final int clientId1 = 1005;
+        final int clientId2 = 1007;
         final int clusterLow = 5;
         final int clusterHigh = 100;
         final int masterPref = 111;
-        final int reason = WifiNanSessionCallback.FAIL_REASON_NO_RESOURCES;
+        final int reason = WifiNanEventCallback.REASON_OTHER;
         final byte[] someMac = HexEncoding.decode("000102030405".toCharArray(), false);
 
         ConfigRequest configRequest1 = new ConfigRequest.Builder().setClusterLow(clusterLow)
                 .setClusterHigh(clusterHigh).setMasterPreference(masterPref)
-                .setEnableIdentityChangeCallback(true).build();
+                .setEnableIdentityChangeCallback(false).build();
 
         ConfigRequest configRequest2 = new ConfigRequest.Builder().setClusterLow(clusterLow)
                 .setClusterHigh(clusterHigh).setMasterPreference(masterPref)
-                .setEnableIdentityChangeCallback(false).build();
+                .setEnableIdentityChangeCallback(true).build();
 
-        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
-        ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        IWifiNanEventCallback mockCallback1 = mock(IWifiNanEventCallback.class);
+        IWifiNanEventCallback mockCallback2 = mock(IWifiNanEventCallback.class);
+        ArgumentCaptor<Short> transactionIdCapture = ArgumentCaptor.forClass(Short.class);
+        InOrder inOrder = inOrder(mockCallback1, mockCallback2, mMockNative);
 
-        mDut.connect(clientId, mockCallback);
-        mDut.requestConfig(clientId, configRequest1);
+        // (1) connect 1st and 2nd clients
+        mDut.connect(clientId1, mockCallback1, configRequest1);
+        mDut.connect(clientId2, mockCallback2, configRequest2);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionIdCapture.capture(),
+                eq(configRequest1), eq(true));
+        short transactionId = transactionIdCapture.getValue();
+        mDut.onConfigSuccessResponse(transactionId);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback1).onConnectSuccess();
+
+        // (2) finish connection of 2nd client
+        inOrder.verify(mMockNative).enableAndConfigure(transactionIdCapture.capture(),
+                eq(configRequest2), eq(false));
+        transactionId = transactionIdCapture.getValue();
+        mDut.onConfigSuccessResponse(transactionId);
+
+        // (3) deliver NAN events
+        mDut.onClusterChangeNotification(WifiNanClientState.CLUSTER_CHANGE_EVENT_STARTED, someMac);
+        mDut.onInterfaceAddressChangeNotification(someMac);
+        mDut.onNanDownNotification(reason);
         mMockLooper.dispatchAll();
 
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest1));
-        short transactionId1 = transactionId.getValue();
 
-        mDut.requestConfig(clientId, configRequest2);
-        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback2).onConnectSuccess();
+        inOrder.verify(mockCallback2, times(2)).onIdentityChanged();
+        inOrder.verify(mockCallback1).onNanDown(reason);
+        inOrder.verify(mockCallback2).onNanDown(reason);
 
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest2));
-        short transactionId2 = transactionId.getValue();
+        validateInternalClientInfoCleanedUp(clientId1);
+        validateInternalClientInfoCleanedUp(clientId2);
 
-        mDut.onClusterChange(WifiNanClientState.CLUSTER_CHANGE_EVENT_JOINED, someMac);
-        mDut.onConfigCompleted(transactionId1);
-        mDut.onConfigCompleted(transactionId2);
-        mDut.onInterfaceAddressChange(someMac);
-        mDut.onNanDown(reason);
-        mMockLooper.dispatchAll();
-
-        inOrder.verify(mockCallback).onConfigCompleted(configRequest1);
-        inOrder.verify(mockCallback).onConfigCompleted(configRequest2);
-        inOrder.verify(mockCallback).onNanDown(reason);
-        verifyNoMoreInteractions(mockCallback);
-
-        validateInternalTransactionInfoCleanedUp(transactionId1);
-        validateInternalTransactionInfoCleanedUp(transactionId2);
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
-     * Validates publish flow: (1) initial publish (2) fail. Expected: only get
-     * a failure callback.
+     * Validates publish flow: (1) initial publish (2) fail. Expected: get a
+     * failure callback.
      */
     @Test
     public void testPublishFail() throws Exception {
         final int clientId = 1005;
-        final int reasonFail = WifiNanSessionCallback.FAIL_REASON_NO_RESOURCES;
+        final int reasonFail = WifiNanSessionCallback.REASON_NO_RESOURCES;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         PublishConfig publishConfig = new PublishConfig.Builder().build();
 
-        IWifiNanSessionCallback mockCallback = mock(IWifiNanSessionCallback.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
-        mDut.connect(clientId, null);
+        // (0) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
+                eq(configRequest), eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
         // (1) initial publish
-        mDut.publish(clientId, publishConfig, mockCallback);
+        mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
 
         // (2) publish failure
-        mDut.onPublishFail(transactionId.getValue(), reasonFail);
+        mDut.onSessionConfigFailResponse(transactionId.getValue(), true, reasonFail);
         mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onSessionConfigFail(reasonFail);
+        validateInternalNoSessions(clientId);
 
-        inOrder.verify(mockCallback).onSessionConfigFail(reasonFail);
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-
-        verifyNoMoreInteractions(mockCallback, mMockNative);
+        verifyNoMoreInteractions(mockCallback, mockSessionCallback, mMockNative);
     }
 
     /**
      * Validates the publish flow: (1) initial publish (2) success (3)
      * termination (e.g. DONE) (4) update session attempt (5) terminateSession
      * (6) update session attempt. Expected: session ID callback + session
-     * cleaned-up + failure with session terminated error code for the first
-     * attempt - but no callback for the second update attempt (since client
-     * already explicitly asked for termination).
+     * cleaned-up.
      */
     @Test
     public void testPublishSuccessTerminated() throws Exception {
@@ -235,47 +193,56 @@
         final int reasonTerminate = WifiNanSessionCallback.TERMINATE_REASON_DONE;
         final int publishId = 15;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         PublishConfig publishConfig = new PublishConfig.Builder().build();
 
-        IWifiNanSessionCallback mockCallback = mock(IWifiNanSessionCallback.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
-        mDut.connect(clientId, null);
+        // (0) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
+                eq(configRequest), eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
         // (1) initial publish
-        mDut.publish(clientId, publishConfig, mockCallback);
+        mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
 
         // (2) publish success
-        mDut.onPublishSuccess(transactionId.getValue(), publishId);
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback).onSessionStarted(sessionId.capture());
+        inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
         // (3) publish termination (from firmware - not app!)
-        mDut.onPublishTerminated(publishId, reasonTerminate);
+        mDut.onSessionTerminatedNotification(publishId, reasonTerminate, true);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onSessionTerminated(reasonTerminate);
+
         // (4) app update session (race condition: app didn't get termination
         // yet)
         mDut.updatePublish(clientId, sessionId.getValue(), publishConfig);
+        mMockLooper.dispatchAll();
+
         // (5) app terminates session
         mDut.terminateSession(clientId, sessionId.getValue());
+        mMockLooper.dispatchAll();
+
         // (6) app updates session (app already knows that terminated - will get
         // a local FAIL).
         mDut.updatePublish(clientId, sessionId.getValue(), publishConfig);
         mMockLooper.dispatchAll();
 
-        inOrder.verify(mockCallback).onSessionTerminated(reasonTerminate);
-        inOrder.verify(mockCallback)
-                .onSessionConfigFail(WifiNanSessionCallback.FAIL_REASON_SESSION_TERMINATED);
-
         validateInternalSessionInfoCleanedUp(clientId, sessionId.getValue());
 
-        verifyNoMoreInteractions(mockCallback, mMockNative);
+        verifyNoMoreInteractions(mockSessionCallback, mMockNative);
     }
 
     /**
@@ -287,54 +254,58 @@
     public void testPublishUpdateFail() throws Exception {
         final int clientId = 2005;
         final int publishId = 15;
-        final int reasonFail = WifiNanSessionCallback.FAIL_REASON_INVALID_ARGS;
+        final int reasonFail = WifiNanSessionCallback.REASON_INVALID_ARGS;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         PublishConfig publishConfig = new PublishConfig.Builder().build();
 
-        IWifiNanSessionCallback mockCallback = mock(IWifiNanSessionCallback.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
-        mDut.connect(clientId, null);
+        // (0) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
         // (1) initial publish
-        mDut.publish(clientId, publishConfig, mockCallback);
+        mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
 
         // (2) publish success
-        mDut.onPublishSuccess(transactionId.getValue(), publishId);
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback).onSessionStarted(sessionId.capture());
+        inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
         // (3) update publish
         mDut.updatePublish(clientId, sessionId.getValue(), publishConfig);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).publish(transactionId.capture(), eq(publishId),
                 eq(publishConfig));
 
         // (4) update fails
-        mDut.onPublishFail(transactionId.getValue(), reasonFail);
+        mDut.onSessionConfigFailResponse(transactionId.getValue(), true, reasonFail);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onSessionConfigFail(reasonFail);
 
         // (5) another update publish
         mDut.updatePublish(clientId, sessionId.getValue(), publishConfig);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback).onSessionConfigFail(reasonFail);
         inOrder.verify(mMockNative).publish(transactionId.capture(), eq(publishId),
                 eq(publishConfig));
 
         // (6) update succeeds
-        mDut.onPublishSuccess(transactionId.getValue(), publishId);
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
         mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onSessionConfigSuccess();
 
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
         inOrder.verifyNoMoreInteractions();
     }
 
@@ -349,76 +320,89 @@
         final int clientId = 2005;
         final int publishId = 15;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         PublishConfig publishConfig = new PublishConfig.Builder().build();
 
-        IWifiNanSessionCallback mockCallback = mock(IWifiNanSessionCallback.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
-        mDut.connect(clientId, null);
+        // (0) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
-        // (1) initial publish & disconnect
-        mDut.publish(clientId, publishConfig, mockCallback);
+        // (1) initial publish
+        mDut.publish(clientId, publishConfig, mockSessionCallback);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+
+        // (2) disconnect (but doesn't get executed until get response for
+        // publish command)
         mDut.disconnect(clientId);
         mMockLooper.dispatchAll();
 
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
-
-        /*
-         * normally transaction isn't cleaned-up at this point (only after
-         * response for this transaction ID). However, since associated with a
-         * disconnected client should be cleaned-up now.
-         */
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-
-        // (2) publish success
-        mDut.onPublishSuccess(transactionId.getValue(), publishId);
+        // (3) publish success
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).stopPublish(transactionId.capture(), eq(publishId));
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
+        inOrder.verify(mMockNative).disable((short) 0);
+
+        validateInternalClientInfoCleanedUp(clientId);
+
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
-     * Validates subscribe flow: (1) initial subscribe (2) fail. Expected: only
-     * get a failure callback.
+     * Validates subscribe flow: (1) initial subscribe (2) fail. Expected: get a
+     * failure callback.
      */
     @Test
     public void testSubscribeFail() throws Exception {
         final int clientId = 1005;
-        final int reasonFail = WifiNanSessionCallback.FAIL_REASON_NO_RESOURCES;
+        final int reasonFail = WifiNanSessionCallback.REASON_NO_RESOURCES;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
 
-        IWifiNanSessionCallback mockCallback = mock(IWifiNanSessionCallback.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
-        mDut.connect(clientId, null);
+        // (0) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
         // (1) initial subscribe
-        mDut.subscribe(clientId, subscribeConfig, mockCallback);
+        mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
 
         // (2) subscribe failure
-        mDut.onSubscribeFail(transactionId.getValue(), reasonFail);
+        mDut.onSessionConfigFailResponse(transactionId.getValue(), false, reasonFail);
         mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onSessionConfigFail(reasonFail);
+        validateInternalNoSessions(clientId);
 
-        inOrder.verify(mockCallback).onSessionConfigFail(reasonFail);
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-
-        verifyNoMoreInteractions(mockCallback, mMockNative);
+        verifyNoMoreInteractions(mockCallback, mockSessionCallback, mMockNative);
     }
 
     /**
      * Validates the subscribe flow: (1) initial subscribe (2) success (3)
      * termination (e.g. DONE) (4) update session attempt (5) terminateSession
      * (6) update session attempt. Expected: session ID callback + session
-     * cleaned-up + failure with session terminated error code for the first
-     * attempt - but no callback for the second update attempt (since client
-     * already explicitly asked for termination).
+     * cleaned-up
      */
     @Test
     public void testSubscribeSuccessTerminated() throws Exception {
@@ -426,47 +410,55 @@
         final int reasonTerminate = WifiNanSessionCallback.TERMINATE_REASON_DONE;
         final int subscribeId = 15;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
 
-        IWifiNanSessionCallback mockCallback = mock(IWifiNanSessionCallback.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
-        mDut.connect(clientId, null);
+        // (0) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
         // (1) initial subscribe
-        mDut.subscribe(clientId, subscribeConfig, mockCallback);
+        mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
 
         // (2) subscribe success
-        mDut.onSubscribeSuccess(transactionId.getValue(), subscribeId);
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback).onSessionStarted(sessionId.capture());
+        inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
         // (3) subscribe termination (from firmware - not app!)
-        mDut.onSubscribeTerminated(subscribeId, reasonTerminate);
+        mDut.onSessionTerminatedNotification(subscribeId, reasonTerminate, false);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onSessionTerminated(reasonTerminate);
+
         // (4) app update session (race condition: app didn't get termination
         // yet)
         mDut.updateSubscribe(clientId, sessionId.getValue(), subscribeConfig);
+        mMockLooper.dispatchAll();
+
         // (5) app terminates session
         mDut.terminateSession(clientId, sessionId.getValue());
-        // (6) app updates session (app already knows that terminated - will get
-        // a local FAIL).
+        mMockLooper.dispatchAll();
+
+        // (6) app updates session
         mDut.updateSubscribe(clientId, sessionId.getValue(), subscribeConfig);
         mMockLooper.dispatchAll();
 
-        inOrder.verify(mockCallback).onSessionTerminated(reasonTerminate);
-        inOrder.verify(mockCallback)
-                .onSessionConfigFail(WifiNanSessionCallback.FAIL_REASON_SESSION_TERMINATED);
-
         validateInternalSessionInfoCleanedUp(clientId, sessionId.getValue());
 
-        verifyNoMoreInteractions(mockCallback, mMockNative);
+        verifyNoMoreInteractions(mockSessionCallback, mMockNative);
     }
 
     /**
@@ -478,54 +470,58 @@
     public void testSubscribeUpdateFail() throws Exception {
         final int clientId = 2005;
         final int subscribeId = 15;
-        final int reasonFail = WifiNanSessionCallback.FAIL_REASON_INVALID_ARGS;
+        final int reasonFail = WifiNanSessionCallback.REASON_INVALID_ARGS;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
 
-        IWifiNanSessionCallback mockCallback = mock(IWifiNanSessionCallback.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
-        mDut.connect(clientId, null);
+        // (0) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
         // (1) initial subscribe
-        mDut.subscribe(clientId, subscribeConfig, mockCallback);
+        mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
 
         // (2) subscribe success
-        mDut.onSubscribeSuccess(transactionId.getValue(), subscribeId);
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback).onSessionStarted(sessionId.capture());
+        inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
         // (3) update subscribe
         mDut.updateSubscribe(clientId, sessionId.getValue(), subscribeConfig);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(subscribeId),
                 eq(subscribeConfig));
 
         // (4) update fails
-        mDut.onSubscribeFail(transactionId.getValue(), reasonFail);
+        mDut.onSessionConfigFailResponse(transactionId.getValue(), false, reasonFail);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onSessionConfigFail(reasonFail);
 
         // (5) another update subscribe
         mDut.updateSubscribe(clientId, sessionId.getValue(), subscribeConfig);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback).onSessionConfigFail(reasonFail);
         inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(subscribeId),
                 eq(subscribeConfig));
 
         // (6) update succeeds
-        mDut.onSubscribeSuccess(transactionId.getValue(), subscribeId);
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onSessionConfigSuccess();
 
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
         inOrder.verifyNoMoreInteractions();
     }
 
@@ -540,34 +536,42 @@
         final int clientId = 2005;
         final int subscribeId = 15;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
 
-        IWifiNanSessionCallback mockCallback = mock(IWifiNanSessionCallback.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
-        mDut.connect(clientId, null);
+        // (0) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
-        // (1) initial subscribe & disconnect
-        mDut.subscribe(clientId, subscribeConfig, mockCallback);
+        // (1) initial subscribe
+        mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+
+        // (2) disconnect (but doesn't get executed until get response for
+        // subscribe command)
         mDut.disconnect(clientId);
         mMockLooper.dispatchAll();
 
-        inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
-
-        /*
-         * normally transaction isn't cleaned-up at this point (only after
-         * response for this transaction ID). However, since associated with a
-         * disconnected client should be cleaned-up now.
-         */
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-
-        // (2) subscribe success
-        mDut.onSubscribeSuccess(transactionId.getValue(), subscribeId);
+        // (3) subscribe success
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).stopSubscribe((short) 0, subscribeId);
+        inOrder.verify(mMockNative).disable((short) 0);
 
-        inOrder.verify(mMockNative).stopSubscribe(transactionId.capture(), eq(subscribeId));
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
+        validateInternalClientInfoCleanedUp(clientId);
+
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
@@ -581,7 +585,7 @@
         final String serviceName = "some-service-name";
         final String ssi = "some much longer and more arbitrary data";
         final int subscribeCount = 7;
-        final int reasonFail = WifiNanSessionCallback.FAIL_REASON_NO_RESOURCES;
+        final int reasonFail = WifiNanSessionCallback.REASON_TX_FAIL;
         final int subscribeId = 15;
         final int requestorId = 22;
         final byte[] peerMac = HexEncoding.decode("060708090A0B".toCharArray(), false);
@@ -590,63 +594,70 @@
         final String peerMsg = "some message from peer";
         final int messageId = 6948;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().setServiceName(serviceName)
                 .setServiceSpecificInfo(ssi)
                 .setSubscribeType(SubscribeConfig.SUBSCRIBE_TYPE_PASSIVE)
                 .setSubscribeCount(subscribeCount).build();
 
-        IWifiNanSessionCallback mockCallback = mock(IWifiNanSessionCallback.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
-        mDut.connect(clientId, null);
-        mDut.subscribe(clientId, subscribeConfig, mockCallback);
+        // (0) connect
+        mDut.connect(clientId, mockCallback, configRequest);
         mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
+                eq(configRequest), eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
+        // (1) subscribe
+        mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
+        mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
-        mDut.onSubscribeSuccess(transactionId.getValue(), subscribeId);
-        mDut.onMatch(subscribeId, requestorId, peerMac, peerSsi.getBytes(), peerSsi.length(),
-                peerMatchFilter.getBytes(), peerMatchFilter.length());
-        mDut.onMessageReceived(subscribeId, requestorId, peerMac, peerMsg.getBytes(),
+        // (2) match
+        mDut.onMatchNotification(subscribeId, requestorId, peerMac, peerSsi.getBytes(),
+                peerSsi.length(), peerMatchFilter.getBytes(), peerMatchFilter.length());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onMatch(requestorId, peerSsi.getBytes(),
+                peerSsi.length(), peerMatchFilter.getBytes(), peerMatchFilter.length());
+
+        // (3) message Rx
+        mDut.onMessageReceivedNotification(subscribeId, requestorId, peerMac, peerMsg.getBytes(),
                 peerMsg.length());
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback).onSessionStarted(sessionId.capture());
-        inOrder.verify(mockCallback).onMatch(requestorId, peerSsi.getBytes(), peerSsi.length(),
-                peerMatchFilter.getBytes(), peerMatchFilter.length());
-        inOrder.verify(mockCallback).onMessageReceived(requestorId, peerMsg.getBytes(),
+        inOrder.verify(mockSessionCallback).onMessageReceived(requestorId, peerMsg.getBytes(),
                 peerMsg.length());
 
+        // (4) message Tx fail
         mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(), ssi.length(),
                 messageId);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
                 eq(requestorId), eq(peerMac), eq(ssi.getBytes()), eq(ssi.length()));
-
-        mDut.onMessageSendFail(transactionId.getValue(), reasonFail);
+        mDut.onMessageSendFailResponse(transactionId.getValue(), reasonFail);
         mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onMessageSendFail(messageId, reasonFail);
 
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback).onMessageSendFail(messageId, reasonFail);
-
+        // (5) message Tx success
         mDut.sendMessage(clientId, sessionId.getValue(), requestorId, ssi.getBytes(), ssi.length(),
                 messageId);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(subscribeId),
                 eq(requestorId), eq(peerMac), eq(ssi.getBytes()), eq(ssi.length()));
-
-        mDut.onMessageSendSuccess(transactionId.getValue());
+        mDut.onMessageSendSuccessResponse(transactionId.getValue());
         mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onMessageSendSuccess(messageId);
 
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback).onMessageSendSuccess(messageId);
-
-        verifyNoMoreInteractions(mockCallback, mMockNative);
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
@@ -671,7 +682,7 @@
         final String msgToPeer2 = "hey there 0506...";
         final int msgToPeerId1 = 546;
         final int msgToPeerId2 = 9654;
-        final int reason = WifiNanSessionCallback.FAIL_REASON_OTHER;
+        final int reason = WifiNanSessionCallback.REASON_OTHER;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().setClusterLow(clusterLow)
                 .setClusterHigh(clusterHigh).setMasterPreference(masterPref).build();
@@ -685,56 +696,55 @@
         IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         InOrder inOrder = inOrder(mMockNative, mockCallback, mockSessionCallback);
 
-        mDut.connect(clientId, mockCallback);
-        mDut.requestConfig(clientId, configRequest);
+        // (1) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        // (2) publish
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest));
-        short transactionIdConfig = transactionId.getValue();
-
         inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
-        short transactionIdPublish = transactionId.getValue();
-
-        mDut.onConfigCompleted(transactionIdConfig);
-        mDut.onPublishSuccess(transactionIdPublish, publishId);
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionIdConfig);
-        validateInternalTransactionInfoCleanedUp(transactionIdPublish);
-        inOrder.verify(mockCallback).onConfigCompleted(configRequest);
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
-        mDut.onMessageReceived(publishId, peerId1, peerMac1, msgFromPeer1.getBytes(),
+        // (3) message received from peers 1 & 2
+        mDut.onMessageReceivedNotification(publishId, peerId1, peerMac1, msgFromPeer1.getBytes(),
                 msgFromPeer1.length());
-        mDut.onMessageReceived(publishId, peerId2, peerMac2, msgFromPeer2.getBytes(),
+        mDut.onMessageReceivedNotification(publishId, peerId2, peerMac2, msgFromPeer2.getBytes(),
                 msgFromPeer2.length());
-        mDut.sendMessage(clientId, sessionId.getValue(), peerId2, msgToPeer2.getBytes(),
-                msgToPeer2.length(), msgToPeerId2);
-        mDut.sendMessage(clientId, sessionId.getValue(), peerId1, msgToPeer1.getBytes(),
-                msgToPeer1.length(), msgToPeerId1);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mockSessionCallback).onMessageReceived(peerId1, msgFromPeer1.getBytes(),
                 msgFromPeer1.length());
         inOrder.verify(mockSessionCallback).onMessageReceived(peerId2, msgFromPeer2.getBytes(),
                 msgFromPeer2.length());
+
+        // (4) sending messages back to same peers: one Tx fails, other succeeds
+        mDut.sendMessage(clientId, sessionId.getValue(), peerId2, msgToPeer2.getBytes(),
+                msgToPeer2.length(), msgToPeerId2);
+        mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId), eq(peerId2),
                 eq(peerMac2), eq(msgToPeer2.getBytes()), eq(msgToPeer2.length()));
-        short transactionIdMsg2 = transactionId.getValue();
+        short transactionIdVal = transactionId.getValue();
+        mDut.onMessageSendSuccessResponse(transactionIdVal);
+
+        mDut.sendMessage(clientId, sessionId.getValue(), peerId1, msgToPeer1.getBytes(),
+                msgToPeer1.length(), msgToPeerId1);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onMessageSendSuccess(msgToPeerId2);
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId), eq(peerId1),
                 eq(peerMac1), eq(msgToPeer1.getBytes()), eq(msgToPeer1.length()));
-        short transactionIdMsg1 = transactionId.getValue();
-
-        mDut.onMessageSendFail(transactionIdMsg1, reason);
-        mDut.onMessageSendSuccess(transactionIdMsg2);
+        transactionIdVal = transactionId.getValue();
+        mDut.onMessageSendFailResponse(transactionIdVal, reason);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionIdMsg1);
-        validateInternalTransactionInfoCleanedUp(transactionIdMsg2);
         inOrder.verify(mockSessionCallback).onMessageSendFail(msgToPeerId1, reason);
-        inOrder.verify(mockSessionCallback).onMessageSendSuccess(msgToPeerId2);
-        verifyNoMoreInteractions(mMockNative, mockCallback, mockSessionCallback);
+
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
@@ -770,62 +780,58 @@
         IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         InOrder inOrder = inOrder(mMockNative, mockCallback, mockSessionCallback);
 
-        mDut.connect(clientId, mockCallback);
-        mDut.requestConfig(clientId, configRequest);
+        // (1) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        // (2) publish
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest));
-        short transactionIdConfig = transactionId.getValue();
-
         inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
-        short transactionIdPublish = transactionId.getValue();
-
-        mDut.onConfigCompleted(transactionIdConfig);
-        mDut.onPublishSuccess(transactionIdPublish, publishId);
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionIdConfig);
-        validateInternalTransactionInfoCleanedUp(transactionIdPublish);
-        inOrder.verify(mockCallback).onConfigCompleted(configRequest);
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
-        mDut.onMessageReceived(publishId, peerId, peerMacOrig, msgFromPeer1.getBytes(),
+        // (3) message received & responded to
+        mDut.onMessageReceivedNotification(publishId, peerId, peerMacOrig, msgFromPeer1.getBytes(),
                 msgFromPeer1.length());
         mDut.sendMessage(clientId, sessionId.getValue(), peerId, msgToPeer1.getBytes(),
                 msgToPeer1.length(), msgToPeerId1);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mockSessionCallback).onMessageReceived(peerId, msgFromPeer1.getBytes(),
                 msgFromPeer1.length());
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId), eq(peerId),
                 eq(peerMacOrig), eq(msgToPeer1.getBytes()), eq(msgToPeer1.length()));
-        short transactionIdMsg = transactionId.getValue();
+        mDut.onMessageSendSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onMessageSendSuccess(msgToPeerId1);
 
-        mDut.onMessageSendSuccess(transactionIdMsg);
-        mDut.onMessageReceived(publishId, peerId, peerMacLater, msgFromPeer2.getBytes(),
+        // (4) message received with same peer ID but different MAC
+        mDut.onMessageReceivedNotification(publishId, peerId, peerMacLater, msgFromPeer2.getBytes(),
                 msgFromPeer2.length());
         mDut.sendMessage(clientId, sessionId.getValue(), peerId, msgToPeer2.getBytes(),
                 msgToPeer2.length(), msgToPeerId2);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionIdMsg);
-        inOrder.verify(mockSessionCallback).onMessageSendSuccess(msgToPeerId1);
         inOrder.verify(mockSessionCallback).onMessageReceived(peerId, msgFromPeer2.getBytes(),
                 msgFromPeer2.length());
         inOrder.verify(mMockNative).sendMessage(transactionId.capture(), eq(publishId), eq(peerId),
                 eq(peerMacLater), eq(msgToPeer2.getBytes()), eq(msgToPeer2.length()));
-        transactionIdMsg = transactionId.getValue();
-
-        mDut.onMessageSendSuccess(transactionIdMsg);
+        mDut.onMessageSendSuccessResponse(transactionId.getValue());
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionIdMsg);
         inOrder.verify(mockSessionCallback).onMessageSendSuccess(msgToPeerId2);
 
-        verifyNoMoreInteractions(mMockNative, mockCallback, mockSessionCallback);
+        inOrder.verifyNoMoreInteractions();
     }
 
+    /**
+     * Validate that get failure (with correct code) when trying to send a
+     * message to an invalid peer ID.
+     */
     @Test
     public void testSendMessageToInvalidPeerId() throws Exception {
         final int clientId = 1005;
@@ -837,40 +843,52 @@
         final String peerMatchFilter = "filter binary array represented as string";
         final int messageId = 6948;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
 
-        IWifiNanSessionCallback mockCallback = mock(IWifiNanSessionCallback.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
-        InOrder inOrder = inOrder(mockCallback, mMockNative);
+        InOrder inOrder = inOrder(mockCallback, mockSessionCallback, mMockNative);
 
-        mDut.connect(clientId, null);
-        mDut.subscribe(clientId, subscribeConfig, mockCallback);
+        // (1) connect
+        mDut.connect(clientId, mockCallback, configRequest);
         mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
+        // (2) subscribe & match
+        mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
+        mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
-
-        mDut.onSubscribeSuccess(transactionId.getValue(), subscribeId);
-        mDut.onMatch(subscribeId, requestorId, peerMac, peerSsi.getBytes(), peerSsi.length(),
-                peerMatchFilter.getBytes(), peerMatchFilter.length());
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
+        mDut.onMatchNotification(subscribeId, requestorId, peerMac, peerSsi.getBytes(),
+                peerSsi.length(), peerMatchFilter.getBytes(), peerMatchFilter.length());
         mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
+        inOrder.verify(mockSessionCallback).onMatch(requestorId, peerSsi.getBytes(),
+                peerSsi.length(), peerMatchFilter.getBytes(), peerMatchFilter.length());
 
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback).onSessionStarted(sessionId.capture());
-        inOrder.verify(mockCallback).onMatch(requestorId, peerSsi.getBytes(), peerSsi.length(),
-                peerMatchFilter.getBytes(), peerMatchFilter.length());
-
+        // (3) send message to invalid peer ID
         mDut.sendMessage(clientId, sessionId.getValue(), requestorId + 5, ssi.getBytes(),
                 ssi.length(), messageId);
         mMockLooper.dispatchAll();
+        inOrder.verify(mockSessionCallback).onMessageSendFail(messageId,
+                WifiNanSessionCallback.REASON_NO_MATCH_SESSION);
 
-        inOrder.verify(mockCallback).onMessageSendFail(messageId,
-                WifiNanSessionCallback.FAIL_REASON_NO_MATCH_SESSION);
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-
-        verifyNoMoreInteractions(mockCallback, mMockNative);
+        inOrder.verifyNoMoreInteractions();
     }
 
+    /**
+     * Test sequence of configuration: (1) config1, (2) config2 - incompatible,
+     * (3) config3 - compatible with config1 (requiring upgrade), (4) disconnect
+     * config3 (should get a downgrade), (5) disconnect config1 (should get a
+     * disable).
+     */
     @Test
     public void testConfigs() throws Exception {
         final int clientId1 = 9999;
@@ -888,13 +906,16 @@
         ArgumentCaptor<ConfigRequest> crCapture = ArgumentCaptor.forClass(ConfigRequest.class);
 
         ConfigRequest configRequest1 = new ConfigRequest.Builder().setClusterLow(clusterLow1)
-                .setClusterHigh(clusterHigh1).setMasterPreference(masterPref1).build();
+                .setClusterHigh(clusterHigh1).setMasterPreference(masterPref1)
+                .setEnableIdentityChangeCallback(false).build();
 
         ConfigRequest configRequest2 = new ConfigRequest.Builder().setSupport5gBand(support5g2)
                 .setClusterLow(clusterLow2).setClusterHigh(clusterHigh2)
                 .setMasterPreference(masterPref2).build();
 
-        ConfigRequest configRequest3 = new ConfigRequest.Builder().build();
+        ConfigRequest configRequest3 = new ConfigRequest.Builder().setClusterLow(clusterLow1)
+                .setClusterHigh(clusterHigh1).setMasterPreference(masterPref1)
+                .setEnableIdentityChangeCallback(true).build();
 
         IWifiNanEventCallback mockCallback1 = mock(IWifiNanEventCallback.class);
         IWifiNanEventCallback mockCallback2 = mock(IWifiNanEventCallback.class);
@@ -902,103 +923,64 @@
 
         InOrder inOrder = inOrder(mMockNative, mockCallback1, mockCallback2, mockCallback3);
 
-        mDut.connect(clientId1, mockCallback1);
-        mDut.requestConfig(clientId1, configRequest1);
+        // (1) config1 (valid)
+        mDut.connect(clientId1, mockCallback1, configRequest1);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
-                crCapture.capture());
-        collector.checkThat("merge: stage 0", configRequest1, equalTo(crCapture.getValue()));
-
-        mDut.onConfigCompleted(transactionId.getValue());
+                crCapture.capture(), eq(true));
+        collector.checkThat("merge: stage 1", crCapture.getValue(), equalTo(configRequest1));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
         mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback1).onConnectSuccess();
 
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback1).onConfigCompleted(configRequest1);
-
-        mDut.connect(clientId2, mockCallback2);
-        mDut.requestConfig(clientId2, configRequest2);
+        // (2) config2 (incompatible with config1)
+        mDut.connect(clientId2, mockCallback2, configRequest2);
         mMockLooper.dispatchAll();
-
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
-                crCapture.capture());
-        collector.checkThat("merge: stage 1: support 5g", crCapture.getValue().mSupport5gBand,
-                equalTo(true));
-        collector.checkThat("merge: stage 1: master pref", crCapture.getValue().mMasterPreference,
-                equalTo(Math.max(masterPref1, masterPref2)));
-        collector.checkThat("merge: stage 1: cluster low", crCapture.getValue().mClusterLow,
-                equalTo(Math.min(clusterLow1, clusterLow2)));
-        collector.checkThat("merge: stage 1: cluster high", crCapture.getValue().mClusterHigh,
-                equalTo(Math.max(clusterHigh1, clusterHigh2)));
-
-        mDut.onConfigCompleted(transactionId.getValue());
-        mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback1).onConfigCompleted(crCapture.getValue());
-
-        mDut.connect(clientId3, mockCallback3);
-        mDut.requestConfig(clientId3, configRequest3);
-        mMockLooper.dispatchAll();
-
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
-                crCapture.capture());
-        collector.checkThat("merge: stage 2: support 5g", crCapture.getValue().mSupport5gBand,
-                equalTo(true));
-        collector.checkThat("merge: stage 2: master pref", crCapture.getValue().mMasterPreference,
-                equalTo(Math.max(masterPref1, masterPref2)));
-        collector.checkThat("merge: stage 2: cluster low", crCapture.getValue().mClusterLow,
-                equalTo(Math.min(clusterLow1, clusterLow2)));
-        collector.checkThat("merge: stage 2: cluster high", crCapture.getValue().mClusterHigh,
-                equalTo(Math.max(clusterHigh1, clusterHigh2)));
-
-        mDut.onConfigCompleted(transactionId.getValue());
-        mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback1).onConfigCompleted(crCapture.getValue());
-
-        mDut.disconnect(clientId2);
-        mMockLooper.dispatchAll();
-
+        inOrder.verify(mockCallback2)
+                .onConnectFail(WifiNanEventCallback.REASON_ALREADY_CONNECTED_INCOMPAT_CONFIG);
         validateInternalClientInfoCleanedUp(clientId2);
+
+        // (3) config3 (compatible with config1 but requires upgrade - i.e. no
+        // OTA changes)
+        mDut.connect(clientId3, mockCallback3, configRequest3);
+        mMockLooper.dispatchAll();
         inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
-                crCapture.capture());
-        collector.checkThat("merge: stage 3", configRequest1, equalTo(crCapture.getValue()));
-
-        mDut.onConfigCompleted(transactionId.getValue());
+                crCapture.capture(), eq(false));
+        collector.checkThat("merge: stage 3: support 5g", crCapture.getValue().mSupport5gBand,
+                equalTo(false));
+        collector.checkThat("merge: stage 3: master pref", crCapture.getValue().mMasterPreference,
+                equalTo(masterPref1));
+        collector.checkThat("merge: stage 3: cluster low", crCapture.getValue().mClusterLow,
+                equalTo(clusterLow1));
+        collector.checkThat("merge: stage 3: cluster high", crCapture.getValue().mClusterHigh,
+                equalTo(clusterHigh1));
+        collector.checkThat("merge: stage 3: enable identity change callback",
+                crCapture.getValue().mEnableIdentityChangeCallback, equalTo(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
         mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback3).onConnectSuccess();
 
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback1).onConfigCompleted(crCapture.getValue());
-
-        mDut.disconnect(clientId1);
-        mMockLooper.dispatchAll();
-
-        validateInternalClientInfoCleanedUp(clientId2);
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
-                crCapture.capture());
-        collector.checkThat("merge: stage 4", configRequest3, equalTo(crCapture.getValue()));
-
-        mDut.onConfigCompleted(transactionId.getValue());
-        mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
-        inOrder.verify(mockCallback3).onConfigCompleted(crCapture.getValue());
-
+        // (4) disconnect config3: want a downgrade
         mDut.disconnect(clientId3);
         mMockLooper.dispatchAll();
+        validateInternalClientInfoCleanedUp(clientId3);
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
+                crCapture.capture(), eq(false));
+        collector.checkThat("merge: stage 4", crCapture.getValue(), equalTo(configRequest1));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
 
-        validateInternalClientInfoCleanedUp(clientId2);
-        inOrder.verify(mMockNative).disable(anyShort());
+        // (5) disconnect config1: disable
+        mDut.disconnect(clientId1);
+        mMockLooper.dispatchAll();
+        validateInternalClientInfoCleanedUp(clientId1);
+        inOrder.verify(mMockNative).disable((short) 0);
 
-        verifyNoMoreInteractions(mMockNative);
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
      * Summary: disconnect a client while there are pending transactions.
-     * Validate that no callbacks are called and that internal state is
-     * cleaned-up.
      */
     @Test
     public void testDisconnectWithPendingTransactions() throws Exception {
@@ -1024,31 +1006,42 @@
         IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         InOrder inOrder = inOrder(mMockNative, mockCallback, mockSessionCallback);
 
-        mDut.connect(clientId, mockCallback);
-        mDut.requestConfig(clientId, configRequest);
-        mDut.publish(clientId, publishConfig, mockSessionCallback);
-        mDut.disconnect(clientId);
+        // (1) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        // (2) publish (no response yet)
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest));
-        short transactionIdConfig = transactionId.getValue();
-
         inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
-        short transactionIdPublish = transactionId.getValue();
+
+        // (3) disconnect (but doesn't get executed until get a RESPONSE to the
+        // previous publish)
+        mDut.disconnect(clientId);
+        mMockLooper.dispatchAll();
+
+        // (4) get successful response to the publish
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).stopPublish((short) 0, publishId);
+        inOrder.verify(mMockNative).disable((short) 0);
 
         validateInternalClientInfoCleanedUp(clientId);
-        validateInternalTransactionInfoCleanedUp(transactionIdPublish);
 
-        mDut.onConfigCompleted(transactionIdConfig);
-        mDut.onPublishSuccess(transactionIdPublish, publishId);
+        // (5) trying to publish on the same client: NOP
+        mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
 
-        mDut.onPublishTerminated(publishId, reason);
+        // (6) got some callback on original publishId - should be ignored
+        mDut.onSessionTerminatedNotification(publishId, reason, true);
         mMockLooper.dispatchAll();
 
-        inOrder.verify(mMockNative).disable(anyShort());
-        verifyZeroInteractions(mockCallback, mockSessionCallback);
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
@@ -1064,8 +1057,6 @@
         final String serviceName = "some-service-name";
         final String ssi = "some much longer and more arbitrary data";
         final int publishCount = 15;
-        final int status = WifiNanSessionCallback.FAIL_REASON_OTHER;
-        final int responseType = 9999;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().setClusterLow(clusterLow)
                 .setClusterHigh(clusterHigh).setMasterPreference(masterPref).build();
@@ -1079,25 +1070,20 @@
         IWifiNanSessionCallback mockPublishSessionCallback = mock(IWifiNanSessionCallback.class);
         InOrder inOrder = inOrder(mMockNative, mockCallback, mockPublishSessionCallback);
 
-        mDut.connect(clientId, mockCallback);
-        mDut.requestConfig(clientId, configRequest);
+        // (1) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+
+        // (2) publish - no response
         mDut.publish(clientId, publishConfig, mockPublishSessionCallback);
         mMockLooper.dispatchAll();
-
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest));
-        short transactionIdConfig = transactionId.getValue();
-
         inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
-        short transactionIdPublish = transactionId.getValue();
 
-        mDut.onUnknownTransaction(responseType, transactionIdConfig, status);
-        mDut.onUnknownTransaction(responseType, transactionIdPublish, status);
-        mMockLooper.dispatchAll();
-
-        verifyNoMoreInteractions(mockCallback);
-        verifyNoMoreInteractions(mockPublishSessionCallback);
-        validateInternalTransactionInfoCleanedUp(transactionIdConfig);
-        validateInternalTransactionInfoCleanedUp(transactionIdPublish);
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
@@ -1109,25 +1095,20 @@
     public void testNoOpTransaction() throws Exception {
         final int clientId = 1294;
 
-        PublishConfig publishConfig = new PublishConfig.Builder().build();
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
 
         ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
         IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
         IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         InOrder inOrder = inOrder(mMockNative, mockCallback, mockSessionCallback);
 
-        mDut.connect(clientId, mockCallback);
-        mDut.publish(clientId, publishConfig, mockSessionCallback);
+        // (1) connect (no response)
+        mDut.connect(clientId, mockCallback, configRequest);
         mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
 
-        inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
-
-        mDut.onNoOpTransaction(transactionId.getValue());
-        mMockLooper.dispatchAll();
-
-        verifyNoMoreInteractions(mockCallback);
-        verifyNoMoreInteractions(mockSessionCallback);
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
@@ -1137,6 +1118,7 @@
      */
     @Test
     public void testInvalidCallbackIdParameters() throws Exception {
+        final int pubSubId = 1235;
         final int clientId = 132;
 
         ConfigRequest configRequest = new ConfigRequest.Builder().build();
@@ -1145,36 +1127,32 @@
         IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
         InOrder inOrder = inOrder(mMockNative, mockCallback);
 
-        mDut.connect(clientId, mockCallback);
-        mDut.requestConfig(clientId, configRequest);
+        // (1) connect and succeed
+        mDut.connect(clientId, mockCallback, configRequest);
         mMockLooper.dispatchAll();
-
-        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest));
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
         short transactionIdConfig = transactionId.getValue();
+        mDut.onConfigSuccessResponse(transactionIdConfig);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
 
-        mDut.onConfigCompleted(transactionIdConfig);
+        // (2) use the same transaction ID to send a bunch of other responses
+        mDut.onConfigSuccessResponse(transactionIdConfig);
+        mDut.onConfigFailedResponse(transactionIdConfig, -1);
+        mDut.onSessionConfigFailResponse(transactionIdConfig, true, -1);
+        mDut.onMessageSendSuccessResponse(transactionIdConfig);
+        mDut.onMessageSendFailResponse(transactionIdConfig, -1);
+        mDut.onSessionConfigFailResponse(transactionIdConfig, false, -1);
+        mDut.onMatchNotification(-1, -1, new byte[0], new byte[0], 0, new byte[0], 0);
+        mDut.onSessionTerminatedNotification(-1, -1, true);
+        mDut.onSessionTerminatedNotification(-1, -1, false);
+        mDut.onMessageReceivedNotification(-1, -1, new byte[0], new byte[0], 0);
+        mDut.onSessionConfigSuccessResponse(transactionIdConfig, true, pubSubId);
+        mDut.onSessionConfigSuccessResponse(transactionIdConfig, false, pubSubId);
         mMockLooper.dispatchAll();
 
-        inOrder.verify(mockCallback).onConfigCompleted(configRequest);
-        validateInternalTransactionInfoCleanedUp(transactionIdConfig);
-
-        mDut.onCapabilitiesUpdate(transactionIdConfig, null);
-        mDut.onConfigCompleted(transactionIdConfig);
-        mDut.onConfigFailed(transactionIdConfig, -1);
-        mDut.onPublishSuccess(transactionIdConfig, -1);
-        mDut.onPublishFail(transactionIdConfig, -1);
-        mDut.onMessageSendSuccess(transactionIdConfig);
-        mDut.onMessageSendFail(transactionIdConfig, -1);
-        mDut.onSubscribeSuccess(transactionIdConfig, -1);
-        mDut.onSubscribeFail(transactionIdConfig, -1);
-        mDut.onUnknownTransaction(-10, transactionIdConfig, -1);
-        mDut.onMatch(-1, -1, null, null, 0, null, 0);
-        mDut.onPublishTerminated(-1, -1);
-        mDut.onSubscribeTerminated(-1, -1);
-        mDut.onMessageReceived(-1, -1, null, null, 0);
-        mMockLooper.dispatchAll();
-
-        verifyNoMoreInteractions(mockCallback);
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
@@ -1185,6 +1163,7 @@
         final int clientId = 188;
         final int publishId = 25;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         PublishConfig publishConfig = new PublishConfig.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
 
@@ -1194,24 +1173,30 @@
         IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         InOrder inOrder = inOrder(mMockNative, mockCallback, mockSessionCallback);
 
-        mDut.connect(clientId, mockCallback);
+        // (1) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(), eq(configRequest),
+                eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        // (2) publish
         mDut.publish(clientId, publishConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
-
-        mDut.onPublishSuccess(transactionId.getValue(), publishId);
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, publishId);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
+        // (3) update-subscribe -> failure
         mDut.updateSubscribe(clientId, sessionId.getValue(), subscribeConfig);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mockSessionCallback)
-                .onSessionConfigFail(WifiNanSessionCallback.FAIL_REASON_OTHER);
-        verifyNoMoreInteractions(mMockNative, mockCallback, mockSessionCallback);
+                .onSessionConfigFail(WifiNanSessionCallback.REASON_OTHER);
+
+        inOrder.verifyNoMoreInteractions();
     }
 
     /**
@@ -1223,6 +1208,7 @@
         final int clientId = 188;
         final int subscribeId = 25;
 
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
         PublishConfig publishConfig = new PublishConfig.Builder().build();
         SubscribeConfig subscribeConfig = new SubscribeConfig.Builder().build();
 
@@ -1232,37 +1218,74 @@
         IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
         InOrder inOrder = inOrder(mMockNative, mockCallback, mockSessionCallback);
 
-        mDut.connect(clientId, mockCallback);
+        // (1) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
+                eq(configRequest), eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        // (2) subscribe
         mDut.subscribe(clientId, subscribeConfig, mockSessionCallback);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mMockNative).subscribe(transactionId.capture(), eq(0), eq(subscribeConfig));
-
-        mDut.onPublishSuccess(transactionId.getValue(), subscribeId);
+        mDut.onSessionConfigSuccessResponse(transactionId.getValue(), false, subscribeId);
         mMockLooper.dispatchAll();
-
-        validateInternalTransactionInfoCleanedUp(transactionId.getValue());
         inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
 
+        // (3) update-publish -> error
         mDut.updatePublish(clientId, sessionId.getValue(), publishConfig);
         mMockLooper.dispatchAll();
-
         inOrder.verify(mockSessionCallback)
-                .onSessionConfigFail(WifiNanSessionCallback.FAIL_REASON_OTHER);
-        verifyNoMoreInteractions(mMockNative, mockCallback, mockSessionCallback);
+                .onSessionConfigFail(WifiNanSessionCallback.REASON_OTHER);
+
+        inOrder.verifyNoMoreInteractions();
     }
 
+    /**
+     * Validate that the session ID increments monotonically
+     */
     @Test
-    public void testTransactionIdIncrement() {
+    public void testSessionIdIncrement() throws Exception {
+        final int clientId = 188;
         int loopCount = 100;
 
-        short prevId = 0;
+        ConfigRequest configRequest = new ConfigRequest.Builder().build();
+        PublishConfig publishConfig = new PublishConfig.Builder().build();
+
+        ArgumentCaptor<Short> transactionId = ArgumentCaptor.forClass(Short.class);
+        ArgumentCaptor<Integer> sessionId = ArgumentCaptor.forClass(Integer.class);
+        IWifiNanEventCallback mockCallback = mock(IWifiNanEventCallback.class);
+        IWifiNanSessionCallback mockSessionCallback = mock(IWifiNanSessionCallback.class);
+        InOrder inOrder = inOrder(mMockNative, mockCallback, mockSessionCallback);
+
+        // (1) connect
+        mDut.connect(clientId, mockCallback, configRequest);
+        mMockLooper.dispatchAll();
+        inOrder.verify(mMockNative).enableAndConfigure(transactionId.capture(),
+                eq(configRequest), eq(true));
+        mDut.onConfigSuccessResponse(transactionId.getValue());
+        mMockLooper.dispatchAll();
+        inOrder.verify(mockCallback).onConnectSuccess();
+
+        int prevId = 0;
         for (int i = 0; i < loopCount; ++i) {
-            short id = mDut.createNextTransactionId();
+            // (2) publish
+            mDut.publish(clientId, publishConfig, mockSessionCallback);
+            mMockLooper.dispatchAll();
+            inOrder.verify(mMockNative).publish(transactionId.capture(), eq(0), eq(publishConfig));
+
+            // (3) publish-success
+            mDut.onSessionConfigSuccessResponse(transactionId.getValue(), true, i + 1);
+            mMockLooper.dispatchAll();
+            inOrder.verify(mockSessionCallback).onSessionStarted(sessionId.capture());
+
             if (i != 0) {
-                assertTrue("Transaction ID incrementing", id > prevId);
+                assertTrue("Session ID incrementing", sessionId.getValue() > prevId);
             }
-            prevId = id;
+            prevId = sessionId.getValue();
         }
     }
 
@@ -1275,19 +1298,6 @@
 
     /**
      * Utility routine used to validate that the internal state is cleaned-up
-     * after the specific transaction ID. To be used in every test which
-     * involves a transaction.
-     *
-     * @param transactionId The transaction ID whose state should be erased.
-     */
-    public void validateInternalTransactionInfoCleanedUp(short transactionId) throws Exception {
-        Object info = getInternalPendingTransactionInfo(mDut, transactionId);
-        collector.checkThat("Transaction record not cleared up for transactionId=" + transactionId,
-                info, nullValue());
-    }
-
-    /**
-     * Utility routine used to validate that the internal state is cleaned-up
      * after a client is disconnected. To be used in every test which terminates
      * a client.
      *
@@ -1297,20 +1307,6 @@
         WifiNanClientState client = getInternalClientState(mDut, clientId);
         collector.checkThat("Client record not cleared up for clientId=" + clientId, client,
                 nullValue());
-
-        Class<?> transactionInfoSessionClass = Class
-                .forName("com.android.server.wifi.nan.WifiNanStateManager$TransactionInfoSession");
-        Field clientField = transactionInfoSessionClass.getField("mClient");
-
-        SparseArray<Object> pending = getInternalPendingTransactions(mDut);
-        for (int i = 0; i < pending.size(); ++i) {
-            Object e = pending.valueAt(i);
-            if (transactionInfoSessionClass.isInstance(e)) {
-                WifiNanClientState clientInTransaction = (WifiNanClientState) clientField.get(e);
-                collector.checkThat("Client transaction not cleaned-up for clientId=" + clientId,
-                        clientId, not(equalTo(clientInTransaction.getClientId())));
-            }
-        }
     }
 
     /**
@@ -1327,26 +1323,27 @@
         WifiNanSessionState session = getInternalSessionState(client, sessionId);
         collector.checkThat("Client record not cleaned-up for sessionId=" + sessionId, session,
                 nullValue());
+    }
 
-        Class<?> transactionInfoSessionClass = Class
-                .forName("com.android.server.wifi.nan.WifiNanStateManager$TransactionInfoSession");
-        Field clientField = transactionInfoSessionClass.getField("mClient");
-        Field sessionField = transactionInfoSessionClass.getField("mSession");
+    /**
+     * Utility routine used to validate that the internal state is cleaned-up
+     * (deleted) correctly. Checks that a specific client has no sessions
+     * attached to it.
+     *
+     * @param clientId The ID of the client which we want to check.
+     */
+    public void validateInternalNoSessions(int clientId) throws Exception {
+        WifiNanClientState client = getInternalClientState(mDut, clientId);
+        collector.checkThat("Client record exists clientId=" + clientId, client, notNullValue());
 
-        SparseArray<Object> pending = getInternalPendingTransactions(mDut);
-        for (int i = 0; i < pending.size(); ++i) {
-            Object e = pending.valueAt(i);
-            if (transactionInfoSessionClass.isInstance(e)) {
-                WifiNanClientState clientInTransaction = (WifiNanClientState) clientField.get(e);
-                WifiNanSessionState sessionInTransaction = (WifiNanSessionState) sessionField
-                        .get(e);
-                if (clientId == clientInTransaction.getClientId()
-                        && sessionId == sessionInTransaction.getSessionId()) {
-                    collector.checkThat("Session record not cleaned-up for clientId=" + clientId
-                            + ", sessionId=" + sessionId, false, equalTo(true));
-                }
-            }
-        }
+        Field field = WifiNanClientState.class.getDeclaredField("mSessions");
+        field.setAccessible(true);
+        @SuppressWarnings("unchecked")
+        SparseArray<WifiNanSessionState> sessions = (SparseArray<WifiNanSessionState>) field
+                .get(client);
+
+        collector.checkThat("No sessions exist for clientId=" + clientId, sessions.size(),
+                equalTo(0));
     }
 
     /*
@@ -1371,21 +1368,6 @@
         field.set(null, obj);
     }
 
-    private static Object getInternalPendingTransactionInfo(WifiNanStateManager dut,
-            short transactionId) throws Exception {
-        return getInternalPendingTransactions(dut).get(transactionId);
-    }
-
-    private static SparseArray<Object> getInternalPendingTransactions(WifiNanStateManager dut)
-            throws Exception {
-        Field field = WifiNanStateManager.class.getDeclaredField("mPendingResponses");
-        field.setAccessible(true);
-        @SuppressWarnings("unchecked")
-        SparseArray<Object> pendingTransactions = (SparseArray<Object>) field.get(dut);
-
-        return pendingTransactions;
-    }
-
     private static WifiNanClientState getInternalClientState(WifiNanStateManager dut, int clientId)
             throws Exception {
         Field field = WifiNanStateManager.class.getDeclaredField("mClients");