1 files changed, 58 insertions, 13 deletions

diff --git a/tools/firmware-utils/src/tplink-safeloader.c b/tools/firmware-utils/src/tplink-safeloader.c
index c7942d1ef6..fe57905fac 100644
--- a/tools/firmware-utils/src/tplink-safeloader.c
+++ b/tools/firmware-utils/src/tplink-safeloader.c
@@ -65,7 +65,7 @@ struct image_partition_entry {
 
 /** A flash partition table entry */
 struct flash_partition_entry {
-	const char *name;
+	char *name;
 	uint32_t base;
 	uint32_t size;
 };
@@ -77,7 +77,7 @@ struct device_info {
 	const char *support_list;
 	char support_trail;
 	const char *soft_ver;
-	const struct flash_partition_entry partitions[MAX_PARTITIONS+1];
+	struct flash_partition_entry partitions[MAX_PARTITIONS+1];
 	const char *first_sysupgrade_partition;
 	const char *last_sysupgrade_partition;
 };
@@ -1223,7 +1223,7 @@ static struct image_partition_entry make_soft_version_from_string(const char *so
 }
 
 /** Generates the support-list partition */
-static struct image_partition_entry make_support_list(const struct device_info *info) {
+static struct image_partition_entry make_support_list(struct device_info *info) {
 	size_t len = strlen(info->support_list);
 	struct image_partition_entry entry = alloc_image_partition("support-list", len + 9);
 
@@ -1236,7 +1236,7 @@ static struct image_partition_entry make_support_list(const struct device_info *
 }
 
 /** Creates a new image partition with an arbitrary name from a file */
-static struct image_partition_entry read_file(const char *part_name, const char *filename, bool add_jffs2_eof) {
+static struct image_partition_entry read_file(const char *part_name, const char *filename, bool add_jffs2_eof, struct flash_partition_entry *file_system_partition) {
 	struct stat statbuf;
 
 	if (stat(filename, &statbuf) < 0)
@@ -1245,7 +1245,10 @@ static struct image_partition_entry read_file(const char *part_name, const char
 	size_t len = statbuf.st_size;
 
 	if (add_jffs2_eof)
-		len = ALIGN(len, 0x10000) + sizeof(jffs2_eof_mark);
+		if (file_system_partition)
+			len = ALIGN(len + file_system_partition->base, 0x10000) + sizeof(jffs2_eof_mark) - file_system_partition->base;
+		else
+			len = ALIGN(len, 0x10000) + sizeof(jffs2_eof_mark);
 
 	struct image_partition_entry entry = alloc_image_partition(part_name, len);
 
@@ -1361,7 +1364,7 @@ static void put_md5(uint8_t *md5, uint8_t *buffer, unsigned int len) {
      1014-1813    Image partition table (2048 bytes, padded with 0xff)
      1814-xxxx    Firmware partitions
 */
-static void * generate_factory_image(const struct device_info *info, const struct image_partition_entry *parts, size_t *len) {
+static void * generate_factory_image(struct device_info *info, const struct image_partition_entry *parts, size_t *len) {
 	*len = 0x1814;
 
 	size_t i;
@@ -1394,7 +1397,7 @@ static void * generate_factory_image(const struct device_info *info, const struc
    should be generalized when TP-LINK starts building its safeloader into hardware with
    different flash layouts.
 */
-static void * generate_sysupgrade_image(const struct device_info *info, const struct image_partition_entry *image_parts, size_t *len) {
+static void * generate_sysupgrade_image(struct device_info *info, const struct image_partition_entry *image_parts, size_t *len) {
 	size_t i, j;
 	size_t flash_first_partition_index = 0;
 	size_t flash_last_partition_index = 0;
@@ -1457,10 +1460,53 @@ static void build_image(const char *output,
 		uint32_t rev,
 		bool add_jffs2_eof,
 		bool sysupgrade,
-		const struct device_info *info) {
+		struct device_info *info) {
+
+	size_t i;
 
 	struct image_partition_entry parts[7] = {};
 
+	struct flash_partition_entry *firmware_partition = NULL;
+	struct flash_partition_entry *os_image_partition = NULL;
+	struct flash_partition_entry *file_system_partition = NULL;
+	size_t firmware_partition_index = 0;
+
+	for (i = 0; info->partitions[i].name; i++) {
+		if (!strcmp(info->partitions[i].name, "firmware"))
+		{
+			firmware_partition = &info->partitions[i];
+			firmware_partition_index = i;
+		}
+	}
+
+	if (firmware_partition)
+	{
+		os_image_partition = &info->partitions[firmware_partition_index];
+		file_system_partition = &info->partitions[firmware_partition_index + 1];
+
+		struct stat kernel;
+		if (stat(kernel_image, &kernel) < 0)
+			error(1, errno, "unable to stat file `%s'", kernel_image);
+
+		if (kernel.st_size > firmware_partition->size)
+			error(1, 0, "kernel overflowed firmware partition\n");
+
+		for (i = MAX_PARTITIONS-1; i >= firmware_partition_index + 1; i--)
+			info->partitions[i+1] = info->partitions[i];
+
+		file_system_partition->name = "file-system";
+		file_system_partition->base = firmware_partition->base + kernel.st_size;
+
+		/* Align partition start to erase blocks for factory images only */
+		if (!sysupgrade)
+			file_system_partition->base = ALIGN(firmware_partition->base + kernel.st_size, 0x10000);
+
+		file_system_partition->size = firmware_partition->size - file_system_partition->base;
+
+		os_image_partition->name = "os-image";
+		os_image_partition->size = kernel.st_size;
+	}
+
 	parts[0] = make_partition_table(info->partitions);
 	if (info->soft_ver)
 		parts[1] = make_soft_version_from_string(info->soft_ver);
@@ -1468,8 +1514,8 @@ static void build_image(const char *output,
 		parts[1] = make_soft_version(rev);
 
 	parts[2] = make_support_list(info);
-	parts[3] = read_file("os-image", kernel_image, false);
-	parts[4] = read_file("file-system", rootfs_image, add_jffs2_eof);
+	parts[3] = read_file("os-image", kernel_image, false, NULL);
+	parts[4] = read_file("file-system", rootfs_image, add_jffs2_eof, file_system_partition);
 
 	/* Some devices need the extra-para partition to accept the firmware */
 	if (strcasecmp(info->id, "ARCHER-C25-V1") == 0 ||
@@ -1500,7 +1546,6 @@ static void build_image(const char *output,
 
 	free(image);
 
-	size_t i;
 	for (i = 0; parts[i].name; i++)
 		free_image_partition(parts[i]);
 }
@@ -1530,7 +1575,7 @@ static void usage(const char *argv0) {
 };
 
 
-static const struct device_info *find_board(const char *id)
+static struct device_info *find_board(const char *id)
 {
 	struct device_info *board = NULL;
 
@@ -1860,7 +1905,7 @@ int main(int argc, char *argv[]) {
 	const char *extract_image = NULL, *output_directory = NULL, *convert_image = NULL;
 	bool add_jffs2_eof = false, sysupgrade = false;
 	unsigned rev = 0;
-	const struct device_info *info;
+	struct device_info *info;
 	set_source_date_epoch();
 
 	while (true) {