Correct .modinfo section parsing

xtldr/protocol.c

@@ -749,100 +749,96 @@ BlpGetModuleInfoStrings(IN PWCHAR SectionData,
     EFI_STATUS Status;
     PWCHAR *Array;
     PWCHAR String;
+    ULONG DataSize;
 
     /* Check input parameters */
     InfoStrings = SectionData;
     if(!InfoStrings || !SectionSize)
     {
         /* Invalid input parameters */
+        *ModInfo = NULL;
+        *InfoCount = 0;
         return STATUS_EFI_INVALID_PARAMETER;
     }
 
-    /* Skip zero padding */
+    /* Calculate the size of the data based on the size of the section */
-    while(InfoStrings[0] == L'\0' && SectionSize > 1)
+    DataSize = SectionSize / sizeof(WCHAR);
+
+    /* Skip zero padding at the beginning */
+    while(DataSize > 0 && *InfoStrings == L'\0')
     {
-        /* Get next character and decrement section size */
         InfoStrings++;
-        SectionSize--;
+        DataSize--;
     }
 
     /* Make sure there is at least one string available */
-    if(SectionSize <= 1)
+    if(DataSize < 1)
     {
         /* No strings found */
+        *ModInfo = NULL;
+        *InfoCount = 0;
         return STATUS_EFI_END_OF_FILE;
     }
 
     /* Count number of strings */
     Index = 0;
     Count = 0;
-    while(Index < SectionSize)
+    while(Index < DataSize)
     {
-        /* Get to the next string */
+        /* Found start of a new string */
-        if(InfoStrings[Index] != L'\0')
+        Count++;
-        {
-            /* Get next character */
-            Index++;
-            continue;
-        }
 
-        /* Skip zero padding */
+        /* Go to the end of the string */
-        while(InfoStrings[Index] == L'\0' && Index < SectionSize)
+        while(Index < DataSize && InfoStrings[Index] != L'\0')
+        {
+            Index++;
+        }
+        /* Skip all null terminators */
+        while(Index < DataSize && InfoStrings[Index] == L'\0')
         {
-            /* Get next character */
             Index++;
         }
-
-        /* New string found, increment counter */
-        Count++;
     }
 
-    /* Make sure there is no missing string */
+    /* Allocate memory for the pointer array and the string data */
-    if(InfoStrings[Index - 1] != L'\0')
+    Status = BlAllocateMemoryPool(sizeof(PWCHAR) * (Count + 1) + (DataSize + 1) * sizeof(WCHAR), (PVOID *)&Array);
-    {
-        /* One more string available */
-        Count++;
-    }
-
-    /* Allocate memory for array of strings */
-    Status = BlAllocateMemoryPool(SectionSize + 1 + sizeof(PWCHAR) * (Count + 1), (PVOID *)&Array);
     if(Status != STATUS_EFI_SUCCESS)
     {
         /* Failed to allocate memory */
         return STATUS_EFI_OUT_OF_RESOURCES;
     }
 
-    /* Allocate memory and copy strings read from '.modinfo' section */
+    /* The string buffer is located right after the pointer array */
-    BlAllocateMemoryPool(SectionSize, (PVOID*)&String);
+    String = (PWCHAR)(Array + Count + 1);
-    RtlCopyMemory(String, InfoStrings, SectionSize);
 
-    /* Make sure last string is NULL-terminated */
+    /* Copy the raw string data */
+    RtlCopyMemory(String, InfoStrings, DataSize * sizeof(WCHAR));
+
+    /* Ensure the entire buffer is null-terminated for safety */
+    String[DataSize] = L'\0';
+
+    /* Set the last element of the pointer array to NULL */
     Array[Count] = NULL;
-    Array[0] = String;
 
-    /* Parse strings into array */
+    /* Populate the array with pointers to the strings within the buffer */
     Index = 0;
-    ArrayIndex = 1;
+    ArrayIndex = 0;
-    while(Index < SectionSize && ArrayIndex < Count)
+    while(Index < DataSize && ArrayIndex < Count)
     {
-        /* Get to the next string */
+        /* Set pointer to the beginning of the string */
-        if(String[Index] != L'\0')
+        Array[ArrayIndex++] = &String[Index];
-        {
-            /* Get next character */
-            Index++;
-            continue;
-        }
 
-        /* Skip zero padding */
+        /* Find the end of the current string */
-        while(InfoStrings[Index] == L'\0' && Index < SectionSize)
+        while(Index < DataSize && String[Index] != L'\0')
         {
-            /* Get next character */
             Index++;
         }
 
-        /* Push string into array */
+        /* Skip all null terminators to find the beginning of the next string */
-        Array[ArrayIndex] = &String[Index];
+        while(Index < DataSize && String[Index] == L'\0')
-        ArrayIndex++;
+        {
+            Index++;
+        }
     }
 
     /* Return array of strings and its size */