diff --git a/net/wireless/mlme.c b/net/wireless/mlme.c
index 097a87d..525e8e2 100644
--- a/net/wireless/mlme.c
+++ b/net/wireless/mlme.c
@@ -61,7 +61,7 @@
 	struct ieee80211_mgmt *mgmt = (struct ieee80211_mgmt *)buf;
 	u8 *ie = mgmt->u.assoc_resp.variable;
 	int i, ieoffs = offsetof(struct ieee80211_mgmt, u.assoc_resp.variable);
-	bool done;
+	struct cfg80211_internal_bss *bss = NULL;
 
 	wdev_lock(wdev);
 
@@ -69,22 +69,32 @@
 
 	nl80211_send_rx_assoc(rdev, dev, buf, len, GFP_KERNEL);
 
-	__cfg80211_connect_result(dev, mgmt->bssid, NULL, 0, ie, len - ieoffs,
-				  status_code,
-				  status_code == WLAN_STATUS_SUCCESS);
-
 	if (status_code == WLAN_STATUS_SUCCESS) {
-		for (i = 0; wdev->current_bss && i < MAX_AUTH_BSSES; i++) {
-			if (wdev->auth_bsses[i] == wdev->current_bss) {
-				cfg80211_unhold_bss(wdev->auth_bsses[i]);
-				cfg80211_put_bss(&wdev->auth_bsses[i]->pub);
+		for (i = 0; i < MAX_AUTH_BSSES; i++) {
+			if (!wdev->auth_bsses[i])
+				continue;
+			if (memcmp(wdev->auth_bsses[i]->pub.bssid, mgmt->bssid,
+				   ETH_ALEN) == 0) {
+				bss = wdev->auth_bsses[i];
 				wdev->auth_bsses[i] = NULL;
-				done = true;
+				/* additional reference to drop hold */
+				cfg80211_ref_bss(bss);
 				break;
 			}
 		}
 
-		WARN_ON(!done);
+		WARN_ON(!bss);
+	}
+
+	/* this consumes one bss reference (unless bss is NULL) */
+	__cfg80211_connect_result(dev, mgmt->bssid, NULL, 0, ie, len - ieoffs,
+				  status_code,
+				  status_code == WLAN_STATUS_SUCCESS,
+				  bss ? &bss->pub : NULL);
+	/* drop hold now, and also reference acquired above */
+	if (bss) {
+		cfg80211_unhold_bss(bss);
+		cfg80211_put_bss(&bss->pub);
 	}
 
 	wdev_unlock(wdev);
@@ -144,7 +154,7 @@
 	} else if (wdev->sme_state == CFG80211_SME_CONNECTING) {
 		__cfg80211_connect_result(dev, mgmt->bssid, NULL, 0, NULL, 0,
 					  WLAN_STATUS_UNSPECIFIED_FAILURE,
-					  false);
+					  false, NULL);
 	}
 }
 
@@ -241,7 +251,7 @@
 	if (wdev->sme_state == CFG80211_SME_CONNECTING)
 		__cfg80211_connect_result(dev, addr, NULL, 0, NULL, 0,
 					  WLAN_STATUS_UNSPECIFIED_FAILURE,
-					  false);
+					  false, NULL);
 
 	for (i = 0; addr && i < MAX_AUTH_BSSES; i++) {
 		if (wdev->authtry_bsses[i] &&
@@ -275,7 +285,7 @@
 	if (wdev->sme_state == CFG80211_SME_CONNECTING)
 		__cfg80211_connect_result(dev, addr, NULL, 0, NULL, 0,
 					  WLAN_STATUS_UNSPECIFIED_FAILURE,
-					  false);
+					  false, NULL);
 
 	for (i = 0; addr && i < MAX_AUTH_BSSES; i++) {
 		if (wdev->auth_bsses[i] &&