[pk] fixed elf loader

1 files changed, 11 insertions, 4 deletions

diff --git a/pk/elf.c b/pk/elf.c
index f61f1da..ca44645 100644
--- a/pk/elf.c
+++ b/pk/elf.c
@@ -14,11 +14,9 @@ long load_elf(const char* fn, int* user64)
 
   char buf[2048]; // XXX
   int header_size = file_read(file, buf, sizeof(buf)).result;
-  if(header_size >= (int)sizeof(Elf64_Ehdr))
-    goto fail;
-
   const Elf64_Ehdr* eh64 = (const Elf64_Ehdr*)buf;
-  if(!(eh64->e_ident[0] == '\177' && eh64->e_ident[1] == 'E' &&
+  if(header_size < (int)sizeof(Elf64_Ehdr) ||
+     !(eh64->e_ident[0] == '\177' && eh64->e_ident[1] == 'E' &&
        eh64->e_ident[2] == 'L'    && eh64->e_ident[3] == 'F'))
     goto fail;
 
@@ -28,6 +26,15 @@ long load_elf(const char* fn, int* user64)
     ph = (typeof(ph))(buf+eh->e_phoff); \
     for(int i = 0; i < eh->e_phnum; i++, ph++) { \
       if(ph->p_type == SHT_PROGBITS && ph->p_memsz) { \
+        extern char _end; \
+        if((char*)(long)ph->p_vaddr < &_end) \
+        { \
+          long diff = &_end - (char*)(long)ph->p_vaddr; \
+          ph->p_vaddr += diff; \
+          ph->p_offset += diff; \
+          ph->p_memsz = diff >= ph->p_memsz ? 0 : ph->p_memsz - diff; \
+          ph->p_filesz = diff >= ph->p_filesz ? 0 : ph->p_filesz - diff; \
+        } \
         if(file_pread(file, (char*)(long)ph->p_vaddr, ph->p_filesz, ph->p_offset).result != ph->p_filesz) \
           goto fail; \
         memset((char*)(long)ph->p_vaddr+ph->p_filesz, 0, ph->p_memsz-ph->p_filesz); \