/*
 * Copyright (c) 2018-2019 - Zephyr Software LLC
 *
 * This file may be used and modified for non-commercial purposes as long as
 * all copyright, permission, and nonwarranty notices are preserved.
 * Redistribution is prohibited without prior written consent from Zephyr
 * Software.
 *
 * Please contact the authors for restrictions applying to commercial use.
 *
 * THIS SOURCE IS PROVIDED "AS IS" AND WITHOUT ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED WARRANTIES OF
 * MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 *
 * Author: Zephyr Software
 * e-mail: jwd@zephyr-software.com
 * URL   : http://www.zephyr-software.com/
 *
 */

#define _GNU_SOURCE 
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <link.h> 
#include <elf.h>
#include <stdlib.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/shm.h>
#include <sys/mman.h>
#include <sys/wait.h>
#include <stdio.h>
#include <string.h>
#include <assert.h>

// these must match config.h in AFL source code
#define SHM_ENV_VAR         "__AFL_SHM_ID"
#define FORKSRV_FD          198
#define MAP_SIZE_POW2       16
#define MAP_SIZE            (1 << MAP_SIZE_POW2)
typedef uint8_t  u8;
typedef int32_t  s32;

// externally visible so that Zipr transformations can access directly
u8* zafl_trace_map = NULL;
unsigned int zafl_prev_id = 123;   
unsigned int zafl_context = 45678;

#define AFL_VARS_OFFSET  (MAP_SIZE + 4096)
#define AFL_VARS_STORAGE_SIZE  1024
#define AFL_VAR_PREVID_OFFSET  32
#define AFL_VAR_CONTEXT_OFFSET 64
s32* zafl_var_map = NULL;

static s32 shm_id;
static int __afl_temp_data;
static pid_t __afl_fork_pid;
static int debug = 0;

static int tracemap_shared_memory_is_setup = 0;
static int aflvars_shared_memory_is_setup = 0;

static void zafl_setupAflTracemapSharedMemory();

#define PRINT_ERROR(string) {int x=write(2, string, my_strlen(string));}
#define PRINT_DEBUG(string) if (debug) {int x=write(1, string, my_strlen(string));}

#ifdef ZAFL_AUTO_INIT_FORK_SERVER
void __attribute__((constructor)) zafl_initAflForkServer();
#else
void __attribute__((constructor)) zafl_setupAflTracemapSharedMemory();
#endif

void __attribute__((destructor)) zafl_dumpTracemap();

static size_t my_strlen(const char *s)
{
	const char *p = s;
	while (*s) ++s;
	return s - p;
}

static char* my_ull_toa(unsigned long long value, char* result, const int base)
{
        if (base < 2 || base > 36)
        {
                *result = '\0';
                return result;
        }
        char* ptr = result, *ptr1 = result, tmp_char;
	unsigned long long tmp_value;

	do 
	{
		tmp_value = value;
		value /= base;
		*ptr++ = "zyxwvutsrqponmlkjihgfedcba9876543210123456789abcdefghijklmnopqrstuvwxyz" [35 + (tmp_value - value * base)];
	} while ( value );

	*ptr-- = '\0';
	while(ptr1 < ptr) 
	{
		tmp_char = *ptr;
		*ptr--= *ptr1;
		*ptr1++ = tmp_char;
	}
	return result;
}

static void zafl_setupAflVariablesMemory(void *vars_addr)
{
	if (aflvars_shared_memory_is_setup)
		return;

	zafl_var_map = (s32*) mmap(vars_addr, AFL_VARS_STORAGE_SIZE, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
	if (zafl_var_map == MAP_FAILED)
	{
		perror("mmap of afl variables failed");
		return;
	}
	else if (zafl_var_map != vars_addr)
	{
		perror("mmap of afl variables failed -- unexpected address");
		return;
	}

	// initialize map
	s32* prev_id = vars_addr + AFL_VAR_PREVID_OFFSET;
	s32* context_id = vars_addr + AFL_VAR_CONTEXT_OFFSET;

	// any values will do here
	*prev_id = 0x234;
	*context_id = 0xabc;

	aflvars_shared_memory_is_setup = 1;

	if (debug)
	{
		char buf[1024];

		my_ull_toa((unsigned long long)vars_addr, buf, 16);
		PRINT_DEBUG("&var_map: 0x"); PRINT_DEBUG(buf); PRINT_DEBUG("\n");
		printf("var_map: 0x%p\n", (void*) vars_addr);

		my_ull_toa((unsigned long long)prev_id, buf, 16);
		PRINT_DEBUG("&prev_id: 0x"); PRINT_DEBUG(buf); PRINT_DEBUG("\n");

		my_ull_toa((unsigned long long)context_id, buf, 16);
		PRINT_DEBUG("&context_id: 0x"); PRINT_DEBUG(buf); PRINT_DEBUG("\n");
	}
}

static int find_map_addr(struct dl_phdr_info *p_info, size_t p_size, void *p_data) 
{
#define close_and_err     \
	{                 \
		close(fd);\
		return 0; \
	}

	void** trace_map_fixed_addr = (void**)p_data;
	*trace_map_fixed_addr = NULL;

	if(p_info->dlpi_name == NULL) return 0;

	PRINT_DEBUG ("Name: ");
	PRINT_DEBUG(p_info->dlpi_name);
	PRINT_DEBUG ("\n");

	const char* the_name = (p_info->dlpi_name[0] == '\0')  ? 
		"/proc/self/exe" : 
		p_info->dlpi_name;


	int fd = open(the_name, O_RDONLY);

	// no such file
	if(fd == -1 ) return 0;

	Elf64_Ehdr file_header =  {};

	const int hdr_size = read(fd, &file_header, sizeof(file_header));	

	// file too small
	if(hdr_size!=sizeof(file_header)) close_and_err;

	// magic number check
	if(file_header.e_ident[0] != ELFMAG0) close_and_err;
	if(file_header.e_ident[1] != ELFMAG1) close_and_err;
	if(file_header.e_ident[2] != ELFMAG2) close_and_err;
	if(file_header.e_ident[3] != ELFMAG3) close_and_err;

	// sanity check section header entry size
	if(file_header.e_shentsize != sizeof(Elf64_Shdr)) close_and_err;

	// sanity check string table index 
	if(file_header.e_shstrndx >= file_header.e_shnum) close_and_err;

	// read section headers
	Elf64_Shdr section_headers[file_header.e_shnum];
	memset(section_headers,0,sizeof(section_headers));

	const int section_header_offset = lseek(fd, file_header.e_shoff, SEEK_SET);

	// sanity check successful seek
	if(section_header_offset != file_header.e_shoff) close_and_err;
	
	const int sec_hdr_size = read(fd, section_headers, sizeof(section_headers));

	// sanity check we read them all
	if(sec_hdr_size != sizeof(section_headers)) close_and_err;

	// sanity check string table type
	if(section_headers[file_header.e_shstrndx].sh_type != SHT_STRTAB) close_and_err;

	// seek to the string section
	int string_sec_offset = lseek(fd, section_headers[file_header.e_shstrndx].sh_offset, SEEK_SET);

	// declare the string table
	char string_table[section_headers[file_header.e_shstrndx].sh_size+1]; // definitely null terminate
	memset(string_table,0,sizeof(string_table));

	// read the string table
	const int string_table_size = read(fd, string_table, section_headers[file_header.e_shstrndx].sh_size);

	// sanity check read
	if(string_table_size != section_headers[file_header.e_shstrndx].sh_size) close_and_err;

	const char* integration_name = "libZaflIntegration";

	int i=0;
	for(i=0; i < file_header.e_shnum; i++)
	{
		const uint64_t sec_name_offset = section_headers[i].sh_name;

		// sanity check name -- bail from function if corrupt
		if(sec_name_offset >= section_headers[file_header.e_shstrndx].sh_size) close_and_err;

		// check if right section, try next if not.
		if(strcmp(integration_name, string_table+sec_name_offset) != 0) continue; 

		// sanity check size
		if(section_headers[i].sh_size != sizeof(void*)) close_and_err;

		// found!  seek to section's data.
		int integration_table_pos = lseek(fd, section_headers[i].sh_offset, SEEK_SET);
		if(integration_table_pos != section_headers[i].sh_offset) close_and_err;

		// now read
		char integration_table[section_headers[i].sh_size];
		memset(integration_table,0,sizeof(integration_table));

		// read the data table
		const int integration_table_size = read(fd, integration_table, section_headers[i].sh_size);

		// sanity check read, bail on function if corrupt.
		if(integration_table_size != section_headers[i].sh_size) close_and_err;

		// extract the value
		*trace_map_fixed_addr = *(void**)(integration_table);
		close(fd);
		return 1; 
	}
	close_and_err;
}




// always setup a trace map so that an instrumented application will run
// even if not running under AFL

static void zafl_setupAflTracemapSharedMemory()
{
	if (getenv("ZAFL_DEBUG")) debug = 1;

	if (tracemap_shared_memory_is_setup)
		return;

	zafl_prev_id = 0;
	zafl_trace_map = NULL;

	void* trace_map_fixed_addr = (void*) NULL;
	dl_iterate_phdr(find_map_addr, (void*)&trace_map_fixed_addr);
	const int do_fixed_addr_optimization = (trace_map_fixed_addr != NULL);

	void* shmat_addr = trace_map_fixed_addr;

	if (do_fixed_addr_optimization)
	{
		// setup the area for storing prev_id and context_id 4K away from the end of the trace map
		void* vars_addr = (void*)(shmat_addr + AFL_VARS_OFFSET);
		zafl_setupAflVariablesMemory(vars_addr);
	}

	char *shm_env_var = getenv(SHM_ENV_VAR);
	if(!shm_env_var) 
	{
		PRINT_DEBUG("Error getting shm environment variable - fake allocate AFL trace map\n");

		// fake allocate until someone calls zafl_initAflForkServer()
		zafl_trace_map = (u8*)mmap(shmat_addr, MAP_SIZE, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
		if(zafl_trace_map == MAP_FAILED)
		{
			perror("Cannot allocate fake map\n");
			exit(1);
		}
		else
		{
			PRINT_DEBUG("Success at mmap!\n");
			if(shmat_addr!=0)
				assert(zafl_trace_map==shmat_addr);
		}

		tracemap_shared_memory_is_setup = 1;
		return;
	}

	shm_id = atoi(shm_env_var);
	zafl_trace_map = (u8*)shmat(shm_id, shmat_addr, 0);
	if(zafl_trace_map == (u8*)-1) 
	{
		perror("shmat_attach");
		return;
	}
	PRINT_DEBUG("libzafl: shared memory segment is setup\n");
	tracemap_shared_memory_is_setup = 1;
}

void zafl_initAflForkServer()
{
	if (getenv("ZAFL_DEBUG")) debug = 1;

	PRINT_DEBUG("libautozafl: auto-initialize fork server\n");

	static int fork_server_initialized = 0;

	if (fork_server_initialized) return;

	zafl_setupAflTracemapSharedMemory();

	if (!zafl_trace_map) 
	{
		zafl_trace_map = (u8*)malloc(MAP_SIZE);
	}

	// force map entry to get past afl sanity checks
	if (debug)
		zafl_trace_map[MAP_SIZE-1] = 1;

	int n = write(FORKSRV_FD+1, &__afl_temp_data,4);
	if(n != 4) 
	{
		return;
	}

	fork_server_initialized = 1;

	// afl forkserver handshake
	while(1) 
	{
		n = read(FORKSRV_FD,&__afl_temp_data,4);
		if(n != 4) 
		{
			perror("Error reading fork server\n");
			return;
		}

		__afl_fork_pid = fork();
		if(__afl_fork_pid < 0) 
		{
		    perror("Error on fork()\n");
		    return;
		}
		if(__afl_fork_pid == 0) 
		{
		    // child
		    close(FORKSRV_FD);
		    close(FORKSRV_FD+1);
		    break;
		} 
		else 
		{
		    // parent
		    n = write(FORKSRV_FD+1,&__afl_fork_pid, 4);
		    pid_t temp_pid = waitpid(__afl_fork_pid,&__afl_temp_data,2);
		    if(temp_pid == 0) 
		    {
			return;
		    }
		    n = write(FORKSRV_FD+1,&__afl_temp_data,4);
		}
	}
}

//
// debugging code below
// 
void zafl_bbInstrument(unsigned short id) 
{
	zafl_trace_map[zafl_prev_id ^ id]++;
	zafl_prev_id = id >> 1;
}

void zafl_dumpTracemap()
{
	int i;
	if (!debug) return;
	if (!zafl_trace_map) return;

	PRINT_DEBUG("\n=== DEBUG output -- trace map\n");
	for (i = 0; i < MAP_SIZE; ++i)
	{
		char buf[1024];
		if (zafl_trace_map[i]!=0)
		{
			PRINT_DEBUG(my_ull_toa(i, buf, 16));
			PRINT_DEBUG(":")
			PRINT_DEBUG(my_ull_toa(zafl_trace_map[i], buf, 10));
			PRINT_DEBUG("\n")
		}
	}
}