Repository: DataBiosphere/toil

Issue:
Add support for running custom init commands before starting appliance
This is similar to #2561, except that the command is run before the appliance starts (i.e. it would be `TOIL_CUSTOM_INIT_COMMAND`). One use case for this is for any custom authentication for pulling the "base" docker images (e.g. we actually use this to pull the Toil docker image itself from a private repo).

I initially hesitated adding this upstream, but given the recent announcement from docker about limiting pulls from anonymous users, I figured it may become useful for other use cases as well (while I realize that we're switching to quay for now, it may still become useful in the future...and also reduces the diffs from our fork :) ).

┆Issue is synchronized with this [Jira Task](https://ucsc-cgl.atlassian.net/browse/TOIL-627)
┆Issue Number: TOIL-627



Discussion:

Resolving patch:
diff --git a/src/toil/__init__.py b/src/toil/__init__.py
--- a/src/toil/__init__.py
+++ b/src/toil/__init__.py
@@ -17,6 +17,7 @@
 import errno
 import logging
 import os
+import re
 import requests
 import socket
 import sys
@@ -220,17 +221,43 @@ def customDockerInitCmd():
     Returns the custom command (if any) provided through the ``TOIL_CUSTOM_DOCKER_INIT_COMMAND``
     environment variable to run prior to running the workers and/or the primary node's services.
     This can be useful for doing any custom initialization on instances (e.g. authenticating to
-    private docker registries). An empty string is returned if the environment variable is not
-    set.
+    private docker registries). Any single quotes are escaped and the command cannot contain a
+    set of blacklisted chars (newline or tab). An empty string is returned if the environment
+    variable is not set.
 
     :rtype: str
     """
     command = lookupEnvVar(name='user-defined custom docker init command',
                            envName='TOIL_CUSTOM_DOCKER_INIT_COMMAND',
                            defaultValue='')
+    _check_custom_bash_cmd(command)
     return command.replace("'", "'\\''")  # Ensure any single quotes are escaped.
 
 
+def customInitCmd():
+    """
+    Returns the custom command (if any) provided through the ``TOIL_CUSTOM_INIT_COMMAND``
+    environment variable to run prior to running Toil appliance itself in workers and/or the
+    primary node (i.e. this is run one stage before ``TOIL_CUSTOM_DOCKER_INIT_COMMAND``).
+    This can be useful for doing any custom initialization on instances (e.g. authenticating to
+    private docker registries). Any single quotes are escaped and the command cannot contain a
+    set of blacklisted chars (newline or tab). An empty string is returned if the environment
+    variable is not set.
+
+    :rtype: str
+    """
+    command = lookupEnvVar(name='user-defined custom init command',
+                           envName='TOIL_CUSTOM_INIT_COMMAND',
+                           defaultValue='')
+    _check_custom_bash_cmd(command)
+    return command.replace("'", "'\\''")  # Ensure any single quotes are escaped.
+
+
+def _check_custom_bash_cmd(cmd_str):
+    """Ensures that the bash command doesn't contain blacklisted characters."""
+    assert not re.search(r'[\n\r\t]', cmd_str), f'"{cmd_str}" contains invalid characters (newline and/or tab).'
+
+
 def lookupEnvVar(name, envName, defaultValue):
     """
     Use this for looking up environment variables that control Toil and are important enough to
diff --git a/src/toil/provisioners/abstractProvisioner.py b/src/toil/provisioners/abstractProvisioner.py
--- a/src/toil/provisioners/abstractProvisioner.py
+++ b/src/toil/provisioners/abstractProvisioner.py
@@ -19,7 +19,7 @@
 import os.path
 
 import subprocess
-from toil import applianceSelf, customDockerInitCmd
+from toil import applianceSelf, customDockerInitCmd, customInitCmd
 from toil.lib.retry import never
 
 a_short_time = 5
@@ -337,6 +337,7 @@ def _setSSH(self):
         Restart=on-failure
         RestartSec=2
         ExecStartPre=-/usr/bin/docker rm toil_{role}
+        ExecStartPre=-/usr/bin/bash -c '{customInitCommand}'
         ExecStart=/usr/bin/docker run \
             --entrypoint={entrypoint} \
             --net=host \
@@ -418,6 +419,7 @@ def _getCloudConfigUserData(self, role, masterPublicKey=None, keyPath=None, pree
                             dockerImage=applianceSelf(),
                             entrypoint=entryPoint,
                             sshKey=masterPublicKey,   # ignored if None
-                            mesosArgs=mesosArgs)
+                            mesosArgs=mesosArgs,
+                            customInitCommand=customInitCmd())
         userData = template.format(**templateArgs)
         return userData


Test patch: