#include <linux/fs.h>
/* #include <asm/current.h> */
#include <asm/uaccess.h>
#include <linux/sched.h>
#include <asm/fcntl.h>
#include <linux/errno.h>
#include <linux/mm.h>
#include <asm-generic/smplock.h>
#include <linux/slab.h>
#include <asm/mman.h>
#include <linux/string.h>

#define NAME_LEN	16	

#define	DEFAULT		0

struct task_struct *target;
	
/*    Making The file_name  */
int make_filename(char *filename, int maxlen){
	char *extend = ".ckpt";
	int i, plen, exlen;
	i=plen=exlen=0;
	
	exlen = strlen(extend)+1;
	plen = strlen(target->comm);
	memcpy(filename, target->comm, plen);	
	
	for(i=0; i < exlen; i++){
		if(plen > maxlen)
			return -1;
		else
			filename[plen++] = extend[i];
	}
	
	return 0;
}

/*  allocate a flag  */

asmlinkage int sys_ckpt(int pid){
	char filename[NAME_LEN];
	size_t n=0;
	struct file *file;
	mm_segment_t	oldfs;
	int vm_area_count=0;
	struct vm_area_struct *tmp;
	struct pt_regs *regs;
	
	ssize_t (*write)(struct file*, const char*, size_t, loff_t*);
	
	 if((target = find_task_by_pid(pid)) == NULL) 
		 return -1;
	/* save the addr_limit of target task */
	/* and set it to KERNEL_DS(0xFFFFFFFF) */
	oldfs = get_fs();
	set_fs(get_ds());

	regs = ((struct pt_regs*)((unsigned long)target + 2*PAGE_SIZE)) -1;

	/* determine filename("process.ckpt") */
	if(make_filename(filename, NAME_LEN)) return -1;

	/* file open */
	file = filp_open(filename, O_CREAT | O_WRONLY | O_TRUNC, 0600);
	if(!file){
		set_fs(oldfs);
		return -1;
	}
	
	/* define write() function */
	if(!(file->f_op) || !(write=file->f_op->write)){
		printk("f_op write error\n");
		goto close_exit;
	}
	
	/* write target task structure */
	n = sizeof(struct task_struct);
	if(write(file, (char*)target, n, &file->f_pos) != n){
		printk("target write error\n");
		goto close_exit;
	}

	/* write the signal_struct */
	n = sizeof(struct signal_struct);
	if(write(file, (char*)target->sig, n, &file->f_pos) != n){
		printk("signal write error\n");
		goto close_exit;
	}

	/* write vm_area_count, that is, the number of vm_area_struct */
	for(tmp=target->mm->mmap; tmp; tmp=tmp->vm_next)
		vm_area_count++;
	n = sizeof(int);
	if(write(file, (char*)&vm_area_count, n, &file->f_pos) != n){
		printk("vm_area_count write error\n");
		goto close_exit;
	}

	/* write the mm_struct */
	n = sizeof(struct mm_struct);
	if(write(file, (char*)target->mm, n, &file->f_pos) != n){
		printk("mm write error\n");
		goto close_exit;
	}
	
	/* write pt_regs */
	n = sizeof(struct pt_regs);
	if(write(file, (char*)regs, n, &file->f_pos) != n){
		printk("regs write error\n");
		goto close_exit;
	}

	/* write each vm_area_struct(segment) of target task */
	for(tmp=target->mm->mmap; tmp; tmp=tmp->vm_next){
		set_fs(KERNEL_DS);

		n = sizeof(struct vm_area_struct);
		if(write(file, (char*)tmp, n, &file->f_pos) != n){
			printk("vm_area_struct write error\n");
			goto close_exit;
		}
		
		set_fs(USER_DS);
		n = tmp->vm_end - tmp->vm_start;
		if(write(file, (char*)tmp->vm_start, n, &file->f_pos) != n){
			printk("pid number : %d\n",pid);
			printk("segment write error\n");
			goto close_exit;
		}
	}

	/* restore addr_limit of target task */
	set_fs(oldfs);
	filp_close(file, target->files);
//	target->should_ckpt = CKPTED;
	return 0;

close_exit:
	set_fs(oldfs);
	filp_close(file, target->files);
	return -1;
}

/* recover implementation */

#define	make_prot(x)	((x) & 0x07)
#define	make_flags(x)	(0x12 | ((x)&0x1900))

extern void *kmalloc(size_t, int);
extern unsigned long do_mmap(struct file*, unsigned long, unsigned long, unsigned long, unsigned long, unsigned long);
void restore_task(struct task_struct *task);
void restore_mm(struct mm_struct *mm);
void restore_sig(struct signal_struct *sig);

asmlinkage int sys_recover(char *name){
	size_t n;
	mm_segment_t	oldfs;
	struct vm_area_struct	tmp;
	struct vm_area_struct	*vma=NULL;
	struct file *file;
	int vm_area_count=0;
	char	*filename;
	unsigned long	old_prot;
	unsigned short	old_flags;
	struct linux_binprm	bprm;
	unsigned long result;
	struct task_struct task;
	struct mm_struct mm;
	struct signal_struct sig;
	struct pt_regs regs;
	struct pt_regs *child;

	ssize_t (*read)(struct file*, char*, size_t, loff_t*);

	/* save the addr_limit of current task
	 * and set it to KERNEL_DS(0xFFFFFFFF) */
	oldfs = get_fs();
	set_fs(get_ds());

	filename = getname(name);
	
	/* file open */
	file = filp_open(filename, O_RDONLY, 0600);
	if(!file){
		set_fs(oldfs);
		return -1;
	}

	/* define read() function */
	if(!(file) || !(read=file->f_op->read)){
		printk("f_op error\n");
		goto error;
	}
	file->f_pos = 0;

	/* read current task structure */
	n = sizeof(struct task_struct);
	if(read(file, (char*)&task, n, &file->f_pos) != n){
		printk("current read error\n");
		goto error;
	}

	/* read the signal_struct */
	n = sizeof(struct signal_struct);
	if(read(file, (char*)&sig, n, &file->f_pos) != n){
		printk("signal read error\n");
		goto error;
	}

	/* read vm_area_count, that is, the number of vm_area_struct */
	n = sizeof(vm_area_count);
	if(read(file, (char*)&vm_area_count, n, &file->f_pos) != n){
		printk("vm_area_count read error\n");
		goto error;
	}

	/* read mm_struct */
	n = sizeof(struct mm_struct);
	if(read(file, (char*)&mm, n, &file->f_pos) != n){
		printk("mm read error\n");
		goto error;
	}

	/* read pt_regs */
	n = sizeof(struct pt_regs);
	if(read(file, (char*)&regs, n, &file->f_pos) != n){
		printk("regs read error\n");
		goto error;
	}

	/* free all the previous segments, mm and so forth... */
	bprm.filename = filename;
	bprm.file = file;
	bprm.e_uid = current->euid;
	bprm.e_gid = current->egid;
	flush_old_exec(&bprm);

	restore_task(&task);
	restore_mm(&mm);
	restore_sig(&sig);
	
	/* read each vm_area_struct(segment) */
	while(vm_area_count--){
		set_fs(KERNEL_DS);
		n = sizeof(struct vm_area_struct);
		if(read(file, (char*)&tmp, n, &file->f_pos) != n){
			printk("vm_area_struct read error\n");
			goto error;
		}
		
		n = tmp.vm_end - tmp.vm_start;
		if(tmp.vm_start != (result=do_mmap(NULL, tmp.vm_start, n, make_prot(tmp.vm_flags), make_flags(tmp.vm_flags), 0))){
			printk("do_mmap error\n");
			goto error;
		}
		
		vma = find_vma(current->mm, tmp.vm_start);
		if(!vma || vma->vm_start > tmp.vm_start){
			printk("find_vma error\n");
			goto error;
		}
		
		old_prot = vma->vm_page_prot.pgprot;
		old_flags = vma->vm_flags;
		
		vma->vm_page_prot.pgprot |= PROT_WRITE;
		vma->vm_flags &= (~VM_DENYWRITE);
		vma->vm_flags |= VM_WRITE;
		
		set_fs(USER_DS);
		if(read(file, (char*)tmp.vm_start, n, &file->f_pos) != n){
			printk("segment read error\n");
			goto error;
		}
		
		vma->vm_page_prot.pgprot = old_prot;
		vma->vm_flags = old_flags;
	}
	
	/* close the file, ".ckpt" */
	filp_close(file, current->files);
	
	child = ((struct pt_regs*)(2*PAGE_SIZE + (unsigned long)current))-1;
	*child = regs;
	child->eax = 0;

	putname(filename);
	return 0;

error:
	set_fs(oldfs);
	filp_close(file, current->files);
	return -1;
}

void restore_task(struct task_struct *task){
	int i;

	current->state = task->state;
	current->counter = task->counter;
	current->nice = task->nice;
	current->blocked = task->blocked;
	current->flags = task->flags;
	for(i=0; i<8; i++) current->thread.debugreg[i] = task->thread.debugreg[i];

	current->exit_code = task->exit_code;
	current->exit_signal = task->exit_signal;
	current->personality = task->personality;

	for(i=0; i<RLIM_NLIMITS; i++) current->rlim[i] = task->rlim[i];
	current->used_math = task->used_math;
	for(i=0; i<16; i++) current->comm[i] = task->comm[i];	

	current->link_count = task->link_count;
}

void restore_mm(struct mm_struct *mm){
	current->mm->start_code = mm->start_code;
	current->mm->end_code = mm->end_code;
	current->mm->start_data = mm->start_data;
	current->mm->end_data = mm->end_data;
	current->mm->start_brk = mm->start_brk;
	current->mm->brk = mm->brk;
	current->mm->start_stack = mm->start_stack;
	current->mm->arg_start = mm->arg_start;
	current->mm->arg_end = mm->arg_end;
	current->mm->env_start = mm->env_start;
	current->mm->env_end = mm->env_end;

	current->mm->rss = 0;
	current->mm->mmap = NULL;
}

void restore_sig(struct signal_struct *sig){
	int i;
	char *src=(char*)sig->action;
	char *dst=(char*)current->sig->action;

	current->sig = kmalloc(sizeof(*current->sig), GFP_KERNEL);
	if (!current->sig) {
		printk("restore_sig: nomem\n");
		return;
	}
	current->sig->count.counter = 1;
	for(i=0; i<sizeof(sig->action); i++)
		dst[i] = src[i];
}