1 files changed, 48 insertions, 4 deletions

diff --git a/src/p2p/p2p_pd.c b/src/p2p/p2p_pd.c
index 007b643..610f0df 100644
--- a/src/p2p/p2p_pd.c
+++ b/src/p2p/p2p_pd.c
@@ -46,8 +46,14 @@ static struct wpabuf * p2p_build_prov_disc_req(struct p2p_data *p2p,
 {
 	struct wpabuf *buf;
 	u8 *len;
+	size_t extra = 0;
 
-	buf = wpabuf_alloc(1000);
+#ifdef CONFIG_WIFI_DISPLAY
+	if (p2p->wfd_ie_prov_disc_req)
+		extra = wpabuf_len(p2p->wfd_ie_prov_disc_req);
+#endif /* CONFIG_WIFI_DISPLAY */
+
+	buf = wpabuf_alloc(1000 + extra);
 	if (buf == NULL)
 		return NULL;
 
@@ -66,17 +72,46 @@ static struct wpabuf * p2p_build_prov_disc_req(struct p2p_data *p2p,
 	/* WPS IE with Config Methods attribute */
 	p2p_build_wps_ie_config_methods(buf, config_methods);
 
+#ifdef CONFIG_WIFI_DISPLAY
+	if (p2p->wfd_ie_prov_disc_req)
+		wpabuf_put_buf(buf, p2p->wfd_ie_prov_disc_req);
+#endif /* CONFIG_WIFI_DISPLAY */
+
 	return buf;
 }
 
 
 static struct wpabuf * p2p_build_prov_disc_resp(struct p2p_data *p2p,
 						u8 dialog_token,
-						u16 config_methods)
+						u16 config_methods,
+						const u8 *group_id,
+						size_t group_id_len)
 {
 	struct wpabuf *buf;
+	size_t extra = 0;
+
+#ifdef CONFIG_WIFI_DISPLAY
+	struct wpabuf *wfd_ie = p2p->wfd_ie_prov_disc_resp;
+	if (wfd_ie && group_id) {
+		size_t i;
+		for (i = 0; i < p2p->num_groups; i++) {
+			struct p2p_group *g = p2p->groups[i];
+			struct wpabuf *ie;
+			if (!p2p_group_is_group_id_match(g, group_id,
+							 group_id_len))
+				continue;
+			ie = p2p_group_get_wfd_ie(g);
+			if (ie) {
+				wfd_ie = ie;
+				break;
+			}
+		}
+	}
+	if (wfd_ie)
+		extra = wpabuf_len(wfd_ie);
+#endif /* CONFIG_WIFI_DISPLAY */
 
-	buf = wpabuf_alloc(100);
+	buf = wpabuf_alloc(100 + extra);
 	if (buf == NULL)
 		return NULL;
 
@@ -85,6 +120,11 @@ static struct wpabuf * p2p_build_prov_disc_resp(struct p2p_data *p2p,
 	/* WPS IE with Config Methods attribute */
 	p2p_build_wps_ie_config_methods(buf, config_methods);
 
+#ifdef CONFIG_WIFI_DISPLAY
+	if (wfd_ie)
+		wpabuf_put_buf(buf, wfd_ie);
+#endif /* CONFIG_WIFI_DISPLAY */
+
 	return buf;
 }
 
@@ -119,6 +159,9 @@ void p2p_process_prov_disc_req(struct p2p_data *p2p, const u8 *sa,
 			        "P2P: Provision Discovery Request add device "
 				"failed " MACSTR, MAC2STR(sa));
 		}
+	} else if (msg.wfd_subelems) {
+		wpabuf_free(dev->info.wfd_subelems);
+		dev->info.wfd_subelems = wpabuf_dup(msg.wfd_subelems);
 	}
 
 	if (!(msg.wps_config_methods &
@@ -164,7 +207,8 @@ void p2p_process_prov_disc_req(struct p2p_data *p2p, const u8 *sa,
 
 out:
 	resp = p2p_build_prov_disc_resp(p2p, msg.dialog_token,
-					reject ? 0 : msg.wps_config_methods);
+					reject ? 0 : msg.wps_config_methods,
+					msg.group_id, msg.group_id_len);
 	if (resp == NULL) {
 		p2p_parse_free(&msg);
 		return;