crosvm/tools/dev_container

#!/usr/bin/env python3
# Copyright 2021 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# Usage:
#
# To get an interactive shell for development:
#   ./tools/dev_container
#
# To run a command in the container, e.g. to run presubmits:
#   ./tools/dev_container ./tools/presubmit
#
# The state of the container (including build artifacts) are preserved between
# calls. To stop the container call:
#   ./tools/dev_container --stop
#
# The dev container can also be called with a fresh container for each call that
# is cleaned up afterwards (e.g. when run by Kokoro):
#
#   ./tools/dev_container --hermetic CMD

import argparse
from argh import arg  # type: ignore
from impl.common import CROSVM_ROOT, cmd, chdir, cros_repo_root, is_cros_repo, quoted, run_main
from typing import Optional, Tuple, List
import getpass
import shutil
import sys
import unittest
import os
import zlib

CONTAINER_NAME = (
    f"crosvm_dev_{getpass.getuser()}_{zlib.crc32(os.path.realpath(__file__).encode('utf-8')):x}"
)
IMAGE_VERSION = (CROSVM_ROOT / "tools/impl/dev_container/version").read_text().strip()
CACHE_DIR = os.environ.get("CROSVM_CONTAINER_CACHE", None)

DOCKER_ARGS = [
    # Share cache dir
    f"--volume {CACHE_DIR}:/cache:rw" if CACHE_DIR else None,
    # Use tmpfs in the container for faster performance.
    "--mount type=tmpfs,destination=/tmp",
    # KVM is required to run a VM for testing.
    "--device /dev/kvm",
    f"--env OUTSIDE_UID={os.getuid()}",
    f"--env OUTSIDE_GID={os.getgid()}",
    f"gcr.io/crosvm-infra/crosvm_dev_user:{IMAGE_VERSION}",
]

PODMAN_ARGS = [
    # Share cache dir
    f"--volume {CACHE_DIR}:/cache:rw" if CACHE_DIR else None,
    # Use tmpfs in the container for faster performance.
    "--mount type=tmpfs,destination=/tmp",
    # KVM is required to run a VM for testing.
    "--device /dev/kvm",
    f"gcr.io/crosvm-infra/crosvm_dev:{IMAGE_VERSION}",
]

PRIVILEGED_ARGS = [
    # Share devices and syslog
    "--volume /dev/log:/dev/log",
    "--device /dev/net/tun",
    "--device /dev/vhost-net",
    "--device /dev/vhost-vsock",
    # For plugin process jail
    "--mount type=tmpfs,destination=/var/empty",
]


PODMAN_IS_DEFAULT = shutil.which("docker") == None


def container_revision(docker: cmd, container_id: str):
    image = docker("container inspect -f {{.Config.Image}}", container_id).stdout()
    parts = image.split(":")
    assert len(parts) == 2, f"Invalid image name {image}"
    return parts[1]


def container_id(docker: cmd):
    return docker(f"ps -a -q -f name={CONTAINER_NAME}").stdout()


def container_is_running(docker: cmd):
    return bool(docker(f"ps -q -f name={CONTAINER_NAME}").stdout())


def delete_container(docker: cmd):
    cid = container_id(docker)
    if cid:
        print(f"Deleting dev-container {cid}.")
        docker("rm -f", cid).fg(quiet=True)
        return True
    return False


def workspace_mount_args():
    """
    Returns arguments for mounting the crosvm sources to /workspace.

    In ChromeOS checkouts the crosvm repo uses a symlink or worktree checkout, which links to a
    different folder in the ChromeOS checkout. So we need to mount the whole CrOS checkout.
    """
    if is_cros_repo():
        return [
            f"--volume {quoted(cros_repo_root())}:/workspace:rw",
            "--workdir /workspace/src/platform/crosvm",
        ]
    else:
        return [
            f"--volume {quoted(CROSVM_ROOT)}:/workspace:rw",
        ]


def ensure_container_is_alive(docker: cmd, docker_args: List[Optional[str]]):
    cid = container_id(docker)
    if cid and not container_is_running(docker):
        print("Existing dev-container is not running.")
        delete_container(docker)
    elif cid and container_revision(docker, cid) != IMAGE_VERSION:
        print(f"New image is available.")
        delete_container(docker)

    if not container_is_running(docker):
        # Run neverending sleep to keep container alive while we 'docker exec' commands.
        docker(f"run --detach --name {CONTAINER_NAME}", *docker_args, "sleep infinity").fg(
            quiet=True
        )
        cid = container_id(docker)
        print(f"Started dev-container ({cid}).")
    else:
        cid = container_id(docker)
        print(f"Using existing dev-container ({cid}).")
    return cid


@arg("command", nargs=argparse.REMAINDER)
def main(
    command: Tuple[str, ...],
    stop: bool = False,
    clean: bool = False,
    hermetic: bool = False,
    interactive: bool = False,
    podman: bool = PODMAN_IS_DEFAULT,
    self_test: bool = False,
    pull: bool = False,
    unprivileged: bool = False,
):
    chdir(CROSVM_ROOT)

    if unprivileged:
        print("WARNING: Running dev_container with --unprivileged is a work in progress.")
        print("Not all tests are expected to pass.")
        print()

    docker_args = [
        *workspace_mount_args(),
        *(PRIVILEGED_ARGS if not unprivileged else []),
    ]
    if podman:
        print("WARNING: Running dev_container with podman is experimental.")
        print("It is strongly recommended to use docker.")
        print()
        docker = cmd("podman")
        docker_args += [*PODMAN_ARGS]
    else:
        docker = cmd("docker")
        docker_args += [
            "--privileged" if not unprivileged else None,
            *DOCKER_ARGS,
        ]

    if self_test:
        TestDevContainer.docker = docker
        suite = unittest.defaultTestLoader.loadTestsFromTestCase(TestDevContainer)
        unittest.TextTestRunner().run(suite)
        return

    if stop:
        if not delete_container(docker):
            print(f"Dev-container is not running.")
        return

    if clean:
        delete_container(docker)

    if pull:
        docker("pull", f"gcr.io/crosvm-infra/crosvm_dev:{IMAGE_VERSION}").fg()
        docker("pull", f"gcr.io/crosvm-infra/crosvm_dev_user:{IMAGE_VERSION}").fg()
        return

    # If a command is provided run non-interactive unless explicitly asked for.
    tty_args = []
    if not command or interactive:
        if not sys.stdin.isatty():
            raise Exception("Trying to run an interactive session in a non-interactive terminal.")
        tty_args = ["--interactive", "--tty"]
    elif sys.stdin.isatty():
        # Even if run non-interactively, we do want to pass along a tty for proper output.
        tty_args = ["--tty"]

    # Start an interactive shell by default
    if hermetic:
        # cmd is passed to entrypoint
        quoted_cmd = list(map(quoted, command))
        docker(f"run --rm", *tty_args, *docker_args, *quoted_cmd).fg()
    else:
        # cmd is executed directly
        cid = ensure_container_is_alive(docker, docker_args)
        if podman:
            if not command:
                command = ("/bin/bash",)
        else:
            if not command:
                command = ("/tools/entrypoint.sh",)
            else:
                command = ("/tools/entrypoint.sh",) + tuple(command)
        quoted_cmd = list(map(quoted, command))
        docker("exec", *tty_args, cid, *quoted_cmd).fg()


class TestDevContainer(unittest.TestCase):
    """
    Runs live tests using the docker service.

    Note: This test is not run by health-check since it cannot be run inside the
    container. It is run by infra/recipes/health_check.py before running health checks.
    """

    docker: cmd
    docker_args = [
        *workspace_mount_args(),
        *DOCKER_ARGS,
    ]

    def setUp(self):
        # Start with a stopped container for each test.
        delete_container(self.docker)

    def test_stopped_container(self):
        # Create but do not run a new container.
        self.docker(f"create --name {CONTAINER_NAME}", *self.docker_args, "sleep infinity").stdout()
        self.assertTrue(container_id(self.docker))
        self.assertFalse(container_is_running(self.docker))

    def test_container_reuse(self):
        cid = ensure_container_is_alive(self.docker, self.docker_args)
        cid2 = ensure_container_is_alive(self.docker, self.docker_args)
        self.assertEqual(cid, cid2)

    def test_handling_of_stopped_container(self):
        cid = ensure_container_is_alive(self.docker, self.docker_args)
        self.docker("kill", cid).fg()

        # Make sure we can get back into a good state and execute commands.
        ensure_container_is_alive(self.docker, self.docker_args)
        self.assertTrue(container_is_running(self.docker))
        main(("true",))


if __name__ == "__main__":
    run_main(main)