Merge driver ops set_wps_beacon_ie and set_wps_probe_resp_ie

set_ap_wps_ie() is not used to set WPS IE for both Beacon and Probe
Response frames with a single call. In addition, struct wpabuf is used
instead of separate u8* and length fields. This avoids duplicated
allocation of the IEs and simplifies code in general.

hostapd/ap_drv_ops.c

@@ -20,16 +20,13 @@
 
 
 static int hostapd_set_ap_wps_ie(struct hostapd_data *hapd,
-				 const u8 *beacon_ie, size_t beacon_ie_len,
-				 const u8 *probe_resp_ie,
-				 size_t probe_resp_ie_len)
+				 const struct wpabuf *beacon,
+				 const struct wpabuf *proberesp)
 {
-	if (hostapd_set_wps_beacon_ie(hapd, hapd->wps_beacon_ie,
-				      hapd->wps_beacon_ie_len) < 0 ||
-	    hostapd_set_wps_probe_resp_ie(hapd, hapd->wps_probe_resp_ie,
-					  hapd->wps_probe_resp_ie_len) < 0)
-		return -1;
-	return 0;
+	if (hapd->driver == NULL || hapd->driver->set_ap_wps_ie == NULL)
+		return 0;
+	return hapd->driver->set_ap_wps_ie(hapd->conf->iface, hapd->drv_priv,
+					   beacon, proberesp);
 }
 
 

hostapd/beacon.c

@@ -309,9 +309,9 @@ void handle_probe_req(struct hostapd_data *hapd,
 
 #ifdef CONFIG_WPS
 	if (hapd->conf->wps_state && hapd->wps_probe_resp_ie) {
-		os_memcpy(pos, hapd->wps_probe_resp_ie,
-			  hapd->wps_probe_resp_ie_len);
-		pos += hapd->wps_probe_resp_ie_len;
+		os_memcpy(pos, wpabuf_head(hapd->wps_probe_resp_ie),
+			  wpabuf_len(hapd->wps_probe_resp_ie));
+		pos += wpabuf_len(hapd->wps_probe_resp_ie);
 	}
 #endif /* CONFIG_WPS */
 
@@ -422,9 +422,9 @@ void ieee802_11_set_beacon(struct hostapd_data *hapd)
 
 #ifdef CONFIG_WPS
 	if (hapd->conf->wps_state && hapd->wps_beacon_ie) {
-		os_memcpy(tailpos, hapd->wps_beacon_ie,
-			  hapd->wps_beacon_ie_len);
-		tailpos += hapd->wps_beacon_ie_len;
+		os_memcpy(tailpos, wpabuf_head(hapd->wps_beacon_ie),
+			  wpabuf_len(hapd->wps_beacon_ie));
+		tailpos += wpabuf_len(hapd->wps_beacon_ie);
 	}
 #endif /* CONFIG_WPS */
 

hostapd/driver_i.h

@@ -454,26 +454,6 @@ hostapd_drv_none(struct hostapd_data *hapd)
 	return hapd->driver && os_strcmp(hapd->driver->name, "none") == 0;
 }
 
-static inline int
-hostapd_set_wps_beacon_ie(struct hostapd_data *hapd, const u8 *ie, size_t len)
-{
-	if (hapd->driver == NULL || hapd->driver->set_wps_beacon_ie == NULL)
-		return 0;
-	return hapd->driver->set_wps_beacon_ie(hapd->conf->iface,
-					       hapd->drv_priv, ie, len);
-}
-
-static inline int
-hostapd_set_wps_probe_resp_ie(struct hostapd_data *hapd, const u8 *ie,
-			      size_t len)
-{
-	if (hapd->driver == NULL ||
-	    hapd->driver->set_wps_probe_resp_ie == NULL)
-		return 0;
-	return hapd->driver->set_wps_probe_resp_ie(hapd->conf->iface,
-						   hapd->drv_priv, ie, len);
-}
-
 static inline int hostapd_driver_scan(struct hostapd_data *hapd,
 				      struct wpa_driver_scan_params *params)
 {

hostapd/hostapd.h

@@ -43,9 +43,8 @@ struct hostapd_rate_data {
 
 struct hostapd_driver_ops {
 	int (*set_ap_wps_ie)(struct hostapd_data *hapd,
-			     const u8 *beacon_ie, size_t beacon_ie_len,
-			     const u8 *probe_resp_ie,
-			     size_t probe_resp_ie_len);
+			     const struct wpabuf *beacon,
+			     const struct wpabuf *probe);
 };
 
 /**
@@ -112,10 +111,8 @@ struct hostapd_data {
 	struct wps_context *wps;
 
 #ifdef CONFIG_WPS
-	u8 *wps_beacon_ie;
-	size_t wps_beacon_ie_len;
-	u8 *wps_probe_resp_ie;
-	size_t wps_probe_resp_ie_len;
+	struct wpabuf *wps_beacon_ie;
+	struct wpabuf *wps_probe_resp_ie;
 	unsigned int ap_pin_failures;
 	struct upnp_wps_device_sm *wps_upnp;
 #endif /* CONFIG_WPS */

hostapd/wps_hostapd.c

@@ -90,46 +90,16 @@ static int hostapd_wps_new_psk_cb(void *ctx, const u8 *mac_addr, const u8 *psk,
 }
 
 
-static int hostapd_wps_set_ie_cb(void *ctx, const u8 *beacon_ie,
-				 size_t beacon_ie_len, const u8 *probe_resp_ie,
-				 size_t probe_resp_ie_len)
+static int hostapd_wps_set_ie_cb(void *ctx, struct wpabuf *beacon_ie,
+				 struct wpabuf *probe_resp_ie)
 {
 	struct hostapd_data *hapd = ctx;
-
-	os_free(hapd->wps_beacon_ie);
-	if (beacon_ie_len == 0) {
-		hapd->wps_beacon_ie = NULL;
-		hapd->wps_beacon_ie_len = 0;
-	} else {
-		hapd->wps_beacon_ie = os_malloc(beacon_ie_len);
-		if (hapd->wps_beacon_ie == NULL) {
-			hapd->wps_beacon_ie_len = 0;
-			return -1;
-		}
-		os_memcpy(hapd->wps_beacon_ie, beacon_ie, beacon_ie_len);
-		hapd->wps_beacon_ie_len = beacon_ie_len;
-	}
-
-	os_free(hapd->wps_probe_resp_ie);
-	if (probe_resp_ie_len == 0) {
-		hapd->wps_probe_resp_ie = NULL;
-		hapd->wps_probe_resp_ie_len = 0;
-	} else {
-		hapd->wps_probe_resp_ie = os_malloc(probe_resp_ie_len);
-		if (hapd->wps_probe_resp_ie == NULL) {
-			hapd->wps_probe_resp_ie_len = 0;
-			return -1;
-		}
-		os_memcpy(hapd->wps_probe_resp_ie, probe_resp_ie,
-			  probe_resp_ie_len);
-		hapd->wps_probe_resp_ie_len = probe_resp_ie_len;
-	}
-	hapd->drv.set_ap_wps_ie(hapd, hapd->wps_beacon_ie,
-				hapd->wps_beacon_ie_len,
-				hapd->wps_probe_resp_ie,
-				hapd->wps_probe_resp_ie_len);
-
-	return 0;
+	wpabuf_free(hapd->wps_beacon_ie);
+	hapd->wps_beacon_ie = beacon_ie;
+	wpabuf_free(hapd->wps_probe_resp_ie);
+	hapd->wps_probe_resp_ie = probe_resp_ie;
+	return hapd->drv.set_ap_wps_ie(hapd, hapd->wps_beacon_ie,
+				       hapd->wps_probe_resp_ie);
 }
 
 
@@ -477,15 +447,13 @@ static void hostapd_wps_event_cb(void *ctx, enum wps_event event,
 
 static void hostapd_wps_clear_ies(struct hostapd_data *hapd)
 {
-	os_free(hapd->wps_beacon_ie);
+	wpabuf_free(hapd->wps_beacon_ie);
 	hapd->wps_beacon_ie = NULL;
-	hapd->wps_beacon_ie_len = 0;
 
-	os_free(hapd->wps_probe_resp_ie);
+	wpabuf_free(hapd->wps_probe_resp_ie);
 	hapd->wps_probe_resp_ie = NULL;
-	hapd->wps_probe_resp_ie_len = 0;
 
-	hapd->drv.set_ap_wps_ie(hapd, NULL, 0, NULL, 0);
+	hapd->drv.set_ap_wps_ie(hapd, NULL, NULL);
 }
 
 

src/drivers/driver.h

@@ -1491,34 +1491,21 @@ struct wpa_driver_ops {
 			     const u8 *ht_oper, size_t ht_oper_len);
 
 	/**
-	 * set_wps_beacon_ie - Add WPS IE into Beacon frames (AP only)
+	 * set_ap_wps_ie - Add WPS IE into Beacon/Probe Response frames (AP)
 	 * @ifname: The interface name (main or virtual BSS)
 	 * @priv: Private driver interface data
-	 * @ie: WPS IE
-	 * @len: Length of the ie buffer in octets
+	 * @beacon: WPS IE for Beacon frames
+	 * @proberesp: WPS IE for Probe Response frames
 	 * Returns: 0 on success, -1 on failure
 	 *
 	 * This is an optional function to add WPS IE in the kernel driver for
-	 * Beacon frames. This can be left undefined (set to %NULL) if the
-	 * driver uses the Beacon template from set_beacon().
+	 * Beacon and Probe Response frames. This can be left undefined (set
+	 * to %NULL) if the driver uses the Beacon template from set_beacon()
+	 * and does not process Probe Request frames.
 	 */
-	int (*set_wps_beacon_ie)(const char *ifname, void *priv,
-				 const u8 *ie, size_t len);
-
-	/**
-	 * set_wps_probe_resp_ie - Add WPS IE into Probe Response frames (AP)
-	 * @ifname: The interface name (main or virtual BSS)
-	 * @priv: Private driver interface data
-	 * @ie: WPS IE
-	 * @len: Length of the ie buffer in octets
-	 * Returns: 0 on success, -1 on failure
-	 *
-	 * This is an optional function to add WPS IE in the kernel driver for
-	 * Beacon frames. This can be left undefined (set to %NULL) if the
-	 * driver does process Probe Request frames.
-	 */
-	int (*set_wps_probe_resp_ie)(const char *ifname, void *priv,
-				     const u8 *ie, size_t len);
+	int (*set_ap_wps_ie)(const char *ifname, void *priv,
+			     const struct wpabuf *beacon,
+			     const struct wpabuf *proberesp);
 
 	/**
 	 * set_supp_port - Set IEEE 802.1X Supplicant Port status

src/drivers/driver_atheros.c

@@ -715,22 +715,21 @@ madwifi_set_wps_ie(void *priv, const u8 *ie, size_t len, u32 frametype)
 }
 
 static int
-madwifi_set_wps_beacon_ie(const char *ifname, void *priv, const u8 *ie,
-			  size_t len)
+madwifi_set_ap_wps_ie(const char *ifname, void *priv,
+		      const struct wpabuf *beacon,
+		      const struct wpabuf *proberesp)
 {
-	return madwifi_set_wps_ie(priv, ie, len, IEEE80211_APPIE_FRAME_BEACON);
-}
-
-static int
-madwifi_set_wps_probe_resp_ie(const char *ifname, void *priv, const u8 *ie,
-			      size_t len)
-{
-	return madwifi_set_wps_ie(priv, ie, len,
+	if (madwifi_set_wps_ie(priv, beacon ? wpabuf_head(beacon) : NULL,
+			       beacon ? wpabuf_len(beacon) : 0,
+			       IEEE80211_APPIE_FRAME_BEACON))
+		return -1;
+	return madwifi_set_wps_ie(priv,
+				  proberesp ? wpabuf_head(proberesp) : NULL,
+				  proberesp ? wpabuf_len(proberesp): 0,
 				  IEEE80211_APPIE_FRAME_PROBE_RESP);
 }
 #else /* CONFIG_WPS */
-#define madwifi_set_wps_beacon_ie NULL
-#define madwifi_set_wps_probe_resp_ie NULL
+#define madwifi_set_ap_wps_ie NULL
 #endif /* CONFIG_WPS */
 
 static int
@@ -1249,6 +1248,5 @@ const struct wpa_driver_ops wpa_driver_atheros_ops = {
 	.set_countermeasures	= madwifi_set_countermeasures,
 	.sta_clear_stats	= madwifi_sta_clear_stats,
 	.commit			= madwifi_commit,
-	.set_wps_beacon_ie	= madwifi_set_wps_beacon_ie,
-	.set_wps_probe_resp_ie	= madwifi_set_wps_probe_resp_ie,
+	.set_ap_wps_ie		= madwifi_set_ap_wps_ie,
 };

src/drivers/driver_hostap.c

@@ -792,30 +792,28 @@ static int hostap_set_generic_elem(const char *ifname, void *priv,
 }
 
 
-static int hostap_set_wps_beacon_ie(const char *ifname, void *priv,
-				    const u8 *ie, size_t len)
-{
-	/* Host AP driver supports only one set of extra IEs, so we need to
-	 * use the ProbeResp IEs also for Beacon frames since they include more
-	 * information. */
-	return 0;
-}
-
-
-static int hostap_set_wps_probe_resp_ie(const char *ifname, void *priv,
-					const u8 *ie, size_t len)
+static int hostap_set_ap_wps_ie(const char *ifname, void *priv,
+				const struct wpabuf *beacon,
+				const struct wpabuf *proberesp)
 {
 	struct hostap_driver_data *drv = priv;
 
+	/*
+	 * Host AP driver supports only one set of extra IEs, so we need to
+	 * use the Probe Response IEs also for Beacon frames since they include
+	 * more information.
+	 */
+
 	os_free(drv->wps_ie);
 	drv->wps_ie = NULL;
 	drv->wps_ie_len = 0;
-	if (ie) {
-		drv->wps_ie = os_malloc(len);
+	if (proberesp) {
+		drv->wps_ie = os_malloc(wpabuf_len(proberesp));
 		if (drv->wps_ie == NULL)
 			return -1;
-		os_memcpy(drv->wps_ie, ie, len);
-		drv->wps_ie_len = len;
+		os_memcpy(drv->wps_ie, wpabuf_head(proberesp),
+			  wpabuf_len(proberesp));
+		drv->wps_ie_len = wpabuf_len(proberesp);
 	}
 
 	return hostapd_ioctl_set_generic_elem(drv);
@@ -1650,8 +1648,7 @@ const struct wpa_driver_ops wpa_driver_hostap_ops = {
 	.get_inact_sec = hostap_get_inact_sec,
 	.sta_clear_stats = hostap_sta_clear_stats,
 	.get_hw_feature_data = hostap_get_hw_feature_data,
-	.set_wps_beacon_ie = hostap_set_wps_beacon_ie,
-	.set_wps_probe_resp_ie = hostap_set_wps_probe_resp_ie,
+	.set_ap_wps_ie = hostap_set_ap_wps_ie,
 #else /* HOSTAPD */
 	.get_bssid = wpa_driver_hostap_get_bssid,
 	.get_ssid = wpa_driver_hostap_get_ssid,

src/drivers/driver_madwifi.c

@@ -818,22 +818,21 @@ madwifi_set_wps_ie(void *priv, const u8 *ie, size_t len, u32 frametype)
 }
 
 static int
-madwifi_set_wps_beacon_ie(const char *ifname, void *priv, const u8 *ie,
-			  size_t len)
+madwifi_set_ap_wps_ie(const char *ifname, void *priv,
+		      const struct wpabuf *beacon,
+		      const struct wpabuf *proberesp)
 {
-	return madwifi_set_wps_ie(priv, ie, len, IEEE80211_APPIE_FRAME_BEACON);
-}
-
-static int
-madwifi_set_wps_probe_resp_ie(const char *ifname, void *priv, const u8 *ie,
-			      size_t len)
-{
-	return madwifi_set_wps_ie(priv, ie, len,
+	if (madwifi_set_wps_ie(priv, beacon ? wpabuf_head(beacon) : NULL,
+			       beacon ? wpabuf_len(beacon) : 0,
+			       IEEE80211_APPIE_FRAME_BEACON) < 0)
+		return -1;
+	return madwifi_set_wps_ie(priv,
+				  proberesp ? wpabuf_head(proberesp) : NULL,
+				  proberesp ? wpabuf_len(proberesp) : 0,
 				  IEEE80211_APPIE_FRAME_PROBE_RESP);
 }
 #else /* CONFIG_WPS */
-#define madwifi_set_wps_beacon_ie NULL
-#define madwifi_set_wps_probe_resp_ie NULL
+#define madwifi_set_ap_wps_ie NULL
 #endif /* CONFIG_WPS */
 
 static int
@@ -1864,8 +1863,7 @@ const struct wpa_driver_ops wpa_driver_madwifi_ops = {
 	.hapd_set_countermeasures	= madwifi_set_countermeasures,
 	.sta_clear_stats        = madwifi_sta_clear_stats,
 	.commit			= madwifi_commit,
-	.set_wps_beacon_ie	= madwifi_set_wps_beacon_ie,
-	.set_wps_probe_resp_ie	= madwifi_set_wps_probe_resp_ie,
+	.set_ap_wps_ie		= madwifi_set_ap_wps_ie,
 #else /* HOSTAPD */
 	.get_bssid		= wpa_driver_madwifi_get_bssid,
 	.get_ssid		= wpa_driver_madwifi_get_ssid,

src/drivers/driver_ndis.c

@@ -3242,8 +3242,7 @@ const struct wpa_driver_ops wpa_driver_ndis_ops = {
 	NULL /* set_radius_acl_auth */,
 	NULL /* set_radius_acl_expire */,
 	NULL /* set_ht_params */,
-	NULL /* set_wps_beacon_ie */,
-	NULL /* set_wps_probe_resp_ie */,
+	NULL /* set_ap_wps_ie */,
 	NULL /* set_supp_port */,
 	NULL /* set_wds_sta */
 };

src/drivers/driver_test.c

@@ -822,64 +822,64 @@ static int test_driver_set_generic_elem(const char *ifname, void *priv,
 }
 
 
-static int test_driver_set_wps_beacon_ie(const char *ifname, void *priv,
-					 const u8 *ie, size_t len)
+static int test_driver_set_ap_wps_ie(const char *ifname, void *priv,
+				     const struct wpabuf *beacon,
+				     const struct wpabuf *proberesp)
 {
 	struct wpa_driver_test_data *drv = priv;
 	struct test_driver_bss *bss;
 
-	wpa_hexdump(MSG_DEBUG, "test_driver: Beacon WPS IE", ie, len);
 	bss = test_driver_get_bss(drv, ifname);
 	if (bss == NULL)
 		return -1;
 
+	if (beacon == NULL)
+		wpa_printf(MSG_DEBUG, "test_driver: Clear Beacon WPS IE");
+	else
+		wpa_hexdump_buf(MSG_DEBUG, "test_driver: Beacon WPS IE",
+				beacon);
+
 	os_free(bss->wps_beacon_ie);
 
-	if (ie == NULL) {
+	if (beacon == NULL) {
 		bss->wps_beacon_ie = NULL;
 		bss->wps_beacon_ie_len = 0;
-		return 0;
-	}
+	} else {
+		bss->wps_beacon_ie = os_malloc(wpabuf_len(beacon));
+		if (bss->wps_beacon_ie == NULL) {
+			bss->wps_beacon_ie_len = 0;
+			return -1;
+		}
 
-	bss->wps_beacon_ie = os_malloc(len);
-	if (bss->wps_beacon_ie == NULL) {
-		bss->wps_beacon_ie_len = 0;
-		return -1;
+		os_memcpy(bss->wps_beacon_ie, wpabuf_head(beacon),
+			  wpabuf_len(beacon));
+		bss->wps_beacon_ie_len = wpabuf_len(beacon);
 	}
 
-	memcpy(bss->wps_beacon_ie, ie, len);
-	bss->wps_beacon_ie_len = len;
-	return 0;
-}
-
-
-static int test_driver_set_wps_probe_resp_ie(const char *ifname, void *priv,
-					     const u8 *ie, size_t len)
-{
-	struct wpa_driver_test_data *drv = priv;
-	struct test_driver_bss *bss;
-
-	wpa_hexdump(MSG_DEBUG, "test_driver: ProbeResp WPS IE", ie, len);
-	bss = test_driver_get_bss(drv, ifname);
-	if (bss == NULL)
-		return -1;
+	if (proberesp == NULL)
+		wpa_printf(MSG_DEBUG, "test_driver: Clear Probe Response WPS "
+			   "IE");
+	else
+		wpa_hexdump_buf(MSG_DEBUG, "test_driver: Probe Response WPS "
+				"IE", proberesp);
 
 	os_free(bss->wps_probe_resp_ie);
 
-	if (ie == NULL) {
+	if (proberesp == NULL) {
 		bss->wps_probe_resp_ie = NULL;
 		bss->wps_probe_resp_ie_len = 0;
-		return 0;
-	}
+	} else {
+		bss->wps_probe_resp_ie = os_malloc(wpabuf_len(proberesp));
+		if (bss->wps_probe_resp_ie == NULL) {
+			bss->wps_probe_resp_ie_len = 0;
+			return -1;
+		}
 
-	bss->wps_probe_resp_ie = os_malloc(len);
-	if (bss->wps_probe_resp_ie == NULL) {
-		bss->wps_probe_resp_ie_len = 0;
-		return -1;
+		os_memcpy(bss->wps_probe_resp_ie, wpabuf_head(proberesp),
+			  wpabuf_len(proberesp));
+		bss->wps_probe_resp_ie_len = wpabuf_len(proberesp);
 	}
 
-	memcpy(bss->wps_probe_resp_ie, ie, len);
-	bss->wps_probe_resp_ie_len = len;
 	return 0;
 }
 
@@ -2429,8 +2429,7 @@ const struct wpa_driver_ops wpa_driver_test_ops = {
 	.set_sta_vlan = test_driver_set_sta_vlan,
 	.sta_add = test_driver_sta_add,
 	.send_ether = test_driver_send_ether,
-	.set_wps_beacon_ie = test_driver_set_wps_beacon_ie,
-	.set_wps_probe_resp_ie = test_driver_set_wps_probe_resp_ie,
+	.set_ap_wps_ie = test_driver_set_ap_wps_ie,
 	.get_bssid = wpa_driver_test_get_bssid,
 	.get_ssid = wpa_driver_test_get_ssid,
 	.set_key = wpa_driver_test_set_key,

src/wps/wps.h

@@ -224,16 +224,15 @@ struct wps_registrar_config {
 	 * set_ie_cb - Callback for WPS IE changes
 	 * @ctx: Higher layer context data (cb_ctx)
 	 * @beacon_ie: WPS IE for Beacon
-	 * @beacon_ie_len: WPS IE length for Beacon
 	 * @probe_resp_ie: WPS IE for Probe Response
-	 * @probe_resp_ie_len: WPS IE length for Probe Response
 	 * Returns: 0 on success, -1 on failure
 	 *
 	 * This callback is called whenever the WPS IE in Beacon or Probe
-	 * Response frames needs to be changed (AP only).
+	 * Response frames needs to be changed (AP only). Callee is responsible
+	 * for freeing the buffers.
 	 */
-	int (*set_ie_cb)(void *ctx, const u8 *beacon_ie, size_t beacon_ie_len,
-			 const u8 *probe_resp_ie, size_t probe_resp_ie_len);
+	int (*set_ie_cb)(void *ctx, struct wpabuf *beacon_ie,
+			 struct wpabuf *probe_resp_ie);
 
 	/**
 	 * pin_needed_cb - Callback for requesting a PIN

src/wps/wps_registrar.c

@@ -99,8 +99,8 @@ struct wps_registrar {
 
 	int (*new_psk_cb)(void *ctx, const u8 *mac_addr, const u8 *psk,
 			  size_t psk_len);
-	int (*set_ie_cb)(void *ctx, const u8 *beacon_ie, size_t beacon_ie_len,
-			 const u8 *probe_resp_ie, size_t probe_resp_ie_len);
+	int (*set_ie_cb)(void *ctx, struct wpabuf *beacon_ie,
+			 struct wpabuf *probe_resp_ie);
 	void (*pin_needed_cb)(void *ctx, const u8 *uuid_e,
 			      const struct wps_device_data *dev);
 	void (*reg_success_cb)(void *ctx, const u8 *mac_addr,
@@ -816,17 +816,13 @@ static void wps_cb_reg_success(struct wps_registrar *reg, const u8 *mac_addr,
 }
 
 
-static int wps_cb_set_ie(struct wps_registrar *reg,
-			 const struct wpabuf *beacon_ie,
-			 const struct wpabuf *probe_resp_ie)
+static int wps_cb_set_ie(struct wps_registrar *reg, struct wpabuf *beacon_ie,
+			 struct wpabuf *probe_resp_ie)
 {
 	if (reg->set_ie_cb == NULL)
 		return 0;
 
-	return reg->set_ie_cb(reg->cb_ctx, wpabuf_head(beacon_ie),
-			      wpabuf_len(beacon_ie),
-			      wpabuf_head(probe_resp_ie),
-			      wpabuf_len(probe_resp_ie));
+	return reg->set_ie_cb(reg->cb_ctx, beacon_ie, probe_resp_ie);
 }
 
 
@@ -884,7 +880,6 @@ static int wps_set_ie(struct wps_registrar *reg)
 {
 	struct wpabuf *beacon;
 	struct wpabuf *probe;
-	int ret;
 
 	wpa_printf(MSG_DEBUG, "WPS: Build Beacon and Probe Response IEs");
 
@@ -950,11 +945,7 @@ static int wps_set_ie(struct wps_registrar *reg)
 		wpabuf_put_data(probe, ms_wps, sizeof(ms_wps));
 	}
 
-	ret = wps_cb_set_ie(reg, beacon, probe);
-	wpabuf_free(beacon);
-	wpabuf_free(probe);
-
-	return ret;
+	return wps_cb_set_ie(reg, beacon, probe);
 }