Part of snapshot taking code in C

GetPhysicalMemoryRanges.c

/**
 * @brief Returns an array of physical memory address ranges on the system.
 *
 * @param This - The pointer to the platform API interface.
 * @param RangeCount - The pointer to receive the number of entries in the returned
 *      pointer on success.
 * @return The pointer to the array of ranges on success. This must be freed with
 *      Freed with FREE_CONTIGUOUS_PAGES. On failure, NULL.
 */
STATIC
SNAPSHOT_MEMORY_RANGE*
MVAPI
GetPhysicalMemoryRanges (
    IN CONST PLATFORM_API_INTERFACE* This,
    OUT UINT64* RangeCount
    )
{
    EFI_STATUS status;
    UINT64 memoryMapSize;
    EFI_MEMORY_DESCRIPTOR* memoryMaps;
    UINT64 descriptorSize;
    SNAPSHOT_MEMORY_RANGE* ranges;
    UINT64 rangeCount;
    //
    // Get the memory map using GetMemoryMap().
    //
    memoryMapSize = 0;
    status = gBS->GetMemoryMap(&memoryMapSize,
                               NULL,
                               NULL,
                               NULL,
                               NULL);
    memoryMaps = This->AllocateContiguousPages(This,
                                               EFI_SIZE_TO_PAGES(memoryMapSize),
                                               AllocationTypeData);
    if (memoryMaps == NULL)
    {
        return NULL;
    }
    memoryMapSize = EFI_PAGES_TO_SIZE(EFI_SIZE_TO_PAGES(memoryMapSize));
    status = gBS->GetMemoryMap(&memoryMapSize,
                               memoryMaps,
                               NULL,
                               &descriptorSize,
                               NULL);
    ASSERT_EFI_ERROR(status);
    //
    // For each entry returned from GetMemoryMap(), save when the entry indicates
    // DRAM backed memory, ie, exclude MMIO regions.
    //
    rangeCount = 0;
    ranges = (SNAPSHOT_MEMORY_RANGE*)memoryMaps;
    for (UINT64 i = 0; i < memoryMapSize / descriptorSize; ++i)
    {
        CONST EFI_MEMORY_DESCRIPTOR* memoryMap;
        memoryMap = (EFI_MEMORY_DESCRIPTOR*)((UINT8*)memoryMaps + descriptorSize * i);
        if ((memoryMap->Type < EfiLoaderCode) ||
            (memoryMap->Type > EfiACPIMemoryNVS))
        {
            continue;
        }
        ranges[rangeCount].PageBase = memoryMap->PhysicalStart;
        ranges[rangeCount].PageCount = memoryMap->NumberOfPages;
        rangeCount++;
    }
    *RangeCount = rangeCount;
    return ranges;
}

Ide.c

/**
 * @file Ide.c
 * @author Satoshi Tanda ([email protected])
 * @brief Write data into the IDE storage device.
 * @date 2022-01-16
 *
 * @copyright Copyright (c) 2022, Satoshi Tanda. All rights reserved.
 *
 */
#include "Ide.h"
#include "HostUtils.h"

//
// See
// https://wiki.osdev.org/PCI_IDE_Controller#Read.2FWrite_From_ATA_Drive
// https://wiki.osdev.org/ATA_PIO_Mode
//
#define IDE_SECTOR_SIZE             512

#define IDE_IO_ADDR_BASE            0x1f0
#define IDE_DATA                    (IDE_IO_ADDR_BASE + 0)
#define IDE_SECCOUNT0               (IDE_IO_ADDR_BASE + 2)
#define IDE_SECCOUNT1               (IDE_IO_ADDR_BASE + 2)
#define IDE_LBA0                    (IDE_IO_ADDR_BASE + 3)
#define IDE_LBA3                    (IDE_IO_ADDR_BASE + 3)
#define IDE_LBA1                    (IDE_IO_ADDR_BASE + 4)
#define IDE_LBA4                    (IDE_IO_ADDR_BASE + 4)
#define IDE_LBA2                    (IDE_IO_ADDR_BASE + 5)
#define IDE_LBA5                    (IDE_IO_ADDR_BASE + 5)
#define IDE_HDDEVSEL                (IDE_IO_ADDR_BASE + 6)
#define IDE_COMMAND                 (IDE_IO_ADDR_BASE + 7)
#define IDE_COMMAND_WRITE_PIO       0x30
#define IDE_COMMAND_CACHE_FLUSH     0xe7
#define IDE_STATUS                  (IDE_IO_ADDR_BASE + 7)
#define IDE_STATUS_BUSY             (1 << 7)

#define IDE_CONTROL                 0x3f6
#define IDE_CONTROL_NIEN            (1 << 1)
#define IDE_CONTROL_SRST            (1 << 2)

/**
 * @brief Wait until the IDE complete its processing and ready for another request.
 */
STATIC
VOID
IdeWait (
    VOID
    )
{
    while (IoRead8(IDE_STATUS) & IDE_STATUS_BUSY);
}

/**
 * @brief Write data into an IDE-related IO port and wait for IDE to complete it.
 *
 * @param Port - The IO port address to write.
 * @param Value - The value to write to the IO port.
 */
STATIC
VOID
IdeWriteRegister (
    IN UINT64 Port,
    IN UINT8 Value
    )
{
    IoWrite8(Port, Value);
    IdeWait();
}

STATIC BOOLEAN mIdeHasBeenReset = FALSE;

/**
 * @brief Writes one sector (512 bytes) at the specified LBA of the master IDE device.
 *
 * @param LogicalBlockAddress - The logical block address to start writing.
 * @param Data - The pointer to the data to write.
 * @param DataSize - The size of the data to write in bytes.
 */
STATIC
VOID
IdeWriteDisk (
    UINT64 LogicalBlockAddress,
    CONST UINT8* Data,
    UINT64 DataSize
    )
{
    MV_HOST_ASSERT((DataSize % 2) == 0);

    //
    // Reset the device on the first use, or when it is left as busy.
    //
    if ((mIdeHasBeenReset == FALSE) ||
        (IoRead8(IDE_STATUS) & IDE_STATUS_BUSY))
    {
        IoWrite8(IDE_CONTROL, IDE_CONTROL_SRST);
        IdeWriteRegister(IDE_CONTROL, 0);
        mIdeHasBeenReset = TRUE;
    }

    //
    // Disable interrupts.
    //
    IdeWriteRegister(IDE_CONTROL, IDE_CONTROL_NIEN);

    //
    // Select master drive with LBA addressing mode.
    //
    IdeWriteRegister(IDE_HDDEVSEL, 0xe0);

    //
    // Set up to write one sector with the LBA specified.
    //
    IdeWriteRegister(IDE_SECCOUNT1, 0);
    IdeWriteRegister(IDE_LBA5, (LogicalBlockAddress >> (8 * 5)) & 0xff);
    IdeWriteRegister(IDE_LBA4, (LogicalBlockAddress >> (8 * 4)) & 0xff);
    IdeWriteRegister(IDE_LBA3, (LogicalBlockAddress >> (8 * 3)) & 0xff);
    IdeWriteRegister(IDE_SECCOUNT0, 1);
    IdeWriteRegister(IDE_LBA2, (LogicalBlockAddress >> (8 * 2)) & 0xff);
    IdeWriteRegister(IDE_LBA1, (LogicalBlockAddress >> (8 * 1)) & 0xff);
    IdeWriteRegister(IDE_LBA0, (LogicalBlockAddress >> (8 * 0)) & 0xff);

    //
    // Send the 28-bit LBA PIO write request command and write data onto the disk.
    // Note that the large IDE devices cannot work with this. Change this to
    // IDE_COMMAND_WRITE_PIO_EX and IDE_COMMAND_CACHE_FLUSH_EX if desired.
    //
    IdeWriteRegister(IDE_COMMAND, IDE_COMMAND_WRITE_PIO);
    for (UINT64 i = 0; i < DataSize; i += sizeof(UINT16))
    {
        IoWrite16(IDE_DATA, *(UINT16*)(Data + i));
    }
    IdeWriteRegister(IDE_COMMAND, IDE_COMMAND_CACHE_FLUSH);
}

VOID
IdeWrite4Kb (
    IN UINT64 Offset,
    IN CONST UINT8* Data
    )
{
    //
    // Write sectors (512 bytes for each) one by one while convert the offset to
    // the LBA.
    //
    for (UINT64 i = 0; i < SIZE_4KB / IDE_SECTOR_SIZE; i++)
    {
        IdeWriteDisk(Offset / IDE_SECTOR_SIZE + i,
                     Data + IDE_SECTOR_SIZE * i,
                     IDE_SECTOR_SIZE);
    }
}

Ide.h

/**
 * @file Ide.h
 * @author Satoshi Tanda ([email protected])
 * @brief Write data into the IDE storage device.
 * @date 2022-01-16
 *
 * @copyright Copyright (c) 2022, Satoshi Tanda. All rights reserved.
 *
 */
#pragma once
#include "Common/Common.h"

/**
 * @brief Writes 4KB of data onto at the specified offset of the IDE master device.
 *
 * @param Offset - The offset in the device to start writing.
 * @param Data - The pointer to 4KB data to write to the device.
 */
VOID
IdeWrite4Kb (
    IN UINT64 Offset,
    IN CONST UINT8* Data
    );

Snapshot.c

/**
 * @file Snapshot.c
 * @author Satoshi Tanda ([email protected])
 * @brief Takes the snapshot of the guest onto a disk.
 * @date 2022-01-16
 *
 * @copyright Copyright (c) 2022, Satoshi Tanda. All rights reserved.
 *
 */
#include "Snapshot.h"
#include "Ide.h"
#include "HostUtils.h"
#include "Common/VtxUtils.h"

CONST SNAPSHOT_MEMORY_RANGE* gMemoryRanges;
UINT64 gRangeCount;

VOID
TakeSnapshot (
    IN CONST VMEXIT_CONTEXT* VmExitContext
    )
{
    SNAPSHOT_HEADER snapshotHeader;
    UINT64 offset;
    UINT64 pageBase;

    //
    // The first 4KB of the snapshot file is a header containing memory ranges
    // and register values. Capture them and write it into the disk.
    //
    offset = 0;
    ZeroMem(&snapshotHeader, sizeof(snapshotHeader));
    snapshotHeader.Magic = SNAPSHOT_SIGNATURE;
    CopyMem(&snapshotHeader.MemoryRanges[0], gMemoryRanges, sizeof(*gMemoryRanges) * gRangeCount);
    snapshotHeader.Registers.Gdtr.Limit = (UINT16)VmxRead(VMCS_GUEST_GDTR_LIMIT);
    snapshotHeader.Registers.Gdtr.Base = VmxRead(VMCS_GUEST_GDTR_BASE);
    snapshotHeader.Registers.Idtr.Limit = (UINT16)VmxRead(VMCS_GUEST_IDTR_LIMIT);
    snapshotHeader.Registers.Idtr.Base = VmxRead(VMCS_GUEST_IDTR_BASE);
    snapshotHeader.Registers.Es = (UINT16)VmxRead(VMCS_GUEST_ES_SELECTOR);
    snapshotHeader.Registers.Cs = (UINT16)VmxRead(VMCS_GUEST_CS_SELECTOR);
    snapshotHeader.Registers.Ss = (UINT16)VmxRead(VMCS_GUEST_SS_SELECTOR);
    snapshotHeader.Registers.Ds = (UINT16)VmxRead(VMCS_GUEST_DS_SELECTOR);
    snapshotHeader.Registers.Fs = (UINT16)VmxRead(VMCS_GUEST_FS_SELECTOR);
    snapshotHeader.Registers.Gs = (UINT16)VmxRead(VMCS_GUEST_GS_SELECTOR);
    snapshotHeader.Registers.Ldtr = (UINT16)VmxRead(VMCS_GUEST_LDTR_SELECTOR);
    snapshotHeader.Registers.Tr = (UINT16)VmxRead(VMCS_GUEST_TR_SELECTOR);
    snapshotHeader.Registers.Efer = VmxRead(VMCS_GUEST_EFER);
    snapshotHeader.Registers.SysenterCs = VmxRead(VMCS_GUEST_SYSENTER_CS);
    snapshotHeader.Registers.Cr0 = VmxRead(VMCS_GUEST_CR0);
    snapshotHeader.Registers.Cr3 = VmxRead(VMCS_GUEST_CR3);
    snapshotHeader.Registers.Cr4 = VmxRead(VMCS_GUEST_CR4);
    snapshotHeader.Registers.FsBase = VmxRead(VMCS_GUEST_FS_BASE);
    snapshotHeader.Registers.GsBase = VmxRead(VMCS_GUEST_GS_BASE);
    snapshotHeader.Registers.LdtrBase = VmxRead(VMCS_GUEST_LDTR_BASE);
    snapshotHeader.Registers.TrBase = VmxRead(VMCS_GUEST_TR_BASE);
    snapshotHeader.Registers.Rsp = VmxRead(VMCS_GUEST_RSP);
    snapshotHeader.Registers.Rip = VmxRead(VMCS_GUEST_RIP);
    snapshotHeader.Registers.Rflags = VmxRead(VMCS_GUEST_RFLAGS);
    snapshotHeader.Registers.SysenterEsp = VmxRead(VMCS_GUEST_SYSENTER_ESP);
    snapshotHeader.Registers.SysenterEip = VmxRead(VMCS_GUEST_SYSENTER_EIP);
    snapshotHeader.Registers.Rax = VmExitContext->Guest.StackSaved->Rax;
    snapshotHeader.Registers.Rbx = VmExitContext->Guest.StackSaved->Rbx;
    snapshotHeader.Registers.Rcx = VmExitContext->Guest.StackSaved->Rcx;
    snapshotHeader.Registers.Rdx = VmExitContext->Guest.StackSaved->Rdx;
    snapshotHeader.Registers.Rdi = VmExitContext->Guest.StackSaved->Rdi;
    snapshotHeader.Registers.Rsi = VmExitContext->Guest.StackSaved->Rsi;
    snapshotHeader.Registers.Rbp = VmExitContext->Guest.StackSaved->Rbp;
    snapshotHeader.Registers.R8 = VmExitContext->Guest.StackSaved->R8;
    snapshotHeader.Registers.R9 = VmExitContext->Guest.StackSaved->R9;
    snapshotHeader.Registers.R10 = VmExitContext->Guest.StackSaved->R10;
    snapshotHeader.Registers.R11 = VmExitContext->Guest.StackSaved->R11;
    snapshotHeader.Registers.R12 = VmExitContext->Guest.StackSaved->R12;
    snapshotHeader.Registers.R13 = VmExitContext->Guest.StackSaved->R13;
    snapshotHeader.Registers.R14 = VmExitContext->Guest.StackSaved->R14;
    snapshotHeader.Registers.R15 = VmExitContext->Guest.StackSaved->R15;
    IdeWrite4Kb(offset, (void*)&snapshotHeader);
    offset += sizeof(snapshotHeader);

    //
    // According with the memory range information, write the contents of memory
    // for each 4KB onto the disk. This is not limited to the guest memory and
    // does include hypervisor-owned memory which is unused for the purpose of
    // fuzzing.
    //
    pageBase = 0;
    for (UINT64 i = 0; i < gRangeCount; ++i)
    {
        CONST SNAPSHOT_MEMORY_RANGE* range;

        range = &gMemoryRanges[i];
        MV_HOST_INFO("%16llx - %16llx : Processing %llx pages",
                     range->PageBase,
                     range->PageBase + range->PageCount * SIZE_4KB,
                     range->PageCount);

        for (UINT64 pageCount = 0; pageCount < range->PageCount; ++pageCount)
        {
            if ((pageCount != 0) && (pageCount % 1000) == 0)
            {
                MV_HOST_INFO("... done with %llu / %llu pages",
                             pageCount,
                             range->PageCount);
            }

            pageBase = range->PageBase + (SIZE_4KB * pageCount);
            IdeWrite4Kb(pageBase + offset, (const void*)pageBase);
        }
    }

    //
    // Write the end marker for the ease of reviewing the contents of the snapshot
    // file.
    //
    IdeWrite4Kb(pageBase + offset + SIZE_4KB, (const void*)"END_OF_FILE");
}

Snapshot.h

/**
 * @file Snapshot.h
 * @author Satoshi Tanda ([email protected])
 * @brief Takes the snapshot of the guest onto a disk.
 * @date 2022-01-16
 *
 * @copyright Copyright (c) 2022, Satoshi Tanda. All rights reserved.
 *
 */
#pragma once
#include "Common/Common.h"
#include "VmExitContext.h"

/**
 * @brief Writes the snapshot at the offset 0 of the IDE master device.
 *
 * @param VmExitContext - The pointer to the VM-exit context.
 */
VOID
TakeSnapshot (
    IN CONST VMEXIT_CONTEXT* VmExitContext
    );

structures.h

//
// Snapshot related definitions
//
#define SNAPSHOT_SIGNATURE  0x544F485350414E53       // 'SNAPSHOT'
#define SNAPSHOT_MAX_MEMORY_RANGE_COUNT 47  // arbitrary
typedef struct
{
    IA32_DESCRIPTOR Gdtr;   // +0x0
    UINT8 Reserved1[0x10 - sizeof(IA32_DESCRIPTOR)];
    IA32_DESCRIPTOR Idtr;   // +0x10
    UINT8 Reserved2[0x10 - sizeof(IA32_DESCRIPTOR)];
    UINT16 Es;      // +0x20
    UINT16 Cs;
    UINT16 Ss;
    UINT16 Ds;
    UINT16 Fs;
    UINT16 Gs;
    UINT16 Ldtr;
    UINT16 Tr;
    UINT64 Efer;    // +0x30
    UINT64 SysenterCs;
    UINT64 Cr0;     // +0x40
    UINT64 Cr3;
    UINT64 Cr4;     // +0x50
    UINT64 FsBase;
    UINT64 GsBase;  // +0x60
    UINT64 LdtrBase;
    UINT64 TrBase;  // +0x70
    UINT64 Rsp;
    UINT64 Rip;     // +0x80
    UINT64 Rflags;
    UINT64 SysenterEsp; // +0x90
    UINT64 SysenterEip;
    UINT64 Rax;     // +0xa0
    UINT64 Rbx;
    UINT64 Rcx;     // +0xb0
    UINT64 Rdx;
    UINT64 Rdi;     // +0xc0
    UINT64 Rsi;
    UINT64 Rbp;     // +0xd0
    UINT64 R8;
    UINT64 R9;      // +0xe0
    UINT64 R10;
    UINT64 R11;     // +0xf0
    UINT64 R12;
    UINT64 R13;     // +0x100
    UINT64 R14;
    UINT64 R15;     // +0x110
} SNAPSHOT_REGISTERS;

typedef struct
{
    UINT64 Magic;
    UINT64 Reserved1;
    SNAPSHOT_MEMORY_RANGE MemoryRanges[SNAPSHOT_MAX_MEMORY_RANGE_COUNT];
    SNAPSHOT_REGISTERS Registers;
    UINT8 Reserved2[SIZE_4KB -
                    sizeof(UINT64) * 2 -
                    sizeof(SNAPSHOT_MEMORY_RANGE) * SNAPSHOT_MAX_MEMORY_RANGE_COUNT -
                    sizeof(SNAPSHOT_REGISTERS)];
} SNAPSHOT_HEADER;
STATIC_ASSERT(sizeof(SNAPSHOT_HEADER) == SIZE_4KB, "Must be 4KB");

typedef struct
{
    UINT64 PageBase;
    UINT64 PageCount;
} SNAPSHOT_MEMORY_RANGE;