From ad5f37f7a8ccc9b9abe61c7374b4d654a72bb2fc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-Fran=C3=A7ois=20Roche?= <jfroche@pyxel.be>
Date: Fri, 2 Jun 2023 23:03:48 +0200
Subject: [PATCH] feat: add optional command timeout

Some template command can last longer than the default 30s. Add option
to override default timeout.

Define template commands using `exec` as `command` is deprecated.
(https://developer.hashicorp.com/vault/docs/agent/template#command)
---
 nix/checks/systemd-vaultd-test.nix |  1 +
 nix/modules/vault-secrets.nix      | 40 +++++++++++++++++++-----------
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/nix/checks/systemd-vaultd-test.nix b/nix/checks/systemd-vaultd-test.nix
index b7dafe3..250e3ae 100644
--- a/nix/checks/systemd-vaultd-test.nix
+++ b/nix/checks/systemd-vaultd-test.nix
@@ -94,6 +94,7 @@
     machine.wait_for_unit("vault.service")
     machine.wait_for_open_port(8200)
     machine.wait_for_unit("setup-vault-agent-approle.service")
+    machine.wait_for_unit("vault-agent-default.service")
 
     out = machine.wait_until_succeeds("grep -q bar /tmp/service1")
 
diff --git a/nix/modules/vault-secrets.nix b/nix/modules/vault-secrets.nix
index 7a30ec7..6748582 100644
--- a/nix/modules/vault-secrets.nix
+++ b/nix/modules/vault-secrets.nix
@@ -29,19 +29,28 @@ let
 
   services = config.systemd.services;
 
+  templateExec = serviceName: vaultConfig: { } //
+    lib.optionalAttrs (vaultConfig.changeAction != null && vaultConfig.changeAction != "none") {
+      exec = [
+        ({
+          command = "systemctl ${
+        if vaultConfig.changeAction == "restart"
+        then "try-restart"
+        else "try-reload-or-restart"
+      } ${lib.escapeShellArg "${serviceName}.service"}";
+        } // lib.optionalAttrs
+          (vaultConfig.command_timeout != null)
+          { timeout = vaultConfig.command_timeout; })
+      ];
+    };
+
   getSecretTemplate = serviceName: vaultConfig:
     {
       contents = vaultConfig.template;
       destination = "/run/systemd-vaultd/secrets/${serviceName}.service.json";
       perms = "0400";
     }
-    // lib.optionalAttrs (vaultConfig.changeAction != null && vaultConfig.changeAction != "none") {
-      command = "systemctl ${
-        if vaultConfig.changeAction == "restart"
-        then "try-restart"
-        else "try-reload-or-restart"
-      } ${lib.escapeShellArg "${serviceName}.service"}";
-    };
+    // templateExec serviceName vaultConfig;
 
   getEnvironmentTemplate = serviceName: vaultConfig:
     {
@@ -49,13 +58,7 @@ let
       destination = "/run/systemd-vaultd/secrets/${serviceName}.service.EnvironmentFile";
       perms = "0400";
     }
-    // lib.optionalAttrs (vaultConfig.changeAction != null) {
-      command = "systemctl ${
-        if vaultConfig.changeAction == "restart"
-        then "try-restart"
-        else "try-reload-or-restart"
-      } ${lib.escapeShellArg "${serviceName}.service"}";
-    };
+    // templateExec serviceName vaultConfig;
 
   vaultTemplates = config:
     (lib.mapAttrsToList
@@ -119,6 +122,15 @@ in
                 some-secret.template = ''{{ with secret "secret/some-secret" }}{{ .Data.data.some-key }}{{ end }}'';
               };
             };
+
+            command_timeout = lib.mkOption {
+              type = lib.types.nullOr lib.types.str;
+              default = null;
+              description = ''
+                Maximum amount of time to wait for the optional command to return.
+              '';
+            };
+
           };
           config =
             let