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
#
# There's an alternative container which can be used to test crosvm in crOS tree.
# It can be launched with:
#   ./tools/dev_container --cros

import argparse
from pathlib import Path
import shutil
from impl.util import (
    add_common_args,
    confirm,
    cros_repo_root,
    CROSVM_ROOT,
    is_cros_repo,
    is_kiwi_repo,
    kiwi_repo_root,
    is_aosp_repo,
    aosp_repo_root,
)
from impl.command import (
    chdir,
    cmd,
    quoted,
)
from typing import Optional, List
import getpass
import sys
import unittest
import os
import zlib

DEV_CONTAINER_NAME = (
    f"crosvm_dev_{getpass.getuser()}_{zlib.crc32(os.path.realpath(__file__).encode('utf-8')):x}"
)
CROS_CONTAINER_NAME = (
    f"crosvm_cros_{getpass.getuser()}_{zlib.crc32(os.path.realpath(__file__).encode('utf-8')):x}"
)

DEV_IMAGE_NAME = "gcr.io/crosvm-infra/crosvm_dev"
CROS_IMAGE_NAME = "gcr.io/crosvm-infra/crosvm_cros_cloudbuild"
DEV_IMAGE_VERSION = (CROSVM_ROOT / "tools/impl/dev_container/version").read_text().strip()

CACHE_DIR = os.environ.get("CROSVM_CONTAINER_CACHE", None)

COMMON_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" if Path("/dev/kvm").is_char_device() else None,
    # Enable terminal colors
    f"--env TERM={os.environ.get('TERM', 'xterm-256color')}",
]

DOCKER_ARGS = [
    *COMMON_ARGS,
]

PODMAN_ARGS = [
    *COMMON_ARGS,
    # Allow access to group permissions of the user (e.g. for kvm access).
    "--group-add keep-groups" if os.name == "posix" else None,
    # Increase number of PIDs the container can spawn (we run a lot of test processes in parallel)
    "--pids-limit=4096" if os.name == "posix" else None,
]

# Environment variables to pass through to the container if they are specified.
ENV_PASSTHROUGH = [
    "NEXTEST_PROFILE",
    "http_proxy",
    "https_proxy",
]


def machine_is_running(docker: cmd):
    machine_state = docker("machine info").stdout()
    return "MachineState: Running" in machine_state


def container_name(cros: bool):
    if cros:
        return CROS_CONTAINER_NAME
    else:
        return DEV_CONTAINER_NAME


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, cros: bool):
    return docker(f"ps -a -q -f name={container_name(cros)}").stdout()


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


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


def workspace_mount_args(cros: bool):
    """
    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 cros:
        return ["--workdir /home/crosvmdev/chromiumos/src/platform/crosvm"]
    elif is_cros_repo():
        return [
            f"--volume {quoted(cros_repo_root())}:/workspace:rw",
            "--workdir /workspace/src/platform/crosvm",
        ]
    elif is_kiwi_repo():
        return [
            f"--volume {quoted(kiwi_repo_root())}:/workspace:rw",
            # We override /scratch because we run out of memory if we use memory to back the
            # `/scratch` mount point.
            f"--volume {quoted(kiwi_repo_root())}/scratch:/scratch/cargo_target:rw",
            "--workdir /workspace/platform/crosvm",
        ]
    elif is_aosp_repo():
        return [
            f"--volume {quoted(aosp_repo_root())}:/workspace:rw",
            "--workdir /workspace/external/crosvm",
        ]
    else:
        return [
            f"--volume {quoted(CROSVM_ROOT)}:/workspace:rw",
        ]


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

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


def validate_podman(podman: cmd):
    graph_driver_name = podman("info --format={{.Store.GraphDriverName}}").stdout()
    config_file_name = podman("info --format={{.Store.ConfigFile}}").stdout()
    if graph_driver_name == "vfs":
        print("You are using vfs as a storage driver. This will be extremely slow.")
        print("Using the overlay driver is strongly recommended.")
        print("Note: This will delete all existing podman images and containers.")
        if confirm(f"Do you want me to update your config in {config_file_name}?"):
            podman("system reset -f").fg()
            with open(config_file_name, "a") as config_file:
                print("[storage]", file=config_file)
                print('driver = "overlay"', file=config_file)

    if os.name == "posix":
        username = os.environ["USER"]
        subuids = Path("/etc/subuid").read_text()
        if not username in subuids:
            print("Rootless podman requires subuid's to be set up for your user.")
            usermod = cmd(
                "sudo usermod --add-subuids 900000-965535 --add-subgids 900000-965535", username
            )
            print("I can fix that by running:", usermod)
            if confirm("Ok?"):
                usermod.fg()
                podman("system migrate").fg()


def main(argv: List[str]):
    parser = argparse.ArgumentParser()
    add_common_args(parser)
    parser.add_argument("--stop", action="store_true")
    parser.add_argument("--clean", action="store_true")
    parser.add_argument("--hermetic", action="store_true")
    parser.add_argument("--no-interactive", action="store_true")
    parser.add_argument("--use-docker", action="store_true")
    parser.add_argument("--self-test", action="store_true")
    parser.add_argument("--pull", action="store_true")
    parser.add_argument("--cros", action="store_true")
    parser.add_argument("command", nargs=argparse.REMAINDER)

    args = parser.parse_args(argv)

    chdir(CROSVM_ROOT)

    if CACHE_DIR:
        Path(CACHE_DIR).mkdir(exist_ok=True)

    has_docker = shutil.which("docker") != None
    has_podman = shutil.which("podman") != None
    if not has_podman and not has_docker:
        raise Exception("Please install podman (or docker) to use the dev container.")

    use_docker = args.use_docker
    if has_docker and not has_podman:
        use_docker = True

    # cros container only works in docker
    if args.cros:
        use_docker = True

    if use_docker:
        print(
            "WARNING: Running dev_container with docker may cause root-owned files to be created."
        )
        print("Use podman to prevent this.")
        print()
        docker = cmd("docker")
        docker_args = [
            *DOCKER_ARGS,
            *workspace_mount_args(args.cros),
        ]
    else:
        docker = cmd("podman")

        # On windows, podman uses wsl vm. start the default podman vm for the rest of the script
        # to work properly.
        if os.name == "nt" and not machine_is_running(docker):
            print("Starting podman default machine.")
            docker("machine start").fg(quiet=True)
        docker_args = [
            *PODMAN_ARGS,
            *workspace_mount_args(args.cros),
        ]
        validate_podman(docker)

    if args.cros:
        docker_args.append("--privileged")  # cros container requires privileged container
        docker_args.append(CROS_IMAGE_NAME)
    else:
        docker_args.append(DEV_IMAGE_NAME + ":" + DEV_IMAGE_VERSION)

    # Add environment variables to command line
    exec_args: List[str] = []
    for key in ENV_PASSTHROUGH:
        value = os.environ.get(key)
        if value is not None:
            exec_args.append("--env")
            exec_args.append(f"{key}={quoted(value)}")

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

    if args.stop:
        if not delete_container(docker, args.cros):
            print(f"container is not running.")
        return

    if args.clean:
        delete_container(docker, args.cros)

    if args.pull:
        if args.cros:
            docker("pull", CROS_IMAGE_NAME).fg()
        else:
            docker("pull", f"gcr.io/crosvm-infra/crosvm_dev:{DEV_IMAGE_VERSION}").fg()
        return

    command = args.command

    # Default to interactive mode if a tty is present.
    tty_args: List[str] = []
    if sys.stdin.isatty():
        tty_args += ["--tty"]
        if not args.no_interactive:
            tty_args += ["--interactive"]

    # Start an interactive shell by default
    if args.hermetic:
        # cmd is passed to entrypoint
        quoted_cmd = list(map(quoted, command))
        docker(f"run --rm", *tty_args, *docker_args, *exec_args, *quoted_cmd).fg()
    else:
        # cmd is executed directly
        cid = ensure_container_is_alive(docker, docker_args, args.cros)
        if not command:
            command = ("/bin/bash",)
        quoted_cmd = list(map(quoted, command))
        docker("exec", *tty_args, *exec_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(cros=False),
        *DOCKER_ARGS,
    ]

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

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

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

    def test_handling_of_stopped_container(self):
        cid = ensure_container_is_alive(self.docker, self.docker_args, cros=False)
        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, cros=False)
        self.assertTrue(container_is_running(self.docker, cros=False))
        main(["true"])


if __name__ == "__main__":
    main(sys.argv[1:])