app: aboot: Add UBI image awareness to fastboot code

Current implementation of flasher in LK just writes the provided image as raw data.
This is not good enough when writing a UBI image since some UBI META data should
be preserved.

Change-Id: I2c85b02d167db9a1a22a6eb6fe8181f7db1c64b2
diff --git a/app/aboot/aboot.c b/app/aboot/aboot.c
index f3b63b7..1ebf2a6 100644
--- a/app/aboot/aboot.c
+++ b/app/aboot/aboot.c
@@ -40,6 +40,7 @@
 #include <arch/ops.h>
 
 #include <dev/flash.h>
+#include <dev/flash-ubi.h>
 #include <lib/ptable.h>
 #include <dev/keys.h>
 #include <dev/fbcon.h>
@@ -118,9 +119,7 @@
 #define DEFAULT_ERASE_SIZE  4096
 #define MAX_PANEL_BUF_SIZE 128
 
-#define UBI_MAGIC      "UBI#"
 #define DISPLAY_DEFAULT_PREFIX "mdss_mdp"
-#define UBI_MAGIC_SIZE 0x04
 #define BOOT_DEV_MAX_LEN  64
 
 #define IS_ARM64(ptr) (ptr->magic_64 == KERNEL64_HDR_MAGIC) ? true : false
@@ -2316,19 +2315,21 @@
 		|| !strcmp(ptn->name, "persist")
 		|| !strcmp(ptn->name, "recoveryfs")
 		|| !strcmp(ptn->name, "modem"))
-	{
-		if (memcmp((void *)data, UBI_MAGIC, UBI_MAGIC_SIZE))
-			extra = 1;
-		else
-			extra = 0;
-	}
+		extra = 1;
 	else
 		sz = ROUND_TO_PAGE(sz, page_mask);
 
 	dprintf(INFO, "writing %d bytes to '%s'\n", sz, ptn->name);
-	if (flash_write(ptn, extra, data, sz)) {
-		fastboot_fail("flash write failure");
-		return;
+	if (!memcmp((void *)data, UBI_MAGIC, UBI_MAGIC_SIZE)) {
+		if (flash_ubi_img(ptn, data, sz)) {
+			fastboot_fail("flash write failure");
+			return;
+		}
+	} else {
+		if (flash_write(ptn, extra, data, sz)) {
+			fastboot_fail("flash write failure");
+			return;
+		}
 	}
 	dprintf(INFO, "partition '%s' updated\n", ptn->name);
 	fastboot_okay("");
diff --git a/include/dev/flash-ubi.h b/include/dev/flash-ubi.h
new file mode 100644
index 0000000..09a7bc0
--- /dev/null
+++ b/include/dev/flash-ubi.h
@@ -0,0 +1,48 @@
+/* Copyright (c) 2014-2015, The Linux Foundation. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ *       copyright notice, this list of conditions and the following
+ *       disclaimer in the documentation and/or other materials provided
+ *       with the distribution.
+ *     * Neither the name of The Linux Foundation nor the names of its
+ *       contributors may be used to endorse or promote products derived
+ *       from this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESS OR IMPLIED
+ * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT
+ * ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
+ * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
+ * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+ * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
+ * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
+ * IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#ifndef _FLASH_UBI_H_
+#define _FLASH_UBI_H_
+
+#include <sys/types.h>
+#include <lib/ptable.h>
+#include <stdint.h>
+
+/* Erase counter header magic number (ASCII "UBI#") */
+#define UBI_EC_HDR_MAGIC  0x55424923
+
+#define UBI_MAGIC      "UBI#"
+#define UBI_MAGIC_SIZE 0x04
+
+inline int flash_ubi_img(struct ptentry *ptn, void *data, unsigned size)
+{
+	dprintf(CRITICAL, "ERROR: flash_ubi_img not yet implemented\n");
+	dprintf(CRITICAL, "falling back to flash_write\n");
+	return flash_write(ptn, 0, data, size);
+}
+#endif
diff --git a/target/init.c b/target/init.c
index c866187..9c850e5 100644
--- a/target/init.c
+++ b/target/init.c
@@ -54,6 +54,11 @@
     return (120 * 1024 * 1024);
 }
 
+__WEAK int flash_ubi_img(void)
+{
+    return 0;
+}
+
 __WEAK int target_is_emmc_boot(void)
 {
 #if _EMMC_BOOT