Use a power button listener service to watch /dev/input/eventX to sense power button press

As the hardened image lacks `systemd-logind`, the new service, power button listener, takes the responsibility of watching /dev/input/eventX which was previously done by `logind`. When it detects a power button press, or VM stop, it triggers systemd to stop services, including the listener itself. Then, it uses the service's ExecStop script to send SIGTERM to all containers.

Author: susienme

launcher/cloudbuild.yaml

@@ -72,6 +72,10 @@ steps:
       export CGO_LDFLAGS="-L/workspace/keymanager/target/release"
       go build -o ../image/launcher -ldflags="-extldflags=-Wl,-z,lazy -X 'main.BuildCommit=${SHORT_SHA}'"
 
+      # Build the power button listener
+      cd /workspace/launcher/power-button-listener
+      go build -o /workspace/launcher/image/power-button-listener
+
 - name: 'gcr.io/cloud-builders/gcloud'
   id: DownloadExpBinary
   entrypoint: 'gcloud'

launcher/image/container-cleanup.service

@@ -1,11 +1,12 @@
 #!/bin/bash
+echo "Container cleanup script started"
 
 SHUTDOWN_TIMEOUT_SEC=15
 
 # Send SIGTERM to all running workloads so that they can shutdown gracefully.
 for ns in $(ctr ns ls -q); do
     tasks=$(ctr -n "$ns" task ls -q)
-    
+
     if [ -n "$tasks" ]; then
         # Send SIGTERM and move on. No waiting, no killing, no deleting.
         # A workload may decide to ignore or not handle SIGTERM.

launcher/image/container-cleanup.sh

@@ -3,7 +3,7 @@
 main() {
   # Copy service files.
   cp /usr/share/oem/confidential_space/container-runner.service /etc/systemd/system/container-runner.service
-  cp /usr/share/oem/confidential_space/container-cleanup.service /etc/systemd/system/container-cleanup.service
+  cp /usr/share/oem/confidential_space/power-button-listener.service /etc/systemd/system/power-button-listener.service
   # Override default fluent-bit config.
   cp /usr/share/oem/confidential_space/fluent-bit-cs.conf /etc/fluent-bit/fluent-bit.conf
 
@@ -16,8 +16,8 @@ main() {
   # Override default kernel-monitor.json for node-problem-detector.
   cp /usr/share/oem/confidential_space/kernel-monitor-cs.json /etc/node_problem_detector/kernel-monitor.json
   systemctl daemon-reload
-  systemctl enable container-runner.service container-cleanup.service
-  systemctl start container-runner.service container-cleanup.service
+  systemctl enable container-runner.service power-button-listener.service
+  systemctl start container-runner.service power-button-listener.service
   systemctl start fluent-bit.service
 }
 

launcher/image/entrypoint.sh

@@ -0,0 +1,16 @@
+[Unit]
+Description=Power Button Listener
+After=containerd.service container-runner.service
+
+[Service]
+Type=simple
+ExecStart=/usr/share/oem/confidential_space/power-button-listener
+ExecStop=/usr/share/oem/confidential_space/container-cleanup.sh
+Restart=on-failure
+RestartSec=5
+TTYPath=/dev/ttyS0
+StandardOutput=tty
+StandardError=tty
+
+[Install]
+WantedBy=basic.target

launcher/image/power-button-listener.service

@@ -19,8 +19,9 @@ setup_launcher_systemd_unit() {
   cp exit_script.sh "${CS_PATH}/exit_script.sh"
 }
 
-setup_container_cleanup_service() {
-  cp container-cleanup.service "${CS_PATH}/container-cleanup.service"
+copy_power_button_listener() {
+  cp power-button-listener "${CS_PATH}/power-button-listener"
+  cp power-button-listener.service "${CS_PATH}/power-button-listener.service"
   cp container-cleanup.sh "${CS_PATH}/container-cleanup.sh"
 }
 
@@ -118,8 +119,8 @@ main() {
   # Install container launcher.
   copy_launcher
   setup_launcher_systemd_unit
-  # Install container cleanup service.
-  setup_container_cleanup_service
+  # Install power button listener.
+  copy_power_button_listener
   # Minimum required COS version for 'e': cos-dev-105-17222-0-0.
   # Minimum required COS version for 'm': cos-dev-113-18203-0-0.
   append_cmdline "cos.protected_stateful_partition=m"

launcher/image/preload.sh

@@ -19,19 +19,31 @@ check_timeout() {
   echo $remaining
 }
 
-echo "Polling for heartbeat..."
+echo "Starting to tail serial port in background..."
+gcloud compute instances tail-serial-port-output $MONITOR_VM --zone $ZONE > /workspace/serial_output.txt &
+TAIL_PID=$!
+
+# Give gcloud a few seconds to establish the connection
+sleep 5
+
+echo "Polling for heartbeat (Workload: $WORKLOAD_VM, Monitor: $MONITOR_VM)..."
 remaining=$(check_timeout "Timeout before heartbeat poll")
-timeout $remaining \
-  gcloud compute instances tail-serial-port-output $MONITOR_VM --zone $ZONE | \
-  grep -q 'Workload heartbeat' || \
-  { echo "failed: Heartbeat not found within timeout" > /workspace/status.txt; exit 0; }
+timeout $remaining bash -c "until grep -q 'Workload heartbeat' /workspace/serial_output.txt; do sleep 1; done" || {
+  echo "failed: Heartbeat not found within timeout" > /workspace/status.txt
+  kill $TAIL_PID
+  exit 0
+}
 
 echo "Stopping workload VM..."
 gcloud compute instances stop $WORKLOAD_VM --zone $ZONE
 
 echo "Polling for graceful exit..."
 remaining=$(check_timeout "Timeout before graceful exit poll")
-timeout $remaining \
-  gcloud compute instances tail-serial-port-output $MONITOR_VM --zone $ZONE | \
-  grep -q 'Workload exiting gracefully' || \
-  { echo "failed: Graceful exit message not found within timeout" > /workspace/status.txt; exit 0; }
+timeout $remaining bash -c "until grep -q 'Workload exiting gracefully' /workspace/serial_output.txt; do sleep 1; done" || {
+  echo "failed: Graceful exit message not found within timeout" > /workspace/status.txt
+  kill $TAIL_PID
+  exit 0
+}
+
+# Success! Clean up the background process
+kill $TAIL_PID

launcher/image/test/scripts/test_workload_cleanup.sh

@@ -10,7 +10,7 @@ substitutions:
   '_IMAGE_PROJECT': ''
   '_CLEANUP': 'true'
   '_IMAGE_REPO': 'us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images'
-  '_WORKLOAD_IMAGE': '${_IMAGE_REPO}/workloadcleanup/workload:latest'
+  '_WORKLOAD_IMAGE': '${_IMAGE_REPO}/workloadcleanup:latest'
   '_ZONE': 'us-west1-a'
 
 steps:
@@ -68,7 +68,7 @@ steps:
     MONITOR_IP=$(cat /workspace/monitor_ip.txt)
     bash create_vm.sh -i $IMAGE_NAME \
       -p $IMAGE_PROJECT \
-      -m "tee-image-reference=$WORKLOAD_IMAGE,tee-container-log-redirect=true,tee-cmd=[\"--server-addr=${MONITOR_IP}:2020\"]" \
+      -m "tee-image-reference=$WORKLOAD_IMAGE,tee-cmd=[\"--server-addr=${MONITOR_IP}:2020\"]" \
       -n cleanup-test-$BUILD_ID \
       -z $ZONE
 

launcher/image/test/test_workloadcleanup_cloudbuild.yaml

@@ -1,12 +1,11 @@
 # From current directory:
 # GOOS=linux GOARCH=amd64 CGO_ENABLED=0 go build -o main .
-# gcloud builds submit --tag us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/workloadcleanup/workload:latest
+# gcloud builds submit --tag us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/workloadcleanup:latest
 
 FROM alpine
 
 COPY main /
 
-LABEL "tee.launch_policy.log_redirect"="always"
 LABEL "tee.launch_policy.allow_cmd_override"="true"
 
 ENTRYPOINT ["/main"]

…oads/workloadcleanup/workload/Dockerfile …testworkloads/workloadcleanup/Dockerfilelauncher/image/testworkloads/workloadcleanup/workload/Dockerfile renamed to launcher/image/testworkloads/workloadcleanup/Dockerfile launcher/image/testworkloads/workloadcleanup/workload/Dockerfile renamed to launcher/image/testworkloads/workloadcleanup/Dockerfile

@@ -1,4 +1,4 @@
-// Package main provides a test workload for cleanup.
+// A test workload for testing workload cleanup.
 package main
 
 import (