From 24a496c861fff6f8633453eedb92079976d3058f Mon Sep 17 00:00:00 2001
From: Lorenz Brun <lorenz@monogon.tech>
Date: Thu, 29 Jun 2023 03:54:01 +0200
Subject: [PATCH 1/2] Implement filename-based A/B slot handling

---
 src/boot/efi/stub.c | 66 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/src/boot/efi/stub.c b/src/boot/efi/stub.c
index 841a0e41bd..267ecbcc81 100644
--- a/src/boot/efi/stub.c
+++ b/src/boot/efi/stub.c
@@ -6,6 +6,8 @@
 #include "cpio.h"
 #include "devicetree.h"
 #include "disk.h"
+#include "efibind.h"
+#include "efidevp.h"
 #include "graphics.h"
 #include "linux.h"
 #include "measure.h"
@@ -15,6 +17,31 @@
 #include "tpm-pcr.h"
 #include "util.h"
 
+// From picolibc under BSD-3-Clause (c) 2018 Arm Ltd.
+/* Small and efficient strstr implementation.  */
+char16_t * strstr (const char16_t *hs, const char16_t *ne)
+{
+  UINTN i;
+  int c = ne[0];
+
+  if (c == 0)
+    return (char16_t*)hs;
+
+  for ( ; hs[0] != '\0'; hs++)
+    {
+      if (hs[0] != c)
+	continue;
+      for (i = 1; ne[i] != 0; i++)
+	if (hs[i] != ne[i])
+	  break;
+      if (ne[i] == '\0')
+	return (char16_t*)hs;
+    }
+
+  return NULL;
+}
+
+
 /* magic string to find in the binary image */
 _used_ _section_(".sdmagic") static const char magic[] = "#### LoaderInfo: systemd-stub " GIT_VERSION " ####";
 
@@ -232,6 +259,45 @@ EFI_STATUS efi_main(EFI_HANDLE image, EFI_SYSTEM_TABLE *sys_table) {
                 mangle_stub_cmdline(cmdline);
         }
 
+        /* Extract last FILEPATH element from image path to check file name */
+        CHAR16 *last_file_path = NULL;
+        EFI_DEVICE_PATH *current_path_elem = loaded_image->FilePath;
+        while (current_path_elem != NULL && !IsDevicePathEnd(current_path_elem)) {
+                if (DevicePathType(current_path_elem) == MEDIA_DEVICE_PATH &&
+                    current_path_elem->SubType == MEDIA_FILEPATH_DP) {
+                        last_file_path = ((FILEPATH_DEVICE_PATH*)current_path_elem)->PathName;
+                }
+                current_path_elem = NextDevicePathNode(current_path_elem);
+        }
+        /* Check slot based on suffix of the last FILE_PATH value */
+        char16_t slot = 'A';
+        const char16_t suffix_a[] = u"boot-a.efi";
+        const char16_t suffix_b[] = u"boot-b.efi";
+        const UINTN suffix_len = (sizeof(suffix_a)/sizeof(char16_t))-1;
+        if (last_file_path != NULL) {
+                UINTN plen = StrLen(last_file_path);
+                if (suffix_len > plen) {
+                        Print(L"File name too short, blindly booting slot A\n");
+                } else if (StriCmp(suffix_a, &last_file_path[plen-suffix_len]) == 0){
+                        slot = 'A';
+                } else if (StriCmp(suffix_b, &last_file_path[plen-suffix_len]) == 0) {
+                        slot = 'B';
+                } else {
+                        Print(L"Unknown file name, blindly booting slot A\n");
+                }
+        }
+        Print(L"Booting into Slot %c\n", slot);
+        /* Replace METROPOLIS-SYSTEM-X with the correct slot */
+        const char16_t slot_identifier[] = u"METROPOLIS-SYSTEM-X\0";
+        const UINTN slot_id_len = (sizeof(slot_identifier)/sizeof(char16_t))-1;
+        if (cmdline != NULL) {
+                char16_t *rest_ptr = cmdline;
+                while((rest_ptr = strstr(rest_ptr, slot_identifier))) {
+                        rest_ptr[slot_id_len-2] = slot;
+                        rest_ptr += slot_id_len;
+                }                
+        }
+
         export_variables(loaded_image);
 
         if (pack_cpio(loaded_image,
-- 
2.47.2