From 1dadb3eea114329cf70ae29531b5451ba82cfc4a Mon Sep 17 00:00:00 2001
From: Gregory Szorc <gps@mozilla.com>
Date: Fri, 29 Jul 2016 20:53:21 -0700
Subject: [PATCH] Bug 1290620 - Implement a run-task wrapper script; r=dustin

Before, we simply executed scripts inside Docker containers. This
frequently resulted in a wall of text with command output. It was
difficult to discern things like the time spent performing certain
actions.

Before, individual tasks had to drop permissions from the default
root user themselves. Dropping permissions isn't exactly a trivial
thing to do and a number of tasks didn't do it or did it wrong.

Before, we had a "checkout-gecko-and-run" script that kinda/sorta
did common activities for us. But it was written as a shell script
and doing advanced things was difficult.

This commit can be treated as a rewrite of "checkout-gecko-and-run"
as a Python script. But it also does a bit more. It prefixes output
with timestamps so we know how long operations took. It features more
robust argument parsing, so we can add new features more easily.

To prove the new wrapper script works, the lint image and all tasks
using it have been converted to use it.

MozReview-Commit-ID: 5d95u5Xebtq

--HG--
extra : rebase_source : 3a1d84782b01d7743e846bd0c04d7867813dd8a3
---
 .../ci/legacy/tasks/tests/eslint-gecko.yml    |   5 +-
 .../legacy/tasks/tests/mozharness-gecko.yml   |   5 +-
 .../ci/legacy/tasks/tests/mozlint-flake8.yml  |   5 +-
 .../ci/legacy/tasks/tests/taskgraph-tests.yml |   5 +-
 testing/docker/lint/Dockerfile                |   4 +-
 testing/docker/recipes/install-mercurial.sh   |   3 +
 testing/docker/recipes/run-task               | 168 ++++++++++++++++++
 7 files changed, 185 insertions(+), 10 deletions(-)
 create mode 100755 testing/docker/recipes/run-task

diff --git a/taskcluster/ci/legacy/tasks/tests/eslint-gecko.yml b/taskcluster/ci/legacy/tasks/tests/eslint-gecko.yml
index df70df678326..0726020d2b2d 100644
--- a/taskcluster/ci/legacy/tasks/tests/eslint-gecko.yml
+++ b/taskcluster/ci/legacy/tasks/tests/eslint-gecko.yml
@@ -20,8 +20,9 @@ task:
         task-reference: "<docker-image>"
 
     command:
-      - /home/worker/bin/checkout-gecko-and-run
-      - /home/worker/checkouts/gecko
+      - /home/worker/bin/run-task
+      - '--vcs-checkout=/home/worker/checkouts/gecko'
+      - '--'
       - bash
       - -cx
       - >
diff --git a/taskcluster/ci/legacy/tasks/tests/mozharness-gecko.yml b/taskcluster/ci/legacy/tasks/tests/mozharness-gecko.yml
index cb05566b8edb..b6a74cc56dc5 100644
--- a/taskcluster/ci/legacy/tasks/tests/mozharness-gecko.yml
+++ b/taskcluster/ci/legacy/tasks/tests/mozharness-gecko.yml
@@ -26,8 +26,9 @@ task:
       level-{{level}}-{{project}}-dotcache: '/home/worker/.cache'
 
     command:
-      - /home/worker/bin/checkout-gecko-and-run
-      - /home/worker/checkouts/gecko
+      - /home/worker/bin/run-task
+      - '--vcs-checkout=/home/worker/checkouts/gecko'
+      - '--'
       - bash
       - -cx
       - >
diff --git a/taskcluster/ci/legacy/tasks/tests/mozlint-flake8.yml b/taskcluster/ci/legacy/tasks/tests/mozlint-flake8.yml
index 05b00a3f6a3c..1b7cce17db9e 100644
--- a/taskcluster/ci/legacy/tasks/tests/mozlint-flake8.yml
+++ b/taskcluster/ci/legacy/tasks/tests/mozlint-flake8.yml
@@ -18,8 +18,9 @@ task:
       taskId:
         task-reference: "<docker-image>"
     command:
-      - /home/worker/bin/checkout-gecko-and-run
-      - /home/worker/checkouts/gecko
+      - /home/worker/bin/run-task
+      - '--vcs-checkout=/home/worker/checkouts/gecko'
+      - '--'
       - bash
       - -cx
       - >
diff --git a/taskcluster/ci/legacy/tasks/tests/taskgraph-tests.yml b/taskcluster/ci/legacy/tasks/tests/taskgraph-tests.yml
index 6a2ab63a9914..a78e4f11450d 100644
--- a/taskcluster/ci/legacy/tasks/tests/taskgraph-tests.yml
+++ b/taskcluster/ci/legacy/tasks/tests/taskgraph-tests.yml
@@ -18,8 +18,9 @@ task:
       taskId:
         task-reference: "<docker-image>"
     command:
-      - /home/worker/bin/checkout-gecko-and-run
-      - /home/worker/checkouts/gecko
+      - /home/worker/bin/run-task
+      - '--vcs-checkout=/home/worker/checkouts/gecko'
+      - '--'
       - bash
       - -cx
       - >
diff --git a/testing/docker/lint/Dockerfile b/testing/docker/lint/Dockerfile
index 26d64ee9d162..4a8bae631b9f 100644
--- a/testing/docker/lint/Dockerfile
+++ b/testing/docker/lint/Dockerfile
@@ -13,8 +13,8 @@ ADD topsrcdir/testing/docker/recipes/install-mercurial.sh /build/install-mercuri
 ADD system-setup.sh /tmp/system-setup.sh
 RUN bash /tmp/system-setup.sh
 
-# %include testing/docker/recipes/checkout-gecko-and-run
-ADD topsrcdir/testing/docker/recipes/checkout-gecko-and-run /home/worker/bin/checkout-gecko-and-run
+# %include testing/docker/recipes/run-task
+ADD topsrcdir/testing/docker/recipes/run-task /home/worker/bin/run-task
 RUN chown -R worker:worker /home/worker/bin && chmod 755 /home/worker/bin/*
 
 # Set variable normally configured at login, by the shells parent process, these
diff --git a/testing/docker/recipes/install-mercurial.sh b/testing/docker/recipes/install-mercurial.sh
index bc1ea620266c..c15d63088efd 100644
--- a/testing/docker/recipes/install-mercurial.sh
+++ b/testing/docker/recipes/install-mercurial.sh
@@ -45,9 +45,12 @@ mkdir -p /etc/mercurial
 cat >/etc/mercurial/hgrc <<EOF
 # By default the progress bar starts after 3s and updates every 0.1s. We
 # change this so it shows and updates every 1.0s.
+# We also tell progress to assume a TTY is present so updates are printed
+# even if there is no known TTY.
 [progress]
 delay = 1.0
 refresh = 1.0
+assume-tty = true
 
 [web]
 cacerts = /etc/ssl/certs/ca-certificates.crt
diff --git a/testing/docker/recipes/run-task b/testing/docker/recipes/run-task
new file mode 100755
index 000000000000..35d7b7f476b6
--- /dev/null
+++ b/testing/docker/recipes/run-task
@@ -0,0 +1,168 @@
+#!/usr/bin/python2.7 -u
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+"""Run a task after performing common actions.
+
+This script is meant to be the "driver" for TaskCluster based tasks.
+It receives some common arguments to control the run-time environment.
+
+It performs actions as requested from the arguments. Then it executes
+the requested process and prints its output, prefixing it with the
+current time to improve log usefulness.
+"""
+
+from __future__ import absolute_import, print_function, unicode_literals
+
+import argparse
+import datetime
+import errno
+import grp
+import os
+import pwd
+import subprocess
+import sys
+
+
+def print_line(prefix, m):
+    now = datetime.datetime.utcnow()
+    print('[%s %sZ] %s' % (prefix, now.isoformat(), m), end='')
+
+
+def run_and_prefix_output(prefix, args):
+    """Runs a process and prefixes its output with the time.
+
+    Returns the process exit code.
+    """
+    print_line(prefix, 'executing %s\n' % args)
+
+    # Note: TaskCluster's stdin is a TTY. This attribute is lost
+    # when we pass sys.stdin to the invoked process. If we cared
+    # to preserve stdin as a TTY, we could make this work. But until
+    # someone needs it, don't bother.
+    p = subprocess.Popen(args,
+                         # Disable buffering because we want to receive output
+                         # as it is generated so timestamps in logs are
+                         # accurate.
+                         bufsize=0,
+                         stdout=subprocess.PIPE,
+                         stderr=subprocess.STDOUT,
+                         stdin=sys.stdin.fileno(),
+                         cwd='/',
+                         # So \r in progress bars are rendered as multiple
+                         # lines, preserving progress indicators.
+                         universal_newlines=True)
+
+    while True:
+        data = p.stdout.readline()
+        if data == '':
+            break
+
+        print_line(prefix, data)
+
+    return p.returncode
+
+
+def vcs_checkout(args):
+    # TODO get VCS parameters from arguments.
+    base_repo = os.environ.get('GECKO_BASE_REPOSITORY')
+
+    # We set the base repository to mozilla-central so tc-vcs doesn't get
+    # confused. Switch to mozilla-unified because robustcheckout works best
+    # with it.
+    if base_repo == 'https://hg.mozilla.org/mozilla-central':
+        base_repo = b'https://hg.mozilla.org/mozilla-unified'
+
+    res = run_and_prefix_output(b'vcs', [
+        b'/usr/bin/hg', b'robustcheckout',
+        b'--sharebase', b'/home/worker/hg-shared',
+        b'--purge',
+        b'--upstream', base_repo,
+        b'--revision', os.environ['GECKO_HEAD_REV'],
+        os.environ['GECKO_HEAD_REPOSITORY'], args.vcs_checkout
+    ])
+
+    if res:
+        sys.exit(res)
+
+
+def main(args):
+    print_line('setup', 'run-task started\n')
+
+    if os.getuid() != 0:
+        print('assertion failed: not running as root')
+        return 1
+
+    # Arguments up to '--' are ours. After are for the main task
+    # to be executed.
+    try:
+        i = args.index('--')
+        our_args = args[0:i]
+        task_args = args[i + 1:]
+    except ValueError:
+        our_args = args
+        task_args = []
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--user', default='worker', help='user to run as')
+    parser.add_argument('--group', default='worker', help='group to run as')
+    parser.add_argument('--vcs-checkout',
+                        help='Directory where Gecko checkout should be created')
+
+    args = parser.parse_args(our_args)
+
+    try:
+        user = pwd.getpwnam(args.user)
+    except KeyError:
+        print('could not find user %s; specify --user to a known user' %
+              args.user)
+        return 1
+    try:
+        group = grp.getgrnam(args.group)
+    except KeyError:
+        print('could not find group %s; specify --group to a known group' %
+              args.group)
+        return 1
+
+    uid = user.pw_uid
+    gid = group.gr_gid
+
+    # Find all groups to which this user is a member.
+    gids = [g.gr_gid for g in grp.getgrall() if args.group in g.gr_mem]
+
+    checkout = args.vcs_checkout
+    if checkout:
+        # Ensure the directory for the source checkout exists.
+        try:
+            os.makedirs(os.path.dirname(checkout))
+        except OSError as e:
+            if e.errno != errno.EEXIST:
+                raise
+
+        # And that it is owned by the appropriate user/group.
+        os.chown('/home/worker/hg-shared', uid, gid)
+        os.chown(os.path.dirname(checkout), uid, gid)
+
+    # Drop permissions to requested user.
+    # This code is modeled after what `sudo` was observed to do in a Docker
+    # container. We do not bother calling setrlimit() because containers have
+    # their own limits.
+    print_line('setup', 'running as %s:%s\n' % (args.user, args.group))
+    os.setgroups(gids)
+    os.umask(022)
+    os.setresgid(gid, gid, gid)
+    os.setresuid(uid, uid, uid)
+
+    if checkout:
+        vcs_checkout(args)
+
+    return run_and_prefix_output('task', task_args)
+
+
+if __name__ == '__main__':
+    # Unbuffer stdio.
+    sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)
+    sys.stderr = os.fdopen(sys.stderr.fileno(), 'w', 0)
+
+    sys.exit(main(sys.argv[1:]))