UBI: Fastmap: Fix memory leak while attaching

Currently we leak a few ubi_ainf_pebs while attaching.

Signed-off-by: Richard Weinberger <richard@nod.at>
diff --git a/drivers/mtd/ubi/attach.c b/drivers/mtd/ubi/attach.c
index b5e1548..ef619c3 100644
--- a/drivers/mtd/ubi/attach.c
+++ b/drivers/mtd/ubi/attach.c
@@ -1301,67 +1301,6 @@
 	return err;
 }
 
-#ifdef CONFIG_MTD_UBI_FASTMAP
-
-/**
- * scan_fastmap - try to find a fastmap and attach from it.
- * @ubi: UBI device description object
- * @ai: attach info object
- *
- * Returns 0 on success, negative return values indicate an internal
- * error.
- * UBI_NO_FASTMAP denotes that no fastmap was found.
- * UBI_BAD_FASTMAP denotes that the found fastmap was invalid.
- */
-static int scan_fast(struct ubi_device *ubi, struct ubi_attach_info *ai)
-{
-	int err, pnum, fm_anchor = -1;
-	unsigned long long max_sqnum = 0;
-
-	err = -ENOMEM;
-
-	ech = kzalloc(ubi->ec_hdr_alsize, GFP_KERNEL);
-	if (!ech)
-		goto out;
-
-	vidh = ubi_zalloc_vid_hdr(ubi, GFP_KERNEL);
-	if (!vidh)
-		goto out_ech;
-
-	for (pnum = 0; pnum < UBI_FM_MAX_START; pnum++) {
-		int vol_id = -1;
-		unsigned long long sqnum = -1;
-		cond_resched();
-
-		dbg_gen("process PEB %d", pnum);
-		err = scan_peb(ubi, ai, pnum, &vol_id, &sqnum);
-		if (err < 0)
-			goto out_vidh;
-
-		if (vol_id == UBI_FM_SB_VOLUME_ID && sqnum > max_sqnum) {
-			max_sqnum = sqnum;
-			fm_anchor = pnum;
-		}
-	}
-
-	ubi_free_vid_hdr(ubi, vidh);
-	kfree(ech);
-
-	if (fm_anchor < 0)
-		return UBI_NO_FASTMAP;
-
-	return ubi_scan_fastmap(ubi, ai, fm_anchor);
-
-out_vidh:
-	ubi_free_vid_hdr(ubi, vidh);
-out_ech:
-	kfree(ech);
-out:
-	return err;
-}
-
-#endif
-
 static struct ubi_attach_info *alloc_ai(const char *slab_name)
 {
 	struct ubi_attach_info *ai;
@@ -1386,6 +1325,72 @@
 	return ai;
 }
 
+#ifdef CONFIG_MTD_UBI_FASTMAP
+
+/**
+ * scan_fastmap - try to find a fastmap and attach from it.
+ * @ubi: UBI device description object
+ * @ai: attach info object
+ *
+ * Returns 0 on success, negative return values indicate an internal
+ * error.
+ * UBI_NO_FASTMAP denotes that no fastmap was found.
+ * UBI_BAD_FASTMAP denotes that the found fastmap was invalid.
+ */
+static int scan_fast(struct ubi_device *ubi, struct ubi_attach_info **ai)
+{
+	int err, pnum, fm_anchor = -1;
+	unsigned long long max_sqnum = 0;
+
+	err = -ENOMEM;
+
+	ech = kzalloc(ubi->ec_hdr_alsize, GFP_KERNEL);
+	if (!ech)
+		goto out;
+
+	vidh = ubi_zalloc_vid_hdr(ubi, GFP_KERNEL);
+	if (!vidh)
+		goto out_ech;
+
+	for (pnum = 0; pnum < UBI_FM_MAX_START; pnum++) {
+		int vol_id = -1;
+		unsigned long long sqnum = -1;
+		cond_resched();
+
+		dbg_gen("process PEB %d", pnum);
+		err = scan_peb(ubi, *ai, pnum, &vol_id, &sqnum);
+		if (err < 0)
+			goto out_vidh;
+
+		if (vol_id == UBI_FM_SB_VOLUME_ID && sqnum > max_sqnum) {
+			max_sqnum = sqnum;
+			fm_anchor = pnum;
+		}
+	}
+
+	ubi_free_vid_hdr(ubi, vidh);
+	kfree(ech);
+
+	if (fm_anchor < 0)
+		return UBI_NO_FASTMAP;
+
+	destroy_ai(*ai);
+	*ai = alloc_ai("ubi_aeb_slab_cache");
+	if (!*ai)
+		return -ENOMEM;
+
+	return ubi_scan_fastmap(ubi, *ai, fm_anchor);
+
+out_vidh:
+	ubi_free_vid_hdr(ubi, vidh);
+out_ech:
+	kfree(ech);
+out:
+	return err;
+}
+
+#endif
+
 /**
  * ubi_attach - attach an MTD device.
  * @ubi: UBI device descriptor
@@ -1413,7 +1418,7 @@
 	if (force_scan)
 		err = scan_all(ubi, ai, 0);
 	else {
-		err = scan_fast(ubi, ai);
+		err = scan_fast(ubi, &ai);
 		if (err > 0) {
 			if (err != UBI_NO_FASTMAP) {
 				destroy_ai(ai);