[PATCH] arm64: kexec_file: try more regions if loading segments fails

[Benjamin Gwin]

It's possible that the first region picked for the new kernel will make
it impossible to fit the other segments in the required 32GB window,
especially if we have a very large initrd.

Instead of giving up, we can keep testing other regions for the kernel
until we find one that works.

Signed-off-by: Benjamin Gwin <bgwin at google.com>
Suggested-by: Ryan O'Leary <ryanoleary at google.com>
---
 arch/arm64/kernel/kexec_image.c        | 41 +++++++++++++++++++-------
 arch/arm64/kernel/machine_kexec_file.c |  9 +++++-
 2 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/arch/arm64/kernel/kexec_image.c b/arch/arm64/kernel/kexec_image.c
index af9987c154ca..e2251d61ea35 100644
--- a/arch/arm64/kernel/kexec_image.c
+++ b/arch/arm64/kernel/kexec_image.c
@@ -43,7 +43,7 @@ static void *image_load(struct kimage *image,
        u64 flags, value;
        bool be_image, be_kernel;
        struct kexec_buf kbuf;
-       unsigned long text_offset;
+       unsigned long text_offset, kernel_segment_number;
        struct kexec_segment *kernel_segment;
        int ret;

@@ -88,11 +88,37 @@ static void *image_load(struct kimage *image,
        /* Adjust kernel segment with TEXT_OFFSET */
        kbuf.memsz += text_offset;

-       ret = kexec_add_buffer(&kbuf);
-       if (ret)
+       kernel_segment_number = image->nr_segments;
+
+       /*
+        * The location of the kernel segment may make it impossible to satisfy
+        * the other segment requirements, so we try repeatedly to find a
+        * location that will work.
+        */
+       while ((ret = kexec_add_buffer(&kbuf)) == 0) {
+               /* Try to load additional data */
+               kernel_segment = &image->segment[kernel_segment_number];
+               ret = load_other_segments(image,
+                                       kernel_segment->mem, kernel_segment->memsz,
+                                       initrd, initrd_len, cmdline);
+               if (!ret)
+                       break;
+
+               /*
+                * We couldn't find space for the other segments; erase the
+                * kernel segment and try the next available hole.
+                * */
+               image->nr_segments -= 1;
+               kbuf.buf_min = kernel_segment->mem + kernel_segment->memsz;
+               kbuf.mem = KEXEC_BUF_MEM_UNKNOWN;
+       }
+
+       if (ret) {
+               pr_err("Could not find any suitable kernel location!");
                return ERR_PTR(ret);
+       }

-       kernel_segment = &image->segment[image->nr_segments - 1];
+       kernel_segment = &image->segment[kernel_segment_number];
        kernel_segment->mem += text_offset;
        kernel_segment->memsz -= text_offset;
        image->start = kernel_segment->mem;
@@ -101,12 +127,7 @@ static void *image_load(struct kimage *image,
                                kernel_segment->mem, kbuf.bufsz,
                                kernel_segment->memsz);

-       /* Load additional data */
-       ret = load_other_segments(image,
-                               kernel_segment->mem, kernel_segment->memsz,
-                               initrd, initrd_len, cmdline);
-
-       return ERR_PTR(ret);
+       return 0;
 }

 #ifdef CONFIG_KEXEC_IMAGE_VERIFY_SIG
diff --git a/arch/arm64/kernel/machine_kexec_file.c b/arch/arm64/kernel/machine_kexec_file.c
index 5b0e67b93cdc..03210f644790 100644
--- a/arch/arm64/kernel/machine_kexec_file.c
+++ b/arch/arm64/kernel/machine_kexec_file.c
@@ -240,6 +240,11 @@ static int prepare_elf_headers(void **addr, unsigned long *sz)
        return ret;
 }

+/*
+ * Tries to add the initrd and DTB to the image. If it is not possible to find
+ * valid locations, this function will undo changes to the image and return non
+ * zero.
+ */
 int load_other_segments(struct kimage *image,
                        unsigned long kernel_load_addr,
                        unsigned long kernel_size,
@@ -248,7 +253,8 @@ int load_other_segments(struct kimage *image,
 {
        struct kexec_buf kbuf;
        void *headers, *dtb = NULL;
-       unsigned long headers_sz, initrd_load_addr = 0, dtb_len;
+       unsigned long headers_sz, initrd_load_addr = 0, dtb_len,
+                     orig_segments = image->nr_segments;
        int ret = 0;

        kbuf.image = image;
@@ -334,6 +340,7 @@ int load_other_segments(struct kimage *image,
        return 0;

 out_err:
+       image->nr_segments = orig_segments;
        vfree(dtb);
        return ret;
 }
--
2.29.1.341.ge80a0c044ae-goog