libcxx/utils/ssh.py

   1 #!/usr/bin/env python
   2 # ===----------------------------------------------------------------------===##
   3 #
   4 # Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   5 # See https://llvm.org/LICENSE.txt for license information.
   6 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   7 #
   8 # ===----------------------------------------------------------------------===##
   9
  10 """
  11 Runs an executable on a remote host.
  12
  13 This is meant to be used as an executor when running the C++ Standard Library
  14 conformance test suite.
  15 """
  16
  17 import argparse
  18 import os
  19 import posixpath
  20 import shlex
  21 import subprocess
  22 import sys
  23 import tarfile
  24 import tempfile
  25
  26
  27 def main():
  28     parser = argparse.ArgumentParser()
  29     parser.add_argument("--host", type=str, required=True)
  30     parser.add_argument("--execdir", type=str, required=True)
  31     parser.add_argument("--tempdir", type=str, required=False, default="/tmp")
  32     parser.add_argument("--extra-ssh-args", type=str, required=False)
  33     parser.add_argument("--extra-scp-args", type=str, required=False)
  34     parser.add_argument("--codesign_identity", type=str, required=False, default=None)
  35     parser.add_argument("--env", type=str, nargs="*", required=False, default=[])
  36     parser.add_argument("--prepend_env", type=str, nargs="*", required=False, default=[])
  37     parser.add_argument("-v", "--verbose", action='store_true')
  38     parser.add_argument("command", nargs=argparse.ONE_OR_MORE)
  39     args = parser.parse_args()
  40     commandLine = args.command
  41
  42     def ssh(command):
  43         cmd = ["ssh", "-oBatchMode=yes"]
  44         if args.extra_ssh_args is not None:
  45             cmd.extend(shlex.split(args.extra_ssh_args))
  46         return cmd + [args.host, command]
  47
  48     def scp(src, dst):
  49         cmd = ["scp", "-q", "-oBatchMode=yes"]
  50         if args.extra_scp_args is not None:
  51             cmd.extend(shlex.split(args.extra_scp_args))
  52         return cmd + [src, "{}:{}".format(args.host, dst)]
  53
  54     def runCommand(command, *args_, **kwargs):
  55         if args.verbose:
  56             print(f"$ {' '.join(command)}", file=sys.stderr)
  57         return subprocess.run(command, *args_, **kwargs)
  58
  59     # Create a temporary directory where the test will be run.
  60     # That is effectively the value of %T on the remote host.
  61     tmp = runCommand(
  62         ssh("mktemp -d {}/libcxx.XXXXXXXXXX".format(args.tempdir)),
  63         universal_newlines=True,
  64         check=True,
  65         capture_output=True,
  66         stdin=subprocess.DEVNULL
  67     ).stdout.strip()
  68
  69     # HACK:
  70     # If an argument is a file that ends in `.tmp.exe`, assume it is the name
  71     # of an executable generated by a test file. We call these test-executables
  72     # below. This allows us to do custom processing like codesigning test-executables
  73     # and changing their path when running on the remote host. It's also possible
  74     # for there to be no such executable, for example in the case of a .sh.cpp
  75     # test.
  76     isTestExe = lambda exe: exe.endswith(".tmp.exe") and os.path.exists(exe)
  77     pathOnRemote = lambda file: posixpath.join(tmp, os.path.basename(file))
  78
  79     try:
  80         # Do any necessary codesigning of test-executables found in the command line.
  81         if args.codesign_identity:
  82             for exe in filter(isTestExe, commandLine):
  83                 codesign = ["codesign", "-f", "-s", args.codesign_identity, exe]
  84                 runCommand(codesign, env={}, check=True, stdin=subprocess.DEVNULL)
  85
  86         # tar up the execution directory (which contains everything that's needed
  87         # to run the test), and copy the tarball over to the remote host.
  88         try:
  89             tmpTar = tempfile.NamedTemporaryFile(suffix=".tar", delete=False)
  90             with tarfile.open(fileobj=tmpTar, mode="w") as tarball:
  91                 tarball.add(args.execdir, arcname=os.path.basename(args.execdir))
  92
  93             # Make sure we close the file before we scp it, because accessing
  94             # the temporary file while still open doesn't work on Windows.
  95             tmpTar.close()
  96             remoteTarball = pathOnRemote(tmpTar.name)
  97             runCommand(scp(tmpTar.name, remoteTarball), check=True, stdin=subprocess.DEVNULL)
  98         finally:
  99             # Make sure we close the file in case an exception happens before
 100             # we've closed it above -- otherwise close() is idempotent.
 101             tmpTar.close()
 102             os.remove(tmpTar.name)
 103
 104         # Untar the dependencies in the temporary directory and remove the tarball.
 105         remoteCommands = [
 106             "tar -xf {} -C {} --strip-components 1".format(remoteTarball, tmp),
 107             "rm {}".format(remoteTarball),
 108         ]
 109
 110         # Make sure all test-executables in the remote command line have 'execute'
 111         # permissions on the remote host. The host that compiled the test-executable
 112         # might not have a notion of 'executable' permissions.
 113         for exe in map(pathOnRemote, filter(isTestExe, commandLine)):
 114             remoteCommands.append("chmod +x {}".format(exe))
 115
 116         # Execute the command through SSH in the temporary directory, with the
 117         # correct environment. We tweak the command line to run it on the remote
 118         # host by transforming the path of test-executables to their path in the
 119         # temporary directory on the remote host.
 120         commandLine = (pathOnRemote(x) if isTestExe(x) else x for x in commandLine)
 121         remoteCommands.append("cd {}".format(tmp))
 122
 123         if args.prepend_env:
 124             # We can't sensibly know the original value of the env vars
 125             # in order to prepend to them, so just overwrite these variables.
 126             args.env.extend(args.prepend_env)
 127
 128         if args.env:
 129             env = list(map(shlex.quote, args.env))
 130             remoteCommands.append("export {}".format(" ".join(args.env)))
 131         remoteCommands.append(subprocess.list2cmdline(commandLine))
 132
 133         # Finally, SSH to the remote host and execute all the commands.
 134         # Make sure to forward stdin to the process so that the test suite
 135         # can pipe stuff into the executor.
 136         rc = runCommand(ssh(" && ".join(remoteCommands))).returncode
 137         return rc
 138
 139     finally:
 140         # Make sure the temporary directory is removed when we're done.
 141         runCommand(ssh("rm -r {}".format(tmp)), check=True)
 142
 143
 144 if __name__ == "__main__":
 145     exit(main())