devtool

#!/usr/bin/env bash

# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

# Firecracker devtool
#
# Use this script to build and test Firecracker.
#
# TL;DR
# Make sure you have Docker installed and properly configured
# (http://docker.com). Then,
#   building: `./devtool build`
#     Then find the binaries under build/debug/
#   testing: `./devtool test`
#     Will run the entire test battery; will take serveral minutes to complete.
#   deep-dive: `./devtool shell`
#     Open a shell prompt inside the container. Then build or test (or do
#     anything, really) manually.
#
# Still TL;DR: have Docker; ./devtool build; ./devtool test; ./devtool help.
#
#
# Both building and testing are done inside a Docker container. Please make sure
# you have Docker up and running on your system (see http:/docker.com) and your
# user has permission to run Docker containers.
#
# The Firecracker sources dir will be bind-mounted inside the development
# container (under /firecracker) and any files generated by the build process
# will show up under the build/ dir.  This includes the final binaries, as well
# as any intermediate or cache files.
#
# By default, all devtool commands run the container transparently, removing
# it after the command completes. Any persisting files will be stored under
# build/.
# If, for any reason, you want to access the container directly, please use
# `devtool shell`. This will perform the initial setup (bind-mounting the
# sources dir, setting privileges) and will then drop into a BASH shell inside
# the container.
#
# Building:
#   Run `./devtool build`.
#   By default, the debug binaries are built and placed under build/debug/.
#   To build the release version, run `./devtool build --release` instead.
#   You can then find the binaries under build/release/.
#
# Testing:
#   Run `./devtool test`.
#   This will run the entire integration test battery. The testing system is
#   based on pytest (http://pytest.org).
#
# Opening a shell prompt inside the development container:
#   Run `./devtool shell`.
#
# Additional information:
#   Run `./devtool help`.
#
#
# TODO:
#   - Cache test binaries, preserving them across `./devtool test` invocations.
#     At the moment, Firecracker is rebuilt everytime a test is run.
#   - List tests by parsing the `pytest --collect-only` output.
#   - Implement test filtering with `./devtool test --filter <filter>`
#   - Find an easier way to run individual tests on existing binaries.
#   - Add a `./devtool run` command to set up and run Firecracker.
#   - Add a `./devtool diag` command to help with troubleshooting, by checking
#     the most common failure conditions.
#   - Add a `./devtool build-devctr` command to build the development container
#     from its Dockerfile.
#   - Look into caching the Cargo registry within the container and if that
#     would help with reproducible builds (in addition to pinning Cargo.lock)

# Development container image (name:tag)
# This should be updated whenever we upgrade the development container.
# (Yet another step on our way to reproducible builds.)
DEVCTR_IMAGE="fcuvm/dev:v7"

# Naming things is hard
MY_NAME="Firecracker $(basename "$0")"

# Full path to the Firecracker tools dir on the host.
FC_TOOLS_DIR=$(cd "$(dirname "$0")" && pwd)

# Full path to the Firecracker sources dir on the host.
FC_ROOT_DIR=$(cd "${FC_TOOLS_DIR}/.." && pwd)

# Full path to the build dir on the host.
FC_BUILD_DIR="${FC_ROOT_DIR}/build"

# Full path to the cargo registry dir on the host. This appears on the host
# because we want to persist the cargo registry across container invocations.
# Otherwise, any rust crates from crates.io would be downloaded again each time
# we build or test.
CARGO_REGISTRY_DIR="${FC_BUILD_DIR}/cargo_registry"

# Full path to the cargo target dir on the host.
CARGO_TARGET_DIR="${FC_BUILD_DIR}/cargo_target"


# Full path to the Firecracker sources dir, as bind-mounted in the container.
CTR_FC_ROOT_DIR="/firecracker"

# Full path to the build dir, as bind-mounted in the container.
CTR_FC_BUILD_DIR="${CTR_FC_ROOT_DIR}/build"

# Full path to the cargo target dir, as bind-mounted in the container.
CTR_CARGO_TARGET_DIR="$CTR_FC_BUILD_DIR/cargo_target"

# Full path to the microVM images cache dir
CTR_MICROVM_IMAGES_DIR="$CTR_FC_BUILD_DIR/img"

# Global options received by $0
# These options are not command-specific, so we store them as global vars
OPT_UNATTENDED=false


# Send a decorated message to stdout, followed by a new line
#
say() {
    [ -t 1 ] && [ -n "$TERM" ] \
        && echo "$(tput setaf 2)[$MY_NAME]$(tput sgr0) $*" \
        || echo "[$MY_NAME] $*"
}

# Send a decorated message to stdout, without a trailing new line
#
say_noln() {
    [ -t 1 ] && [ -n "$TERM" ] \
        && echo -n "$(tput setaf 2)[$MY_NAME]$(tput sgr0) $*" \
        || echo "[$MY_NAME] $*"
}

# Send a text message to stderr
#
say_err() {
    [ -t 2 ] && [ -n "$TERM" ] \
        && echo "$(tput setaf 1)[$MY_NAME] $*$(tput sgr0)" 1>&2 \
        || echo "[$MY_NAME] $*" 1>&2
}

# Send a warning-highlighted text to stdout
say_warn() {
    [ -t 1 ] && [ -n "$TERM" ] \
        && echo "$(tput setaf 3)[$MY_NAME] $*$(tput sgr0)" \
        || echo "[$MY_NAME] $*"
}

# Exit with an error message and (optional) code
# Usage: die [-c <error code>] <error message>
#
die() {
    code=1
    [[ "$1" = "-c" ]] && {
        code="$2"
        shift 2
    }
    say_err "$@"
    exit $code
}

# Exit with an error message if the last exit code is not 0
#
ok_or_die() {
    code=$?
    [[ $code -eq 0 ]] || die -c $code "$@"
}

# Check if Docker is available and exit if it's not.
# Upon returning from this call, the caller can be certain Docker is available.
#
ensure_docker() {
    which docker > /dev/null 2>&1
    ok_or_die "Docker not found. Aborting." \
        "Please make sure you have Docker (http://docker.com) installed" \
        "and properly configured."

    docker ps > /dev/null 2>&1
    ok_or_die "Error accessing Docker. Please make sure the Docker daemon" \
        "is running and that you are part of the docker group." \
        "For more information, see" \
        "https://docs.docker.com/install/linux/linux-postinstall/"
}

# Attempt to download our Docker image. Exit if that fails.
# Upon returning from this call, the caller can be certain our Docker image is
# available on this system.
#
ensure_devctr() {

    # We depend on having Docker present.
    ensure_docker

    # Check if we have the container image available locally. Attempt to
    # download it, if we don't.
    [[ $(docker images -q "$DEVCTR_IMAGE" | wc -l) -gt 0 ]] || {
        say "About to pull docker image $DEVCTR_IMAGE"
        get_user_confirmation || die "Aborted."

        docker pull "${DEVCTR_IMAGE}"

        ok_or_die "Error pulling docker image. Aborting."
    }
}

# Check if /dev/kvm exists. Exit if it doesn't.
# Upon returning from this call, the caller can be certain /dev/kvm is
# available.
#
ensure_kvm() {
    [[ -c /dev/kvm ]] || die "/dev/kvm not found. Aborting."
}

# Make sure the build/ dirs are available. Exit if we can't create them.
# Upon returning from this call, the caller can be certain the build/ dirs exist.
#
ensure_build_dir() {
    for dir in "$FC_BUILD_DIR" "$CARGO_TARGET_DIR" "$CARGO_REGISTRY_DIR"; do
        mkdir -p "$dir" || die "Error: cannot create dir $dir"
        [ -x "$dir" ] && [ -w "$dir" ] || \
            {
                say "Wrong permissions for $dir. Attempting to fix them ..."
                chmod +x+w "$dir"
            } || \
            die "Error: wrong permissions for $dir. Should be +x+w"
    done
}

# Fix build/ dir permissions after a privileged container run.
# Since the privileged cotainer runs as root, any files it creates will be
# owned by root. This fixes that by recursively changing the ownership of build/
# to the current user.
#
fix_build_dir_perms() {
    # Yes, running Docker to get elevated privileges, just to chown some files
    # is a dirty hack.
    run_devctr \
        -- \
        chown -R "$(id -u):$(id -g)" "$CTR_FC_BUILD_DIR"
}

# Prompt the user for confirmation before proceeding.
# Args:
#   $1  prompt text.
#       Default: Continue? (y/n)
#   $2  confirmation input.
#       Default: y
# Return:
#   exit code 0 for successful confirmation
#   exit code != 0 if the user declined
#
get_user_confirmation() {

    # Pass if running unattended
    [[ "$OPT_UNATTENDED" = true ]] && return 0

    # Fail if STDIN is not a terminal (there's no user to confirm anything)
    [[ -t 0 ]] || return 1

    # Otherwise, ask the user
    #
    msg=$([ -n "$1" ] && echo -n "$1" || echo -n "Continue? (y/n) ")
    yes=$([ -n "$2" ] && echo -n "$2" || echo -n "y")
    say_noln "$msg"
    read c && [ "$c" = "$yes" ] && return 0
    return 1
}

# Validate the user supplied version number.
# It must be composed of 3 groups of integers separated by dot.
#
validate_version() {
    declare version_regex="^([0-9]+.){2}[0-9]+$"
    version="$1"

    if [ -z "$version" ]; then
        die "Version cannot be empty."
    elif [[ ! "$version" =~ $version_regex ]]; then
        die "Invalid version number: $version (expected: \$Major.\$Minor.\$Build)."
    fi

}

# Compose the text for a new release tag using the information in the changelog,
# between the two specified releases.
# The following transformations are applied:
# * `-` is replaced with `*` for unnumbered lists.
# * section headers (`###`) are removed.
#
# Args:
#   $1  previous version.
#   $2  new version.
#
compose_tag_text() {
    declare prev_ver="$1"
    declare curr_ver="$2"
    declare changelog="$FC_ROOT_DIR/CHANGELOG.md"

    validate_version "$prev_ver"
    validate_version "$curr_ver"

    # Patterns for the sections in the changelog corresponding to the versions.
    pat_prev="^##\s\[$prev_ver\]"
    pat_curr="^##\s\[$curr_ver\]"

    # Extract the section enclosed between the 2 headers and strip off the first
    # 2 and last 2 lines (one is blank and one contains the header `## [A.B.C]`).
    # Then, replace `-` with `*` and remove section headers.
    sed "/$pat_curr/,/$pat_prev/!d" "$changelog" \
      | sed '1,2d;$d' \
      | sed "s/^-/*/g" \
      | sed "s/^###\s//g"
}

# Helper function to run the dev container.
# Usage: run_devctr <docker args> -- <container args>
# Example: run_devctr --privileged -- bash -c "echo 'hello world'"
run_devctr() {
    docker_args=()
    ctr_args=()
    docker_args_done=false
    while [[ $# -gt 0 ]]; do
        [[ "$1" = "--" ]] && {
            docker_args_done=true
            shift
            continue
        }
        [[ $docker_args_done = true ]] && ctr_args+=("$1") || docker_args+=("$1")
        shift
    done

    # If we're running in a terminal, pass the terminal to Docker and run
    # the container interactively
    [[ -t 0 ]] && docker_args+=("-i")
    [[ -t 1 ]] && docker_args+=("-t")

    # Try to pass these environments from host into container for network proxies
    proxies=(http_proxy HTTP_PROXY https_proxy HTTPS_PROXY no_proxy NO_PROXY)
    for i in "${proxies[@]}"; do
        if [[ ! -z ${!i} ]]; then
            docker_args+=("--env") && docker_args+=("$i=${!i}")
        fi
    done

    # Finally, run the dev container
    # Use 'z' on the --volume parameter for docker to automatically relabel the
    # content and allow sharing between containers.
    docker run "${docker_args[@]}" \
        --init \
        --rm \
        --volume "$FC_ROOT_DIR:$CTR_FC_ROOT_DIR:z" \
        --env OPT_LOCAL_IMAGES_PATH="$(dirname "$CTR_MICROVM_IMAGES_DIR")" \
        --env PYTHONDONTWRITEBYTECODE=1 \
        "$DEVCTR_IMAGE" "${ctr_args[@]}"
}

# `$0 help`
# Show the detailed devtool usage information.
#
cmd_help() {
    echo ""
    echo "Firecracker $(basename $0)"
    echo "Usage: $(basename $0) [<args>] <command> [<command args>]"
    echo ""
    echo "Global arguments"
    echo "    -y, --unattended         Run unattended. Assume the user would always"
    echo "                             answer \"yes\" to any confirmation prompt."
    echo ""
    echo "Available commands:"
    echo ""
    echo "    build [--debug|--release] [-l|--libc musl|gnu] [-- [<cargo args>]]"
    echo "        Build the Firecracker binaries."
    echo "        Firecracker is built using the Rust build system (cargo). All arguments after --"
    echo "        will be passed through to cargo."
    echo "        --debug               Build the debug binaries. This is the default."
    echo "        --release             Build the release binaries."
    echo "        -l, --libc musl|gnu   Choose the libc flavor against which Firecracker will"
    echo "                              be linked. Default is musl."
    echo ""
    echo "    fmt"
    echo "        Auto-format all Rust source files, to match the Firecracker requirements."
    echo "        This should be used as the last step in every commit, to ensure that the"
    echo "        Rust style tests pass."
    echo ""
    echo "    help"
    echo "        Display this help message."
    echo ""
    echo "    prepare_release <version>"
    echo "        Prepare a new Firecracker release by updating the version number, crate "
    echo "        dependencies and credits."
    echo ""
    echo "    shell [--privileged]"
    echo "        Launch the development container and open an interactive BASH shell."
    echo "        -p, --privileged    Run the container as root, in privileged mode."
    echo "                            Running Firecracker via the jailer requires elevated"
    echo "                            privileges, though the build phase does not."
    echo ""
    echo "    tag <version>"
    echo "        Create a git tag for the specified version. The tag message will contain "
    echo "        the contents of CHANGELOG.md enclosed between the header corresponding to "
    echo "        the specified version and the one corresponding to the previous version."
    echo ""
    echo "    test [-- [<pytest args>]]"
    echo "        Run the Firecracker integration tests."
    echo "        The Firecracker testing system is based on pytest. All arguments after --"
    echo "        will be passed through to pytest."
    echo ""
    echo "    checkenv"
    echo "        Performs prerequisites checks needed to execute firecracker."
    echo ""
}

# `$0 build` - build Firecracker
# Please see `$0 help` for more information.
#
cmd_build() {

    # By default, we'll build the debug binaries.
    profile="debug"
    libc="musl"

    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")  { cmd_help; exit 1;     } ;;
            "--debug")      { profile="debug";      } ;;
            "--release")    { profile="release";    } ;;
            "-l"|"--libc")
                shift
                [[ "$1" =~ ^(musl|gnu)$ ]] || \
                    die "Invalid libc: $1. Valid options are \"musl\" and \"gnu\"."
                libc="$1"
                ;;
            "--")           { shift; break;         } ;;
            *)
                die "Unknown argument: $1. Please use --help for help."
            ;;
        esac
        shift
    done

    target="$(uname -m)-unknown-linux-${libc}"

    # Check prerequisites
    ensure_devctr
    ensure_build_dir

    say "Starting build ($profile, $libc) ..."

    # Cargo uses the debug profile by default. If we're building the release
    # binaries, we need to pass an extra argument to cargo.
    cargo_args=("$@")

    [ $profile = "release" ] && cargo_args+=("--release")

    # Run the cargo build process inside the container.
    # We don't need any special privileges for the build phase, so we run the
    # container as the current user/group.
    #
    # Note: we need to pass an explicit TARGET_CC to cargo, because the `cc`
    # crate gets confused for aarch64 musl.
    run_devctr \
        --user "$(id -u):$(id -g)" \
        --workdir "$CTR_FC_ROOT_DIR" \
        --env TARGET_CC=musl-gcc \
        -- \
        cargo build \
            --target-dir "$CTR_CARGO_TARGET_DIR" \
            --target "$target" \
            "${cargo_args[@]}"
    ret=$?

    # If `cargo build` was successful, let's copy the binaries to a more
    # accessible location.
    [ $ret -eq 0 ] && {
        cargo_bin_dir="$CARGO_TARGET_DIR/$target/$profile"
        friendly_bin_dir="$FC_BUILD_DIR/$profile-$libc"
        say "Build successful."

        mkdir -p "$friendly_bin_dir"
        [ $? -eq 0 ] || {
            say_warn "Unable to create dir $friendly_bin_dir."
            say_warn "Binaries placed under $cargo_bin_dir/"
            return $ret
        }
        cp -f "$cargo_bin_dir"/{firecracker,jailer} "$friendly_bin_dir/"
        say "Build complete."
        say "Binaries placed under $friendly_bin_dir"
    }

    return $ret
}

# `$0 test` - run integration tests
# Please see `$0 help` for more information.
#
cmd_test() {

    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")      { cmd_help; exit 1; } ;;
            "--")               { shift; break;     } ;;
            *)
                die "Unknown argument: $1. Please use --help for help."
            ;;
        esac
        shift
    done

    # Check prerequisites.
    ensure_kvm
    ensure_devctr
    ensure_build_dir


    # If we got to here, we've got all we need to continue.
    say "Starting test run ..."

    # Testing (running Firecracker via the jailer) needs root access,
    # in order to set-up the Firecracker jail (manipulating cgroups, net
    # namespaces, etc).
    # We need to run a privileged container to get that kind of access.
    run_devctr \
        --privileged \
        --security-opt seccomp=unconfined \
        --ulimit core=0 \
        --workdir "$CTR_FC_ROOT_DIR/tests" \
        -- \
        pytest "$@"

    ret=$?

    # Running as root would have created some root-owned files under the build
    # dir. Let's fix that.
    fix_build_dir_perms

    return $ret
}

# `$0 shell` - drop to a shell prompt inside the dev container
# Please see `$0 help` for more information.
#
cmd_shell() {

    # By default, we run the container as the current user.
    privileged=false

    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")          { cmd_help; exit 1; } ;;
            "-p"|"--privileged")    { privileged=true;  } ;;
            *)
                die "Unknown argument: $1. Please use --help for help."
            ;;
        esac
        shift
    done

    # Make sure we have what we need to continue.
    ensure_devctr
    ensure_build_dir

    if [[ $privileged = true ]]; then
        # If requested, spin up a privileged container.
        #
        say "Dropping to a privileged shell prompt ..."
        say "Note: $FC_ROOT_DIR is bind-mounted under $CTR_FC_ROOT_DIR"
        say_warn "You are running as root; any files that get created under" \
            "$CTR_FC_ROOT_DIR will be owned by root."
        run_devctr \
            --privileged \
            --security-opt seccomp=unconfined \
            --workdir "$CTR_FC_ROOT_DIR" \
            -- \
            bash
        ret=$?

        # Running as root may have created some root-owned files under the build
        # dir. Let's fix that.
        #
        fix_build_dir_perms
    else
        say "Dropping to shell prompt as user $(whoami) ..."
        say "Note: $FC_ROOT_DIR is bind-mounted under $CTR_FC_ROOT_DIR"
        say_warn "You won't be able to run Firecracker via the jailer," \
            "but you can still build it."
        say "You can use \`$0 shell --privileged\` to get a root shell."

        [ -w /dev/kvm ] || \
            say_warn "WARNING: user $(whoami) doesn't have permission to" \
                "access /dev/kvm. You won't be able to run Firecracker."

        run_devctr \
            --user "$(id -u):$(id -g)" \
            --device=/dev/kvm:/dev/kvm \
            --workdir "$CTR_FC_ROOT_DIR" \
            --env PS1="$(whoami)@\h:\w\$ " \
            -- \
            bash --norc
        ret=$?
    fi

    return $ret
}


# Auto-format all source code, to match the Firecracker requirements. For the
# moment, this is just a wrapper over `cargo fmt --all`
# Example: `devtool fmt`
#
cmd_fmt() {

    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")      { cmd_help; exit 1; } ;;
            *)
                die "Unknown argument: $1. Please use --help for help."
            ;;
        esac
        shift
    done

    ensure_devctr

    say "Applying rustfmt ..."
    run_devctr \
        --user "$(id -u):$(id -g)" \
        --workdir "$CTR_FC_ROOT_DIR" \
        -- \
        cargo fmt --all
}


# Prepare a Firecracker release by updating the version, crate dependencies
# and credits.
# Example: `devtool release 0.42.0`
#
cmd_prepare_release() {

    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")      { cmd_help; exit 1;    } ;;
            *)                  { version="$1"; break; } ;;
        esac
        shift
    done

    validate_version "$version"

    # We'll be needing the dev container later on.
    ensure_devctr

    # The cargo registry dir needs to be there for `cargo update`.
    ensure_build_dir

    # Get current version from the swagger spec.
    swagger="$FC_ROOT_DIR/api_server/swagger/firecracker.yaml"
    curr_ver=$(grep "version: " "$swagger" | awk -F : '{print $2}' | tr -d ' ')

    say "Updating from $curr_ver to $version ..."
    get_user_confirmation || die "Aborted."

    # Update version in files.
    files_to_change=("$swagger"                         \
                     "$FC_ROOT_DIR/Cargo.toml"          \
                     "$FC_ROOT_DIR/jailer/Cargo.toml")
    say "Updating source files:"
    for file in "${files_to_change[@]}"; do
        say "- $file"
        # Dirty hack to make this work on both macOS/BSD and Linux.
        sed -i="" "s/$curr_ver/$version/g" "$file"
        rm -f "${file}="
    done

    # Update crate dependencies.
    say "Updating crate dependencies..."
    run_devctr \
        --user "$(id -u):$(id -g)" \
        --workdir "$CTR_FC_ROOT_DIR" \
        -- \
        cargo update
    ok_or_die "cargo update failed."

    # Update credits.
    say "Updating credits..."
    "$FC_TOOLS_DIR/update-credits.sh"

    # Update changelog.
    say "Updating changelog..."
    sed -i="" "s/\[Unreleased\]/\[$version\]/g" "$FC_ROOT_DIR/CHANGELOG.md"
    rm -f "$FC_ROOT_DIR/CHANGELOG.md="
}

# Create a tag for the specified release.
# The tag text will be composed from the changelog contents enclosed between the
# specified release number and the previous one.
# Args:
#   $1  release number.
# Example: `devtool tag 0.42.0`
#
cmd_tag() {

    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")      { cmd_help; exit 1;    } ;;
            *)                  { version="$1"; break; } ;;
        esac
        shift
    done

    validate_version "$version"

    declare pat_release="^## \[([0-9]+\.){2}[0-9]+\]"
    declare changelog="$FC_ROOT_DIR/CHANGELOG.md"

    grep -q "\[$version\]" "$changelog"
    ok_or_die "No changelog entry for release $version can be found."

    # We work with the assumption that the changelog has already been updated
    # and contains a header (and a corresponding section) for the new release.

    # Step 1: Get all release numbers.
    all_releases=($(grep -E "$pat_release" "$changelog"))

    # Step 2: Trim out headers (`##`).
    all_releases=(${all_releases[@]//##*})

    # Step 3: Walk the array until we come across the desired release number,
    # then pick up the next one. Since the latest releases are at the top of the
    # changelog, the next one in line will be the previous one chronologically.
    # The array now contains all the release numbers in the changelog, enclosed
    # in square brackets.
    found=
    for release in "${all_releases[@]}"; do
        if [ ! -z "$found" ]; then
            # Trim out square brackets.
            prev_version=$(echo "$release" | awk -F"[][]" "{print \$2}")
            break
        elif [ "$release" == "[$version]" ]; then
            found=1
        fi
    done

    # Create tag.
    tag_text=$(compose_tag_text "$prev_version" "$version")
    say "Preparing to create tag..."
    say "Tag: v$version"
    say "Tag text:"
    echo "$tag_text"
    say "Continue with tag creation?"
    get_user_confirmation || die "Tag not created."

    git tag -a "v$version" -m "$tag_text"
    ok_or_die "Tag v$version not created."
    say "Tag v$version created."
}

# Check if able to run firecracker.
# ../docs/getting-started.md#prerequisites

ensure_kvm_rw () {
    [[ -c /dev/kvm && -w /dev/kvm && -r /dev/kvm ]] || \
        say_err "FAILED: user $(whoami) doesn't have permission to" \
                "access /dev/kvm."
}

check_kernver () {
    KERN_MAJOR=4
    KERN_MINOR=14
    (uname -r | awk -v MAJOR=$KERN_MAJOR -v MINOR=$KERN_MINOR '{ split($0,kver,".");
    if( (kver[1] + (kver[2] / 100) ) <  MAJOR + (MINOR/100) )
    {
      exit 1;
    } }') ||
	say_err "FAILED: Kernel version must be >= $KERN_MAJOR.$KERN_MINOR"
}

# Check Production Host Setup
# ../docs/prod-host-setup.md

check_SMT () {
    (grep -q "^forceoff$\|^notsupported$" \
	  /sys/devices/system/cpu/smt/control) ||
	say_warn "WARNING: Hyperthreading ENABLED."
}

check_KPTI () {
    (grep -q "^Mitigation: PTI$" \
	  /sys/devices/system/cpu/vulnerabilities/meltdown) || \
	say_warn "WARNING: KPTI NOT SUPPORTED"
}

check_KSM () {
    (grep -q "^0$" /sys/kernel/mm/ksm/run) || \
	say_warn "WARNING: KSM ENABLED"
}

check_IBPB_IBRS () {
    (grep -q "^Mitigation: Full generic retpoline, IBPB, IBRS_FW$"\
	  /sys/devices/system/cpu/vulnerabilities/spectre_v2) || \
	say_warn "WARNING: retpoline, IBPB, IBRS: DISABLED."
}

check_L1TF () {
    declare -a CONDITIONS=("Mitigation: PTE Inversion" "VMX: cache flushes")
    for cond in "${CONDITIONS[@]}";
    do (grep -q "$cond" /sys/devices/system/cpu/vulnerabilities/l1tf) ||
	   say_warn "WARNING: $cond: DISABLED";
    done
}

check_swap () {
    (grep -q "swap.img" /proc/swaps ) && \
	say_warn "WARNING: SWAP enabled"
}

cmd_checkenv() {
    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")      { cmd_help; exit 1; } ;;
            *)
                die "Unknown argument: $1. Please use --help for help."
		;;
        esac
        shift
    done
    PROD_DOC="../docs/prod-host-setup.md"
    QUICKSTART="../docs/getting-started.md#prerequisites"
    say "Checking prerequisites for running Firecracker."
    say "Please check $QUICKSTART in case of any error."
    ensure_kvm_rw
    check_kernver
    say "Checking Host Security Configuration."
    say "Please check $PROD_DOC in case of any error."
    check_KSM
    check_IBPB_IBRS
    check_L1TF
    check_SMT
    check_swap
}

main() {

    if [ $# = 0 ]; then
	die "No command provided. Please use \`$0 help\` for help."
    fi

    # Parse main command line args.
    #
    while [ $# -gt 0 ]; do
        case "$1" in
            -h|--help)              { cmd_help; exit 1;     } ;;
            -y|--unattended)        { OPT_UNATTENDED=true;  } ;;
            -*)
                die "Unknown arg: $1. Please use \`$0 help\` for help."
            ;;
            *)
                break
            ;;
        esac
        shift
    done

    # $1 is now a command name. Check if it is a valid command and, if so,
    # run it.
    #
    declare -f "cmd_$1" > /dev/null
    ok_or_die "Unknown command: $1. Please use \`$0 help\` for help."

    cmd=cmd_$1
    shift

    # $@ is now a list of command-specific args
    #
    $cmd "$@"
}

main "$@"