[libcxx][lit] Add support for custom ssh/scp flags in ssh.py

[arichardson]

Summary

In our CHERI Jenkins CI we need to pass -F <custom_config_file> to each
ssh/scp command to set various arguments such as the localhost port, usage
of controlmaster, etc. to speed up connections to our emulated QEMU systems.

For our specific use-case I could have also added a single --ssh-config-file
argument that can be used for both the scp and ssh commands, but being able
to pass arbitrary extra flags for both commands seems more flexible.

[arichardson]

rebase

[arichardson]

ping?

[ldionne]

LGTM, but based on D84421, do you actually need the extra scp args?

[arichardson]

In D84097#2271738, @ldionne wrote:

LGTM, but based on D84421, do you actually need the extra scp args?

Unfortunately, in some configuration variants we can't use shared mounts since they are (currently) unreliable.

libcxx/utils/ssh.py

 
 This is meant to be used as an executor when running the C++ Standard Library
 conformance test suite.
 """
 
 import argparse
 import os
 import posixpath
+import shlex
 import subprocess
 import sys
 import tarfile
 import tempfile
 
+def ssh(args, command):
+    cmd = ['ssh', '-oBatchMode=yes']
+    if args.extra_ssh_args is not None:
+        cmd.extend(shlex.split(args.extra_ssh_args))
+    return cmd + [args.host, command]
+
+
+def scp(args, src, dst):
+    cmd = ['scp', '-q', '-oBatchMode=yes']
+    if args.extra_scp_args is not None:
+        cmd.extend(shlex.split(args.extra_scp_args))
+    return cmd + [src, '{}:{}'.format(args.host, dst)]
+
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--host', type=str, required=True)
     parser.add_argument('--execdir', type=str, required=True)
+    parser.add_argument('--extra-ssh-args', type=str, required=False)
+    parser.add_argument('--extra-scp-args', type=str, required=False)
     parser.add_argument('--codesign_identity', type=str, required=False, default=None)
     parser.add_argument('--env', type=str, nargs='*', required=False, default=dict())
     parser.add_argument("command", nargs=argparse.ONE_OR_MORE)
     args = parser.parse_args()
     commandLine = args.command
 
-    ssh = lambda command: ['ssh', '-oBatchMode=yes', args.host, command]
-    scp = lambda src, dst: ['scp', '-q', '-oBatchMode=yes', src, '{}:{}'.format(args.host, dst)]
-
     # Create a temporary directory where the test will be run.
     # That is effectively the value of %T on the remote host.
-    tmp = subprocess.check_output(ssh('mktemp -d /tmp/libcxx.XXXXXXXXXX'), universal_newlines=True).strip()
+    tmp = subprocess.check_output(ssh(args, 'mktemp -d /tmp/libcxx.XXXXXXXXXX'), universal_newlines=True).strip()
 
     # HACK:
     # If an argument is a file that ends in `.tmp.exe`, assume it is the name
     # of an executable generated by a test file. We call these test-executables
     # below. This allows us to do custom processing like codesigning test-executables
     # and changing their path when running on the remote host. It's also possible
     # for there to be no such executable, for example in the case of a .sh.cpp
     # test.
try:
             tmpTar = tempfile.NamedTemporaryFile(suffix='.tar', delete=False)
             with tarfile.open(fileobj=tmpTar, mode='w') as tarball:
                 tarball.add(args.execdir, arcname=os.path.basename(args.execdir))
 
             # Make sure we close the file before we scp it, because accessing
             # the temporary file while still open doesn't work on Windows.
             tmpTar.close()
             remoteTarball = pathOnRemote(tmpTar.name)
-            subprocess.check_call(scp(tmpTar.name, remoteTarball))
+            subprocess.check_call(scp(args, tmpTar.name, remoteTarball))
         finally:
             # Make sure we close the file in case an exception happens before
             # we've closed it above -- otherwise close() is idempotent.
             tmpTar.close()
             os.remove(tmpTar.name)
 
         # Untar the dependencies in the temporary directory and remove the tarball.
         remoteCommands = [
try:
         # temporary directory on the remote host.
         commandLine = (pathOnRemote(x) if isTestExe(x) else x for x in commandLine)
         remoteCommands.append('cd {}'.format(tmp))
         if args.env:
             remoteCommands.append('export {}'.format(' '.join(args.env)))
         remoteCommands.append(subprocess.list2cmdline(commandLine))
 
         # Finally, SSH to the remote host and execute all the commands.
-        rc = subprocess.call(ssh(' && '.join(remoteCommands)))
+        rc = subprocess.call(ssh(args, ' && '.join(remoteCommands)))
         return rc
 
     finally:
         # Make sure the temporary directory is removed when we're done.
-        subprocess.check_call(ssh('rm -r {}'.format(tmp)))
+        subprocess.check_call(ssh(args, 'rm -r {}'.format(tmp)))
 
 
 if __name__ == '__main__':
     exit(main())