OSDev.org

Code: Select all

	/* allocating the buffer for the memory map */
#define STUPID_EXTRA_BEHIND_DESCRIPTOR	8	/* FW is known to put extra garbage behind the descr. */

	DescriptorSize = (sizeof(EFI_MEMORY_DESCRIPTOR) + STUPID_EXTRA_BEHIND_DESCRIPTOR);
	MemoryMapSize = DescriptorSize << 7;	/* let it be 128 entries first */
	AddendumShift = 3;
A:	Status = pbs->AllocatePool(EfiLoaderData, MemoryMapSize, &MemoryMap);
	if(Status != EFI_SUCCESS){
		pout->OutputString(pout, L"ERROR:AllocatePool()\r\n");
		MemoryMap = NULL;
		goto YYY;
	}

	/* getting the memory map */
	Status = pbs->GetMemoryMap(&MemoryMapSize, MemoryMap, &MapKey, &DescriptorSize, &DescriptorVersion);
	if(Status == EFI_SUCCESS)goto B;
	else if(Status ==  EFI_BUFFER_TOO_SMALL){
		if(AddendumShift < 10){
			AddendumShift++;
		}else{
			pout->OutputString(pout, L"ERROR:GetMemoryMap() too many attempts to guess the buffer size\r\n");
			goto YYY;
		}

		pbs->FreePool(MemoryMap);	/* freeing the small buffer */

		if(DescriptorSize){
			MemoryMapSize += DescriptorSize << AddendumShift;
		}else{
			MemoryMapSize += (sizeof(EFI_MEMORY_DESCRIPTOR) + STUPID_EXTRA_BEHIND_DESCRIPTOR) << AddendumShift;
		}
		goto A;	/* retrying */
	}else{
		pout->OutputString(pout, L"ERROR:GetMemoryMap() failed\r\n");
		goto YYY;
	}

B: /* doing further work */

/* we return to the FW, but the loader doesn't have to */
YYY:	if(MemoryMap)pbs->FreePool(MemoryMap);

Code: Select all

void EfiBoot::Exit(MemoryMap& memoryMap)
{
    UINTN size = 0;
    EFI_MEMORY_DESCRIPTOR* descriptors = nullptr;
    UINTN memoryMapKey = 0;
    UINTN descriptorSize = 0;
    UINT32 descriptorVersion = 0;

    // 1) Retrieve the memory map from the firmware
    EFI_STATUS status;

    std::vector<char> buffer;
    while ((status = g_efiBootServices->GetMemoryMap(&size, descriptors, &memoryMapKey, &descriptorSize, &descriptorVersion)) == EFI_BUFFER_TOO_SMALL)
    {
        // Add some extra space. There are few reasons for this:
        // a) Allocating memory for the buffer can increase the size of the memory map itself.
        //    Adding extra space will prevent an infinite loop.
        // b) We want to try to prevent a "partial shutdown" when calling ExitBootServices().
        //    See comment below about what a "partial shutdown" is.
        // c) If a "partial shutdown" does happen, we won't be able to allocate more memory!
        //    Having some extra space now should mitigate the issue.
        size += descriptorSize * 10;

        buffer.resize(size);
        descriptors = (EFI_MEMORY_DESCRIPTOR*)buffer.data();
    }

    if (EFI_ERROR(status))
    {
        Fatal("Failed to retrieve the EFI memory map: %p\n", status);
    }

    // 2) Exit boot services - it is possible for the firmware to modify the memory map
    // during a call to ExitBootServices(). A so-called "partial shutdown".
    // When that happens, ExitBootServices() will return EFI_INVALID_PARAMETER.
    while ((status = g_efiBootServices->ExitBootServices(g_efiImage, memoryMapKey)) == EFI_INVALID_PARAMETER)
    {
        // Memory map changed during ExitBootServices(), the only APIs we are allowed to
        // call at this point are GetMemoryMap() and ExitBootServices().
        size = buffer.size(); // Probably not needed, but let's play safe since EFI could change that value behind our back (you never know!)
        status = g_efiBootServices->GetMemoryMap(&size, descriptors, &memoryMapKey, &descriptorSize, &descriptorVersion);
        if (EFI_ERROR(status))
        {
            Fatal("Failed to retrieve the EFI memory map: %p\n", status);
        }
    }

    if (EFI_ERROR(status))
    {
        Fatal("Failed to exit boot services: %p\n", status);
    }

    BuildMemoryMap(memoryMap, descriptors, size / descriptorSize, descriptorSize);
}

Code: Select all

Print(L"Getting Memory Map and Exiting Boot Services... ");
EFI_MEMORY_DESCRIPTOR* memMap = NULL;
UINTN memMapSize = 0, memMapKey = 0, descriptorSize = 0;
UINT32 descriptorVer = 0;
{
	while ((status = SystemTable->BootServices->GetMemoryMap(&memMapSize, memMap, &memMapKey, &descriptorSize, &descriptorVer)) == EFI_BUFFER_TOO_SMALL) {
		memMapSize += descriptorSize * 10;

		if (memMap != NULL) SystemTable->BootServices->FreePool(memMap);
		SystemTable->BootServices->AllocatePool(EfiLoaderData, memMapSize, (void**)&memMap);
	}

	if (EFI_ERROR(status)) {
		Print(L"FATAL - Unable to retrieve EFI memory map [%d]\n\r", status);
		while (1);
	}

	while ((status = SystemTable->BootServices->ExitBootServices(ImageHandle, memMapKey)) == EFI_INVALID_PARAMETER) {
		status = SystemTable->BootServices->GetMemoryMap(&memMapSize, memMap, &memMapKey, &descriptorSize, &descriptorVer);
		if (EFI_ERROR(status)) break;
	}

	if (EFI_ERROR(status)) {
		Print(L"FATAL - Unable to exit boot services [%d]\n\r", status);
		while (1);
	}
}

Code: Select all

typedef struct {
  UINT32                Type;
  UINT32                Padding;        // Padding to ensure the next field is 64-bits aligned as per MS packing rules
  EFI_PHYSICAL_ADDRESS  PhysicalStart;
  EFI_VIRTUAL_ADDRESS   VirtualStart;
  UINT64                NumberOfPages;
} EFI_MEMORY_DESCRIPTOR;


// Intel's UEFI header does not properly define EFI_MEMORY_DESCRIPTOR for GCC.
// This check ensures that it is.
static_assert(offsetof(EFI_MEMORY_DESCRIPTOR, PhysicalStart) == 8);

Code: Select all

struct EFI_MEMORY_DESCRIPTOR {
    uint32_t type;
    void* physAddress, *virtAddress;
    uint64_t pageCount, attributes;
};

Code: Select all

// efidef.h

typedef struct {
    UINT32                          Type;           // Field size is 32 bits followed by 32 bit pad
    UINT32                          Pad;
    EFI_PHYSICAL_ADDRESS            PhysicalStart;  // Field size is 64 bits
    EFI_VIRTUAL_ADDRESS             VirtualStart;   // Field size is 64 bits
    UINT64                          NumberOfPages;  // Field size is 64 bits
    UINT64                          Attribute;      // Field size is 64 bits
} EFI_MEMORY_DESCRIPTOR;

Code: Select all

uint64_t GetMemorySize(EFI_MEMORY_DESCRIPTOR* memMap, uint64_t memMapEntries, uint64_t memMapDescSize) {
    static uint64_t memSizeBytes = 0;
    if (memSizeBytes > 0) return memSizeBytes;

    for (int i = 0; i < memMapEntries; i++) {
        EFI_MEMORY_DESCRIPTOR* desc = (EFI_MEMORY_DESCRIPTOR*)((uint64_t)memMap + (i * memMapDescSize));
        memSizeBytes += desc->pageCount * 4096;
    }

    return memSizeBytes;
}