llvm-project

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 LLVM and the Libraries
14
tests on a target.
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
import re
26

27

28
def ssh(args, command):
29
    cmd = ["ssh", "-oBatchMode=yes"]
30
    if args.extra_ssh_args is not None:
31
        cmd.extend(shlex.split(args.extra_ssh_args))
32
    return cmd + [args.host, command]
33

34

35
def scp(args, src, dst):
36
    cmd = ["scp", "-q", "-oBatchMode=yes"]
37
    if args.extra_scp_args is not None:
38
        cmd.extend(shlex.split(args.extra_scp_args))
39
    return cmd + [src, "{}:{}".format(args.host, dst)]
40

41

42
def main():
43
    parser = argparse.ArgumentParser()
44
    parser.add_argument("--host", type=str, required=True)
45
    parser.add_argument("--execdir", type=str, required=False)
46
    parser.add_argument("--extra-ssh-args", type=str, required=False)
47
    parser.add_argument("--extra-scp-args", type=str, required=False)
48
    parser.add_argument("--codesign_identity", type=str, required=False, default=None)
49
    parser.add_argument("--env", type=str, nargs="*", required=False, default=dict())
50

51
    # Note: The default value is for the backward compatibility with a hack in
52
    # libcxx test suite.
53
    # If an argument is a file that ends in `.tmp.exe`, assume it is the name
54
    # of an executable generated by a test file. We call these test-executables
55
    # below. This allows us to do custom processing like codesigning test-executables
56
    # and changing their path when running on the remote host. It's also possible
57
    # for there to be no such executable, for example in the case of a .sh.cpp
58
    # test.
59
    parser.add_argument(
60
        "--exec-pattern",
61
        type=str,
62
        required=False,
63
        default=".*",
64
        help="The name regex pattern of the executables generated by \
65
                              a test file. Specifying it allows us to do custom \
66
                              processing like codesigning test-executables \
67
                              and changing their path when running on \
68
                              the remote host. It's also possible for there \
69
                              to be no such executable, for example in \
70
                              the case of a .sh.cpp test.",
71
    )
72

73
    parser.add_argument("command", nargs=argparse.ONE_OR_MORE)
74
    args = parser.parse_args()
75
    commandLine = args.command
76

77
    execdir = args.execdir
78
    if execdir == ".":
79
        # Retrieve the exec directory from the command line.
80
        execdir, _ = os.path.split(commandLine[0])
81
        if execdir == "":
82
            # Get the current directory in that case.
83
            execdir = os.getcwd()
84
    arcname = os.path.basename(execdir) if execdir else None
85

86
    # Create a temporary directory where the test will be run.
87
    # That is effectively the value of %T on the remote host.
88
    tmp = subprocess.check_output(
89
        ssh(args, "mktemp -d /tmp/llvm.XXXXXXXXXX"), universal_newlines=True
90
    ).strip()
91

92
    isExecutable = lambda exe: re.match(args.exec_pattern, exe) and os.path.exists(exe)
93
    pathOnRemote = lambda file: posixpath.join(tmp, os.path.basename(file))
94

95
    remoteCommands = []
96

97
    try:
98
        # Do any necessary codesigning of test-executables found in the command line.
99
        if args.codesign_identity:
100
            for exe in filter(isExecutable, commandLine):
101
                subprocess.check_call(
102
                    ["xcrun", "codesign", "-f", "-s", args.codesign_identity, exe],
103
                    env={},
104
                )
105

106
        # tar up the execution directory (which contains everything that's needed
107
        # to run the test), and copy the tarball over to the remote host.
108
        if execdir:
109
            try:
110
                tmpTar = tempfile.NamedTemporaryFile(suffix=".tar", delete=False)
111
                with tarfile.open(fileobj=tmpTar, mode="w") as tarball:
112
                    tarball.add(execdir, arcname=arcname)
113

114
                # Make sure we close the file before we scp it, because accessing
115
                # the temporary file while still open doesn't work on Windows.
116
                tmpTar.close()
117
                remoteTarball = pathOnRemote(tmpTar.name)
118
                subprocess.check_call(scp(args, tmpTar.name, remoteTarball))
119
            finally:
120
                # Make sure we close the file in case an exception happens before
121
                # we've closed it above -- otherwise close() is idempotent.
122
                tmpTar.close()
123
                os.remove(tmpTar.name)
124

125
            # Untar the dependencies in the temporary directory and remove the tarball.
126
            remoteCommands.extend(
127
                [
128
                    "tar -xf {} -C {} --strip-components 1".format(remoteTarball, tmp),
129
                    "rm {}".format(remoteTarball),
130
                ]
131
            )
132
        else:
133
            # Copy only the files, which are specified in the command line.
134
            # Copy them to remote host one by one.
135
            for x in commandLine:
136
                _, f = os.path.split(x)
137
                subprocess.check_call(scp(args, x, pathOnRemote(f)))
138

139
        # Make sure all executables in the remote command line have 'execute'
140
        # permissions on the remote host. The host that compiled the test-executable
141
        # might not have a notion of 'executable' permissions.
142
        for exe in filter(isExecutable, commandLine):
143
            remoteCommands.append("chmod +x {}".format(pathOnRemote(exe)))
144

145
        # Execute the command through SSH in the temporary directory, with the
146
        # correct environment. We tweak the command line to run it on the remote
147
        # host by transforming the path of test-executables to their path in the
148
        # temporary directory on the remote host.
149
        for i, x in enumerate(commandLine):
150
            if isExecutable(x):
151
                commandLine[i] = pathOnRemote(x)
152
        remoteCommands.append("cd {}".format(tmp))
153
        if args.env:
154
            remoteCommands.append("export {}".format(" ".join(args.env)))
155
        remoteCommands.append(subprocess.list2cmdline(commandLine))
156

157
        # Finally, SSH to the remote host and execute all the commands.
158
        rc = subprocess.call(ssh(args, " && ".join(remoteCommands)))
159
        return rc
160

161
    finally:
162
        # Make sure the temporary directory is removed when we're done.
163
        subprocess.check_call(ssh(args, "rm -r {}".format(tmp)))
164

165

166
if __name__ == "__main__":
167
    exit(main())
168