Merge git://git.denx.de/u-boot-net

[u-boot] / drivers / mtd / ubi / upd.c

diff --git a/drivers/mtd/ubi/upd.c b/drivers/mtd/ubi/upd.c
index 220c120515ac508f58eb9e1459cb8aed4d3eeb0e..e0caf8eed23dee7daf11b8499d1d29e671c58e06 100644 (file)
--- a/drivers/mtd/ubi/upd.c
+++ b/drivers/mtd/ubi/upd.c
@@ -26,7 +26,6 @@
  * transaction with a roll-back capability.
  */
 
-#define __UBOOT__
 #ifndef __UBOOT__
 #include <linux/uaccess.h>
 #else
@@ -128,6 +127,10 @@ int ubi_start_update(struct ubi_device *ubi, struct ubi_volume *vol,
        ubi_assert(!vol->updating && !vol->changing_leb);
        vol->updating = 1;
 
+       vol->upd_buf = vmalloc(ubi->leb_size);
+       if (!vol->upd_buf)
+               return -ENOMEM;
+
        err = set_update_marker(ubi, vol);
        if (err)
                return err;
@@ -147,14 +150,12 @@ int ubi_start_update(struct ubi_device *ubi, struct ubi_volume *vol,
                err = clear_update_marker(ubi, vol, 0);
                if (err)
                        return err;
+
+               vfree(vol->upd_buf);
                vol->updating = 0;
                return 0;
        }
 
-       vol->upd_buf = vmalloc(ubi->leb_size);
-       if (!vol->upd_buf)
-               return -ENOMEM;
-
        vol->upd_ebs = div_u64(bytes + vol->usable_leb_size - 1,
                               vol->usable_leb_size);
        vol->upd_bytes = bytes;