/*
 *  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; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Library General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 */
/*
Test-code for Kernel-Modules with GRUB - ELF-Relocation
  by Soeren Bleikertz <sb@osdev.de>
  
This is just for demonstration and learning.
I won't explain how relocation works. Please read the ELF-specifications!

http://sac.cc || http://osdev.de || http://soeren.geekgate.org

works more or less. not fully tested!
any questions or comments?
*/

#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdlib.h>
#include <string.h>
#include <elf.h>
#include <sys/mman.h>

//FATAL ERROR
#define FAT_ERROR(s) \
  do { \
    perror(s); \
    exit(-1); \
  } while(1);
  
#define ESUCCESS 0
#define EFAILURE -1
#define ELF_MAGIC "\x7f""ELF"

#define MAX_SYMS 255
  
#define ENTRY_FUNC "mod_init"
  
/* TEST-FUNCTION */
  

/* test.c - I use this small source for relocation-testing. just compile it with
  gcc -c test.c
  
int lala(int *a) {
  *a = 2;
  return 0;
}

static void mod_init(void)
{  
  char *bla = "MOD_INIT!!!!\n";
  foo();
  bar(bla);
}

void hoo(void) {
  int bla;
  bla = 2;
  bla *= 2;
}  
  
 EOF */
 
  
void foo(void)
{
  printf("\nfoo() said: Hello\n");
}
  
void bar(char *b)
{
  printf("\nbar() said: %s(%p)\n", b, b);
}


/* ENTRY-STUFF */

typedef struct entry {
  //Elf32_Shdr *shdr;
  ulong addr;
} entry_t;

entry_t main_entry;


/* SYMBOL-STUFF */
  
typedef struct sym_ent {
  char *name;
  ulong addr;
} sym_t;

// symbol table from module
sym_t sym_tab_mod[MAX_SYMS];
int sym_count;

// global symbol-table
sym_t sym_tab[] = {
  { "foo", (ulong)foo },
  { "bar", (ulong)bar } //test
};

// get addr for a symbol
int sym_lookup(char *sym, ulong *sym_val)
{
  int sym_nr = sizeof(sym_tab)/sizeof(sym_t);
  int i;
  
  for (i=0; i<sym_nr; i++) {
    if (!strcmp(sym_tab[i].name, sym)) {
      *sym_val = sym_tab[i].addr;
      return ESUCCESS;
    }
  }
  return EFAILURE;
}

//only local/global functions with size>0
int sym_valid_name(Elf32_Sym *symt)
{
  if ((symt->st_name != 0) &&
  ((ELF32_ST_BIND(symt->st_info) == STB_LOCAL) || (ELF32_ST_BIND(symt->st_info) == STB_GLOBAL))
  && (symt->st_size > 0) && (ELF32_ST_TYPE(symt->st_info) == STT_FUNC))
    return ESUCCESS;
  return EFAILURE;
}
  

// create symbol-table from mod
int sym_tab_create(char *data)
{
  Elf32_Ehdr *elfhdr = (Elf32_Ehdr*)data;
  int nr_entry, i, nr_shdr = elfhdr->e_shnum;
  Elf32_Shdr *shdr, *off_shdr = (Elf32_Shdr*)(data + elfhdr->e_shoff);
  Elf32_Sym *symt;
  char *strtab;
  ulong text_ofs;
  
  sym_count = 0;
  
  //search Symbol-Table
  for (i=0; i < nr_shdr; i++) {
    if (off_shdr[i].sh_type == SHT_SYMTAB)
      shdr = &off_shdr[i];
  }
  
  printf("Debug: symtab: %p\n", shdr->sh_offset);
  symt = (Elf32_Sym*)(data+shdr->sh_offset);
  nr_entry = shdr->sh_size/sizeof(Elf32_Sym);
  
  //search String-Table
  for (i=0; i < nr_shdr; i++) {
    if ((off_shdr[i].sh_type == SHT_STRTAB) &&
      !(off_shdr[i].sh_flags) && (elfhdr->e_shstrndx != i))
      shdr = &off_shdr[i];
  }
  printf("Debug: strtab: %p\n", shdr->sh_offset);
  strtab = (char*)(data+shdr->sh_offset);
  
  //search .text
  for (i=0;i<nr_shdr; i++) {
    if (!(off_shdr[i].sh_flags & 0x0004))
      continue;
    text_ofs = (ulong)off_shdr[i].sh_offset;
  }
  
  
  printf("Debug: .text ofs: %p\n", text_ofs);
  
  // create symbol-table of module
  for (i=0; i<nr_entry; i++) {
    if ((sym_valid_name(&symt[i]) == ESUCCESS) && (sym_count <= MAX_SYMS)) {
      sym_tab_mod[sym_count].name = &strtab[symt[i].st_name];
      sym_tab_mod[sym_count].addr = (ulong)(symt[i].st_value + data + text_ofs);
      printf("Debug: Symbol: %s (%p)\n",
        sym_tab_mod[sym_count], sym_tab_mod[sym_count].addr);
      sym_count++;
      }
  }

  return ESUCCESS;
}
  

/* KMOD-STUFF */

// validate ELF
int kmod_check(char *data)
{
  Elf32_Ehdr *elfhdr= (Elf32_Ehdr*)data;

  if (strncmp(&elfhdr->e_ident[EI_MAG0], ELF_MAGIC, 4) ||
    (elfhdr->e_ident[EI_CLASS] != ELFCLASS32) ||
    (elfhdr->e_ident[EI_DATA] != ELFDATA2LSB) ||
    (elfhdr->e_type != ET_REL) ||
    (elfhdr->e_machine != EM_386) ||
    (elfhdr->e_version != 1))
      return EFAILURE;
    
  return ESUCCESS;
}

int kmod_load(char *data, char *modname)
{
  if (kmod_check(data) == EFAILURE) {
    printf("Error: loading of '%s' failed!\n", modname);
    return EFAILURE;
  }
  printf("Debug: Valid ELF\n");

  return ESUCCESS;
}

int kmod_elf_sym(char *data, uint symidx, ulong *sym_val, ulong symtab_sect)
{
  Elf32_Ehdr *elfhdr;
  Elf32_Shdr *shdr, *sym_sec, *tmp;
  Elf32_Sym *symtab;
  char *sym, *strtab;
  int i;
  
  elfhdr = (Elf32_Ehdr*)data;
  
  shdr = (Elf32_Shdr*)(data + elfhdr->e_shoff);
  
  //search strtab
  for (i=0; i < elfhdr->e_shnum; i++) {
    if ((shdr[i].sh_type == SHT_STRTAB) && !(shdr[i].sh_flags) && (elfhdr->e_shstrndx != i))
      break;      
  }
  tmp = &shdr[i];
  
  strtab = (char*)(data+tmp->sh_offset);
  
  printf("Debug: StrTab-ofs: %p\n", tmp->sh_offset);
  
  sym_sec = (Elf32_Shdr*)&shdr[symtab_sect];
  
  printf("Debug: SymTab-Sect: %lu, SymTab-IDX: %d\n", symtab_sect, symidx);
  
  if (symidx > sym_sec->sh_size/sym_sec->sh_entsize)
    return EFAILURE;
  
  if (sym_sec->sh_type != SHT_SYMTAB) {
    printf("Error: Not a SymTab!\n");
    return EFAILURE;
  }
  
  symtab = (Elf32_Sym*)(data + sym_sec->sh_offset);
  
  printf("Debug: Symtab-ofs: %p\n", sym_sec->sh_offset);
  
  
  if (!symtab[symidx].st_shndx) {
    //external symbol
    printf("Debug: External Symbol\n");
    sym = &strtab[symtab[symidx].st_name];
    printf("Debug: Symbol: \"%s\"\n", sym);
    if (sym_lookup(sym, sym_val) == EFAILURE) {
      printf("Error: Unknown Symbol!\n");
      return EFAILURE;
    }
    printf("Debug: ext. Symbol-Addr: %p\n", *sym_val);
  } else {
    //internal symbol
    printf("Debug: Internal Symbol\n");
    shdr = (Elf32_Shdr*)(data + elfhdr->e_shoff + elfhdr->e_shentsize * symidx);
    *sym_val = symtab->st_value + (ulong)(data + shdr->sh_offset);
    printf("Debug: int. Symbol-Addr: %p, Symtab-Value: %d\n", *sym_val, symtab->st_value);  
  }
  return ESUCCESS;
}
  

int kmod_elf_reloc_do(char *data, Elf32_Rel *rel, Elf32_Shdr *shdr)
{
  ulong *rel_addr;
  Elf32_Shdr *rel_sect; //section for relocation
  Elf32_Ehdr *elfhdr;
  ulong sym_val;
  
  //ELf-hdr
  elfhdr = (Elf32_Ehdr*)data;
  
  
  //section for relocation
  rel_sect = (Elf32_Shdr*)(data + elfhdr->e_shoff + elfhdr->e_shentsize * shdr->sh_info);
  
  //relocation addr
  rel_addr = (ulong*)(data + rel_sect->sh_offset + rel->r_offset);
  printf("Debug: rel_offset: 0x%x, SYM: 0x%x, TYPE: 0x%x\n",
    rel->r_offset, ELF32_R_SYM(rel->r_info), ELF32_R_TYPE(rel->r_info));
  
  
  printf("Debug: sect_for_rel_ofs: %p, SymTab-Sect-IDX: %d\n", rel_sect->sh_offset, shdr->sh_link);
  
  // get addr of symbol
  if (kmod_elf_sym(data, ELF32_R_SYM(rel->r_info), &sym_val, shdr->sh_link) == EFAILURE)
    return EFAILURE;
  // omg..
  switch (ELF32_R_TYPE(rel->r_info)) {
    case R_386_32:
      *rel_addr = sym_val + *rel_addr;
      printf("Debug: R_386_32: rel_addr: %p(%p)\n", rel_addr, *rel_addr);
    break;
    case R_386_PC32:
      *rel_addr = sym_val + *rel_addr - (ulong)rel_addr;
      printf("Debug: R_386_PC32: rel_addr: %p(%p)\n", rel_addr, *rel_addr);
    break;
    default:
      printf("Error: wrong Reloc-Type\n");
      return EFAILURE;
  }
  
  return ESUCCESS;  
}


int kmod_elf_reloc(char *data)
{
  Elf32_Shdr *shdr;
  Elf32_Ehdr *ehdr;
  Elf32_Rel *rel;
  uint shdr_nr, i=0, shdr_ent_sz, rel_sz=0, j;
  ulong entry=0;
  char *bss;
  
  ehdr = (Elf32_Ehdr*)data;
  shdr_nr = ehdr->e_shnum;
  shdr_ent_sz = ehdr->e_shentsize;
  
  shdr = (Elf32_Shdr*)(data+ehdr->e_shoff);  
  
  // search BBS
  while((shdr[i].sh_type != SHT_NOBITS) && (i<shdr_nr))
    i++;
  if (shdr[i].sh_type != SHT_NOBITS)
    printf("No BSS found!\n");
  printf("Debug: Found BSS in Section %d\n", i);
  bss = malloc(shdr[i].sh_size);
  printf("Debug: new BSS: 0x%x\n", (ulong)bss - (ulong)data);
  shdr[i].sh_offset = (ulong)bss - (ulong)data;
  
  // search relocation-sections
  for(i=0; i<shdr_nr; i++, rel_sz=0) {
    if ((shdr[i].sh_type != SHT_RELA) && (shdr[i].sh_type != SHT_REL))
      continue;
    rel_sz = shdr[i].sh_entsize;
    rel = (Elf32_Rel*)(data+shdr[i].sh_offset);
    printf("Debug: rel_section: %d(%p), rel_ent_size: %d -> rel_type: %d, rel_ent_nr: %d\n",
      i, shdr[i].sh_offset, rel_sz, shdr[i].sh_type, shdr[i].sh_size/rel_sz);
    for(j=0;j<shdr[i].sh_size/rel_sz;j++) {
      printf("\nDebug: Relocation-Nr: %d/%d\n", j+1, shdr[i].sh_size/rel_sz);
      printf("Debug: rel: %p, shdr: %p\n",
        (ulong)&rel[j]-(ulong)data, (ulong)&shdr[i]-(ulong)data);
      // do relocation
      if(kmod_elf_reloc_do(data, &rel[j], &shdr[i]) == EFAILURE)
        return EFAILURE;
    }
  }
  
  // search for ENTRY_FUNC, set main-entry-addr
  for (i=0; i<sym_count; i++) {
    if (!strcmp(sym_tab_mod[i].name, ENTRY_FUNC)) {
      main_entry.addr = sym_tab_mod[i].addr;
      break;
    }
  }
  
  // ENTRY_FUNC not found, set main-entry-addr to beginning of .text
  if (main_entry.addr == 0) {
    printf("Debug: No Main-Entry! Set entry to beginning of .text!\n");
    for (i=0;i<shdr_nr; i++) {
      if (!(shdr[i].sh_flags & 0x0004))
        continue;
      main_entry.addr = (ulong)(data + shdr[i].sh_offset);
      //main_entry.shdr = &shdr[i];
      printf("Debug: Entry: %p\n", main_entry.addr);
      break;
    }
  }  
  
  return ESUCCESS;
}
  


/* OTHER STUFF */

int main(int argc, char **argv)
{
  ulong elf_sz, entry;
  char *elf_ptr;
  int fd;
  struct stat st;
  void (*ble)();
  
  if (argc != 2)
    return -1;
  
  //open file
  if ((fd = open(argv[1], O_RDONLY, 0)) == -1)
    FAT_ERROR("open failed");

  //getting file-infos
  if (fstat(fd, &st) == -1)
    FAT_ERROR("stat failed");

  elf_sz = st.st_size;
  if ((elf_ptr = malloc(elf_sz)) == NULL)
    FAT_ERROR("malloc() failed");
  //read file
  read(fd, elf_ptr, elf_sz);
  
  close(fd);
  
  printf("KMOD-Testing:\n");
  
  printf("Modul: %s:\n---\n", argv[1]);
  printf("[Load Mod]\n");
  kmod_load(elf_ptr, argv[1]);
  printf("> Done!\n\n");
  
  printf("[Create SymTab]\n");
  sym_tab_create(elf_ptr);
  printf("> Done!\n\n");
  
  //relocation
  bzero(&main_entry, sizeof(entry_t));
  printf("[Relocation]\n");
  if (kmod_elf_reloc(elf_ptr)==EFAILURE) {
    printf("Error: Relocation failed!\n");
    return EFAILURE;
  } else
    printf("> Done!\n\n");
  
  printf("[Start]\n");
  ble = (void*)main_entry.addr;
  printf("Debug: entry: %p\n", main_entry.addr);
  printf("<output>\n");
  ble();
  printf("</output\n");
  printf("> Done!\n");
  free(elf_ptr);
  return 0;
}