From bbe7e19554b7bba487d836b49864297947c55cc9 Mon Sep 17 00:00:00 2001
From: Zakaria <zbenmo1@wgu.edu>
Date: Thu, 15 Oct 2026 10:25:56 -0400
Subject: [PATCH] small fixes

---
 .../services/ec2/Ec2ContainerManager.java     | 38 ++++++++++++++++++-
 .../services/ec2/Ec2ContainerManagerTest.java | 21 ++++++++--
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/src/main/java/io/github/hectorvent/floci/services/ec2/Ec2ContainerManager.java b/src/main/java/io/github/hectorvent/floci/services/ec2/Ec2ContainerManager.java
index b91b557..47c0469 100644
--- a/src/main/java/io/github/hectorvent/floci/services/ec2/Ec2ContainerManager.java
+++ b/src/main/java/io/github/hectorvent/floci/services/ec2/Ec2ContainerManager.java
@@ -2036,7 +2036,43 @@ public class Ec2ContainerManager {
     }
 
     static String[] userDataExecutionCommand() {
-        return new String[]{USER_DATA_SCRIPT_PATH};
+        String command = """
+                policy=/usr/sbin/policy-rc.d
+                backup=/var/lib/floci-policy-rc.d
+                actions=/var/lib/floci-service-actions
+                rm -f "$actions"
+                if [ -f "$policy" ]; then mv "$policy" "$backup"; fi
+                printf '%s\n' '#!/bin/sh' \
+                  'while [ "$#" -gt 0 ]; do case "$1" in --*) shift ;; *) break ;; esac; done' \
+                  'service=$1' \
+                  'action=$2' \
+                  'action=${action#(}' \
+                  'action=${action%)}' \
+                  'case "$action" in start|restart|try-restart|reload|force-reload)' \
+                  'printf "%s %s\\n" "$service" "$action" >> /var/lib/floci-service-actions ;; esac' \
+                  'exit 101' > "$policy"
+                chmod 755 "$policy"
+                restore_policy_and_start_services() {
+                  status=$?
+                  rm -f "$policy"
+                  if [ -f "$backup" ]; then mv "$backup" "$policy"; fi
+                  if [ -f "$actions" ]; then
+                    while read -r service action; do
+                      if command -v service >/dev/null 2>&1; then
+                        service "$service" "$action" || true
+                      elif [ -x "/etc/init.d/$service" ]; then
+                        "/etc/init.d/$service" "$action" || true
+                      fi
+                    done < "$actions"
+                    rm -f "$actions"
+                  fi
+                  trap - EXIT
+                  exit "$status"
+                }
+                trap restore_policy_and_start_services EXIT
+                /var/lib/user-data.sh
+                """;
+        return new String[]{"sh", "-c", command};
     }
 
     /**
diff --git a/src/test/java/io/github/hectorvent/floci/services/ec2/Ec2ContainerManagerTest.java b/src/test/java/io/github/hectorvent/floci/services/ec2/Ec2ContainerManagerTest.java
index 01a3938..c6d03f8 100644
--- a/src/test/java/io/github/hectorvent/floci/services/ec2/Ec2ContainerManagerTest.java
+++ b/src/test/java/io/github/hectorvent/floci/services/ec2/Ec2ContainerManagerTest.java
@@ -444,7 +444,20 @@ class Ec2ContainerManagerTest {
 
     @Test
     void userDataExecutionCommandRunsScriptDirectlySoShebangIsHonored() {
-        assertArrayEquals(new String[]{"/var/lib/user-data.sh"}, Ec2ContainerManager.userDataExecutionCommand());
+        String[] command = Ec2ContainerManager.userDataExecutionCommand();
+
+        assertArrayEquals(new String[]{"sh", "-c"}, Arrays.copyOf(command, 2));
+        assertTrue(command[2].contains("/var/lib/user-data.sh"));
+        assertFalse(command[2].contains("sh /var/lib/user-data.sh"));
+    }
+
+    @Test
+    void userDataExecutionCommandRestoresPolicyAndStartsPackageServices() {
+        String command = Ec2ContainerManager.userDataExecutionCommand()[2];
+
+        assertTrue(command.contains("mv \"$backup\" \"$policy\""));
+        assertTrue(command.contains("service \"$service\" \"$action\""));
+        assertTrue(command.contains("exit \"$status\""));
     }
 
     @Test
@@ -1320,7 +1333,7 @@ class Ec2ContainerManagerTest {
         verify(copy).withRemotePath("/var/lib");
         verify(copy, never()).withRemotePath("/tmp");
         assertTrue(harness.executedCommands.stream()
-                .anyMatch(command -> Arrays.equals(command, new String[]{"/var/lib/user-data.sh"})));
+                .anyMatch(command -> Arrays.equals(command, Ec2ContainerManager.userDataExecutionCommand())));
     }
 
     @Test
@@ -1780,7 +1793,7 @@ class Ec2ContainerManagerTest {
         });
         when(execCreate.exec()).thenAnswer(invocation -> {
             String[] command = currentCommand.get();
-            return command != null && command.length == 1 && "/var/lib/user-data.sh".equals(command[0])
+            return Arrays.equals(command, Ec2ContainerManager.userDataExecutionCommand())
                     ? userDataExec : metadataExec;
         });
 
@@ -1969,7 +1982,7 @@ class Ec2ContainerManagerTest {
             });
             when(execCreate.exec()).thenAnswer(invocation -> {
                 String[] command = currentCommand.get();
-                if (command != null && command.length == 1 && "/var/lib/user-data.sh".equals(command[0])) {
+                if (Arrays.equals(command, Ec2ContainerManager.userDataExecutionCommand())) {
                     return userDataExec;
                 }
                 return metadataExec;