Clean up AP mode extra IE construction

Make it easier to add more IEs into the buffers.

src/ap/ap_drv_ops.c

@@ -42,118 +42,103 @@ u32 hostapd_sta_flags_to_drv(u32 flags)
 
 
 int hostapd_build_ap_extra_ies(struct hostapd_data *hapd,
-			       struct wpabuf **beacon,
-			       struct wpabuf **proberesp,
-			       struct wpabuf **assocresp)
+			       struct wpabuf **beacon_ret,
+			       struct wpabuf **proberesp_ret,
+			       struct wpabuf **assocresp_ret)
 {
-	*beacon = hapd->wps_beacon_ie;
-	*proberesp = hapd->wps_probe_resp_ie;
-	*assocresp = NULL;
+	struct wpabuf *beacon = NULL, *proberesp = NULL, *assocresp = NULL;
+
+	*beacon_ret = *proberesp_ret = *assocresp_ret = NULL;
+
+	if (hapd->wps_beacon_ie) {
+		if (wpabuf_resize(&beacon, wpabuf_len(hapd->wps_beacon_ie)) <
+		    0)
+			goto fail;
+		wpabuf_put_buf(beacon, hapd->wps_beacon_ie);
+	}
+
+	if (hapd->wps_probe_resp_ie) {
+		if (wpabuf_resize(&proberesp,
+				  wpabuf_len(hapd->wps_probe_resp_ie)) < 0)
+			goto fail;
+		wpabuf_put_buf(proberesp, hapd->wps_probe_resp_ie);
+	}
 
 #ifdef CONFIG_P2P
-	if (hapd->wps_beacon_ie == NULL && hapd->p2p_beacon_ie == NULL)
-		*beacon = NULL;
-	else {
-		*beacon = wpabuf_alloc((hapd->wps_beacon_ie ?
-					wpabuf_len(hapd->wps_beacon_ie) : 0) +
-				       (hapd->p2p_beacon_ie ?
-					wpabuf_len(hapd->p2p_beacon_ie) : 0));
-		if (*beacon == NULL)
-			return -1;
-		if (hapd->wps_beacon_ie)
-			wpabuf_put_buf(*beacon, hapd->wps_beacon_ie);
-		if (hapd->p2p_beacon_ie)
-			wpabuf_put_buf(*beacon, hapd->p2p_beacon_ie);
+	if (hapd->p2p_beacon_ie) {
+		if (wpabuf_resize(&beacon, wpabuf_len(hapd->p2p_beacon_ie)) <
+		    0)
+			goto fail;
+		wpabuf_put_buf(beacon, hapd->p2p_beacon_ie);
 	}
 
-	if (hapd->wps_probe_resp_ie == NULL && hapd->p2p_probe_resp_ie == NULL)
-		*proberesp = NULL;
-	else {
-		*proberesp = wpabuf_alloc(
-			(hapd->wps_probe_resp_ie ?
-			 wpabuf_len(hapd->wps_probe_resp_ie) : 0) +
-			(hapd->p2p_probe_resp_ie ?
-			 wpabuf_len(hapd->p2p_probe_resp_ie) : 0));
-		if (*proberesp == NULL) {
-			wpabuf_free(*beacon);
-			return -1;
-		}
-		if (hapd->wps_probe_resp_ie)
-			wpabuf_put_buf(*proberesp, hapd->wps_probe_resp_ie);
-		if (hapd->p2p_probe_resp_ie)
-			wpabuf_put_buf(*proberesp, hapd->p2p_probe_resp_ie);
+	if (hapd->p2p_probe_resp_ie) {
+		if (wpabuf_resize(&beacon, wpabuf_len(hapd->p2p_probe_resp_ie))
+		    < 0)
+			goto fail;
+		wpabuf_put_buf(beacon, hapd->p2p_probe_resp_ie);
 	}
 #endif /* CONFIG_P2P */
 
 #ifdef CONFIG_P2P_MANAGER
 	if (hapd->conf->p2p & P2P_MANAGE) {
-		struct wpabuf *a;
-
-		a = wpabuf_alloc(100 + (*beacon ? wpabuf_len(*beacon) : 0));
-		if (a) {
+		if (wpabuf_resize(&beacon, 100) == 0) {
 			u8 *start, *p;
-			if (*beacon)
-				wpabuf_put_buf(a, *beacon);
-			if (*beacon != hapd->wps_beacon_ie)
-				wpabuf_free(*beacon);
-			start = wpabuf_put(a, 0);
+			start = wpabuf_put(beacon, 0);
 			p = hostapd_eid_p2p_manage(hapd, start);
-			wpabuf_put(a, p - start);
-			*beacon = a;
+			wpabuf_put(beacon, p - start);
 		}
 
-		a = wpabuf_alloc(100 + (*proberesp ? wpabuf_len(*proberesp) :
-					0));
-		if (a) {
+		if (wpabuf_resize(&proberesp, 100) == 0) {
 			u8 *start, *p;
-			if (*proberesp)
-				wpabuf_put_buf(a, *proberesp);
-			if (*proberesp != hapd->wps_probe_resp_ie)
-				wpabuf_free(*proberesp);
-			start = wpabuf_put(a, 0);
+			start = wpabuf_put(proberesp, 0);
 			p = hostapd_eid_p2p_manage(hapd, start);
-			wpabuf_put(a, p - start);
-			*proberesp = a;
+			wpabuf_put(proberesp, p - start);
 		}
 	}
 #endif /* CONFIG_P2P_MANAGER */
 
 #ifdef CONFIG_WPS2
-	if (hapd->conf->wps_state)
-		*assocresp = wps_build_assoc_resp_ie();
+	if (hapd->conf->wps_state) {
+		struct wpabuf *a = wps_build_assoc_resp_ie();
+		if (a && wpabuf_resize(&assocresp, wpabuf_len(a)) == 0)
+			wpabuf_put_buf(assocresp, a);
+		wpabuf_free(a);
+	}
 #endif /* CONFIG_WPS2 */
 
 #ifdef CONFIG_P2P_MANAGER
 	if (hapd->conf->p2p & P2P_MANAGE) {
-		struct wpabuf *a;
-		a = wpabuf_alloc(100 + (*assocresp ? wpabuf_len(*assocresp) :
-					0));
-		if (a) {
+		if (wpabuf_resize(&assocresp, 100) == 0) {
 			u8 *start, *p;
-			start = wpabuf_put(a, 0);
+			start = wpabuf_put(assocresp, 0);
 			p = hostapd_eid_p2p_manage(hapd, start);
-			wpabuf_put(a, p - start);
-			if (*assocresp) {
-				wpabuf_put_buf(a, *assocresp);
-				wpabuf_free(*assocresp);
-			}
-			*assocresp = a;
+			wpabuf_put(assocresp, p - start);
 		}
 	}
 #endif /* CONFIG_P2P_MANAGER */
 
+	*beacon_ret = beacon;
+	*proberesp_ret = proberesp;
+	*assocresp_ret = assocresp;
+
 	return 0;
+
+fail:
+	wpabuf_free(beacon);
+	wpabuf_free(proberesp);
+	wpabuf_free(assocresp);
+	return -1;
 }
 
 
-void hostapd_free_ap_extra_ies(struct hostapd_data *hapd, struct wpabuf *beacon,
+void hostapd_free_ap_extra_ies(struct hostapd_data *hapd,
+			       struct wpabuf *beacon,
 			       struct wpabuf *proberesp,
 			       struct wpabuf *assocresp)
 {
-	if (beacon != hapd->wps_beacon_ie)
-		wpabuf_free(beacon);
-	if (proberesp != hapd->wps_probe_resp_ie)
-		wpabuf_free(proberesp);
+	wpabuf_free(beacon);
+	wpabuf_free(proberesp);
 	wpabuf_free(assocresp);
 }