Merge "wifi: MBO-OCE feature support (phase 1)"

9 files changed, 483 insertions, 4 deletions

diff --git a/service/java/com/android/server/wifi/ClientModeImpl.java b/service/java/com/android/server/wifi/ClientModeImpl.java
index 87c7e049a..720b3dccb 100644
--- a/service/java/com/android/server/wifi/ClientModeImpl.java
+++ b/service/java/com/android/server/wifi/ClientModeImpl.java
@@ -221,6 +221,7 @@ public class ClientModeImpl extends StateMachine {
     private final PasspointManager mPasspointManager;
     private final WifiDataStall mWifiDataStall;
     private final LinkProbeManager mLinkProbeManager;
+    private final MboOceController mMboOceController;
 
     private final McastLockManagerFilterController mMcastLockManagerFilterController;
     private final ActivityManager mActivityManager;
@@ -709,7 +710,8 @@ public class ClientModeImpl extends StateMachine {
                             SarManager sarManager, WifiTrafficPoller wifiTrafficPoller,
                             LinkProbeManager linkProbeManager,
                             BatteryStatsManager batteryStatsManager,
-                            SupplicantStateTracker supplicantStateTracker) {
+                            SupplicantStateTracker supplicantStateTracker,
+                            MboOceController mboOceController) {
         super(TAG, looper);
         mWifiInjector = wifiInjector;
         mWifiMetrics = mWifiInjector.getWifiMetrics();
@@ -725,6 +727,7 @@ public class ClientModeImpl extends StateMachine {
         mSarManager = sarManager;
         mWifiTrafficPoller = wifiTrafficPoller;
         mLinkProbeManager = linkProbeManager;
+        mMboOceController = mboOceController;
 
         mNetworkInfo = new NetworkInfo(ConnectivityManager.TYPE_WIFI, 0, NETWORKTYPE, "");
         mBatteryStatsManager = batteryStatsManager;
@@ -1118,6 +1121,7 @@ public class ClientModeImpl extends StateMachine {
         mPasspointManager.enableVerboseLogging(verbose);
         mNetworkFactory.enableVerboseLogging(verbose);
         mLinkProbeManager.enableVerboseLogging(mVerboseLoggingEnabled);
+        mMboOceController.enableVerboseLogging(mVerboseLoggingEnabled);
     }
 
     private static final String SYSTEM_PROPERTY_LOG_CONTROL_WIFIHAL = "log.tag.WifiHAL";
@@ -3303,6 +3307,8 @@ public class ClientModeImpl extends StateMachine {
         mWifiDiagnostics.startPktFateMonitoring(mInterfaceName);
         mWifiDiagnostics.startLogging(mInterfaceName);
 
+        mMboOceController.enable();
+
         /**
          * Enable bluetooth coexistence scan mode when bluetooth connection is active.
          * When this mode is on, some of the low-level scan parameters used by the
@@ -3339,6 +3345,7 @@ public class ClientModeImpl extends StateMachine {
         // exiting supplicant started state is now only applicable to client mode
         mWifiDiagnostics.stopLogging(mInterfaceName);
 
+        mMboOceController.disable();
         if (mIpClient != null && mIpClient.shutdown()) {
             // Block to make sure IpClient has really shut down, lest cleanup
             // race with, say, bringup code over in tethering.
diff --git a/service/java/com/android/server/wifi/MboOceController.java b/service/java/com/android/server/wifi/MboOceController.java
new file mode 100644
index 000000000..9104a42b7
--- /dev/null
+++ b/service/java/com/android/server/wifi/MboOceController.java
@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.wifi;
+
+import static android.net.wifi.WifiManager.WIFI_FEATURE_MBO;
+import static android.net.wifi.WifiManager.WIFI_FEATURE_OCE;
+
+import android.telephony.PhoneStateListener;
+import android.telephony.TelephonyManager;
+import android.util.Log;
+
+/**
+ * MboOceController is responsible for controlling MBO and OCE operations.
+ */
+public class MboOceController {
+    private static final String TAG = "MboOceController";
+
+    /** State of MBO/OCE module. */
+    private boolean mEnabled = false;
+    private boolean mIsMboSupported = false;
+    private boolean mIsOceSupported = false;
+    private boolean mVerboseLoggingEnabled = false;
+
+    private final WifiNative mWifiNative;
+    private final TelephonyManager mTelephonyManager;
+
+    /**
+     * Create new instance of MboOceController.
+     */
+    public MboOceController(TelephonyManager telephonyManager,
+                            WifiNative wifiNative) {
+        mTelephonyManager = telephonyManager;
+        mWifiNative = wifiNative;
+    }
+
+    /**
+     * Enable MBO and OCE functionality.
+     */
+    public void enable() {
+        String iface = mWifiNative.getClientInterfaceName();
+        if (iface == null) {
+            return;
+        }
+        mIsMboSupported = (mWifiNative.getSupportedFeatureSet(iface) & WIFI_FEATURE_MBO) != 0;
+        mIsOceSupported = (mWifiNative.getSupportedFeatureSet(iface) & WIFI_FEATURE_OCE) != 0;
+        mEnabled = true;
+        if (mVerboseLoggingEnabled) {
+            Log.d(TAG, "Enable MBO-OCE MBO support: " + mIsMboSupported
+                    + " OCE support: " + mIsOceSupported);
+        }
+        if (mIsMboSupported) {
+            // Register for data connection state change events (Cellular).
+            mTelephonyManager.listen(mDataConnectionStateListener,
+                    PhoneStateListener.LISTEN_DATA_CONNECTION_STATE);
+        }
+    }
+
+    /**
+     * Disable MBO and OCE functionality.
+     */
+    public void disable() {
+        if (mVerboseLoggingEnabled) {
+            Log.d(TAG, "Disable MBO-OCE");
+        }
+        if (mIsMboSupported) {
+            // Un-register for data connection state change events (Cellular).
+            mTelephonyManager.listen(mDataConnectionStateListener, PhoneStateListener.LISTEN_NONE);
+        }
+        mEnabled = false;
+    }
+
+    /**
+     * Enable/Disable verbose logging.
+     *
+     * @param verbose true to enable and false to disable.
+     */
+    public void enableVerboseLogging(boolean verbose) {
+        mVerboseLoggingEnabled = verbose;
+    }
+
+    /**
+     * Listen for changes to the data connection state (Cellular).
+     */
+    private PhoneStateListener mDataConnectionStateListener = new PhoneStateListener(){
+        public void onDataConnectionStateChanged(int state, int networkType) {
+            boolean dataAvailable;
+
+            String iface = mWifiNative.getClientInterfaceName();
+            if (iface == null) {
+                return;
+            }
+            if (!mEnabled) {
+                Log.e(TAG, "onDataConnectionStateChanged called when MBO is disabled!!");
+                return;
+            }
+            if (state == TelephonyManager.DATA_CONNECTED) {
+                dataAvailable = true;
+            } else if (state == TelephonyManager.DATA_DISCONNECTED) {
+                dataAvailable = false;
+            } else {
+                Log.e(TAG, "onDataConnectionStateChanged unexpected State: " + state);
+                return;
+            }
+            if (mVerboseLoggingEnabled) {
+                Log.d(TAG, "Cell Data: " + dataAvailable);
+            }
+            mWifiNative.setMboCellularDataStatus(iface, dataAvailable);
+        }
+    };
+}
diff --git a/service/java/com/android/server/wifi/SupplicantStaIfaceHal.java b/service/java/com/android/server/wifi/SupplicantStaIfaceHal.java
index 64a9c1f4e..f0baa61d5 100644
--- a/service/java/com/android/server/wifi/SupplicantStaIfaceHal.java
+++ b/service/java/com/android/server/wifi/SupplicantStaIfaceHal.java
@@ -16,6 +16,8 @@
 package com.android.server.wifi;
 
 import static android.net.wifi.WifiManager.WIFI_FEATURE_DPP;
+import static android.net.wifi.WifiManager.WIFI_FEATURE_MBO;
+import static android.net.wifi.WifiManager.WIFI_FEATURE_OCE;
 import static android.net.wifi.WifiManager.WIFI_FEATURE_OWE;
 import static android.net.wifi.WifiManager.WIFI_FEATURE_WPA3_SAE;
 import static android.net.wifi.WifiManager.WIFI_FEATURE_WPA3_SUITE_B;
@@ -34,6 +36,7 @@ import android.hardware.wifi.supplicant.V1_0.SupplicantStatusCode;
 import android.hardware.wifi.supplicant.V1_0.WpsConfigMethods;
 import android.hardware.wifi.supplicant.V1_3.ConnectionCapabilities;
 import android.hardware.wifi.supplicant.V1_3.WifiTechnology;
+import android.hardware.wifi.supplicant.V1_3.WpaDriverCapabilitiesMask;
 import android.hidl.manager.V1_0.IServiceManager;
 import android.hidl.manager.V1_0.IServiceNotification;
 import android.net.wifi.ScanResult;
@@ -2696,6 +2699,64 @@ public class SupplicantStaIfaceHal {
         return keyMgmtMask.value;
     }
 
+    /**
+     * Get the driver supported features through supplicant.
+     *
+     * @param ifaceName Name of the interface.
+     * @return bitmask defined by WifiManager.WIFI_FEATURE_*.
+     */
+    public long getWpaDriverFeatureSet(@NonNull String ifaceName) {
+        final String methodStr = "getWpaDriverFeatureSet";
+        MutableInt drvCapabilitiesMask = new MutableInt(0);
+        long featureSet = 0;
+
+        if (isV1_3()) {
+            ISupplicantStaIface iface = checkSupplicantStaIfaceAndLogFailure(ifaceName, methodStr);
+            if (iface == null) {
+                return 0;
+            }
+            // Get a v1.3 supplicant STA Interface
+            android.hardware.wifi.supplicant.V1_3.ISupplicantStaIface staIfaceV13 =
+                    getStaIfaceMockableV1_3(iface);
+            if (staIfaceV13 == null) {
+                Log.e(TAG, methodStr
+                        + ": SupplicantStaIface is null, cannot get wpa driver features");
+                return 0;
+            }
+
+            try {
+                staIfaceV13.getWpaDriverCapabilities(
+                        (SupplicantStatus statusInternal, int drvCapabilities) -> {
+                            if (statusInternal.code == SupplicantStatusCode.SUCCESS) {
+                                drvCapabilitiesMask.value = drvCapabilities;
+                            }
+                            checkStatusAndLogFailure(statusInternal, methodStr);
+                        });
+            } catch (RemoteException e) {
+                handleRemoteException(e, methodStr);
+            }
+        } else {
+            Log.i(TAG, "Method " + methodStr + " is not supported in existing HAL");
+            return 0;
+        }
+
+        if ((drvCapabilitiesMask.value & WpaDriverCapabilitiesMask.MBO) != 0) {
+            featureSet |= WIFI_FEATURE_MBO;
+            if (mVerboseLoggingEnabled) {
+                Log.v(TAG, methodStr + ": MBO supported");
+            }
+            if ((drvCapabilitiesMask.value
+                    & WpaDriverCapabilitiesMask.OCE) != 0) {
+                featureSet |= WIFI_FEATURE_OCE;
+                if (mVerboseLoggingEnabled) {
+                    Log.v(TAG, methodStr + ": OCE supported");
+                }
+            }
+        }
+
+        return featureSet;
+    }
+
     private @ScanResult.WifiStandard int getWifiStandardFromCap(ConnectionCapabilities capa) {
         switch(capa.technology) {
             case WifiTechnology.HE:
@@ -2988,4 +3049,44 @@ public class SupplicantStaIfaceHal {
     protected DppEventCallback getDppCallback() {
         return mDppCallback;
     }
+
+   /**
+     * Set MBO cellular data availability.
+     *
+     * @param ifaceName Name of the interface.
+     * @param available true means cellular data available, false otherwise.
+     * @return None.
+     */
+    public boolean setMboCellularDataStatus(@NonNull String ifaceName, boolean available) {
+        final String methodStr = "setMboCellularDataStatus";
+
+        if (isV1_3()) {
+            ISupplicantStaIface iface = checkSupplicantStaIfaceAndLogFailure(ifaceName, methodStr);
+            if (iface == null) {
+                return false;
+            }
+
+            // Get a v1.3 supplicant STA Interface
+            android.hardware.wifi.supplicant.V1_3.ISupplicantStaIface staIfaceV13 =
+                    getStaIfaceMockableV1_3(iface);
+            if (staIfaceV13 == null) {
+                Log.e(TAG, methodStr
+                        + ": SupplicantStaIface is null, cannot update cell status");
+                return false;
+            }
+
+            try {
+                SupplicantStatus status = staIfaceV13.setMboCellularDataStatus(available);
+                return checkStatusAndLogFailure(status, methodStr);
+            } catch (RemoteException e) {
+                handleRemoteException(e, methodStr);
+            }
+        } else {
+            Log.e(TAG, "Method " + methodStr + " is not supported in existing HAL");
+            return false;
+        }
+
+        return false;
+    }
+
 }
diff --git a/service/java/com/android/server/wifi/WifiInjector.java b/service/java/com/android/server/wifi/WifiInjector.java
index abf9dfab1..13f6a8267 100644
--- a/service/java/com/android/server/wifi/WifiInjector.java
+++ b/service/java/com/android/server/wifi/WifiInjector.java
@@ -153,6 +153,7 @@ public class WifiInjector {
     private final WifiThreadRunner mWifiThreadRunner;
     private BssidBlocklistMonitor mBssidBlocklistMonitor;
     private final MacAddressUtil mMacAddressUtil;
+    private final MboOceController mMboOceController;
 
     public WifiInjector(Context context) {
         if (context == null) {
@@ -306,12 +307,13 @@ public class WifiInjector {
                 mFrameworkFacade, wifiHandler, mContext);
         SupplicantStateTracker supplicantStateTracker = new SupplicantStateTracker(
                 mContext, mWifiConfigManager, mBatteryStats, wifiHandler);
+        mMboOceController = new MboOceController(makeTelephonyManager(), mWifiNative);
         mClientModeImpl = new ClientModeImpl(mContext, mFrameworkFacade,
                 wifiLooper, mUserManager,
                 this, mBackupManagerProxy, mCountryCode, mWifiNative,
                 new WrongPasswordNotifier(mContext, mFrameworkFacade),
                 mSarManager, mWifiTrafficPoller, mLinkProbeManager, mBatteryStats,
-                supplicantStateTracker);
+                supplicantStateTracker, mMboOceController);
         mActiveModeWarden = new ActiveModeWarden(this, wifiLooper,
                 mWifiNative, new DefaultModeManager(mContext), mBatteryStats, mWifiDiagnostics,
                 mContext, mClientModeImpl, mSettingsStore, mFrameworkFacade, mWifiPermissionsUtil);
diff --git a/service/java/com/android/server/wifi/WifiNative.java b/service/java/com/android/server/wifi/WifiNative.java
index 2b3e50ed5..eeef34fc3 100644
--- a/service/java/com/android/server/wifi/WifiNative.java
+++ b/service/java/com/android/server/wifi/WifiNative.java
@@ -2725,7 +2725,8 @@ public class WifiNative {
      */
     private long getSupportedFeatureSetInternal(@NonNull String ifaceName) {
         return mSupplicantStaIfaceHal.getAdvancedKeyMgmtCapabilities(ifaceName)
-                | mWifiVendorHal.getSupportedFeatureSet(ifaceName);
+                | mWifiVendorHal.getSupportedFeatureSet(ifaceName)
+                | mSupplicantStaIfaceHal.getWpaDriverFeatureSet(ifaceName);
     }
 
     /**
@@ -3346,6 +3347,18 @@ public class WifiNative {
         return mWifiVendorHal.selectTxPowerScenario(sarInfo);
     }
 
+    /**
+     * Set MBO cellular data status
+     *
+     * @param ifaceName Name of the interface.
+     * @param available cellular data status,
+     *        true means cellular data available, false otherwise.
+     */
+    public void setMboCellularDataStatus(@NonNull String ifaceName, boolean available) {
+        mSupplicantStaIfaceHal.setMboCellularDataStatus(ifaceName, available);
+        return;
+    }
+
     /********************************************************
      * JNI operations
      ********************************************************/
diff --git a/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java b/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
index 2efb7b77a..7f1c778ba 100644
--- a/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
@@ -387,6 +387,7 @@ public class ClientModeImplTest extends WifiBaseTest {
     @Mock CarrierNetworkConfig mCarrierNetworkConfig;
     @Mock Handler mNetworkAgentHandler;
     @Mock BatteryStatsManager mBatteryStatsManager;
+    @Mock MboOceController mMboOceController;
 
     final ArgumentCaptor<WifiConfigManager.OnNetworkUpdateListener> mConfigUpdateListenerCaptor =
             ArgumentCaptor.forClass(WifiConfigManager.OnNetworkUpdateListener.class);
@@ -543,7 +544,8 @@ public class ClientModeImplTest extends WifiBaseTest {
         mCmi = new ClientModeImpl(mContext, mFrameworkFacade, mLooper.getLooper(),
                 mUserManager, mWifiInjector, mBackupManagerProxy, mCountryCode, mWifiNative,
                 mWrongPasswordNotifier, mSarManager, mWifiTrafficPoller,
-                mLinkProbeManager, mBatteryStatsManager, mSupplicantStateTracker);
+                mLinkProbeManager, mBatteryStatsManager, mSupplicantStateTracker,
+                mMboOceController);
         mCmi.start();
         mWifiCoreThread = getCmiHandlerThread(mCmi);
 
@@ -3825,4 +3827,16 @@ public class ClientModeImplTest extends WifiBaseTest {
 
         assertEquals("ConnectedState", getCurrentState().getName());
     }
+
+    /**
+     * Verify that MboOce initialization/Deinitialization methods are called in ClientMode.
+     */
+    @Test
+    public void verifyMboOceInitAndDeinitInClientMode() throws Exception {
+        startSupplicantAndDispatchMessages();
+        verify(mMboOceController).enable();
+        mCmi.setOperationalMode(ClientModeImpl.DISABLED_MODE, null);
+        mLooper.dispatchAll();
+        verify(mMboOceController).disable();
+    }
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/MboOceControllerTest.java b/tests/wifitests/src/com/android/server/wifi/MboOceControllerTest.java
new file mode 100644
index 000000000..239fea95f
--- /dev/null
+++ b/tests/wifitests/src/com/android/server/wifi/MboOceControllerTest.java
@@ -0,0 +1,138 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.wifi;
+
+import static org.junit.Assert.assertNotNull;
+import static org.mockito.Mockito.*;
+
+import android.net.wifi.WifiManager;
+import android.os.test.TestLooper;
+import android.telephony.PhoneStateListener;
+import android.telephony.TelephonyManager;
+
+import androidx.test.filters.SmallTest;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.InOrder;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+
+/**
+ * unit tests for {@link com.android.server.wifi.MboOceController}.
+ */
+@SmallTest
+public class MboOceControllerTest extends WifiBaseTest {
+    private static final String TAG = "MboOceControllerTest";
+    private static final String INTERFACE_NAME = "wlan0";
+
+    private MboOceController mMboOceController;
+    private TestLooper mLooper;
+    @Mock WifiNative mWifiNative;
+    @Mock TelephonyManager mTelephonyManager;
+
+    /**
+     * Initializes common state (e.g. mocks) needed by test cases.
+     */
+    @Before
+    public void setUp() throws Exception {
+        /* Ensure Looper exists */
+        mLooper = new TestLooper();
+        MockitoAnnotations.initMocks(this);
+
+        mMboOceController = new MboOceController(mTelephonyManager, mWifiNative);
+
+        when(mWifiNative.getClientInterfaceName()).thenReturn(INTERFACE_NAME);
+    }
+
+    /**
+     * Helper function to initialize mboOceController
+     */
+    private PhoneStateListener enableMboOceController(boolean isMboEnabled, boolean isOceEnabled) {
+        long featureSet = 0;
+        PhoneStateListener dataConnectionStateListener = null;
+
+        if (isMboEnabled) {
+            featureSet |= WifiManager.WIFI_FEATURE_MBO;
+        }
+        if (isOceEnabled) {
+            featureSet |= WifiManager.WIFI_FEATURE_OCE;
+        }
+        when(mWifiNative.getSupportedFeatureSet(INTERFACE_NAME))
+                .thenReturn((long) featureSet);
+
+        mMboOceController.enable();
+
+        if (isMboEnabled) {
+            /* Capture the PhoneStateListener */
+            ArgumentCaptor<PhoneStateListener> phoneStateListenerCaptor =
+                    ArgumentCaptor.forClass(PhoneStateListener.class);
+            verify(mTelephonyManager).listen(phoneStateListenerCaptor.capture(),
+                    eq(PhoneStateListener.LISTEN_DATA_CONNECTION_STATE));
+            dataConnectionStateListener = phoneStateListenerCaptor.getValue();
+            assertNotNull(dataConnectionStateListener);
+        }
+
+        return dataConnectionStateListener;
+
+    }
+
+    /**
+     * Test that we do not register for cellular Data state change
+     * events when MBO is disabled.
+     */
+    @Test
+    public void testMboDisabledDoNotRegisterCellularDataStateChangeEvents()
+            throws Exception {
+        enableMboOceController(false, false);
+        verify(mTelephonyManager, never()).listen(any(), anyInt());
+    }
+
+    /**
+     * Test that we register for cellular Data state change
+     * events and update cellular data status changes when MBO is enabled.
+     */
+    @Test
+    public void testMboEnabledUpdateCellularDataStateChangeEvents() throws Exception {
+        InOrder inOrder = inOrder(mWifiNative);
+        PhoneStateListener dataConnectionStateListener;
+        dataConnectionStateListener = enableMboOceController(true, false);
+        dataConnectionStateListener.onDataConnectionStateChanged(TelephonyManager.DATA_CONNECTED,
+                TelephonyManager.NETWORK_TYPE_LTE);
+        verify(mWifiNative).setMboCellularDataStatus(eq(INTERFACE_NAME), eq(true));
+        dataConnectionStateListener.onDataConnectionStateChanged(
+                TelephonyManager.DATA_DISCONNECTED, TelephonyManager.NETWORK_TYPE_LTE);
+        verify(mWifiNative).setMboCellularDataStatus(eq(INTERFACE_NAME), eq(false));
+    }
+
+    /**
+     * Test that we unregister data connection state events
+     * when disable mMboOceController is called.
+     */
+    @Test
+    public void testDisableMboOceControllerUnRegisterCellularDataStateChangeEvents()
+            throws Exception {
+        enableMboOceController(true, false);
+        mMboOceController.disable();
+        ArgumentCaptor<PhoneStateListener> phoneStateListenerCaptor =
+                ArgumentCaptor.forClass(PhoneStateListener.class);
+        verify(mTelephonyManager).listen(phoneStateListenerCaptor.capture(),
+                eq(PhoneStateListener.LISTEN_NONE));
+    }
+}
diff --git a/tests/wifitests/src/com/android/server/wifi/SupplicantStaIfaceHalTest.java b/tests/wifitests/src/com/android/server/wifi/SupplicantStaIfaceHalTest.java
index fc8ae0cf0..b071786d3 100644
--- a/tests/wifitests/src/com/android/server/wifi/SupplicantStaIfaceHalTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/SupplicantStaIfaceHalTest.java
@@ -16,6 +16,8 @@
 package com.android.server.wifi;
 
 import static android.net.wifi.WifiManager.WIFI_FEATURE_DPP;
+import static android.net.wifi.WifiManager.WIFI_FEATURE_MBO;
+import static android.net.wifi.WifiManager.WIFI_FEATURE_OCE;
 import static android.net.wifi.WifiManager.WIFI_FEATURE_OWE;
 import static android.net.wifi.WifiManager.WIFI_FEATURE_WPA3_SAE;
 import static android.net.wifi.WifiManager.WIFI_FEATURE_WPA3_SUITE_B;
@@ -2361,4 +2363,70 @@ public class SupplicantStaIfaceHalTest extends WifiBaseTest {
         }).when(mSupplicantStaNetworkMock)
                 .setPmkCache(any(ArrayList.class));
     }
+
+    private class GetWpaDriverCapabilitiesAnswer extends MockAnswerUtil.AnswerWithArguments {
+        private int mWpaDriverCapabilities;
+
+        GetWpaDriverCapabilitiesAnswer(int wpaDriverCapabilities) {
+            mWpaDriverCapabilities = wpaDriverCapabilities;
+        }
+
+        public void answer(android.hardware.wifi.supplicant.V1_3.ISupplicantStaIface
+                .getWpaDriverCapabilitiesCallback cb) {
+            cb.onValues(mStatusSuccess, mWpaDriverCapabilities);
+        }
+    }
+
+    /**
+     * Test To get wpa driver capabilities API on old HAL, should
+     * return 0 (not supported)
+     */
+    @Test
+    public void tetGetWpaDriverCapabilitiesOldHal() throws Exception {
+        setupMocksForHalV1_2();
+
+        executeAndValidateInitializationSequenceV1_2();
+
+        assertEquals(0, mDut.getWpaDriverFeatureSet(WLAN0_IFACE_NAME));
+    }
+
+    /**
+     * Test Multi Band operation support (MBO).
+     */
+    @Test
+    public void testGetWpaDriverCapabilitiesMbo() throws Exception {
+        setupMocksForHalV1_3();
+
+        executeAndValidateInitializationSequenceV1_3();
+
+        doAnswer(new GetWpaDriverCapabilitiesAnswer(android.hardware.wifi.supplicant.V1_3
+                .WpaDriverCapabilitiesMask.MBO))
+                .when(mISupplicantStaIfaceMockV13).getWpaDriverCapabilities(any(
+                android.hardware.wifi.supplicant.V1_3.ISupplicantStaIface
+                        .getWpaDriverCapabilitiesCallback.class));
+
+        assertEquals(WIFI_FEATURE_MBO, mDut.getWpaDriverFeatureSet(WLAN0_IFACE_NAME));
+    }
+
+    /**
+     * Test Optimized Connectivity support (OCE).
+     */
+    @Test
+    public void testGetWpaDriverCapabilitiesOce() throws Exception {
+        setupMocksForHalV1_3();
+
+        executeAndValidateInitializationSequenceV1_3();
+
+        doAnswer(new GetWpaDriverCapabilitiesAnswer(android.hardware.wifi.supplicant.V1_3
+                .WpaDriverCapabilitiesMask.MBO
+                | android.hardware.wifi.supplicant.V1_3
+                .WpaDriverCapabilitiesMask.OCE))
+                .when(mISupplicantStaIfaceMockV13).getWpaDriverCapabilities(any(
+                android.hardware.wifi.supplicant.V1_3.ISupplicantStaIface
+                        .getWpaDriverCapabilitiesCallback.class));
+
+        assertEquals(WIFI_FEATURE_MBO | WIFI_FEATURE_OCE,
+                mDut.getWpaDriverFeatureSet(WLAN0_IFACE_NAME));
+    }
+
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiNativeInterfaceManagementTest.java b/tests/wifitests/src/com/android/server/wifi/WifiNativeInterfaceManagementTest.java
index ccd81856b..68c59081a 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiNativeInterfaceManagementTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiNativeInterfaceManagementTest.java
@@ -452,6 +452,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mNwManagementService).getInterfaceConfig(IFACE_NAME_0);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(IFACE_NAME_0);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(IFACE_NAME_0);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(IFACE_NAME_0);
 
         // Execute a teardown of the interface to ensure that the new iface removal works.
         executeAndValidateTeardownSoftApInterface(false, false, IFACE_NAME_0, mIfaceCallback1,
@@ -504,6 +505,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mNwManagementService).disableIpv6(IFACE_NAME_0);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(IFACE_NAME_0);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(IFACE_NAME_0);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(IFACE_NAME_0);
 
         // Execute a teardown of the interface to ensure that the new iface removal works.
         executeAndValidateTeardownClientInterface(false, false, IFACE_NAME_0, mIfaceCallback1,
@@ -715,6 +717,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mNwManagementService).getInterfaceConfig(IFACE_NAME_0);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(IFACE_NAME_0);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(IFACE_NAME_0);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(IFACE_NAME_0);
 
         // Step (c) - Iface up on old iface, ignored!
         mNetworkObserverCaptor0.getValue().interfaceLinkStateChanged(IFACE_NAME_0, true);
@@ -1174,6 +1177,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mNwManagementService).disableIpv6(IFACE_NAME_0);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(IFACE_NAME_0);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(IFACE_NAME_0);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(IFACE_NAME_0);
 
         // Now setup an AP interface.
         assertEquals(IFACE_NAME_0, mWifiNative.setupInterfaceForSoftApMode(mIfaceCallback1));
@@ -1199,6 +1203,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mNwManagementService).getInterfaceConfig(IFACE_NAME_0);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(IFACE_NAME_0);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(IFACE_NAME_0);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(IFACE_NAME_0);
     }
 
     /**
@@ -1227,6 +1232,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mNwManagementService).getInterfaceConfig(IFACE_NAME_0);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(IFACE_NAME_0);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(IFACE_NAME_0);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(IFACE_NAME_0);
 
         // Now setup a STA interface.
         assertEquals(IFACE_NAME_0,
@@ -1257,6 +1263,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mNwManagementService).disableIpv6(IFACE_NAME_0);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(IFACE_NAME_0);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(IFACE_NAME_0);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(IFACE_NAME_0);
     }
 
     /**
@@ -1295,6 +1302,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mSupplicantStaIfaceHal).terminate();
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(IFACE_NAME_0);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(IFACE_NAME_0);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(IFACE_NAME_0);
     }
 
     /**
@@ -1326,6 +1334,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mSupplicantStaIfaceHal).setupIface(IFACE_NAME_0);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(IFACE_NAME_0);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(IFACE_NAME_0);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(IFACE_NAME_0);
     }
 
     /**
@@ -1400,6 +1409,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mNwManagementService).disableIpv6(ifaceName);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(ifaceName);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(ifaceName);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(ifaceName);
     }
 
     private void executeAndValidateTeardownClientInterface(
@@ -1472,6 +1482,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mNwManagementService).getInterfaceConfig(ifaceName);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(ifaceName);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(ifaceName);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(ifaceName);
     }
 
     private void executeAndValidateTeardownClientInterfaceForScan(
@@ -1544,6 +1555,7 @@ public class WifiNativeInterfaceManagementTest extends WifiBaseTest {
         mInOrder.verify(mNwManagementService).getInterfaceConfig(ifaceName);
         mInOrder.verify(mSupplicantStaIfaceHal).getAdvancedKeyMgmtCapabilities(ifaceName);
         mInOrder.verify(mWifiVendorHal).getSupportedFeatureSet(ifaceName);
+        mInOrder.verify(mSupplicantStaIfaceHal).getWpaDriverFeatureSet(ifaceName);
     }
 
     private void executeAndValidateTeardownSoftApInterface(