4 files changed, 16 insertions, 38 deletions

diff --git a/drivers/staging/rtl8723au/core/rtw_ieee80211.c b/drivers/staging/rtl8723au/core/rtw_ieee80211.c
index 780631fd3b6d..a48ab25a7d8a 100644
--- a/drivers/staging/rtl8723au/core/rtw_ieee80211.c
+++ b/drivers/staging/rtl8723au/core/rtw_ieee80211.c
@@ -1496,45 +1496,23 @@ void rtw_wlan_bssid_ex_remove_p2p_attr23a(struct wlan_bssid_ex *bss_ex, u8 attr_
 int rtw_get_wfd_ie(u8 *in_ie, int in_len, u8 *wfd_ie, uint *wfd_ielen)
 {
         int match;
-        uint cnt = 0;
-        u8 eid, wfd_oui[4] = {0x50, 0x6F, 0x9A, 0x0A};
+        const u8 *ie;
 
-        match = false;
+        match = 0;
 
-        if (in_len < 0) {
+        if (in_len < 0)
                 return match;
-        }
-
-        while (cnt < in_len)
-        {
-                eid = in_ie[cnt];
 
-                if ((eid == _VENDOR_SPECIFIC_IE_) &&
-                    !memcmp(&in_ie[cnt+2], wfd_oui, 4)) {
-                        if (wfd_ie != NULL) {
-                                memcpy(wfd_ie, &in_ie[cnt], in_ie[cnt + 1] + 2);
-
-                        } else {
-                                if (wfd_ielen != NULL) {
-                                        *wfd_ielen = 0;
-                                }
-                        }
-
-                        if (wfd_ielen != NULL) {
-                                *wfd_ielen = in_ie[cnt + 1] + 2;
-                        }
-
-                        cnt += in_ie[cnt + 1] + 2;
-
-                        match = true;
-                        break;
-                } else {
-                        cnt += in_ie[cnt + 1] +2; /* goto next */
-                }
-        }
+        ie = cfg80211_find_vendor_ie(0x506F9A, 0x0A, in_ie, in_len);
+        if (ie && (ie[1] <= (MAX_WFD_IE_LEN - 2))) {
+                if (wfd_ie) {
+                        *wfd_ielen = ie[1] + 2;
+                        memcpy(wfd_ie, ie, ie[1] + 2);
+                } else
+                        if (wfd_ielen)
+                                *wfd_ielen = 0;
 
-        if (match == true) {
-                match = cnt;
+                match = 1;
         }
 
         return match;
diff --git a/drivers/staging/rtl8723au/core/rtw_mlme_ext.c b/drivers/staging/rtl8723au/core/rtw_mlme_ext.c
index 4c753639ea5a..1f3e8a0aece4 100644
--- a/drivers/staging/rtl8723au/core/rtw_mlme_ext.c
+++ b/drivers/staging/rtl8723au/core/rtw_mlme_ext.c
@@ -1281,7 +1281,7 @@ unsigned int OnAssocReq23a(struct rtw_adapter *padapter, struct recv_frame *prec
         u8 p2p_status_code = P2P_STATUS_SUCCESS;
         u8 *p2pie;
         u32 p2pielen = 0;
-        u8      wfd_ie[ 128 ] = { 0x00 };
+        u8      wfd_ie[MAX_WFD_IE_LEN] = { 0x00 };
         u32     wfd_ielen = 0;
 #endif /* CONFIG_8723AU_P2P */
 
diff --git a/drivers/staging/rtl8723au/core/rtw_p2p.c b/drivers/staging/rtl8723au/core/rtw_p2p.c
index 27a6cc76973d..1a961e3f3a55 100644
--- a/drivers/staging/rtl8723au/core/rtw_p2p.c
+++ b/drivers/staging/rtl8723au/core/rtw_p2p.c
@@ -2535,7 +2535,7 @@ u8 process_p2p_group_negotation_req23a(struct wifidirect_info *pwdinfo, u8 *pfra
         u16             wps_devicepassword_id = 0x0000;
         uint    wps_devicepassword_id_len = 0;
 #ifdef CONFIG_8723AU_P2P
-        u8      wfd_ie[ 128 ] = { 0x00 };
+        u8      wfd_ie[MAX_WFD_IE_LEN] = { 0x00 };
         u32     wfd_ielen = 0;
 #endif /*  CONFIG_8723AU_P2P */
 
@@ -2741,7 +2741,7 @@ u8 process_p2p_group_negotation_resp23a(struct wifidirect_info *pwdinfo, u8 *pfr
         u32 ies_len;
         u8 * p2p_ie;
 #ifdef CONFIG_8723AU_P2P
-        u8      wfd_ie[ 128 ] = { 0x00 };
+        u8      wfd_ie[MAX_WFD_IE_LEN] = { 0x00 };
         u32     wfd_ielen = 0;
 #endif /*  CONFIG_8723AU_P2P */
 
diff --git a/drivers/staging/rtl8723au/core/rtw_wlan_util.c b/drivers/staging/rtl8723au/core/rtw_wlan_util.c
index 0dfcfbce3b52..e743b053b8a2 100644
--- a/drivers/staging/rtl8723au/core/rtw_wlan_util.c
+++ b/drivers/staging/rtl8723au/core/rtw_wlan_util.c
@@ -570,7 +570,7 @@ void flush_all_cam_entry23a(struct rtw_adapter *padapter)
 int WFD_info_handler(struct rtw_adapter *padapter, struct ndis_802_11_var_ies * pIE)
 {
         struct wifidirect_info  *pwdinfo;
-        u8      wfd_ie[128] = {0x00};
+        u8      wfd_ie[MAX_WFD_IE_LEN] = {0x00};
         u32     wfd_ielen = 0;
 
         pwdinfo = &padapter->wdinfo;