src/sys/pkg/tests/system-tests/bin/create-emu

#!/usr/bin/env fuchsia-vendored-python
# Copyright 2022 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import asyncio
import json
import os
import pathlib
import shlex
import shutil
import subprocess
import sys
import tempfile
import time

async def print_qemu_logs(qemu_reader, qemu_writer):
  while True:
    line = await qemu_reader.readline()
    if not line:
      break

    print(line.decode('utf-8'), end='')

  qemu_writer.close()
  await qemu_writer.wait_closed()


async def poll_process(process):
  try:
    await asyncio.wait_for(process.wait(), 1)
  except asyncio.TimeoutError:
    pass
  return process.returncode


async def main():
  parser = argparse.ArgumentParser(
      formatter_class=argparse.RawDescriptionHelpFormatter,
      description='''
This command creates a Fuchsia image in the specified directory. As part of
this process, it starts up the image in QEMU so that we can copy in our
ssh authorized keys file. We do this through QEMU because there is no
host-side tooling that lets us inject this file into minfs or fxfs.

Once complete, this image will be ready to be used with:

    $ fx qemu --uefi -D path/to/dir/fuchsia.efi

This command delegates to /run-zircon. Flags
are documented in that script, and can be discovered by passing -h or
--help.''')
  parser.add_argument(
    '--arch',
    default='x64',
    help='emulator architecture (arm64, riscv64, or x64)')
  parser.add_argument(
    '--fuchsia-dir',
    type=pathlib.Path,
    default=os.getenv('FUCHSIA_DIR'),
    help='path to the fuchsia directory (default $FUCHSIA_DIR)')
  parser.add_argument(
    '--fuchsia-build-dir',
    type=pathlib.Path,
    default=os.environ.get('FUCHSIA_BUILD_DIR'),
    help='path to the fuchsia build directory (default $FUCHSIA_BUILD_DIR)')
  parser.add_argument(
    '--edk2-dir',
    type=pathlib.Path,
    default=os.environ.get('PREBUILT_EDK2_DIR'),
    help='path to the fuchsia build directory (default $PREBUILT_EDK2_DIR)')
  parser.add_argument(
    '--image-size',
    type=int,
    default=20 * 1024 * 1024 * 1024,
    help='total image size (default 20 GiB)')
  parser.add_argument(
    '--abr-size',
    type=int,
    help='ABR slot size (default 256 MiB)')
  parser.add_argument(
    '--system-disk-size',
    type=int,
    help='system disk size (default 63 MiB)')
  parser.add_argument(
    '--disk-type',
    type=str,
    help='mount the disk with this type in qemu')
  parser.add_argument(
    '--product-bundle-name',
    required=True,
    help='The name of the product bundle to use to create the image')
  parser.add_argument(
    '--filesystem',
    choices=['fxfs', 'blobfs'],
    default='fxfs',
    help='Use fxfs to create the image')
  parser.add_argument(
    'image',
    type=pathlib.Path,
    help='create the image into this path')

  args = parser.parse_args()

  if args.fuchsia_dir is None or not args.fuchsia_dir.exists():
    print(
      '--fuchsia-dir or FUCHSIA_DIR must exist',
      file=sys.stderr)
    return 1

  if args.fuchsia_build_dir is None or not args.fuchsia_build_dir.exists():
    print(
      '--fuchsia-build-dir or FUCHSIA_BUILD_DIR must exist',
      file=sys.stderr)
    return 1

  if os.path.exists(args.image):
    print(
      f'image {args.image} already exists, you should remove it, or run it with `fx qemu --uefi -D {args.image}`',
      file=sys.stderr)

    return 1

  # Check that the product bundle exists.
  with open(args.fuchsia_build_dir / 'product_bundles.json') as f:
    product_bundles = json.load(f)

    for product_bundle in product_bundles:
      if product_bundle['name'] == args.product_bundle_name:
        break
    else:

      print(f'Unknown product bundle "{args.product_bundle_name}", must be one of:',
            file=sys.stderr)

      for product_bundle in product_bundles:
        print(f'  - {product_bundle["name"]}',
              file=sys.stderr)

      return 1

  print(f'Creating {args.image}')

  with tempfile.TemporaryDirectory() as tmp:
    tmp = pathlib.Path(tmp)
    cmdline_path = tmp / 'cmdline'

    with open(cmdline_path, 'w') as f:
      # Ensure that the output is sent to the serial, and that we boot into
      # zedboot.
      f.write(
        'kernel.serial=legacy\n'
        'TERM=xterm-256color\n'
        'kernel.halt-on-panic=true\n'
        'bootloader.default=default\n'
        'bootloader.timeout=5\n')

    ## Construct a Fuchsia image.
    make_fuchsia_vol_args = [
      args.fuchsia_build_dir / 'host-tools/make-fuchsia-vol',
      '--fuchsia-build-dir', args.fuchsia_build_dir,
      '--arch', args.arch,
      '--cmdline', cmdline_path,
      '--resize', str(args.image_size),
      '--product-bundle-name', args.product_bundle_name,
      args.image.resolve(),
    ]

    if args.system_disk_size is not None:
      make_fuchsia_vol_args.extend(['--system-disk-size', str(args.system_disk_size)])

    if args.abr_size is not None:
      make_fuchsia_vol_args.extend(['--abr-size', str(args.abr_size)])

    if args.filesystem == 'fxfs':
      make_fuchsia_vol_args.append('--use-fxfs')

    print(f'Running: {shlex.join(str(s) for s in make_fuchsia_vol_args)}')

    process = await asyncio.create_subprocess_exec(*make_fuchsia_vol_args)
    return_code = await process.wait()
    if return_code != 0:
      print('Failed to create fuchsia image', file=sys.stderr)
      return 1

    print(f'Created {args.image}')

    # Emulators always have the same mac address, so they always have this
    # device name.
    device_name = 'fuchsia-5254-0063-5e7a'

    qemu_socket_path = tmp / 'sock'

    print('Launching QEMU to install the SSH authorized keys. ')

    run_zircon_args = [
      args.fuchsia_dir / 'zircon/scripts/run-zircon',
      '-a', args.arch,
      '-N',
      '--uefi',
      '-S', f'unix:{qemu_socket_path.resolve()},server,nowait',
      '-M', 'null',
      '-D', args.image.resolve(),
      '-q', args.fuchsia_dir.resolve() / f'prebuilt/third_party/qemu/linux-x64/bin',
    ]

    if args.disk_type:
      run_zircon_args.append(f'--disktype={args.disk_type}')

    print(f'Running: {shlex.join(str(s) for s in run_zircon_args)}')

    qemu_process = await asyncio.create_subprocess_exec(
      *run_zircon_args,
      stdin=subprocess.DEVNULL,
      close_fds=True,
    )

    try:
      print('Waiting to connect to qemu serial...')

      for i in range(1000):
        print(f'check {qemu_socket_path}: {qemu_socket_path.exists()}')
        if await poll_process(qemu_process) is not None:
          print('qemu exited early', file=sys.stderr)
          return 1

        if qemu_socket_path.exists():
          break
        else:
          await asyncio.sleep(0.05)
      else:
        print('Failed to connect to serial port', file=sys.stderr)
        return 1

      print('Connecting to socket...')

      qemu_reader, qemu_writer = await asyncio.open_unix_connection(str(qemu_socket_path))

      print('Connected to socket')

      # Enter fastboot.
      while True:
        line = await qemu_reader.readline()
        if not line:
          break

        print(line.decode('utf-8'), end='')

        if line == b'Press f to enter fastboot.\r\n':
          print('Entering fastboot...')
          qemu_writer.write(b'f\r\n')
          await qemu_writer.drain()
          break

      log_task = asyncio.create_task(print_qemu_logs(qemu_reader, qemu_writer))

      print('Flashing device...')

      # Flash the device.
      ffx_path = args.fuchsia_build_dir / 'host-tools/ffx'
      flash_args = [
          ffx_path,
          '--target', device_name,
          '--timeout', '60',
          'target', 'flash',
          '--manifest', args.fuchsia_build_dir / 'flash.json',
      ]
      print(f'Running: {shlex.join(str(s) for s in flash_args)}')

      # Sleep a little bit for ffx to notice the device is in fastboot
      await asyncio.sleep(5)

      process = await asyncio.create_subprocess_exec(*flash_args)
      return_code = await process.wait()
      if return_code != 0:
        print('Failed to flash target', file=sys.stderr)
        return return_code

      print('Flashed device.')
      print('Connecting to device...')

      show_args = [ffx_path, '--target', device_name, '--timeout', '60', 'target', 'wait']
      print(f'Running: {shlex.join(str(s) for s in show_args)}')

      process = await asyncio.create_subprocess_exec(*show_args)
      return_code = await process.wait()
      if return_code != 0:
        print('Target failed to come online', file=sys.stderr)
        return return_code

      print('Connected to device.')

      print('Signalling qemu to shut down...')

      shutdown_args = [ffx_path, '--target', device_name, 'target', 'off']
      print(f'Running: {shlex.join(str(s) for s in shutdown_args)}')

      process = await asyncio.create_subprocess_exec(*shutdown_args)
      return_code = await process.wait()
      if return_code != 0:
        print('Target failed to turn off', file=sys.stderr)
        return return_code

      print('Waiting for qemu to shutdown...')

      for i in range(100):
        if await poll_process(qemu_process) is not None:
          break
        else:
          await asyncio.sleep(0.5)
      else:
        print('Timed out waiting for qemu to shut down', file=sys.stderr)

      print(f'Successfully created Fuchsia image {args.image}')
    finally:
      if await poll_process(qemu_process) is None:
        print('Killing qemu...', file=sys.stderr)
        qemu_process.kill()
        await qemu_process.wait()

if __name__ == '__main__':
  sys.exit(asyncio.run(main()))