crosperf: Add support to run tests with ssh forwarding

The change allows running crosperf with local DUT with port forwarding.

Experiment file can define a local DUT with:
remote: 127.0.0.1:<local-port>

BUG=chromium:1064988
TEST=Tested with crosperf on local and remote DUT.

Change-Id: I62446821f83e8ae39d610ca9ed60c9b37af9852c
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/2158122
Reviewed-by: Denis Nikitin <denik@chromium.org>
Reviewed-by: Tiancong Wang <tcwang@google.com>
Tested-by: Denis Nikitin <denik@chromium.org>

1 files changed, 35 insertions, 23 deletions

diff --git a/cros_utils/command_executer.py b/cros_utils/command_executer.py
index b9acbe31..ee1df95c 100755
--- a/cros_utils/command_executer.py
+++ b/cros_utils/command_executer.py
@@ -229,9 +229,11 @@ class CommandExecuter(object):
     kwargs['return_output'] = True
     return self.RunCommandGeneric(*args, **kwargs)
 
-  def RemoteAccessInitCommand(self, chromeos_root, machine):
+  def RemoteAccessInitCommand(self, chromeos_root, machine, port=None):
     command = ''
     command += '\nset -- --remote=' + machine
+    if port:
+      command += ' --ssh_port=' + port
     command += '\n. ' + chromeos_root + '/src/scripts/common.sh'
     command += '\n. ' + chromeos_root + '/src/scripts/remote_access.sh'
     command += '\nTMP=$(mktemp -d)'
@@ -285,12 +287,16 @@ class CommandExecuter(object):
         sys.exit(1)
     chromeos_root = os.path.expanduser(chromeos_root)
 
+    port = None
+    if ':' in machine:
+      machine, port = machine.split(':')
     # Write all commands to a file.
     command_file = self.WriteToTempShFile(cmd)
     retval = self.CopyFiles(
         command_file,
         command_file,
         dest_machine=machine,
+        dest_port=port,
         command_terminator=command_terminator,
         chromeos_root=chromeos_root,
         dest_cros=True,
@@ -302,7 +308,7 @@ class CommandExecuter(object):
                              ' Is the machine up?')
       return (retval, '', '')
 
-    command = self.RemoteAccessInitCommand(chromeos_root, machine)
+    command = self.RemoteAccessInitCommand(chromeos_root, machine, port)
     command += '\nremote_sh bash %s' % command_file
     command += '\nl_retval=$?; echo "$REMOTE_OUT"; exit $l_retval'
     retval = self.RunCommandGeneric(
@@ -313,8 +319,8 @@ class CommandExecuter(object):
         terminated_timeout=terminated_timeout,
         print_to_console=print_to_console)
     if return_output:
-      connect_signature = (
-          'Initiating first contact with remote host\n' + 'Connection OK\n')
+      connect_signature = ('Initiating first contact with remote host\n' +
+                           'Connection OK\n')
       connect_signature_re = re.compile(connect_signature)
       modded_retval = list(retval)
       modded_retval[1] = connect_signature_re.sub('', retval[1])
@@ -449,7 +455,9 @@ class CommandExecuter(object):
                 src,
                 dest,
                 src_machine=None,
+                src_port=None,
                 dest_machine=None,
+                dest_port=None,
                 src_user=None,
                 dest_user=None,
                 recursive=True,
@@ -475,30 +483,34 @@ class CommandExecuter(object):
         sys.exit(1)
       if src_cros:
         cros_machine = src_machine
+        cros_port = src_port
+        host_machine = dest_machine
+        host_user = dest_user
       else:
         cros_machine = dest_machine
-
-      command = self.RemoteAccessInitCommand(chromeos_root, cros_machine)
-      ssh_command = (
-          'ssh -o StrictHostKeyChecking=no' + ' -o UserKnownHostsFile=$(mktemp)'
-          + ' -i $TMP_PRIVATE_KEY')
+        cros_port = dest_port
+        host_machine = src_machine
+        host_user = src_user
+
+      command = self.RemoteAccessInitCommand(chromeos_root, cros_machine,
+                                             cros_port)
+      ssh_command = ('ssh -o StrictHostKeyChecking=no' +
+                     ' -o UserKnownHostsFile=$(mktemp)' +
+                     ' -i $TMP_PRIVATE_KEY')
+      if cros_port:
+        ssh_command += ' -p %s' % cros_port
       rsync_prefix = '\nrsync -r -e "%s" ' % ssh_command
       if dest_cros:
-        command += rsync_prefix + '%s root@%s:%s' % (src, dest_machine, dest)
-        return self.RunCommand(
-            command,
-            machine=src_machine,
-            username=src_user,
-            command_terminator=command_terminator,
-            print_to_console=print_to_console)
+        command += rsync_prefix + '%s root@%s:%s' % (src, cros_machine, dest)
       else:
-        command += rsync_prefix + 'root@%s:%s %s' % (src_machine, src, dest)
-        return self.RunCommand(
-            command,
-            machine=dest_machine,
-            username=dest_user,
-            command_terminator=command_terminator,
-            print_to_console=print_to_console)
+        command += rsync_prefix + 'root@%s:%s %s' % (cros_machine, src, dest)
+
+      return self.RunCommand(
+          command,
+          machine=host_machine,
+          username=host_user,
+          command_terminator=command_terminator,
+          print_to_console=print_to_console)
 
     if dest_machine == src_machine:
       command = 'rsync -a %s %s' % (src, dest)