--- /dev/null
+/* SPDX-License-Identifier: GPL-2.0-only */
+/*
+ * MTD splitter for ELF loader firmware partitions
+ *
+ * Copyright (C) 2020 Sander Vanheule <sander@svanheule.net>
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License as published by the Free
+ * Software Foundation; version 2.
+ *
+ * To parse the ELF kernel loader, a small ELF parser is used that can
+ * handle both ELF32 or ELF64 class loaders. The splitter assumes that the
+ * kernel is always located before the rootfs, whether it is embedded in the
+ * loader or not.
+ *
+ * The kernel image is preferably embedded inside the ELF loader, so the end
+ * of the loader equals the end of the kernel partition. This is due to the
+ * way mtd_find_rootfs_from searches for the the rootfs:
+ * - if the kernel image is embedded in the loader, the appended rootfs may
+ * follow the loader immediately, within the same erase block.
+ * - if the kernel image is not embedded in the loader, but placed at some
+ * offset behind the loader (OKLI-style loader), the rootfs must be
+ * aligned to an erase-block after the loader and kernel image.
+ */
+
+#include <linux/module.h>
+#include <linux/init.h>
+#include <linux/kernel.h>
+#include <linux/slab.h>
+#include <linux/mtd/mtd.h>
+#include <linux/mtd/partitions.h>
+#include <linux/of.h>
+#include <linux/byteorder/generic.h>
+
+#include "mtdsplit.h"
+
+#define ELF_NR_PARTS 2
+
+#define ELF_MAGIC 0x7f454c46 /* 0x7f E L F */
+#define ELF_CLASS_32 1
+#define ELF_CLASS_64 2
+
+struct elf_header_ident {
+ uint32_t magic;
+ uint8_t class;
+ uint8_t data;
+ uint8_t version;
+ uint8_t osabi;
+ uint8_t abiversion;
+ uint8_t pad[7];
+};
+
+struct elf_header_32 {
+ uint16_t type;
+ uint16_t machine;
+ uint32_t version;
+ uint32_t entry;
+ uint32_t phoff;
+ uint32_t shoff;
+ uint32_t flags;
+ uint16_t ehsize;
+ uint16_t phentsize;
+ uint16_t phnum;
+ uint16_t shentsize;
+ uint16_t shnum;
+ uint16_t shstrndx;
+};
+
+struct elf_header_64 {
+ uint16_t type;
+ uint16_t machine;
+ uint32_t version;
+ uint64_t entry;
+ uint64_t phoff;
+ uint64_t shoff;
+ uint32_t flags;
+ uint16_t ehsize;
+ uint16_t phentsize;
+ uint16_t phnum;
+ uint16_t shentsize;
+ uint16_t shnum;
+ uint16_t shstrndx;
+};
+
+struct elf_header {
+ struct elf_header_ident ident;
+ union {
+ struct elf_header_32 elf32;
+ struct elf_header_64 elf64;
+ };
+};
+
+struct elf_program_header_32 {
+ uint32_t type;
+ uint32_t offset;
+ uint32_t vaddr;
+ uint32_t paddr;
+ uint32_t filesize;
+ uint32_t memsize;
+ uint32_t flags;
+};
+
+struct elf_program_header_64 {
+ uint32_t type;
+ uint32_t flags;
+ uint64_t offset;
+ uint64_t vaddr;
+ uint64_t paddr;
+ uint64_t filesize;
+ uint64_t memsize;
+};
+
+
+static int mtdsplit_elf_read_mtd(struct mtd_info *mtd, size_t offset,
+ uint8_t *dst, size_t len)
+{
+ size_t retlen;
+ int ret;
+
+ ret = mtd_read(mtd, offset, len, &retlen, dst);
+ if (ret) {
+ pr_debug("read error in \"%s\"\n", mtd->name);
+ return ret;
+ }
+
+ if (retlen != len) {
+ pr_debug("short read in \"%s\"\n", mtd->name);
+ return -EIO;
+ }
+
+ return 0;
+}
+
+static int elf32_determine_size(struct mtd_info *mtd, struct elf_header *hdr,
+ size_t *size)
+{
+ struct elf_header_32 *hdr32 = &(hdr->elf32);
+ int err;
+ size_t section_end, ph_table_end, ph_entry;
+ struct elf_program_header_32 ph;
+
+ *size = 0;
+
+ if (hdr32->shoff > 0) {
+ *size = hdr32->shoff + hdr32->shentsize * hdr32->shnum;
+ return 0;
+ }
+
+ ph_entry = hdr32->phoff;
+ ph_table_end = hdr32->phoff + hdr32->phentsize * hdr32->phnum;
+
+ while (ph_entry < ph_table_end) {
+ err = mtdsplit_elf_read_mtd(mtd, ph_entry, (uint8_t *)(&ph),
+ sizeof(ph));
+ if (err)
+ return err;
+
+ section_end = ph.offset + ph.filesize;
+ if (section_end > *size)
+ *size = section_end;
+
+ ph_entry += hdr32->phentsize;
+ }
+
+ return 0;
+}
+
+static int elf64_determine_size(struct mtd_info *mtd, struct elf_header *hdr,
+ size_t *size)
+{
+ struct elf_header_64 *hdr64 = &(hdr->elf64);
+ int err;
+ size_t section_end, ph_table_end, ph_entry;
+ struct elf_program_header_64 ph;
+
+ *size = 0;
+
+ if (hdr64->shoff > 0) {
+ *size = hdr64->shoff + hdr64->shentsize * hdr64->shnum;
+ return 0;
+ }
+
+ ph_entry = hdr64->phoff;
+ ph_table_end = hdr64->phoff + hdr64->phentsize * hdr64->phnum;
+
+ while (ph_entry < ph_table_end) {
+ err = mtdsplit_elf_read_mtd(mtd, ph_entry, (uint8_t *)(&ph),
+ sizeof(ph));
+ if (err)
+ return err;
+
+ section_end = ph.offset + ph.filesize;
+ if (section_end > *size)
+ *size = section_end;
+
+ ph_entry += hdr64->phentsize;
+ }
+
+ return 0;
+}
+
+static int mtdsplit_parse_elf(struct mtd_info *mtd,
+ const struct mtd_partition **pparts,
+ struct mtd_part_parser_data *data)
+{
+ struct elf_header hdr;
+ size_t loader_size, rootfs_offset;
+ enum mtdsplit_part_type type;
+ struct mtd_partition *parts;
+ int err;
+
+ err = mtdsplit_elf_read_mtd(mtd, 0, (uint8_t *)&hdr, sizeof(hdr));
+ if (err)
+ return err;
+
+ if (be32_to_cpu(hdr.ident.magic) != ELF_MAGIC) {
+ pr_debug("invalid ELF magic %08x\n",
+ be32_to_cpu(hdr.ident.magic));
+ return -EINVAL;
+ }
+
+ switch (hdr.ident.class) {
+ case ELF_CLASS_32:
+ err = elf32_determine_size(mtd, &hdr, &loader_size);
+ break;
+ case ELF_CLASS_64:
+ err = elf64_determine_size(mtd, &hdr, &loader_size);
+ break;
+ default:
+ pr_debug("invalid ELF class %i\n", hdr.ident.class);
+ err = -EINVAL;
+ }
+
+ if (err)
+ return err;
+
+ err = mtd_find_rootfs_from(mtd, loader_size, mtd->size,
+ &rootfs_offset, &type);
+ if (err)
+ return err;
+
+ if (rootfs_offset == mtd->size) {
+ pr_debug("no rootfs found in \"%s\"\n", mtd->name);
+ return -ENODEV;
+ }
+
+ parts = kzalloc(ELF_NR_PARTS * sizeof(*parts), GFP_KERNEL);
+ if (!parts)
+ return -ENOMEM;
+
+ parts[0].name = KERNEL_PART_NAME;
+ parts[0].offset = 0;
+ parts[0].size = rootfs_offset;
+
+ if (type == MTDSPLIT_PART_TYPE_UBI)
+ parts[1].name = UBI_PART_NAME;
+ else
+ parts[1].name = ROOTFS_PART_NAME;
+ parts[1].offset = rootfs_offset;
+ parts[1].size = mtd->size - rootfs_offset;
+
+ *pparts = parts;
+ return ELF_NR_PARTS;
+}
+
+static const struct of_device_id mtdsplit_elf_of_match_table[] = {
+ { .compatible = "openwrt,elf" },
+ {},
+};
+MODULE_DEVICE_TABLE(of, mtdsplit_elf_of_match_table);
+
+static struct mtd_part_parser mtdsplit_elf_parser = {
+ .owner = THIS_MODULE,
+ .name = "elf-loader-fw",
+ .of_match_table = mtdsplit_elf_of_match_table,
+ .parse_fn = mtdsplit_parse_elf,
+ .type = MTD_PARSER_TYPE_FIRMWARE,
+};
+
+static int __init mtdsplit_elf_init(void)
+{
+ register_mtd_parser(&mtdsplit_elf_parser);
+
+ return 0;
+}
+
+subsys_initcall(mtdsplit_elf_init);