Fix GetCapabilities to return the right size

Author: elicn

qiling/os/uefi/protocols/EfiSmmAccess2Protocol.py

@@ -10,7 +10,7 @@
 from ..ProcessorBind import *
 from ..UefiBaseType import *
 from ..PiMultiPhase import *
-from ..utils import write_int64, read_int64
+from .. import utils
 
 # @see: MdePkg\Include\Pi\PiMultiPhase.h
 class EFI_MMRAM_DESCRIPTOR(STRUCT):
@@ -108,8 +108,23 @@ def hook_GetCapabilities(ql: Qiling, address: int, params):
 	size = len(chunks) * EFI_SMRAM_DESCRIPTOR.sizeof()
 	MmramMapSize = params["MmramMapSize"]
 
-	if read_int64(ql, MmramMapSize) < size:
-		write_int64(ql, MmramMapSize, size)
+	if utils.read_int64(ql, MmramMapSize) < size:
+		# since the caller cannot predict how much memory would be required for storing
+		# the memory map, this method is normally called twice. the first one passes a
+		# zero size only to determine the expected size, then the caller allocates the
+		# required amount of memory and call it again.
+		#
+		# our memory map is managed differently from the real one, and memory allocations
+		# are likely to generate an additional "map block" (or two, if allocated somewhere
+		# in the last free heap chunk). because the caller allocates a new memory chunk
+		# between the two calls, that would cause the second call to always complain the
+		# buffer is too small.
+		#
+		# to work around that, we have the first call return a larger number than it should
+		# have, to compensate on the coming allocation.
+		extra = 2 * EFI_SMRAM_DESCRIPTOR.sizeof()
+
+		utils.write_int64(ql, MmramMapSize, size + extra)
 		return EFI_BUFFER_TOO_SMALL
 
 	MmramMap = params["MmramMap"]