5 files changed, 38 insertions, 6 deletions

diff --git a/net/mac80211/ieee80211_i.h b/net/mac80211/ieee80211_i.h
index 4947c91c6c86..892bac64a189 100644
--- a/net/mac80211/ieee80211_i.h
+++ b/net/mac80211/ieee80211_i.h
@@ -1197,6 +1197,7 @@ struct ieee802_11_elems {
         u8 *pwr_constr_elem;
         u8 *quiet_elem; /* first quite element */
         u8 *timeout_int;
+        u8 *opmode_notif;
 
         /* length of them, respectively */
         u8 ssid_len;
@@ -1435,7 +1436,7 @@ enum ieee80211_sta_rx_bandwidth ieee80211_sta_cur_vht_bw(struct sta_info *sta);
 void ieee80211_sta_set_rx_nss(struct sta_info *sta);
 void ieee80211_vht_handle_opmode(struct ieee80211_sub_if_data *sdata,
                                  struct sta_info *sta, u8 opmode,
-                                 enum ieee80211_band band);
+                                 enum ieee80211_band band, bool nss_only);
 
 /* Spectrum management */
 void ieee80211_process_measurement_req(struct ieee80211_sub_if_data *sdata,
diff --git a/net/mac80211/mlme.c b/net/mac80211/mlme.c
index e6da6bf8cc27..f2c6f7794f35 100644
--- a/net/mac80211/mlme.c
+++ b/net/mac80211/mlme.c
@@ -2217,6 +2217,21 @@ static bool ieee80211_assoc_success(struct ieee80211_sub_if_data *sdata,
                                                   elems.ht_operation,
                                                   cbss->bssid, false);
 
+        /*
+         * If an operating mode notification IE is present, override the
+         * NSS calculation (that would be done in rate_control_rate_init())
+         * and use the # of streams from that element.
+         */
+        if (elems.opmode_notif &&
+            !(*elems.opmode_notif & IEEE80211_OPMODE_NOTIF_RX_NSS_TYPE_BF)) {
+                u8 nss;
+
+                nss = *elems.opmode_notif & IEEE80211_OPMODE_NOTIF_RX_NSS_MASK;
+                nss >>= IEEE80211_OPMODE_NOTIF_RX_NSS_SHIFT;
+                nss += 1;
+                sta->sta.rx_nss = nss;
+        }
+
         rate_control_rate_init(sta);
 
         if (ifmgd->flags & IEEE80211_STA_MFP_ENABLED)
@@ -2489,6 +2504,7 @@ static void ieee80211_rx_mgmt_beacon(struct ieee80211_sub_if_data *sdata,
         struct ieee80211_local *local = sdata->local;
         struct ieee80211_chanctx_conf *chanctx_conf;
         struct ieee80211_channel *chan;
+        struct sta_info *sta;
         u32 changed = 0;
         bool erp_valid;
         u8 erp_value = 0;
@@ -2728,14 +2744,18 @@ static void ieee80211_rx_mgmt_beacon(struct ieee80211_sub_if_data *sdata,
                         le16_to_cpu(mgmt->u.beacon.capab_info),
                         erp_valid, erp_value);
 
-
         mutex_lock(&local->sta_mtx);
+        sta = sta_info_get(sdata, bssid);
+
         if (elems.ht_cap_elem && elems.ht_operation && elems.wmm_param &&
             !(ifmgd->flags & IEEE80211_STA_DISABLE_HT))
-                changed |= ieee80211_config_ht_tx(sdata,
-                                                  sta_info_get(sdata, bssid),
+                changed |= ieee80211_config_ht_tx(sdata, sta,
                                                   elems.ht_operation,
                                                   bssid, true);
+
+        if (sta && elems.opmode_notif)
+                ieee80211_vht_handle_opmode(sdata, sta, *elems.opmode_notif,
+                                            rx_status->band, true);
         mutex_unlock(&local->sta_mtx);
 
         if (elems.country_elem && elems.pwr_constr_elem &&
diff --git a/net/mac80211/rx.c b/net/mac80211/rx.c
index 1617e0bd4ca6..30f1ba6de8f3 100644
--- a/net/mac80211/rx.c
+++ b/net/mac80211/rx.c
@@ -2459,7 +2459,8 @@ ieee80211_rx_h_action(struct ieee80211_rx_data *rx)
                         opmode = mgmt->u.action.u.vht_opmode_notif.operating_mode;
 
                         ieee80211_vht_handle_opmode(rx->sdata, rx->sta,
-                                                    opmode, status->band);
+                                                    opmode, status->band,
+                                                    false);
                         goto handled;
                 }
                 default:
diff --git a/net/mac80211/util.c b/net/mac80211/util.c
index b231bc2ed740..e24ff38606a9 100644
--- a/net/mac80211/util.c
+++ b/net/mac80211/util.c
@@ -787,6 +787,12 @@ u32 ieee802_11_parse_elems_crc(u8 *start, size_t len,
                         else
                                 elem_parse_failed = true;
                         break;
+                case WLAN_EID_OPMODE_NOTIF:
+                        if (elen > 0)
+                                elems->opmode_notif = pos;
+                        else
+                                elem_parse_failed = true;
+                        break;
                 case WLAN_EID_MESH_ID:
                         elems->mesh_id = pos;
                         elems->mesh_id_len = elen;
diff --git a/net/mac80211/vht.c b/net/mac80211/vht.c
index 0951f74e7ff5..a9549fcc5a04 100644
--- a/net/mac80211/vht.c
+++ b/net/mac80211/vht.c
@@ -141,7 +141,7 @@ void ieee80211_sta_set_rx_nss(struct sta_info *sta)
 
 void ieee80211_vht_handle_opmode(struct ieee80211_sub_if_data *sdata,
                                  struct sta_info *sta, u8 opmode,
-                                 enum ieee80211_band band)
+                                 enum ieee80211_band band, bool nss_only)
 {
         struct ieee80211_local *local = sdata->local;
         struct ieee80211_supported_band *sband;
@@ -164,6 +164,9 @@ void ieee80211_vht_handle_opmode(struct ieee80211_sub_if_data *sdata,
                 changed |= IEEE80211_RC_NSS_CHANGED;
         }
 
+        if (nss_only)
+                goto change;
+
         switch (opmode & IEEE80211_OPMODE_NOTIF_CHANWIDTH_MASK) {
         case IEEE80211_OPMODE_NOTIF_CHANWIDTH_20MHZ:
                 sta->cur_max_bandwidth = IEEE80211_STA_RX_BW_20;
@@ -185,6 +188,7 @@ void ieee80211_vht_handle_opmode(struct ieee80211_sub_if_data *sdata,
                 changed |= IEEE80211_RC_NSS_CHANGED;
         }
 
+ change:
         if (changed)
                 rate_control_rate_update(local, sband, sta, changed);
 }