• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2021 The ChromiumOS Authors
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import argparse
6import itertools
7import json
8import os
9import shutil
10import socket
11import subprocess
12import sys
13import time
14import urllib.request as request
15from contextlib import closing
16from pathlib import Path
17from typing import Dict, Iterable, List, Literal, Optional, Tuple
18
19from .common import CACHE_DIR, download_file
20
21USAGE = """%(prog)s {command} [options]
22
23Manages VMs for testing crosvm.
24
25Can run an x86_64 and an aarch64 vm via `./tools/x86vm` and `./tools/aarch64vm`.
26The VM image will be downloaded and initialized on first use.
27
28The easiest way to use the VM is:
29
30  $ ./tools/aarch64vm ssh
31
32Which will initialize and boot the VM, then wait for SSH to be available and
33opens an SSH session. The VM will stay alive between calls.
34
35Alternatively, you can set up an SSH config to connect to the VM. First ensure
36the VM ready:
37
38  $ ./tools/aarch64vm wait
39
40Then append the VMs ssh config to your SSH config:
41
42  $ ./tools/aarch64vm ssh_config >> ~/.ssh/config
43
44And connect as usual:
45
46  $ ssh crosvm_$ARCH
47
48Commands:
49
50  build: Download base image and create rootfs overlay.
51  up: Ensure that the VM is running in the background.
52  run: Run the VM in the foreground process for debugging.
53  wait: Boot the VM if it's offline and wait until it's available.
54  ssh: SSH into the VM. Boot and wait if it's not available.
55  ssh_config: Prints the ssh config needed to connnect to the VM.
56  stop: Gracefully shutdown the VM.
57  kill: Kill the QEMU process. Might damage the image file.
58  clean: Stop all VMs and delete all data.
59"""
60
61KVM_SUPPORT = os.access("/dev/kvm", os.W_OK)
62
63Arch = Literal["x86_64", "aarch64"]
64
65SCRIPT_DIR = Path(__file__).parent.resolve()
66SRC_DIR = SCRIPT_DIR.joinpath("testvm")
67ID_RSA = SRC_DIR.joinpath("id_rsa")
68BASE_IMG_VERSION = open(SRC_DIR.joinpath("version"), "r").read().strip()
69
70IMAGE_DIR_URL = "https://storage.googleapis.com/crosvm/testvm"
71
72
73def cargo_target_dir():
74    # Do not call cargo if we have the environment variable specified. This
75    # allows the script to be used when cargo is not available but the target
76    # dir is known.
77    env_target = os.environ.get("CARGO_TARGET_DIR")
78    if env_target:
79        return Path(env_target)
80    text = subprocess.run(
81        ["cargo", "metadata", "--no-deps", "--format-version=1"],
82        check=True,
83        capture_output=True,
84        text=True,
85    ).stdout
86    metadata = json.loads(text)
87    return Path(metadata["target_directory"])
88
89
90def data_dir(arch: Arch):
91    return CACHE_DIR.joinpath("crosvm_tools").joinpath(arch)
92
93
94def pid_path(arch: Arch):
95    return data_dir(arch).joinpath("pid")
96
97
98def base_img_name(arch: Arch):
99    return f"base-{arch}-{BASE_IMG_VERSION}.qcow2"
100
101
102def base_img_url(arch: Arch):
103    return f"{IMAGE_DIR_URL}/{base_img_name(arch)}"
104
105
106def base_img_path(arch: Arch):
107    return data_dir(arch).joinpath(base_img_name(arch))
108
109
110def rootfs_img_path(arch: Arch):
111    return data_dir(arch).joinpath(f"rootfs-{arch}-{BASE_IMG_VERSION}.qcow2")
112
113
114# List of ports to use for SSH for each architecture
115SSH_PORTS: Dict[Arch, int] = {
116    "x86_64": 9000,
117    "aarch64": 9001,
118}
119
120# QEMU arguments shared by all architectures
121SHARED_ARGS: List[Tuple[str, str]] = [
122    ("-display", "none"),
123    ("-device", "virtio-net-pci,netdev=net0"),
124    ("-smp", "8"),
125    ("-m", "4G"),
126]
127
128# Arguments to QEMU for each architecture
129ARCH_TO_QEMU: Dict[Arch, Tuple[str, List[Iterable[str]]]] = {
130    # arch: (qemu-binary, [(param, value), ...])
131    "x86_64": (
132        "qemu-system-x86_64",
133        [
134            ("-cpu", "host"),
135            ("-netdev", f"user,id=net0,hostfwd=tcp::{SSH_PORTS['x86_64']}-:22"),
136            *([("-enable-kvm",)] if KVM_SUPPORT else []),
137            *SHARED_ARGS,
138        ],
139    ),
140    "aarch64": (
141        "qemu-system-aarch64",
142        [
143            ("-M", "virt"),
144            ("-machine", "virt,virtualization=true,gic-version=3"),
145            ("-cpu", "cortex-a57"),
146            ("-bios", "/usr/share/qemu-efi-aarch64/QEMU_EFI.fd"),
147            (
148                "-netdev",
149                f"user,id=net0,hostfwd=tcp::{SSH_PORTS['aarch64']}-:22",
150            ),
151            *SHARED_ARGS,
152        ],
153    ),
154}
155
156
157def ssh_opts(arch: Arch) -> Dict[str, str]:
158    return {
159        "Port": str(SSH_PORTS[arch]),
160        "User": "crosvm",
161        "StrictHostKeyChecking": "no",
162        "UserKnownHostsFile": "/dev/null",
163        "LogLevel": "ERROR",
164        "IdentityFile": str(ID_RSA),
165    }
166
167
168def ssh_cmd_args(arch: Arch):
169    return [f"-o{k}={v}" for k, v in ssh_opts(arch).items()]
170
171
172def ssh_exec(arch: Arch, cmd: Optional[str] = None):
173    subprocess.run(
174        [
175            "ssh",
176            "localhost",
177            *ssh_cmd_args(arch),
178            *(["-T", cmd] if cmd else []),
179        ],
180    ).check_returncode()
181
182
183def ping_vm(arch: Arch):
184    os.chmod(ID_RSA, 0o600)
185    return (
186        subprocess.run(
187            [
188                "ssh",
189                "localhost",
190                *ssh_cmd_args(arch),
191                "-oConnectTimeout=1",
192                "-T",
193                "exit",
194            ],
195            capture_output=True,
196        ).returncode
197        == 0
198    )
199
200
201def write_pid_file(arch: Arch, pid: int):
202    with open(pid_path(arch), "w") as pid_file:
203        pid_file.write(str(pid))
204
205
206def read_pid_file(arch: Arch):
207    if not pid_path(arch).exists():
208        return None
209
210    with open(pid_path(arch), "r") as pid_file:
211        return int(pid_file.read())
212
213
214def run_qemu(
215    arch: Arch,
216    hda: Path,
217    background: bool = False,
218):
219    if not is_ssh_port_available(arch):
220        print(f"Port {SSH_PORTS[arch]} is occupied, but is required for the {arch} vm to run.")
221        print(f"You may be running the {arch}vm in another place and need to kill it.")
222        sys.exit(1)
223
224    (binary, arch_args) = ARCH_TO_QEMU[arch]
225    qemu_args = [*arch_args, ("-hda", str(hda))]
226    if background:
227        qemu_args.append(("-serial", f"file:{data_dir(arch).joinpath('vm_log')}"))
228    else:
229        qemu_args.append(("-serial", "stdio"))
230
231    # Flatten list of tuples into flat list of arguments
232    qemu_cmd = [binary, *itertools.chain(*qemu_args)]
233    process = subprocess.Popen(qemu_cmd, start_new_session=background)
234    write_pid_file(arch, process.pid)
235    if not background:
236        process.wait()
237
238
239def run_vm(arch: Arch, background: bool = False):
240    run_qemu(
241        arch,
242        rootfs_img_path(arch),
243        background=background,
244    )
245
246
247def is_running(arch: Arch):
248    pid = read_pid_file(arch)
249    if pid is None:
250        return False
251
252    # Send signal 0 to check if the process is alive
253    try:
254        os.kill(pid, 0)
255    except OSError:
256        return False
257    return True
258
259
260def kill_vm(arch: Arch):
261    pid = read_pid_file(arch)
262    if pid:
263        os.kill(pid, 9)
264
265
266def build_if_needed(arch: Arch, reset: bool = False):
267    if reset and is_running(arch):
268        print("Killing existing VM...")
269        kill_vm(arch)
270        time.sleep(1)
271
272    data_dir(arch).mkdir(parents=True, exist_ok=True)
273
274    base_img = base_img_path(arch)
275    if not base_img.exists():
276        print(f"Downloading base image ({base_img_url(arch)})...")
277        download_file(base_img_url(arch), base_img_path(arch))
278
279    rootfs_img = rootfs_img_path(arch)
280    if not rootfs_img.exists() or reset:
281        # The rootfs is backed by the base image generated above. So we can
282        # easily reset to a clean VM by rebuilding an empty rootfs image.
283        print("Creating rootfs overlay...")
284        subprocess.run(
285            [
286                "qemu-img",
287                "create",
288                "-f",
289                "qcow2",
290                "-F",
291                "qcow2",
292                "-b",
293                base_img,
294                rootfs_img,
295                "4G",
296            ]
297        ).check_returncode()
298
299
300def is_ssh_port_available(arch: Arch):
301    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
302        return sock.connect_ex(("127.0.0.1", SSH_PORTS[arch])) != 0
303
304
305def up(arch: Arch):
306    if is_running(arch):
307        return
308
309    print("Booting VM...")
310    run_qemu(
311        arch,
312        rootfs_img_path(arch),
313        background=True,
314    )
315
316
317def run(arch: Arch):
318    if is_running(arch):
319        raise Exception("VM is already running")
320    run_qemu(
321        arch,
322        rootfs_img_path(arch),
323        background=False,
324    )
325
326
327def wait(arch: Arch, timeout: int = 120):
328    if not is_running(arch):
329        up(arch)
330    elif ping_vm(arch):
331        return
332
333    print("Waiting for VM")
334    start_time = time.time()
335    while (time.time() - start_time) < timeout:
336        time.sleep(1)
337        sys.stdout.write(".")
338        sys.stdout.flush()
339        if ping_vm(arch):
340            print()
341            return
342    raise Exception("Timeout while waiting for VM")
343
344
345def ssh(arch: Arch, timeout: int):
346    wait(arch, timeout)
347    ssh_exec(arch)
348
349
350def ssh_config(arch: Arch):
351    print(f"Host crosvm_{arch}")
352    print(f"    HostName localhost")
353    for opt, value in ssh_opts(arch).items():
354        print(f"    {opt} {value}")
355
356
357def stop(arch: Arch):
358    if not is_running(arch):
359        print("VM is not running.")
360        return
361    ssh_exec(arch, "sudo poweroff")
362
363
364def kill(arch: Arch):
365    if not is_running(arch):
366        print("VM is not running.")
367        return
368    kill_vm(arch)
369
370
371def clean(arch: Arch):
372    if is_running(arch):
373        kill(arch)
374    if data_dir(arch).exists():
375        shutil.rmtree(data_dir(arch))
376
377
378def main(arch: Arch, argv: List[str]):
379    COMMANDS = [
380        "build",
381        "up",
382        "run",
383        "wait",
384        "ssh",
385        "ssh_config",
386        "stop",
387        "kill",
388        "clean",
389    ]
390    parser = argparse.ArgumentParser(usage=USAGE)
391    parser.add_argument("command", choices=COMMANDS)
392    parser.add_argument(
393        "--reset",
394        action="store_true",
395        help="Reset VM image to a fresh state. Removes all user modifications.",
396    )
397    parser.add_argument(
398        "--timeout",
399        type=int,
400        default=60,
401        help="Timeout in seconds while waiting for the VM to come up.",
402    )
403    args = parser.parse_args(argv)
404
405    if args.command == "clean":
406        clean(arch)
407        return
408
409    if args.command == "ssh_config":
410        ssh_config(arch)
411        return
412
413    # Ensure the images are built regardless of which command we execute.
414    build_if_needed(arch, reset=args.reset)
415
416    if args.command == "build":
417        return  # Nothing left to do.
418    elif args.command == "run":
419        run(arch)
420    elif args.command == "up":
421        up(arch)
422    elif args.command == "ssh":
423        ssh(arch, args.timeout)
424    elif args.command == "stop":
425        stop(arch)
426    elif args.command == "kill":
427        kill(arch)
428    elif args.command == "wait":
429        wait(arch, args.timeout)
430    else:
431        print(f"Unknown command {args.command}")
432