From db5b547b2542d01661ad602b437d88e3c75a8606 Mon Sep 17 00:00:00 2001
From: Matej Urbas <matej.urbas@gmail.com>
Date: Sun, 18 Apr 2021 10:19:06 +0100
Subject: [PATCH] nixos/amazon-init: add user-data shell script support

---
 nixos/modules/virtualisation/amazon-init.nix | 10 +++++
 nixos/tests/all-tests.nix                    |  1 +
 nixos/tests/amazon-init-shell.nix            | 40 ++++++++++++++++++++
 3 files changed, 51 insertions(+)
 create mode 100644 nixos/tests/amazon-init-shell.nix

diff --git a/nixos/modules/virtualisation/amazon-init.nix b/nixos/modules/virtualisation/amazon-init.nix
index be83607c0af7..4f2f8df90eb8 100644
--- a/nixos/modules/virtualisation/amazon-init.nix
+++ b/nixos/modules/virtualisation/amazon-init.nix
@@ -16,6 +16,16 @@ let
 
     userData=/etc/ec2-metadata/user-data
 
+    # Check if user-data looks like a shell script and execute it with the
+    # runtime shell if it does. Otherwise treat it as a nixos configuration
+    # expression
+    if IFS= LC_ALL=C read -rN2 shebang < $userData && [ "$shebang" = '#!' ]; then
+      # NB: we cannot chmod the $userData file, this is why we execute it via
+      # `pkgs.runtimeShell`. This means we have only limited support for shell
+      # scripts compatible with the `pkgs.runtimeShell`.
+      exec ${pkgs.runtimeShell} $userData
+    fi
+
     if [ -s "$userData" ]; then
       # If the user-data looks like it could be a nix expression,
       # copy it over. Also, look for a magic three-hash comment and set
diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index fb45ec1a310c..d4fdc668e768 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -24,6 +24,7 @@ in
   _3proxy = handleTest ./3proxy.nix {};
   acme = handleTest ./acme.nix {};
   agda = handleTest ./agda.nix {};
+  amazon-init-shell = handleTest ./amazon-init-shell.nix {};
   ammonite = handleTest ./ammonite.nix {};
   atd = handleTest ./atd.nix {};
   avahi = handleTest ./avahi.nix {};
diff --git a/nixos/tests/amazon-init-shell.nix b/nixos/tests/amazon-init-shell.nix
new file mode 100644
index 000000000000..f9268b2f3a00
--- /dev/null
+++ b/nixos/tests/amazon-init-shell.nix
@@ -0,0 +1,40 @@
+# This test verifies that the amazon-init service can treat the `user-data` ec2
+# metadata file as a shell script. If amazon-init detects that `user-data` is a
+# script (based on the presence of the shebang #! line) it executes it and
+# exits.
+# Note that other tests verify that amazon-init can treat user-data as a nixos
+# configuration expression.
+
+{ system ? builtins.currentSystem,
+  config ? {},
+  pkgs ? import ../.. { inherit system config; }
+}:
+
+with import ../lib/testing-python.nix { inherit system pkgs; };
+with pkgs.lib;
+
+makeTest {
+  name = "amazon-init";
+  meta = with maintainers; {
+    maintainers = [ urbas ];
+  };
+  machine = { ... }:
+  {
+    imports = [ ../modules/profiles/headless.nix ../modules/virtualisation/amazon-init.nix ];
+    services.openssh.enable = true;
+    networking.hostName = "";
+    environment.etc."ec2-metadata/user-data" = {
+      text = ''
+        #!/usr/bin/bash
+
+        echo successful > /tmp/evidence
+      '';
+    };
+  };
+  testScript = ''
+    # To wait until amazon-init terminates its run
+    unnamed.wait_for_unit("amazon-init.service")
+
+    unnamed.succeed("grep -q successful /tmp/evidence")
+  '';
+}