From 7787fad6f6017911db3342ca031cb4e6e760ad77 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20=C5=BBygowski?= <michal.zygowski@3mdeb.com>
Date: Sat, 16 Dec 2023 13:01:29 +0100
Subject: [PATCH] simple_file: Compact the array of filesystem entries
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

If HandleProtocol or OpenVolume fails, the entries array will
become non-contiguous, i.e. will have NULL pointers between valid
volume names in the array. Because of that count_lines may return
a lower number of entries than expected. As a result one may not
browse all valid filesystems in the file explorer.

Compact the array by locating the holes in filesystem entries array
and fill them with valid entries from higher indices, if any.
As a result, count_lines should return proper length and there
won't be any accesses to invalid entries or lost partitions.

Signed-off-by: Michał Żygowski <michal.zygowski@3mdeb.com>
---
 include/console.h | 13 +++++++++++++
 lib/simple_file.c | 25 ++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/include/console.h b/include/console.h
index 063bc20f1..923ba9796 100644
--- a/include/console.h
+++ b/include/console.h
@@ -63,6 +63,19 @@ count_lines(CHAR16 *str_arr[])
 	return i;
 }
 
+static inline int
+find_first_existing_line(CHAR16 *str_arr[], UINTN start, UINTN end)
+{
+	UINTN i = start;
+
+	while (i < end) {
+		if (str_arr[i])
+			break;
+		i++;
+	}
+	return i;
+}
+
 #define NOSEL 0x7fffffff
 
 typedef struct _EFI_CONSOLE_CONTROL_PROTOCOL   EFI_CONSOLE_CONTROL_PROTOCOL;
diff --git a/lib/simple_file.c b/lib/simple_file.c
index fc082bed1..530a20c5e 100644
--- a/lib/simple_file.c
+++ b/lib/simple_file.c
@@ -170,7 +170,7 @@ simple_file_write_all(EFI_FILE *file, UINTN size, void *buffer)
 EFI_STATUS
 simple_volume_selector(CHAR16 **title, CHAR16 **selected, EFI_HANDLE *h)
 {
-	UINTN count, i;
+	UINTN count, i, j;
 	EFI_HANDLE *vol_handles = NULL;
 	EFI_STATUS efi_status;
 	CHAR16 **entries;
@@ -225,6 +225,29 @@ simple_volume_selector(CHAR16 **title, CHAR16 **selected, EFI_HANDLE *h)
 	}
 	entries[i] = NULL;
 
+	/*
+	 * If HandleProtocol or OpenVolume fails, we will end up with non-contiguous array
+	 * and lose some valid partitions to browse because of how count_lines works.
+	 * Compact the array if needed.
+	 */
+	if ((UINTN)count_lines(entries) < i) {
+		UINTN temp;
+		for (j = (UINTN)count_lines(entries); j < i; j++) {
+			/* Skip contiguous valid entries */
+			if (entries[j])
+				continue;
+
+			temp = (UINTN)find_first_existing_line(entries, j, i);
+			/* If reached the end, there is nothing to compact anymore. */
+			if (temp >= i)
+				break;
+
+			/* Move the entry to the first free slot and remove old entry record. */
+			entries[j] = entries[temp];
+			entries[temp] = NULL;
+		}
+	}
+
 	val = console_select(title, entries, 0);
 
 	if (val >= 0) {