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