#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/shm.h>
#include <sys/types.h>
#include <sys/mman.h>
#include <errno.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/wait.h>
#include <setjmp.h>

#define WORK_DIR "/mnt/hugepages"
#define MMAP_DIR "/mnt/hugepages/mmap"
#define MMAP_FILE "/mnt/hugepages/mmap/mmap_file"

char buff[2000];

int get_htlb_mem_stat(int *total_pages, int *free_pages, long *page_size)
{
        int fd=0;
        int size1=0;
        char *start=0;

        buff[0] = '\0';
        if((fd=open("/proc/meminfo", O_RDONLY)) < 0)
        {
                printf("Current kernel does not support Huge TLB!\n");
                exit(-1);
        }

        size1 = pread(fd, buff, sizeof(buff), 0);
        if(size1>=sizeof(buff))
        {
                size1 = sizeof(buff)-1;
        }
        if(size1<0)
                size1 = 0;

        buff[size1] = '\0';
	start = buff;
	while(*start) {
		if(isupper(*start))
			*start = tolower(*start);
		start ++;
	}

	start = strstr(buff, "hugepagesize:");
	if(start==NULL)
	{
                printf("Current kernel does not support Huge TLB!\n");
                exit(-1);
        }
        start += sizeof("hugepagesize:");
        if(sscanf(start,"%ld", page_size) != 1)
        {
                printf("Current kernel does not support Huge TLB!\n");
                exit(-1);
        }
        *page_size *= 1024; /*Convert to BYTE from KB*/

        close(fd);

        return 0;
}


int mmap_open_mmap(int *fd, void **addr, size_t length, int prot, int flags)
{
	*fd = open(MMAP_FILE, O_CREAT|O_RDWR, 0755);
	if (*fd < 0) {
		/*perror("Open failed");*/
		exit(errno);
	}
	*addr = mmap(NULL, length, prot, flags, *fd, 0);
	if( *addr == (void *)-1 ) {
		/*perror("mmap failed");*/
		close(*fd);
		*fd = -1;
		return -1;
	}

	return 0;
}

int mmap_unmap(void *start, size_t length)
{
	int     result=0;

	result = munmap(start, length);
	if(result == -1 ) {
		/*perror("mmap_unmap failed");*/
		return -1;
	}

	return 0;
}

int test_mmap_mprotect()
{
	int result = -1;
	void *addr=0;
	long tt001;
	int	fd=-1;

	int	total_pages;
	int	free_pages;
	long	LPAGE_SIZE;

	printf("Test mmap mprotect ...");
	get_htlb_mem_stat(&total_pages, &free_pages, &LPAGE_SIZE);
	result = mount("none", MMAP_DIR, "hugetlbfs", 0, NULL);
	if(result != 0) {
		perror("Mount fail:");
		goto out1;
	}

	result = mmap_open_mmap(&fd, &addr, LPAGE_SIZE*3, PROT_READ|PROT_WRITE, MAP_SHARED);
	if(result < 0)
		goto out2;

	munmap(addr, LPAGE_SIZE*3);
	close(fd);

	result = mmap_open_mmap(&fd, &addr, LPAGE_SIZE*3, PROT_NONE, MAP_SHARED);
	if(result < 0)
		goto out2;

	if(mprotect(addr, LPAGE_SIZE*3, PROT_READ)) {
		printf("fail!\n");
		goto out2;
	}

	tt001 = ((long *)addr)[0];

	if(mprotect(addr+LPAGE_SIZE, LPAGE_SIZE, PROT_READ|PROT_WRITE)) {
		printf("fail!\n");
		goto out2;
	}
	((char *)addr)[LPAGE_SIZE+100] = tt001;

	if(mprotect(addr+LPAGE_SIZE, LPAGE_SIZE, PROT_READ)) {
		printf("fail!\n");
		goto out2;
	}

	printf("pass!\n");

out3:
	if(result < 0)
		mmap_unmap(addr, LPAGE_SIZE*3);
	else
		result = mmap_unmap(addr, LPAGE_SIZE*3);
out2:
	close(fd);
	if(result < 0)
		umount(MMAP_DIR);
	else
		result = umount(MMAP_DIR);
out1:
	return result;
}

void init()
{
	mkdir(WORK_DIR, 0777);
	mkdir(MMAP_DIR, 0755);
}

void uninit()
{
	rmdir(MMAP_DIR);
	rmdir(WORK_DIR);
}

int main(int argc, char * argv)
{
	int	total_pages=0;
	int	free_pages=0;
	int	result=0;

	init();

	test_mmap_mprotect();

	uninit();
	return 0;
}