1 files changed, 6 insertions, 4 deletions

diff --git a/drivers/mtd/ubi/upd.c b/drivers/mtd/ubi/upd.c
index ec2c2dc1c1ca..2a1b6e037e1a 100644
--- a/drivers/mtd/ubi/upd.c
+++ b/drivers/mtd/ubi/upd.c
@@ -133,6 +133,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;
@@ -152,14 +156,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;