#!/usr/bin/perl -w

#
#
# Wrap a kernel module (.ko) so that all undefined functions
# are resolved at compile time to wrapper that have the template
# as given here
#
# TODO: parameterize hypercall numbers
# TODO: remove printk special case
#

use Getopt::Long;
use File::Basename;

$user=0;
$kern=0;

$V3_BORDER_OUT_CALL_NR = 0x6003;
$V3_BORDER_IN_RET_NR   = 0x6004;

&GetOptions("user"=>\$user, "kern"=>\$kern);

$#ARGV==0 or die "christo.pl [--user|--kern] <object-file>\n";

$of=shift; chomp($of);
($ofstem) = split(/\./,fileparse($of)); 

if ($user) { 
  
  print "Wrapping $of to produce $ofstem-wrapped\n";
  
  # First, find all undefined symbols that are text
  @funcs = ExtractUndefinedFunctions($of);
  
  print "Wrapping the following functions\n";
  print join("\n", @funcs),"\n";
  
  print "Generating wrapper $ofstem\_wrapper.c\n";
  
  open(WRAPPER, ">$ofstem\_wrapper.c");
  GenerateUserWrapper(WRAPPER, @funcs);
  close(WRAPPER);

  print "Compiling wrapper\n";
  
  CompileUserWrapper($ofstem);


  print "Linking with your supplied file $of to form $ofstem\_wrapped\n";

  LinkUserWrapped($of,"$ofstem\_wrapper.o","$ofstem\_wrapped",@funcs);

  print "Done.\n"

} else {

  print "Wrapping $of to produce $ofstem-wrapped\n";

  @funcs = KernUndefFuncs($of);

  
  print "Wrapping the following functions\n";
  print join("\n", @funcs), "\n";

  print "Generating wrapper $ofstem\_wrapper.S\n";

  open(WRAPPER,">$ofstem\_wrapper.S");
  GenerateKernWrapper($ofstem, WRAPPER, @funcs);
  close(WRAPPER);

}

sub CompileUserWrapper {
  my $stem=shift;
  
  return system("gcc -c -fPIC $stem.c -o $stem.o");
}


sub LinkUserWrapped {
  my ($of,$wrapper,$wrapped,@funcs) = @_;
  
  $cmd = "gcc ".join(" ", 
		     join(" ",
			  map {" -Xlinker --wrap=$_ "} @funcs
			 ), 
		     $of, $wrapper)." -o $wrapped";
  
  print $cmd;
  return system $cmd;

}

sub KernUndefFuncs {
    my $of = shift;
    my $line;
    my @funcs;
    my @lines;
    my @final_funcs;
    my %sysmap;

    print STDERR "Acquiring all undefined symbols from file\n";

    open(OBJDUMP, "objdump -t $of |") or die "Cannot open file $of\n";

    while ($line=<OBJDUMP>) {
        if ($line =~ /.*\*UND\*\s+\S+\s+(\S+)$/) {
            push @funcs, $1;
        }
    }

    close(OBJDUMP);

    print STDERR "Finding all functions listed in System.map\n";

    open(MAP,"System.map") or die "Cannot open System.map\n";

    while ($line=<MAP>) { 
	chomp($line);
	my ($addr,$type,$name) = split(/\s+/,$line);
	if (($type eq "t") or ($type eq "T")) { 
	    $sysmap{$name}=1;
	}
    }

    close(MAP);

    foreach $func (@funcs) { 
	if (defined($sysmap{$func}) and ($func ne "printk") and ($func ne "mcount")
        and ($func ne "__cyg_profile_func_enter") and ($func ne "__cyg_profile_func_exit")) {
	    print STDERR "ACCEPT  $func\n";
	    push @final_funcs, $func;
	} else {
	    print STDERR "DISCARD $func\n";
	}
    }

    return @final_funcs;
}

sub ExtractUndefinedFunctions {
  my $of=shift;
  my $line;
  my @funcs;

  open(OBJDUMP,"objdump -t $of |") or die "Cannot open file $of\n";

  while ($line=<OBJDUMP>) { 
    if ($line =~ /.*\*UND\*\s+\S+\s+(\S+)$/) { 
      push @funcs, $1;
    }
  }

  close(OBJDUMP);

  return @funcs;

}

sub GenerateUserWrapper  {
  my $file=shift;
  

  while ($func=shift) {
    print $file <<ENDF;

int __wrap_$func(void *a1, 
		 void *a2,
		 void *a3,
		 void *a4,
		 void *a5,
		 void *a6,
		 void *a7,
		 void *a8) 
{
 asm volatile("vmmcall"
        :"=a" (ret)
        :"0" ($V3_GUARD_EXIT_HCALL_NR));
  return __real_$func(a1,a2,a3,a4,a5,a6,a7,a8);
 asm volatile("vmmcall"
        :"=a" (ret)
        :"0" ($V3_GUARD_ENTER_HCALL_NR));
}

ENDF

  }
}

sub GenerateKernWrapper {
    my $fname=shift;
    my $file=shift;
    my @funcs=@_;

  open(MKFILE,">Makefile");
  open(CFILE, ">wrap_impl.c");

    print MKFILE "### THIS FILE IS AUTOMATICALLY GENERATED ###\n";
    print MKFILE "EXTRA_LDFLAGS = ";

    print $file "/* THIS FILE IS AUTOMATICALLY GENERATED */\n";
    print $file ".text\n.align 4\n";
    print $file map { ".globl __wrap_$_\n" } @funcs;
    print $file map { ".globl __wrap_impl_$_\n" } @funcs;
 
    print CFILE "/* THIS FILE IS AUTOMATICALLY GENERATED */\n";
    print CFILE "#include <linux/kernel.h>\n\n";
    #print CFILE "extern int wrapper_count;\n\n";
    #print CFILE "static int border_count = 0;\n\n";
    print CFILE map { "void __wrap_impl_$_(void);\n\n" } @funcs;

    foreach $func (@funcs) {
#__wrap_$func:
        #nop
        #nop
        #nop
        #nop
        #nop
        #nop
        #nop
        #pushq %rbp
        #movq  %rsp, %rbp
        #pushq %rax
        #pushq %rbx
        #pushq %rcx
        #pushq %rdx
        #pushq %rdi
        #pushq %rsi
        #pushq %r8
        #pushq %r9
        #pushq %r10
        #pushq %r11
        #pushq %r12
        #pushq %r13
        #pushq %r14
        #pushq %r15
        #callq __wrap_impl_$func
        #popq %r15
        #popq %r14
        #popq %r13
        #popq %r12
        #popq %r11 
        #popq %r10
        #popq %r9
        #popq %r8 
        #popq %rsi
        #popq %rdi
        #popq %rdx
        #popq %rcx
        #popq %rbx
        #popq %rax
        #popq %rbp
        #jmp __real_$func
        #nop
        #nop
        #nop
        #nop
        #nop

    print $file <<ENDF;

__wrap_$func:
        popq  %r11
        pushq %rax
        movq  \$$V3_BORDER_OUT_CALL_NR, %rax
        vmmcall
        popq  %rax
        callq __real_$func
        pushq %rax
        movq \$$V3_BORDER_IN_RET_NR, %rax
        vmmcall
        popq  %rax
        pushq %r11
        ret
ENDF

    print MKFILE "--wrap $func ";

    print CFILE <<ENDF;
void __wrap_impl_$func(void) {

        //wrapper_count++;
        //printk("exit: $func\\n");
        //printk("wc: %d (%s)\\n", wrapper_count,__func__);
}

ENDF

  }
  print MKFILE "\n\nldflags-y := \$(EXTRA_LDFLAGS)\n";

# TODO: change this to not be hardcoded!
  my $kerndir = "/v-test/guard_modules/gm_guest/kyle_gl/";
  print MKFILE <<ENDF;
KERNDIR=$kerndir
  
obj-m := $fname\_wrapped.o
    
$fname\_wrapped-objs := $fname.ko $fname\_wrapper.o wrap_impl.o

$fname\_wrapped.ko: $fname.ko $fname\_wrapper.o wrap_impl.o
\tmake -C \$(KERNDIR) M=\$(PWD) modules

clean:
\tmake -C \$(KERNDIR) M=\$(PWD) clean

ENDF

  close(MKFILE);
  close(CFILE);

}