diff --git a/drivers/staging/rtl8723au/core/rtw_ieee80211.c b/drivers/staging/rtl8723au/core/rtw_ieee80211.c
index 780631fd3b6d4a8f9bbbb9b7662853f7ba5d9608..a48ab25a7d8a251cd93a30cb04349a09ce3cdf1b 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 be593989aec751ba5d898b396ba5d986f1a3cc5c..07493bc7053952cf01d374d81473cec1f4fb9446 100644
--- a/drivers/staging/rtl8723au/core/rtw_mlme_ext.c
+++ b/drivers/staging/rtl8723au/core/rtw_mlme_ext.c
@@ -1280,7 +1280,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 ba851cf4f9921af8e9a42fb7e16bd6ad30f96786..a3bf031ee738929ace4059d4ffecc75ce34e4204 100644
--- a/drivers/staging/rtl8723au/core/rtw_p2p.c
+++ b/drivers/staging/rtl8723au/core/rtw_p2p.c
@@ -2536,7 +2536,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 */
 
@@ -2742,7 +2742,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 e364d7afc08d6497e4699ac1bbd48e3bbd9113a1..f54021ceba5cc9dd566edabf6ee33cea3eef5c29 100644
--- a/drivers/staging/rtl8723au/core/rtw_wlan_util.c
+++ b/drivers/staging/rtl8723au/core/rtw_wlan_util.c
@@ -502,7 +502,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;