Clean up AP mode extra IE construction

Make it easier to add more IEs into the buffers.

diff --git a/src/ap/ap_drv_ops.c b/src/ap/ap_drv_ops.c
index 729cc946aea901139783d5f722ab4fc808d20e10..77d6834f47247ec0adc93d51fb1f57b732da0a53 100644 (file)
--- a/src/ap/ap_drv_ops.c
+++ b/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);
 }