Fix bugs: symbol table size, base vaddr calc, sections flags check

upicine · Apr 11, 2020 · 614079c · 614079c
1 parent ff3ca9e
commit 614079c
Show file tree

Hide file tree

Showing 8 changed files with 90 additions and 56 deletions.
diff --git a/main.c b/main.c
@@ -1,13 +1,10 @@
 #include <stdio.h>
 #include <stdlib.h>
-#include <elf.h>
 #include <unistd.h>
 
 #include "utils.h"
 #include "postlinker.h"
 #include "relocator.h"
-#include "parser.h"
-
 
 int main(int argc, char** argv) {
     if (argc != 4) {
@@ -21,6 +18,8 @@ int main(int argc, char** argv) {
     Elf_Data* exec_elf =  init_exec_elf(argv[1]);
     Elf_Data* rel_elf = init_rel_elf(argv[2]);
 
+    set_base_vaddr(exec_elf);
+
     incr_segments_off(exec_elf);
     rearrange_vaddr(exec_elf);
     copy_file(exec_elf->fd, output_fd, 0x0);
@@ -36,7 +35,11 @@ int main(int argc, char** argv) {
 
     free_elf(&exec_elf);
     free_elf(&rel_elf);
-    close(output_fd);
+    int ret = close(output_fd);
+    if (ret < 0) {
+        perror("Error closing output file");
+        exit(1);
+    }
 
     return 0;
 }
diff --git a/parser.c b/parser.c
@@ -1,7 +1,3 @@
-//
-// Created by michal on 05.04.20.
-//
-
 #include <elf.h>
 #include <stdlib.h>
 #include <zconf.h>
@@ -40,7 +36,12 @@ Elf_Data* init_rel_elf(const char* filename) {
 
 
 void free_elf(Elf_Data** elf) {
-    close((*elf)->fd);
+    int ret = close((*elf)->fd);
+    if (ret < 0) {
+        perror("Error closing elf file");
+        exit(1);
+
+    }
     free((*elf)->ehdr);
     free((*elf)->phdr);
     free((*elf)->shdr);
@@ -68,7 +69,7 @@ Elf64_Ehdr* read_elf_header(int fd) {
 Elf64_Shdr* read_sections_header(int fd, const Elf64_Ehdr *ehdr) {
     size_t sect_size = ehdr->e_shentsize * ehdr->e_shnum;
     Elf64_Shdr* sections = malloc(sect_size);
-    if (!sections) {
+    if (!sections && sect_size > 0) {
         perror("Error reading sections header");
         exit(1);
     }
@@ -86,7 +87,7 @@ Elf64_Shdr* read_sections_header(int fd, const Elf64_Ehdr *ehdr) {
 Elf64_Phdr* read_program_header(int fd, const Elf64_Ehdr* ehdr) {
     size_t phdr_size = ehdr->e_phentsize * ehdr->e_phnum;
     Elf64_Phdr* phdr = malloc(phdr_size);
-    if (!phdr) {
+    if (!phdr && phdr_size > 0) {
         perror("Error reading program header");
         exit(1);
     }
@@ -104,7 +105,7 @@ Elf64_Phdr* read_program_header(int fd, const Elf64_Ehdr* ehdr) {
 char* read_shstr_tab(int fd, const Elf64_Shdr* shdr, const Elf64_Ehdr* ehdr) {
     Elf64_Shdr shstr_hdr = shdr[ehdr->e_shstrndx];
     char* shstr_tab = malloc(shstr_hdr.sh_size);
-    if (!shstr_tab) {
+    if (!shstr_tab && shstr_hdr.sh_size > 0) {
         perror("Error reading shstr tab");
         exit(1);
     }
@@ -154,7 +155,7 @@ Elf64_Sym* read_sym_tab(int fd, const Elf64_Shdr* sym_shdr) {
 
 char* read_str_tab(int fd, const Elf64_Shdr* shdr) {
     char* str_tab = malloc(shdr->sh_size);
-    if (!str_tab) {
+    if (!str_tab && shdr->sh_size > 0) {
         perror("Error reading shstr tab");
         exit(1);
     }

diff --git a/parser.h b/parser.h
@@ -1,8 +1,9 @@
 #ifndef POSTLINKER_PARSER_H_H
 #define POSTLINKER_PARSER_H_H
 
+#include <elf.h>
 
-typedef struct Elf_File {
+typedef struct Elf_Data {
     int fd;
     Elf64_Ehdr* ehdr;
     Elf64_Phdr* phdr;

diff --git a/postlinker.c b/postlinker.c
@@ -3,35 +3,53 @@
 #include <stdlib.h>
 #include <stdio.h>
 #include <string.h>
+#include <stdint.h>
 
 #include "postlinker.h"
 #include "utils.h"
 
-static const Elf64_Xword sect_flags[] = {
-        SHF_ALLOC + SHF_EXECINSTR + SHF_WRITE,
-        SHF_ALLOC + SHF_EXECINSTR,
-        SHF_ALLOC + SHF_WRITE,
-        SHF_ALLOC
-};
-
-static const size_t sect_flags_sz = 4;
-
-
-static Elf64_Word shflag_to_phflag(Elf64_Xword shflag) {
-    switch (shflag) {
-        case SHF_ALLOC + SHF_EXECINSTR + SHF_WRITE:
-            return PF_R + PF_W + PF_X;
-        case SHF_ALLOC + SHF_EXECINSTR:
-            return PF_R + PF_X;
-        case SHF_ALLOC + SHF_WRITE:
-            return PF_R + PF_W;
-        case SHF_ALLOC:
+
+typedef enum {RWX, RX, RW, R} Sect_Flag;
+
+
+static const unsigned int sect_flag_num = 4;
+
+static Elf64_Addr base_vaddr;
+
+
+static Elf64_Word shflag_to_phflag(Sect_Flag flag) {
+    switch (flag) {
+        case RWX:
+            return PF_R | PF_W | PF_X;
+        case RX:
+            return PF_R | PF_X;
+        case RW:
+            return PF_R | PF_W;
+        case R:
             return PF_R;
     }
 }
 
 
-static off_t calc_alignment(off_t off, uint64_t alignment) {
+static int has_flag(Elf64_Xword sect_flag, Sect_Flag flag) {
+    if (sect_flag & SHF_ALLOC) {
+        if (flag == RWX && sect_flag & SHF_EXECINSTR && sect_flag & SHF_WRITE)
+            return 1;
+        if (flag == RW && sect_flag & SHF_WRITE
+            && !(sect_flag & SHF_EXECINSTR))
+            return 1;
+        if (flag == RX && sect_flag & SHF_EXECINSTR
+            && !(sect_flag & SHF_WRITE))
+            return 1;
+        if (flag == R && !(sect_flag & SHF_EXECINSTR)
+            && !(sect_flag & SHF_WRITE))
+            return 1;
+    }
+    return 0;
+}
+
+
+static Elf64_Off calc_alignment(Elf64_Off off, uint64_t alignment) {
     if (off % alignment != 0)
         return alignment - (off % alignment);
 
@@ -48,19 +66,29 @@ static Elf64_Off next_avail_off(int fd) {
         exit(1);
     }
 
-    return (Elf64_Off) file_sz + calc_alignment(file_sz, PAGE_SIZE);
+    return (Elf64_Off) file_sz + calc_alignment((Elf64_Off)file_sz, PAGE_SIZE);
+}
+
+
+void set_base_vaddr(const Elf_Data* elf) {
+    base_vaddr = UINT64_MAX;
+    for (int i = 0; i < elf->ehdr->e_phnum; i++) {
+        if (elf->phdr[i].p_vaddr != 0x0 && elf->phdr[i].p_vaddr < base_vaddr) {
+            base_vaddr = elf->phdr[i].p_vaddr;
+        }
+    }
 }
 
 
-static size_t merge_content(char** content, Elf_Data* elf, Elf64_Xword flag,
+static size_t merge_content(char** content, Elf_Data* elf, Sect_Flag flag,
                             Elf64_Addr vaddr, off_t base_off) {
     Elf64_Ehdr* ehdr = elf->ehdr;
     Elf64_Shdr* shdr = elf->shdr;
     size_t content_sz = 0;
-    off_t off = 0, addr_align = 0;
+    Elf64_Off off = 0, addr_align = 0;
 
     for (int i = 0; i < ehdr->e_shnum; i++) {
-        if (shdr[i].sh_size != 0 && shdr[i].sh_flags == flag) {
+        if (shdr[i].sh_size != 0 && has_flag(shdr[i].sh_flags, flag)) {
             addr_align = calc_alignment(off, shdr[i].sh_addralign);
 
             content_sz += addr_align + shdr[i].sh_size;
@@ -93,16 +121,15 @@ uint16_t merge_sections(int out_fd, Elf_Data* exec_elf, Elf_Data* rel_elf,
     uint16_t j = 0;
     Elf64_Off off = next_avail_off(exec_elf->fd);
 
-    for (int i = 0; i < sect_flags_sz; i++) {
-        Elf64_Addr vaddr = BASE_ADDR + off;
-        size_t content_sz = merge_content(&content, rel_elf, sect_flags[i],
-                                          vaddr, off);
+    for (Sect_Flag f = RWX; f <= R; f++) {
+        Elf64_Addr vaddr = base_vaddr + off;
+        size_t content_sz = merge_content(&content, rel_elf, f, vaddr, off);
 
         if (content_sz == 0)
             continue;
 
         new_phdr[j].p_type = PT_LOAD;
-        new_phdr[j].p_flags = shflag_to_phflag(sect_flags[i]);
+        new_phdr[j].p_flags = shflag_to_phflag(f);
         new_phdr[j].p_offset = off + PAGE_SIZE;
         new_phdr[j].p_memsz = content_sz;
         new_phdr[j].p_filesz = content_sz;
@@ -127,7 +154,7 @@ uint16_t merge_sections(int out_fd, Elf_Data* exec_elf, Elf_Data* rel_elf,
 
 
 void add_new_segments(int out_fd, Elf_Data* rel_elf, Elf_Data** exec_elf) {
-    Elf64_Phdr new_phdr[sect_flags_sz];
+    Elf64_Phdr new_phdr[sect_flag_num];
     Elf64_Ehdr* exec_ehdr = (*exec_elf)->ehdr;
     uint16_t new_seg_num = merge_sections(out_fd, *exec_elf, rel_elf, new_phdr);
     uint32_t new_phdr_sz = ((*exec_elf)->ehdr->e_phnum + new_seg_num)

diff --git a/postlinker.h b/postlinker.h
@@ -11,4 +11,6 @@ void add_new_segments(int out_fd, Elf_Data* rel_elf, Elf_Data** exec_elf);
 
 void incr_sections_off(Elf_Data* elf);
 
+void set_base_vaddr(const Elf_Data* elf);
+
 #endif //POSTLINKER_POSTLINEKR_H
diff --git a/relocator.c b/relocator.c
@@ -7,7 +7,7 @@
 #include "utils.h"
 
 
-static Elf64_Shdr* search_shdr(const char* shdr_name, Elf_Data* elf,
+static Elf64_Shdr* search_shdr(const char* shdr_name, const Elf_Data* elf,
                                char* shstr_tab) {
     Elf64_Shdr* shdr = elf->shdr;
 
@@ -19,19 +19,19 @@ static Elf64_Shdr* search_shdr(const char* shdr_name, Elf_Data* elf,
 }
 
 
-static size_t get_sym_tab(int fd, Elf64_Sym** sym_tab, Elf_Data *elf) {
+static size_t get_sym_tab(int fd, Elf64_Sym** sym_tab, const Elf_Data *elf) {
     Elf64_Shdr* shdr = elf->shdr;
 
     for (int i = 0; i < elf->ehdr->e_shnum; i++) {
         if (shdr[i].sh_type == SHT_SYMTAB) {
             *sym_tab = read_sym_tab(fd, shdr + i);
-            return shdr[i].sh_size;
+            return shdr[i].sh_size / sizeof(Elf64_Sym);
         }
     }
 }
 
 
-static char* get_str_tab(int fd, Elf_Data *elf, char* shstr_tab) {
+static char* get_str_tab(int fd, const Elf_Data *elf, char* shstr_tab) {
     Elf64_Shdr* shdr = elf->shdr;
 
     for (int i = 0; i < elf->ehdr->e_shnum; i++) {
@@ -51,6 +51,8 @@ static Elf64_Sym* search_symbol(const char* sym_name, Elf64_Sym* sym_tab,
         if (strcmp(act_sym_name, sym_name) == 0)
             return sym_tab + i;
     }
+
+    return NULL;
 }
 
 
@@ -107,11 +109,9 @@ void relocate(int fd, Elf_Data* rel_elf, Elf_Data* exec_elf) {
                     } else {
                         sym = search_symbol(sym_name, exec_sym_tab, exec_str,
                                             exec_sym_sz);
-//                        sym_shdr = exec_elf->shdr + sym->st_shndx;
                         sym_addr = sym->st_value;
                     }
                 } else {
-                    char* sym_name = rel_str + rel_sym_tab[symndx].st_name;
                     sym_shdr = rel_elf->shdr + rel_sym_tab[symndx].st_shndx;
                     sym_addr = sym_shdr->sh_addr + rel_sym_tab[symndx].st_value;
                 }
@@ -125,8 +125,10 @@ void relocate(int fd, Elf_Data* rel_elf, Elf_Data* exec_elf) {
 
     Elf64_Sym* start_sym = search_symbol(_START_NAME, rel_sym_tab, rel_str,
                                          rel_sym_sz);
-    Elf64_Shdr* start_hdr = rel_shdr + start_sym->st_shndx;
-    exec_elf->ehdr->e_entry = start_hdr->sh_addr + start_sym->st_value;
+    if (start_sym) {
+        Elf64_Shdr *start_hdr = rel_shdr + start_sym->st_shndx;
+        exec_elf->ehdr->e_entry = start_hdr->sh_addr + start_sym->st_value;
+    }
 
     free(exec_sym_tab);
     free(rel_sym_tab);

diff --git a/utils.c b/utils.c
@@ -71,7 +71,7 @@ void copy_file(int fd_from, int fd_to, off_t start) {
         }
 
         if (rret != 0) {
-            wret = pwrite(fd_to, buffer, rret, off + PAGE_SIZE);
+            wret = pwrite(fd_to, buffer, (size_t)rret, off + PAGE_SIZE);
             if (wret < 0) {
                 perror("Error copy file");
                 exit(1);

diff --git a/utils.h b/utils.h
@@ -5,8 +5,6 @@
 
 #define HEADER_SIZE 0x40
 #define PAGE_SIZE 0x1000
-#define BASE_ADDR 0x400000
-
 
 int create_file(const char* filename);
 
@@ -22,4 +20,4 @@ void write_sheader(int fd, const Elf_Data* elf);
 
 void write_addr(int fd, const void* addr, size_t addr_sz, off_t off);
 
-#endif //POSTLINKER_ELFUTILS_H
+#endif //POSTLINKER_UTILS_H
Original file line number	Diff line number	Diff line change
Expand Up		@@ -11,4 +11,6 @@ void add_new_segments(int out_fd, Elf_Data* rel_elf, Elf_Data** exec_elf);

		void incr_sections_off(Elf_Data* elf);

		void set_base_vaddr(const Elf_Data* elf);

		#endif //POSTLINKER_POSTLINEKR_H