Merge "Cap BW estimation sample by Tx/Rx link speed" into sc-mainline-prod
diff --git a/framework/java/android/net/wifi/ScanResult.java b/framework/java/android/net/wifi/ScanResult.java
index efdf09b..fb831d9 100644
--- a/framework/java/android/net/wifi/ScanResult.java
+++ b/framework/java/android/net/wifi/ScanResult.java
@@ -908,6 +908,8 @@
         @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.R, trackingBug = 170729553)
         public static final int EID_TIM = 5;
         /** @hide */
+        public static final int EID_COUNTRY = 7;
+        /** @hide */
         @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.R, trackingBug = 170729553)
         public static final int EID_BSS_LOAD = 11;
         /** @hide */
diff --git a/service/java/com/android/server/wifi/ClientModeImpl.java b/service/java/com/android/server/wifi/ClientModeImpl.java
index e85db36..a3ee4df 100644
--- a/service/java/com/android/server/wifi/ClientModeImpl.java
+++ b/service/java/com/android/server/wifi/ClientModeImpl.java
@@ -3600,6 +3600,7 @@
                         boolean isIpClientStarted = startIpClient(config, true);
                         if (isIpClientStarted) {
                             mIpClientWithPreConnection = true;
+                            transitionTo(mL2ConnectingState);
                             break;
                         }
                     }
@@ -6329,6 +6330,7 @@
                     .withPreconnection()
                     .withApfCapabilities(
                     mWifiNative.getApfCapabilities(mInterfaceName))
+                    .withDisplayName(config.SSID)
                     .withLayer2Information(layer2Info);
             if (isUsingMacRandomization) {
                 // Use EUI64 address generation for link-local IPv6 addresses.
diff --git a/service/java/com/android/server/wifi/SupplicantStaIfaceCallbackImpl.java b/service/java/com/android/server/wifi/SupplicantStaIfaceCallbackImpl.java
index 4949972..55d83ce 100644
--- a/service/java/com/android/server/wifi/SupplicantStaIfaceCallbackImpl.java
+++ b/service/java/com/android/server/wifi/SupplicantStaIfaceCallbackImpl.java
@@ -318,8 +318,13 @@
                         .getCandidateSecurityParams();
                 if (params != null
                         && params.getSecurityType() == WifiConfiguration.SECURITY_TYPE_SAE) {
-                    mStaIfaceHal.logCallback("SAE incorrect password");
-                    isWrongPwd = true;
+                    // If this is ever connected, the password should be correct.
+                    isWrongPwd = !curConfiguration.getNetworkSelectionStatus().hasEverConnected();
+                    if (isWrongPwd) {
+                        mStaIfaceHal.logCallback("SAE incorrect password");
+                    } else {
+                        mStaIfaceHal.logCallback("SAE association rejection");
+                    }
                 }
             } else if (assocRejectInfo.statusCode == StatusCode.CHALLENGE_FAIL
                     && WifiConfigurationUtil.isConfigForWepNetwork(curConfiguration)) {
diff --git a/service/java/com/android/server/wifi/WifiCountryCode.java b/service/java/com/android/server/wifi/WifiCountryCode.java
index 0398200..6edf9f1 100644
--- a/service/java/com/android/server/wifi/WifiCountryCode.java
+++ b/service/java/com/android/server/wifi/WifiCountryCode.java
@@ -21,6 +21,7 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
+import android.net.wifi.WifiInfo;
 import android.os.SystemProperties;
 import android.telephony.TelephonyManager;
 import android.text.TextUtils;
@@ -50,6 +51,9 @@
 public class WifiCountryCode {
     private static final String TAG = "WifiCountryCode";
     private static final String BOOT_DEFAULT_WIFI_COUNTRY_CODE = "ro.boot.wificountrycode";
+    private static final int PKT_COUNT_HIGH_PKT_PER_SEC = 16;
+    private static final int DISCONNECT_WIFI_COUNT_MAX = 1;
+    private static final String COUNTRY_CODE_WORLD = "00";
     private final Context mContext;
     private final TelephonyManager mTelephonyManager;
     private final ActiveModeWarden mActiveModeWarden;
@@ -79,6 +83,7 @@
     private String mTelephonyCountryTimestamp = null;
     private String mDriverCountryTimestamp = null;
     private String mReadyTimestamp = null;
+    private int mDisconnectWifiToForceUpdateCount = 0;
 
     private class ModeChangeCallbackInternal implements ActiveModeWarden.ModeChangeCallback {
         @Override
@@ -315,7 +320,7 @@
         // Empty country code.
         if (TextUtils.isEmpty(countryCode)) {
             if (mContext.getResources()
-                        .getBoolean(R.bool.config_wifi_revert_country_code_on_cellular_loss)) {
+                    .getBoolean(R.bool.config_wifi_revert_country_code_on_cellular_loss)) {
                 Log.d(TAG, "Received empty country code, reset to default country code");
                 mTelephonyCountryCode = null;
             }
@@ -343,11 +348,54 @@
             updateCountryCode();
         } else {
             Log.d(TAG, "skip update supplicant not ready yet");
+            disconnectWifiToForceUpdateIfNeeded();
         }
 
         return true;
     }
 
+    private void disconnectWifiToForceUpdateIfNeeded() {
+        if (shouldDisconnectWifiToForceUpdate()) {
+            Log.d(TAG, "Disconnect wifi to force update");
+            for (ClientModeManager cmm :
+                    mActiveModeWarden.getInternetConnectivityClientModeManagers()) {
+                if (!cmm.isConnected()) {
+                    continue;
+                }
+                cmm.disconnect();
+            }
+            mDisconnectWifiToForceUpdateCount++;
+        }
+    }
+
+    private boolean shouldDisconnectWifiToForceUpdate() {
+        if (mTelephonyCountryCode == null || mTelephonyCountryCode.equals(mDriverCountryCode)) {
+            return false;
+        }
+
+        if (mDisconnectWifiToForceUpdateCount >= DISCONNECT_WIFI_COUNT_MAX) {
+            return false;
+        }
+
+        if (mDriverCountryCode != null
+                && !mDriverCountryCode.equalsIgnoreCase(COUNTRY_CODE_WORLD)) {
+            return false;
+        }
+
+        for (ClientModeManager cmm :
+                mActiveModeWarden.getInternetConnectivityClientModeManagers()) {
+            if (!cmm.isConnected()) {
+                continue;
+            }
+            WifiInfo wifiInfo = cmm.syncRequestConnectionInfo();
+            if (wifiInfo.getSuccessfulTxPacketsPerSecond() < PKT_COUNT_HIGH_PKT_PER_SEC
+                    && wifiInfo.getSuccessfulRxPacketsPerSecond() < PKT_COUNT_HIGH_PKT_PER_SEC) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     /**
      * Method to get the Country Code that was sent to wpa_supplicant.
      *
@@ -399,12 +447,12 @@
     }
 
     /**
-     * Method to dump the current state of this WifiCounrtyCode object.
+     * Method to dump the current state of this WifiCountryCode object.
      */
     public synchronized void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
         pw.println("mRevertCountryCodeOnCellularLoss: "
                 + mContext.getResources().getBoolean(
-                        R.bool.config_wifi_revert_country_code_on_cellular_loss));
+                R.bool.config_wifi_revert_country_code_on_cellular_loss));
         pw.println("DefaultCountryCode(system property): " + getOemDefaultCountryCode());
         pw.println("DefaultCountryCode(config store): "
                 + mSettingsConfigStore.get(WIFI_DEFAULT_COUNTRY_CODE));
@@ -416,6 +464,7 @@
         pw.println("mReadyTimestamp: " + mReadyTimestamp);
         pw.println("isReady: " + isReady());
         pw.println("mAmmToReadyForChangeMap: " + mAmmToReadyForChangeMap);
+        pw.println("mDisconnectWifiToForceUpdateCount: " + mDisconnectWifiToForceUpdateCount);
     }
 
     private void updateCountryCode() {
@@ -426,7 +475,7 @@
         // There are two reasons:
         // 1. Wpa supplicant may silently modify the country code.
         // 2. If Wifi restarted therefore wpa_supplicant also restarted,
-        // the country code counld be reset to '00' by wpa_supplicant.
+        // the country code could be reset to '00' by wpa_supplicant.
         if (country != null) {
             setCountryCodeNative(country);
         }
@@ -497,4 +546,3 @@
         }
     }
 }
-
diff --git a/service/java/com/android/server/wifi/WifiMetrics.java b/service/java/com/android/server/wifi/WifiMetrics.java
index 3445b76..3affdef 100644
--- a/service/java/com/android/server/wifi/WifiMetrics.java
+++ b/service/java/com/android/server/wifi/WifiMetrics.java
@@ -210,6 +210,9 @@
     public static final int MIN_WIFI_GOOD_USABILITY_STATS_PERIOD_MS = 1000 * 3600; // 1 hour
     public static final int PASSPOINT_DEAUTH_IMMINENT_SCOPE_ESS = 0;
     public static final int PASSPOINT_DEAUTH_IMMINENT_SCOPE_BSS = 1;
+    public static final int COUNTRY_CODE_CONFLICT_WIFI_SCAN = -1;
+    public static final int COUNTRY_CODE_CONFLICT_WIFI_SCAN_TELEPHONY = -2;
+    public static final int MAX_COUNTRY_CODE_COUNT = 4;
     // Histogram for WifiConfigStore IO duration times. Indicates the following 5 buckets (in ms):
     //   < 50
     //   [50, 100)
@@ -237,6 +240,7 @@
     private RttMetrics mRttMetrics;
     private final PnoScanMetrics mPnoScanMetrics = new PnoScanMetrics();
     private final WifiLinkLayerUsageStats mWifiLinkLayerUsageStats = new WifiLinkLayerUsageStats();
+    private final TelephonyManager mTelephonyManager;
     /** Mapping of radio id values to RadioStats objects. */
     private final SparseArray<RadioStats> mRadioStats = new SparseArray<>();
     private final ExperimentValues mExperimentValues = new ExperimentValues();
@@ -394,6 +398,7 @@
     private final SparseIntArray mObservedHotspotR3ApsPerEssInScanHistogram = new SparseIntArray();
 
     private final SparseIntArray mObserved80211mcApInScanHistogram = new SparseIntArray();
+    private final IntCounter mCountryCodeScanHistogram = new IntCounter();
 
     // link probing stats
     private final IntCounter mLinkProbeSuccessRssiCounts = new IntCounter(-85, -65);
@@ -1525,6 +1530,7 @@
         setScreenState(context.getSystemService(PowerManager.class).isInteractive());
 
         mScanMetrics = new ScanMetrics(context, clock);
+        mTelephonyManager = context.getSystemService(TelephonyManager.class);
     }
 
     /** Sets internal ScoringParams member */
@@ -3409,6 +3415,8 @@
                 mWifiLogProto.partialAllSingleScanListenerResults++;
                 return;
             }
+            updateCountryCodeScanStats(scanDetails);
+
             Set<ScanResultMatchInfo> ssids = new HashSet<ScanResultMatchInfo>();
             int bssids = 0;
             Set<ScanResultMatchInfo> openSsids = new HashSet<ScanResultMatchInfo>();
@@ -3547,6 +3555,35 @@
         }
     }
 
+    private void updateCountryCodeScanStats(List<ScanDetail> scanDetails) {
+        String countryCode = null;
+        int countryCodeCount = 0;
+        for (ScanDetail scanDetail : scanDetails) {
+            NetworkDetail networkDetail = scanDetail.getNetworkDetail();
+            String countryCodeCurr = networkDetail.getCountryCode();
+            if (countryCodeCurr == null) {
+                continue;
+            }
+            if (countryCode == null) {
+                countryCode = countryCodeCurr;
+                countryCodeCount = 1;
+                continue;
+            }
+            if (!countryCodeCurr.equalsIgnoreCase(countryCode)) {
+                mCountryCodeScanHistogram.increment(COUNTRY_CODE_CONFLICT_WIFI_SCAN);
+                return;
+            }
+            countryCodeCount++;
+        }
+        String countryCodeTelephony = mTelephonyManager.getNetworkCountryIso();
+        if (countryCodeCount > 0 && !TextUtils.isEmpty(countryCodeTelephony)
+                && !countryCodeTelephony.equalsIgnoreCase(countryCode)) {
+            mCountryCodeScanHistogram.increment(COUNTRY_CODE_CONFLICT_WIFI_SCAN_TELEPHONY);
+            return;
+        }
+        mCountryCodeScanHistogram.increment(Math.min(countryCodeCount, MAX_COUNTRY_CODE_COUNT));
+    }
+
     /** Increments the occurence of a "Connect to Network" notification. */
     public void incrementConnectToNetworkNotification(String notifierTag, int notificationType) {
         synchronized (mLock) {
@@ -4094,6 +4131,8 @@
 
                 pw.println("mWifiLogProto.observed80211mcSupportingApsInScanHistogram"
                         + mObserved80211mcApInScanHistogram);
+                pw.println("mWifiLogProto.CountryCodeScanHistogram="
+                        + mCountryCodeScanHistogram.toString());
                 pw.println("mWifiLogProto.bssidBlocklistStats:");
                 pw.println(mBssidBlocklistStats.toString());
 
@@ -5057,6 +5096,7 @@
             mWifiLogProto.passpointDeauthImminentScope = mPasspointDeauthImminentScope.toProto();
             mWifiLogProto.recentFailureAssociationStatus =
                     mRecentFailureAssociationStatus.toProto();
+            mWifiLogProto.countryCodeScanHistogram = mCountryCodeScanHistogram.toProto();
         }
     }
 
@@ -5216,6 +5256,7 @@
             mObservedHotspotR1ApsPerEssInScanHistogram.clear();
             mObservedHotspotR2ApsPerEssInScanHistogram.clear();
             mObservedHotspotR3ApsPerEssInScanHistogram.clear();
+            mCountryCodeScanHistogram.clear();
             mSoftApEventListTethered.clear();
             mSoftApEventListLocalOnly.clear();
             mWifiWakeMetrics.clear();
diff --git a/service/java/com/android/server/wifi/coex/CoexManager.java b/service/java/com/android/server/wifi/coex/CoexManager.java
index 326bb85..2bfcf13 100644
--- a/service/java/com/android/server/wifi/coex/CoexManager.java
+++ b/service/java/com/android/server/wifi/coex/CoexManager.java
@@ -309,6 +309,11 @@
             Log.e(TAG, "setCoexUnsafeChannels called with undefined restriction flags");
             return;
         }
+        if (new HashSet(mCurrentCoexUnsafeChannels).equals(new HashSet(coexUnsafeChannels))
+                && mCoexRestrictions == coexRestrictions) {
+            // Do not update if the unsafe channels haven't changed since the last time
+            return;
+        }
         mCurrentCoexUnsafeChannels.clear();
         mCurrentCoexUnsafeChannels.addAll(coexUnsafeChannels);
         mCoexRestrictions = coexRestrictions;
diff --git a/service/java/com/android/server/wifi/hotspot2/NetworkDetail.java b/service/java/com/android/server/wifi/hotspot2/NetworkDetail.java
index 22e8554..b43a00b 100644
--- a/service/java/com/android/server/wifi/hotspot2/NetworkDetail.java
+++ b/service/java/com/android/server/wifi/hotspot2/NetworkDetail.java
@@ -115,6 +115,7 @@
     private final int mAnqpOICount;
     private final long[] mRoamingConsortiums;
     private int mDtimInterval = -1;
+    private String mCountryCode;
 
     private final InformationElementUtil.ExtendedCapabilities mExtendedCapabilities;
 
@@ -167,6 +168,9 @@
         InformationElementUtil.ExtendedCapabilities extendedCapabilities =
                 new InformationElementUtil.ExtendedCapabilities();
 
+        InformationElementUtil.Country country =
+                new InformationElementUtil.Country();
+
         InformationElementUtil.TrafficIndicationMap trafficIndicationMap =
                 new InformationElementUtil.TrafficIndicationMap();
 
@@ -212,6 +216,9 @@
                     case ScanResult.InformationElement.EID_EXTENDED_CAPS:
                         extendedCapabilities.from(ie);
                         break;
+                    case ScanResult.InformationElement.EID_COUNTRY:
+                        country.from(ie);
+                        break;
                     case ScanResult.InformationElement.EID_TIM:
                         trafficIndicationMap.from(ie);
                         break;
@@ -348,6 +355,10 @@
         mCenterfreq0 = centerFreq0;
         mCenterfreq1 = centerFreq1;
 
+        if (country.isValid()) {
+            mCountryCode = country.getCountryCode();
+        }
+
         // If trafficIndicationMap is not valid, mDtimPeriod will be negative
         if (trafficIndicationMap.isValid()) {
             mDtimInterval = trafficIndicationMap.mDtimPeriod;
@@ -426,6 +437,7 @@
         mCenterfreq0 = base.mCenterfreq0;
         mCenterfreq1 = base.mCenterfreq1;
         mDtimInterval = base.mDtimInterval;
+        mCountryCode = base.mCountryCode;
         mWifiMode = base.mWifiMode;
         mMaxRate = base.mMaxRate;
         mMaxNumberSpatialStreams = base.mMaxNumberSpatialStreams;
@@ -552,6 +564,10 @@
         return mDtimInterval;
     }
 
+    public String getCountryCode() {
+        return mCountryCode;
+    }
+
     public boolean is80211McResponderSupport() {
         return mExtendedCapabilities.is80211McRTTResponder();
     }
diff --git a/service/java/com/android/server/wifi/util/InformationElementUtil.java b/service/java/com/android/server/wifi/util/InformationElementUtil.java
index 92e5d35..ebcf379 100644
--- a/service/java/com/android/server/wifi/util/InformationElementUtil.java
+++ b/service/java/com/android/server/wifi/util/InformationElementUtil.java
@@ -36,6 +36,7 @@
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.List;
+import java.util.Locale;
 
 public class InformationElementUtil {
     private static final String TAG = "InformationElementUtil";
@@ -1704,4 +1705,58 @@
             return sbuf.toString();
         }
     }
+
+    /**
+     * This util class determines country related information in beacon/probe response
+     */
+    public static class Country {
+        private boolean mValid = false;
+        public String mCountryCode = "00";
+
+        /**
+         * Parse the Information Element Country Information field. Note that element ID and length
+         * fields are already removed.
+         *
+         * Country IE format (size unit: byte)
+         *
+         * ElementID | Length | country string | triplet | padding
+         *      1          1          3            Q*x       0 or 1
+         * First two bytes of country string are country code
+         * See 802.11 spec dot11CountryString definition.
+         */
+        public void from(InformationElement ie) {
+            mValid = false;
+            if (ie == null || ie.bytes == null || ie.bytes.length < 3) return;
+            ByteBuffer data = ByteBuffer.wrap(ie.bytes).order(ByteOrder.LITTLE_ENDIAN);
+            try {
+                char letter1 = (char) (data.get() & Constants.BYTE_MASK);
+                char letter2 = (char) (data.get() & Constants.BYTE_MASK);
+                char letter3 = (char) (data.get() & Constants.BYTE_MASK);
+                // See 802.11 spec dot11CountryString definition.
+                // ' ', 'O', 'I' are for all operation, outdoor, indoor environments, respectively.
+                mValid = (letter3 == ' ' || letter3 == 'O' || letter3 == 'I')
+                        && Character.isLetterOrDigit((int) letter1)
+                        && Character.isLetterOrDigit((int) letter2);
+                if (mValid) {
+                    mCountryCode = (String.valueOf(letter1) + letter2).toUpperCase(Locale.US);
+                }
+            } catch (BufferUnderflowException e) {
+                return;
+            }
+        }
+
+        /**
+         * Is this a valid country information element.
+         */
+        public boolean isValid() {
+            return mValid;
+        }
+
+        /**
+         * @return country code indicated in beacon/probe response frames
+         */
+        public String getCountryCode() {
+            return mCountryCode;
+        }
+    }
 }
diff --git a/service/proto/src/metrics.proto b/service/proto/src/metrics.proto
index 35db206..cc8b154 100644
--- a/service/proto/src/metrics.proto
+++ b/service/proto/src/metrics.proto
@@ -767,6 +767,13 @@
   // Histogram corresponding to the number of times recent connection failure status are reported
   // per WifiConfiguration.RecentFailureReason
   repeated Int32Count recent_failure_association_status = 218;
+
+  // Histogram of country codes observed from scan results.
+  // Bucket WifiMetrics.COUNTRY_CODE_CONFLICT_WIFI_SCAN for the code conflict within scan results.
+  // Bucket WifiMetrics.COUNTRY_CODE_CONFLICT_WIFI_SCAN_TELEPHONY for the code conflict between wifi
+  // and telephony.
+  // Bucket value is capped to WifiMetrics.MAX_COUNTRY_CODE_COUNT.
+  repeated Int32Count country_code_scan_histogram = 219;
 }
 
 // Information that gets logged for every WiFi connection.
diff --git a/service/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java b/service/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
index 3d6fa8f..6d913a1 100644
--- a/service/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
+++ b/service/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
@@ -5171,6 +5171,7 @@
         mCmi.sendMessage(ClientModeImpl.CMD_START_FILS_CONNECTION, 0, 0,
                 Collections.singletonList(l2Packet));
         mLooper.dispatchAll();
+        assertEquals("L2ConnectingState", mCmi.getCurrentState().getName());
     }
 
     /**
@@ -5278,6 +5279,7 @@
 
         mCmi.sendMessage(ClientModeImpl.CMD_START_CONNECT, 0, 0, TEST_BSSID_STR);
         mLooper.dispatchAll();
+        assertEquals("L2ConnectingState", mCmi.getCurrentState().getName());
 
         prepareFilsHlpPktAndSendStartConnect();
 
diff --git a/service/tests/wifitests/src/com/android/server/wifi/SupplicantStaIfaceHalTest.java b/service/tests/wifitests/src/com/android/server/wifi/SupplicantStaIfaceHalTest.java
index 7d429b3..6ef8b54 100644
--- a/service/tests/wifitests/src/com/android/server/wifi/SupplicantStaIfaceHalTest.java
+++ b/service/tests/wifitests/src/com/android/server/wifi/SupplicantStaIfaceHalTest.java
@@ -1266,6 +1266,42 @@
     }
 
     /**
+     * Tests the handling of association rejection for WPA3-Personal networks
+     */
+    @Test
+    public void testWpa3AuthRejectionEverConnected() throws Exception {
+        executeAndValidateInitializationSequence();
+        assertNotNull(mISupplicantStaIfaceCallback);
+
+        WifiConfiguration config = executeAndValidateConnectSequenceWithKeyMgmt(
+                SUPPLICANT_NETWORK_ID, false,
+                WifiConfiguration.SECURITY_TYPE_SAE, null, true);
+        mISupplicantStaIfaceCallback.onStateChanged(
+                ISupplicantStaIfaceCallback.State.ASSOCIATING,
+                NativeUtil.macAddressToByteArray(BSSID),
+                SUPPLICANT_NETWORK_ID,
+                NativeUtil.decodeSsid(SUPPLICANT_SSID));
+        int statusCode = ISupplicantStaIfaceCallback.StatusCode.UNSPECIFIED_FAILURE;
+        mISupplicantStaIfaceCallback.onAssociationRejected(
+                NativeUtil.macAddressToByteArray(BSSID), statusCode, false);
+        verify(mWifiMonitor, never()).broadcastAuthenticationFailureEvent(eq(WLAN0_IFACE_NAME),
+                anyInt(), anyInt());
+        ArgumentCaptor<AssocRejectEventInfo> assocRejectEventInfoCaptor =
+                ArgumentCaptor.forClass(AssocRejectEventInfo.class);
+        verify(mWifiMonitor).broadcastAssociationRejectionEvent(
+                eq(WLAN0_IFACE_NAME), assocRejectEventInfoCaptor.capture());
+        AssocRejectEventInfo assocRejectEventInfo =
+                (AssocRejectEventInfo) assocRejectEventInfoCaptor.getValue();
+        assertNotNull(assocRejectEventInfo);
+        assertEquals(SUPPLICANT_SSID, assocRejectEventInfo.ssid);
+        assertEquals(BSSID, assocRejectEventInfo.bssid);
+        assertEquals(statusCode, assocRejectEventInfo.statusCode);
+        assertFalse(assocRejectEventInfo.timedOut);
+        assertNull(assocRejectEventInfo.oceRssiBasedAssocRejectInfo);
+        assertNull(assocRejectEventInfo.mboAssocDisallowedInfo);
+    }
+
+    /**
      * Tests the handling of incorrect network passwords for WEP networks.
      */
     @Test
@@ -2926,11 +2962,12 @@
      * @param haveExistingNetwork Removes the existing network.
      * @param securityType The security type.
      * @param wepKey if configurations are for a WEP network else null.
+     * @param hasEverConnected indicate that this configuration is ever connected or not.
      * @return the WifiConfiguration object of the new network to connect.
      */
     private WifiConfiguration executeAndValidateConnectSequenceWithKeyMgmt(
             final int newFrameworkNetworkId, final boolean haveExistingNetwork,
-            int securityType, String wepKey) throws Exception {
+            int securityType, String wepKey, boolean hasEverConnected) throws Exception {
         setupMocksForConnectSequence(haveExistingNetwork);
         WifiConfiguration config = new WifiConfiguration();
         config.setSecurityParams(securityType);
@@ -2940,6 +2977,7 @@
         WifiConfiguration.NetworkSelectionStatus networkSelectionStatus =
                 new WifiConfiguration.NetworkSelectionStatus();
         networkSelectionStatus.setCandidateSecurityParams(config.getSecurityParams(securityType));
+        networkSelectionStatus.setHasEverConnected(hasEverConnected);
         config.setNetworkSelectionStatus(networkSelectionStatus);
         assertTrue(mDut.connectToNetwork(WLAN0_IFACE_NAME, config));
         validateConnectSequence(haveExistingNetwork, 1);
@@ -2947,6 +2985,23 @@
     }
 
     /**
+     * Helper function to execute all the actions to perform connection to the network.
+     *
+     * @param newFrameworkNetworkId Framework Network Id of the new network to connect.
+     * @param haveExistingNetwork Removes the existing network.
+     * @param securityType The security type.
+     * @param wepKey if configurations are for a WEP network else null.
+     * @return the WifiConfiguration object of the new network to connect.
+     */
+    private WifiConfiguration executeAndValidateConnectSequenceWithKeyMgmt(
+            final int newFrameworkNetworkId, final boolean haveExistingNetwork,
+            int securityType, String wepKey) throws Exception {
+        return executeAndValidateConnectSequenceWithKeyMgmt(
+                newFrameworkNetworkId, haveExistingNetwork,
+                securityType, wepKey, false);
+    }
+
+    /**
      * Setup mocks for roam sequence.
      */
     private void setupMocksForRoamSequence(String roamBssid) throws Exception {
diff --git a/service/tests/wifitests/src/com/android/server/wifi/WifiCountryCodeTest.java b/service/tests/wifitests/src/com/android/server/wifi/WifiCountryCodeTest.java
index fda21db..c833bf4 100644
--- a/service/tests/wifitests/src/com/android/server/wifi/WifiCountryCodeTest.java
+++ b/service/tests/wifitests/src/com/android/server/wifi/WifiCountryCodeTest.java
@@ -26,6 +26,7 @@
 
 import android.app.test.MockAnswerUtil.AnswerWithArguments;
 import android.content.Context;
+import android.net.wifi.WifiInfo;
 import android.telephony.TelephonyManager;
 
 import androidx.test.filters.SmallTest;
@@ -41,6 +42,8 @@
 
 import java.io.PrintWriter;
 import java.io.StringWriter;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Locale;
 
 /**
@@ -62,7 +65,9 @@
     @Mock ClientModeImplMonitor mClientModeImplMonitor;
     @Mock WifiNative mWifiNative;
     @Mock WifiSettingsConfigStore mSettingsConfigStore;
+    @Mock WifiInfo mWifiInfo;
     private WifiCountryCode mWifiCountryCode;
+    private List<ClientModeManager> mClientManagerList;
 
     @Captor
     private ArgumentCaptor<ActiveModeWarden.ModeChangeCallback> mModeChangeCallbackCaptor;
@@ -80,8 +85,16 @@
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
+        mClientManagerList = new ArrayList<>();
+        mClientManagerList.add(mClientModeManager);
+        when(mActiveModeWarden.getInternetConnectivityClientModeManagers())
+                .thenReturn(mClientManagerList);
         when(mClientModeManager.getRole()).thenReturn(ROLE_CLIENT_PRIMARY);
         when(mClientModeManager.setCountryCode(anyString())).thenReturn(true);
+        when(mClientModeManager.isConnected()).thenReturn(true);
+        when(mClientModeManager.syncRequestConnectionInfo()).thenReturn(mWifiInfo);
+        when(mWifiInfo.getSuccessfulTxPacketsPerSecond()).thenReturn(10.0);
+        when(mWifiInfo.getSuccessfulRxPacketsPerSecond()).thenReturn(5.0);
         when(mContext.getSystemService(Context.TELEPHONY_SERVICE))
                 .thenReturn(mTelephonyManager);
 
@@ -233,17 +246,34 @@
      */
     @Test
     public void telephonyCountryCodeChangeAfterL2Connected() throws Exception {
+        mWifiCountryCode.setDefaultCountryCode("00");
         // Supplicant starts.
         mModeChangeCallbackCaptor.getValue().onActiveModeManagerAdded(mClientModeManager);
         // Wifi get L2 connected.
         mClientModeImplListenerCaptor.getValue().onConnectionStart(mClientModeManager);
+
+        // Wifi traffic is high
+        when(mWifiInfo.getSuccessfulTxPacketsPerSecond()).thenReturn(20.0);
         // Telephony country code arrives.
         mWifiCountryCode.setTelephonyCountryCodeAndUpdate(mTelephonyCountryCode);
-        // Telephony coutry code won't be applied at this time.
-        assertEquals(mDefaultCountryCode, mWifiCountryCode.getCountryCodeSentToDriver());
+        // Telephony country code won't be applied at this time.
+        assertEquals("00", mWifiCountryCode.getCountryCodeSentToDriver());
+        // Wifi is not forced to disconnect
+        verify(mClientModeManager, times(0)).disconnect();
+
+        // Wifi traffic is low
+        when(mWifiInfo.getSuccessfulTxPacketsPerSecond()).thenReturn(10.0);
+        // Telephony country code arrives for multiple times
+        for (int i = 0; i < 3; i++) {
+            mWifiCountryCode.setTelephonyCountryCodeAndUpdate(mTelephonyCountryCode);
+        }
+        // Telephony country code still won't be applied.
+        assertEquals("00", mWifiCountryCode.getCountryCodeSentToDriver());
+        // Wifi is forced to disconnect
+        verify(mClientModeManager, times(1)).disconnect();
 
         mClientModeImplListenerCaptor.getValue().onConnectionEnd(mClientModeManager);
-        // Telephony coutry is applied after supplicant is ready.
+        // Telephony country is applied after supplicant is ready.
         verify(mClientModeManager, times(2)).setCountryCode(anyString());
         assertEquals(mTelephonyCountryCode, mWifiCountryCode.getCountryCodeSentToDriver());
     }
diff --git a/service/tests/wifitests/src/com/android/server/wifi/WifiMetricsTest.java b/service/tests/wifitests/src/com/android/server/wifi/WifiMetricsTest.java
index abb08fa..6084a6c 100644
--- a/service/tests/wifitests/src/com/android/server/wifi/WifiMetricsTest.java
+++ b/service/tests/wifitests/src/com/android/server/wifi/WifiMetricsTest.java
@@ -196,6 +196,7 @@
     @Mock PowerManager mPowerManager;
     @Mock WifiMonitor mWifiMonitor;
     @Mock ActiveModeWarden mActiveModeWarden;
+    @Mock TelephonyManager mTelephonyManager;
 
     @Captor ArgumentCaptor<BroadcastReceiver> mBroadcastReceiverCaptor;
     @Captor ArgumentCaptor<ActiveModeWarden.ModeChangeCallback> mModeChangeCallbackArgumentCaptor;
@@ -211,6 +212,8 @@
         when(mContext.getResources()).thenReturn(mResources);
         when(mContext.getSystemService(PowerManager.class)).thenReturn(mPowerManager);
         when(mPowerManager.isInteractive()).thenReturn(true);
+        when(mTelephonyManager.getNetworkCountryIso()).thenReturn("US");
+        when(mContext.getSystemService(TelephonyManager.class)).thenReturn(mTelephonyManager);
         mWifiMetrics = new WifiMetrics(mContext, mFacade, mClock, mTestLooper.getLooper(),
                 new WifiAwareMetrics(mClock), new RttMetrics(mClock), mWifiPowerMetrics,
                 mWifiP2pMetrics, mDppMetrics, mWifiMonitor);
@@ -3218,6 +3221,46 @@
     }
 
     /**
+     * Test that country code stats are collected correctly
+     */
+    @Test
+    public void testCountryCodeStats() throws Exception {
+        ScanDetail mockScanDetailUs = mock(ScanDetail.class);
+        ScanDetail mockScanDetailNonUs = mock(ScanDetail.class);
+        NetworkDetail mockNetworkDetailUs = mock(NetworkDetail.class);
+        NetworkDetail mockNetworkDetailNonUs = mock(NetworkDetail.class);
+        when(mockNetworkDetailUs.getCountryCode()).thenReturn("US");
+        when(mockNetworkDetailNonUs.getCountryCode()).thenReturn("CA");
+        ScanResult mockScanResult = mock(ScanResult.class);
+        mockScanResult.capabilities = "";
+        when(mockScanDetailUs.getScanResult()).thenReturn(mockScanResult);
+        when(mockScanDetailNonUs.getScanResult()).thenReturn(mockScanResult);
+        when(mockScanDetailUs.getNetworkDetail()).thenReturn(mockNetworkDetailUs);
+        when(mockScanDetailNonUs.getNetworkDetail()).thenReturn(mockNetworkDetailNonUs);
+        List<ScanDetail> scan = new ArrayList<ScanDetail>();
+
+        for (int i = 0; i < (WifiMetrics.MAX_COUNTRY_CODE_COUNT + 1); i++) {
+            scan.add(mockScanDetailUs);
+        }
+        mWifiMetrics.incrementAvailableNetworksHistograms(scan, true);
+        when(mTelephonyManager.getNetworkCountryIso()).thenReturn("CA");
+        mWifiMetrics.incrementAvailableNetworksHistograms(scan, true);
+        scan.add(mockScanDetailNonUs);
+        mWifiMetrics.incrementAvailableNetworksHistograms(scan, true);
+
+        dumpProtoAndDeserialize();
+
+        Int32Count[] expectedCountryCodeScanHistogram = {
+                buildInt32Count(WifiMetrics.COUNTRY_CODE_CONFLICT_WIFI_SCAN_TELEPHONY, 1),
+                buildInt32Count(WifiMetrics.COUNTRY_CODE_CONFLICT_WIFI_SCAN, 1),
+                buildInt32Count(WifiMetrics.MAX_COUNTRY_CODE_COUNT, 1),
+        };
+        assertKeyCountsEqual(expectedCountryCodeScanHistogram,
+                mDecodedProto.countryCodeScanHistogram);
+    }
+
+
+    /**
      * Test Open Network Notification blocklist size and feature state are not cleared when proto
      * is dumped.
      */
diff --git a/service/tests/wifitests/src/com/android/server/wifi/coex/CoexManagerTest.java b/service/tests/wifitests/src/com/android/server/wifi/coex/CoexManagerTest.java
index fbddee8..18ee415 100644
--- a/service/tests/wifitests/src/com/android/server/wifi/coex/CoexManagerTest.java
+++ b/service/tests/wifitests/src/com/android/server/wifi/coex/CoexManagerTest.java
@@ -867,4 +867,34 @@
         // The remote callback has an extra call since it was notified on registration.
         verify(remoteCallback, times(4)).onCoexUnsafeChannelsChanged(any(), anyInt());
     }
+
+    /**
+     * Verify that multiple calls to setCoexUnsafeChannels with the same unsafe channels will only
+     * update the driver the first time.
+     */
+    @Test
+    public void testSetCoexUnsafeChannels_multipleSameChannels_updatesDriverOnce() {
+        CoexManager coexManager = createCoexManager();
+
+        coexManager.setCoexUnsafeChannels(new ArrayList<>(), 0);
+        coexManager.setCoexUnsafeChannels(new ArrayList<>(), 0);
+        coexManager.setCoexUnsafeChannels(new ArrayList<>(), 0);
+        // Default state after initialization is no channels and no restrictions, so no update to
+        // the driver is necessary
+        verify(mMockWifiNative, times(0)).setCoexUnsafeChannels(any(), anyInt());
+
+        coexManager.setCoexUnsafeChannels(
+                Arrays.asList(new CoexUnsafeChannel(WIFI_BAND_24_GHZ, 6)), COEX_RESTRICTION_SOFTAP);
+        coexManager.setCoexUnsafeChannels(
+                Arrays.asList(new CoexUnsafeChannel(WIFI_BAND_24_GHZ, 6)), COEX_RESTRICTION_SOFTAP);
+        coexManager.setCoexUnsafeChannels(
+                Arrays.asList(new CoexUnsafeChannel(WIFI_BAND_24_GHZ, 6)), COEX_RESTRICTION_SOFTAP);
+        // Driver should be updated only once for the same unsafe channels
+        verify(mMockWifiNative, times(1)).setCoexUnsafeChannels(any(), anyInt());
+
+        // Change in restrictions with same unsafe channels should trigger an update
+        coexManager.setCoexUnsafeChannels(
+                Arrays.asList(new CoexUnsafeChannel(WIFI_BAND_24_GHZ, 6)), 0);
+        verify(mMockWifiNative, times(2)).setCoexUnsafeChannels(any(), anyInt());
+    }
 }
diff --git a/service/tests/wifitests/src/com/android/server/wifi/util/InformationElementUtilTest.java b/service/tests/wifitests/src/com/android/server/wifi/util/InformationElementUtilTest.java
index 942adbc..43169f1 100644
--- a/service/tests/wifitests/src/com/android/server/wifi/util/InformationElementUtilTest.java
+++ b/service/tests/wifitests/src/com/android/server/wifi/util/InformationElementUtilTest.java
@@ -1836,5 +1836,22 @@
                       2412, 72000000, true, true, true, false));
     }
 
+    /**
+     * Verify that the country code is parsed correctly from country IE
+     */
+    @Test
+    public void getCountryCodeWithCountryIE() throws Exception {
+        InformationElement ie = new InformationElement();
+        ie.id = InformationElement.EID_EXTENSION_PRESENT;
+        /** Country IE format (size unit: byte)
+         *
+         * |ElementID | Length | country string | triplet | padding
+         *      1          1          3            Q*x       0 or 1
+         */
+        ie.bytes = new byte[]{(byte) 0x75, (byte) 0x73, (byte) 0x49};
+        InformationElementUtil.Country country = new InformationElementUtil.Country();
+        country.from(ie);
+        assertEquals("US", country.getCountryCode());
+    }
     // TODO: SAE, OWN, SUITE_B
 }