}
 EXPORT_SYMBOL(cfg80211_is_element_inherited);
 
-static size_t cfg80211_gen_new_ie(const u8 *ie, size_t ielen,
-                                 const u8 *subelement, size_t subie_len,
-                                 u8 *new_ie, gfp_t gfp)
+static size_t cfg80211_copy_elem_with_frags(const struct element *elem,
+                                           const u8 *ie, size_t ie_len,
+                                           u8 **pos, u8 *buf, size_t buf_len)
 {
-       u8 *pos, *tmp;
-       const u8 *tmp_old, *tmp_new;
-       const struct element *non_inherit_elem;
-       u8 *sub_copy;
+       if (WARN_ON((u8 *)elem < ie || elem->data > ie + ie_len ||
+                   elem->data + elem->datalen > ie + ie_len))
+               return 0;
 
-       /* copy subelement as we need to change its content to
-        * mark an ie after it is processed.
-        */
-       sub_copy = kmemdup(subelement, subie_len, gfp);
-       if (!sub_copy)
+       if (elem->datalen + 2 > buf + buf_len - *pos)
                return 0;
 
-       pos = &new_ie[0];
+       memcpy(*pos, elem, elem->datalen + 2);
+       *pos += elem->datalen + 2;
 
-       /* set new ssid */
-       tmp_new = cfg80211_find_ie(WLAN_EID_SSID, sub_copy, subie_len);
-       if (tmp_new) {
-               memcpy(pos, tmp_new, tmp_new[1] + 2);
-               pos += (tmp_new[1] + 2);
+       /* Finish if it is not fragmented  */
+       if (elem->datalen != 255)
+               return *pos - buf;
+
+       ie_len = ie + ie_len - elem->data - elem->datalen;
+       ie = (const u8 *)elem->data + elem->datalen;
+
+       for_each_element(elem, ie, ie_len) {
+               if (elem->id != WLAN_EID_FRAGMENT)
+                       break;
+
+               if (elem->datalen + 2 > buf + buf_len - *pos)
+                       return 0;
+
+               memcpy(*pos, elem, elem->datalen + 2);
+               *pos += elem->datalen + 2;
+
+               if (elem->datalen != 255)
+                       break;
        }
 
-       /* get non inheritance list if exists */
-       non_inherit_elem =
-               cfg80211_find_ext_elem(WLAN_EID_EXT_NON_INHERITANCE,
-                                      sub_copy, subie_len);
+       return *pos - buf;
+}
 
-       /* go through IEs in ie (skip SSID) and subelement,
-        * merge them into new_ie
+static size_t cfg80211_gen_new_ie(const u8 *ie, size_t ielen,
+                                 const u8 *subie, size_t subie_len,
+                                 u8 *new_ie, size_t new_ie_len)
+{
+       const struct element *non_inherit_elem, *parent, *sub;
+       u8 *pos = new_ie;
+       u8 id, ext_id;
+       unsigned int match_len;
+
+       non_inherit_elem = cfg80211_find_ext_elem(WLAN_EID_EXT_NON_INHERITANCE,
+                                                 subie, subie_len);
+
+       /* We copy the elements one by one from the parent to the generated
+        * elements.
+        * If they are not inherited (included in subie or in the non
+        * inheritance element), then we copy all occurrences the first time
+        * we see this element type.
         */
-       tmp_old = cfg80211_find_ie(WLAN_EID_SSID, ie, ielen);
-       tmp_old = (tmp_old) ? tmp_old + tmp_old[1] + 2 : ie;
-
-       while (tmp_old + 2 - ie <= ielen &&
-              tmp_old + tmp_old[1] + 2 - ie <= ielen) {
-               if (tmp_old[0] == 0) {
-                       tmp_old++;
+       for_each_element(parent, ie, ielen) {
+               if (parent->id == WLAN_EID_FRAGMENT)
                        continue;
+
+               if (parent->id == WLAN_EID_EXTENSION) {
+                       if (parent->datalen < 1)
+                               continue;
+
+                       id = WLAN_EID_EXTENSION;
+                       ext_id = parent->data[0];
+                       match_len = 1;
+               } else {
+                       id = parent->id;
+                       match_len = 0;
                }
 
-               if (tmp_old[0] == WLAN_EID_EXTENSION)
-                       tmp = (u8 *)cfg80211_find_ext_ie(tmp_old[2], sub_copy,
-                                                        subie_len);
-               else
-                       tmp = (u8 *)cfg80211_find_ie(tmp_old[0], sub_copy,
-                                                    subie_len);
+               /* Find first occurrence in subie */
+               sub = cfg80211_find_elem_match(id, subie, subie_len,
+                                              &ext_id, match_len, 0);
 
-               if (!tmp) {
-                       const struct element *old_elem = (void *)tmp_old;
+               /* Copy from parent if not in subie and inherited */
+               if (!sub &&
+                   cfg80211_is_element_inherited(parent, non_inherit_elem)) {
+                       if (!cfg80211_copy_elem_with_frags(parent,
+                                                          ie, ielen,
+                                                          &pos, new_ie,
+                                                          new_ie_len))
+                               return 0;
 
-                       /* ie in old ie but not in subelement */
-                       if (cfg80211_is_element_inherited(old_elem,
-                                                         non_inherit_elem)) {
-                               memcpy(pos, tmp_old, tmp_old[1] + 2);
-                               pos += tmp_old[1] + 2;
-                       }
-               } else {
-                       /* ie in transmitting ie also in subelement,
-                        * copy from subelement and flag the ie in subelement
-                        * as copied (by setting eid field to WLAN_EID_SSID,
-                        * which is skipped anyway).
-                        * For vendor ie, compare OUI + type + subType to
-                        * determine if they are the same ie.
-                        */
-                       if (tmp_old[0] == WLAN_EID_VENDOR_SPECIFIC) {
-                               if (tmp_old[1] >= 5 && tmp[1] >= 5 &&
-                                   !memcmp(tmp_old + 2, tmp + 2, 5)) {
-                                       /* same vendor ie, copy from
-                                        * subelement
-                                        */
-                                       memcpy(pos, tmp, tmp[1] + 2);
-                                       pos += tmp[1] + 2;
-                                       tmp[0] = WLAN_EID_SSID;
-                               } else {
-                                       memcpy(pos, tmp_old, tmp_old[1] + 2);
-                                       pos += tmp_old[1] + 2;
-                               }
-                       } else {
-                               /* copy ie from subelement into new ie */
-                               memcpy(pos, tmp, tmp[1] + 2);
-                               pos += tmp[1] + 2;
-                               tmp[0] = WLAN_EID_SSID;
-                       }
+                       continue;
                }
 
-               if (tmp_old + tmp_old[1] + 2 - ie == ielen)
-                       break;
+               /* Already copied if an earlier element had the same type */
+               if (cfg80211_find_elem_match(id, ie, (u8 *)parent - ie,
+                                            &ext_id, match_len, 0))
+                       continue;
 
-               tmp_old += tmp_old[1] + 2;
+               /* Not inheriting, copy all similar elements from subie */
+               while (sub) {
+                       if (!cfg80211_copy_elem_with_frags(sub,
+                                                          subie, subie_len,
+                                                          &pos, new_ie,
+                                                          new_ie_len))
+                               return 0;
+
+                       sub = cfg80211_find_elem_match(id,
+                                                      sub->data + sub->datalen,
+                                                      subie_len + subie -
+                                                      (sub->data +
+                                                       sub->datalen),
+                                                      &ext_id, match_len, 0);
+               }
        }
 
-       /* go through subelement again to check if there is any ie not
-        * copied to new ie, skip ssid, capability, bssid-index ie
+       /* The above misses elements that are included in subie but not in the
+        * parent, so do a pass over subie and append those.
+        * Skip the non-tx BSSID caps and non-inheritance element.
         */
-       tmp_new = sub_copy;
-       while (tmp_new + 2 - sub_copy <= subie_len &&
-              tmp_new + tmp_new[1] + 2 - sub_copy <= subie_len) {
-               if (!(tmp_new[0] == WLAN_EID_NON_TX_BSSID_CAP ||
-                     tmp_new[0] == WLAN_EID_SSID)) {
-                       memcpy(pos, tmp_new, tmp_new[1] + 2);
-                       pos += tmp_new[1] + 2;
+       for_each_element(sub, subie, subie_len) {
+               if (sub->id == WLAN_EID_NON_TX_BSSID_CAP)
+                       continue;
+
+               if (sub->id == WLAN_EID_FRAGMENT)
+                       continue;
+
+               if (sub->id == WLAN_EID_EXTENSION) {
+                       if (sub->datalen < 1)
+                               continue;
+
+                       id = WLAN_EID_EXTENSION;
+                       ext_id = sub->data[0];
+                       match_len = 1;
+
+                       if (ext_id == WLAN_EID_EXT_NON_INHERITANCE)
+                               continue;
+               } else {
+                       id = sub->id;
+                       match_len = 0;
                }
-               if (tmp_new + tmp_new[1] + 2 - sub_copy == subie_len)
-                       break;
-               tmp_new += tmp_new[1] + 2;
+
+               /* Processed if one was included in the parent */
+               if (cfg80211_find_elem_match(id, ie, ielen,
+                                            &ext_id, match_len, 0))
+                       continue;
+
+               if (!cfg80211_copy_elem_with_frags(sub, subie, subie_len,
+                                                  &pos, new_ie, new_ie_len))
+                       return 0;
        }
 
-       kfree(sub_copy);
        return pos - new_ie;
 }
 
                        new_ie_len = cfg80211_gen_new_ie(ie, ielen,
                                                         profile,
                                                         profile_len, new_ie,
-                                                        gfp);
+                                                        IEEE80211_MAX_DATA_LEN);
                        if (!new_ie_len)
                                continue;