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 several 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:
#   - List tests by parsing the `pytest --collect-only` output.
#   - 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.
#   - 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 (without tag)
DEVCTR_IMAGE_NO_TAG="public.ecr.aws/firecracker/fcuvm"

# Development container tag
DEVCTR_IMAGE_TAG=${DEVCTR_IMAGE_TAG:-v79}

# 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="${DEVCTR_IMAGE_NO_TAG}:${DEVCTR_IMAGE_TAG}"

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

# 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 devctr dir on the host.
FC_DEVCTR_DIR="${FC_ROOT_DIR}/tools/devctr"

# Path to the linux kernel directory on the host.
KERNEL_DIR="${FC_ROOT_DIR}/.kernel"

# 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 git registry on the host. This serves the same purpose
# as CARGO_REGISTRY_DIR, for crates downloaded from GitHub repos instead of
# crates.io.
CARGO_GIT_REGISTRY_DIR="${FC_BUILD_DIR}/cargo_git_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"
CTR_TEST_RESULTS_DIR="${CTR_FC_ROOT_DIR}/test_results"

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

# Path to the microVM images cache dir
MICROVM_IMAGES_DIR="build/img"

# Full path to the public key mapping on the guest
PUB_KEY_PATH=/root/.ssh/id_rsa.pub

# Full path to the private key mapping on the guest
PRIV_KEY_PATH=/root/.ssh/id_rsa

# Path to the linux kernel directory, as bind-mounted in the container.
CTR_KERNEL_DIR="${CTR_FC_ROOT_DIR}/.kernel"

# Get the target prefix to avoid repeated calls to uname -m
TARGET_PREFIX="$(uname -m)-unknown-linux-"

# Container path to directory where we store built CI artifacts.
CTR_CI_ARTIFACTS_PATH="${CTR_FC_ROOT_DIR}/resources/$(uname -m)"

# 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() {
    NEWLINE=$'\n'
    output=$(which docker 2>&1)
    ok_or_die "Docker not found. Aborting." \
        "Please make sure you have Docker (http://docker.com) installed" \
        "and properly configured.${NEWLINE}" \
        "Error: $?, command output: ${output}"

    output=$(docker ps 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.${NEWLINE}" \
        "Error: $?, command output: ${output}${NEWLINE}" \
        "For more information, see" \
        "https://docs.docker.com/install/linux/linux-postinstall/"
}

# Run a command and retry multiple times if it fails. Once it stops
# failing return to normal execution. If there are "retry count"
# failures, set the last error code.
# $1 - command
# $2 - retry count
# $3 - sleep interval between retries
retry_cmd() {
    command=$1
    retry_cnt=$2
    sleep_int=$3

    {
        $command
    } || {
        # Command failed, substract one from retry_cnt
        retry_cnt=$((retry_cnt - 1))

        # If retry_cnt is larger than 0, sleep and call again
        if [ "$retry_cnt" -gt 0 ]; then
            echo "$command failed, retrying..."
            sleep "$sleep_int"
            retry_cmd "$command" "$retry_cnt" "$sleep_int"
        fi
    }
}

# 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"

        # Run docker pull 5 times in case it fails - sleep 3 seconds
        # between attempts
        retry_cmd "docker pull $DEVCTR_IMAGE" 5 3

        ok_or_die "Error pulling docker image. 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" "$CARGO_GIT_REGISTRY_DIR"; do
        create_dir "$dir"
    done
}

build_bin_path() {
    target="$1"
    profile="$2"
    binary="$3"
    echo "$CARGO_TARGET_DIR/$target/$profile/$binary"
}

# Fix build/ dir permissions after a privileged container run.
# Since the privileged container 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.
#
cmd_fix_perms() {
    # Yes, running Docker to get elevated privileges, just to chown some files
    # is a dirty hack.
    run_devctr \
        --workdir "$CTR_FC_ROOT_DIR" \
        -- \
        chown -f -R "$(id -u):$(id -g)" "$CTR_FC_BUILD_DIR" "$CTR_TEST_RESULTS_DIR" "$CTR_CI_ARTIFACTS_PATH" $@
}

# Builds the development container from its Dockerfile.
#
cmd_build_devctr() {
    docker_file_name=$FC_DEVCTR_DIR/Dockerfile
    build_args="--build-arg ARCH=$(uname -m)"

    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

    docker build -t "$DEVCTR_IMAGE_NO_TAG" -f "$docker_file_name" $build_args .
}


# Validate the user supplied kernel version number.
# It must be composed of 2 groups of integers separated by dot, with an optional third group.
validate_kernel_version() {
    local version_regex="^([0-9]+.)[0-9]+(.[0-9]+)?$"
    version="$1"

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

}


# 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[@]}" \
        --rm \
        --volume /dev:/dev \
        --volume "$FC_ROOT_DIR:$CTR_FC_ROOT_DIR:z" \
        --volume "$FC_ROOT_DIR/build/cargo_registry:/usr/local/rust/registry:z" \
        --volume "$FC_ROOT_DIR/build/cargo_git_registry:/usr/local/rust/git:z" \
        --tmpfs /srv:exec,dev,size=32G \
        -v /boot:/boot \
        --env PYTHONDONTWRITEBYTECODE=1 \
        "$DEVCTR_IMAGE" "${ctr_args[@]}"
}

# Helper function to test that the argument provided is a valid path to a SSH key.
#
test_key() {
    ssh-keygen -lf "$1" &>/dev/null
    ret=$?
    [ $ret -ne 0 ] && die "$1 is not a valid key file."
}

create_dir() {
    # Create a dir for the provided path.
    dir="$1"
    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"
}

# `$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]"
    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 "        --ssh-keys            Provide the paths to the public and private SSH keys on the host"
    echo "                              (in this particular order) required for the git authentication."
    echo "                              It is mandatory that both keys are specified."
    echo ""
    echo "    build_devctr"
    echo "        Builds the development container from its Dockerfile."
    echo ""
    echo "    checkenv"
    echo "        Performs prerequisites checks needed to execute firecracker."
    echo ""
    echo "    distclean"
    echo "        Clean up the build tree and remove the docker container."
    echo ""
    echo "    fix_perms"
    echo "        Fixes permissions when devtool dies in the middle of a privileged session."
    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 "    install [-p|--path] [--debug|--release]"
    echo "      Install firecracker, jailer and seccomp binaries to /usr/local/bin or a given path."
    echo "      Only the musl linked binaries are supported."
    echo "        --path                Install binaries to a specified path."
    echo "        --debug               Install the debug binaries."
    echo "        --release             Install the release binaries. This is the default."
    echo ""
    echo "    help"
    echo "        Display this help message."
    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 "    sh CMD..."
    echo "        Launch the development container and run a command."
    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 "        -c, --cpuset-cpus cpulist    Set a dedicated cpulist to be used by the tests."
    echo "        -m, --cpuset-mems memlist    Set a dedicated memlist to be used by the tests."
    echo "            --performance            Tweak various setting of the host running the tests (such as C- and P-states)"
    echo "                                     to achieve consistent performance. Used for running performance tests in CI."
    echo ""
    echo "    build_ci_artifacts [all|rootfs|kernels]"
    echo "        Builds the rootfs and guest kernel artifacts we use for our CI."
    echo "        Run './tools/devtool build_ci_artifacts help' for more details about the available commands."
    echo ""
    echo "    download_ci_artifacts [--force]"
    echo "        Downloads the CI artifacts used for testing from our S3 bucket. If --force is passed, purges any existing"
    echo "        artifacts first. Useful for refreshing local artifacts after an update, or if something got messed up."
    echo ""

    cat <<EOF
    test_debug [-- [<pytest args>]]
        Run tests in a debugging environment

    sandbox
        Run Firecracker in an IPython REPL (in devctr)

    sandbox_native
        Run Firecracker in an IPython REPL (AL2023/Ubuntu)

    mkdocs
        Use 'cargo doc' to generate rustdoc documentation

    checkstyle
        Run style checks
EOF
}


# `$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";    } ;;
            "--rev")       { shift; revision=$1; } ;;
            "--ssh-keys")
                shift
                [[ -z "$1" ]] && \
                    die "Please provide the path to the public SSH key."
                [[ ! -f "$1" ]]  && die "The public key file does not exist: $1."
                test_key "$1"
                host_pub_key_path="$1"
                shift
                [[ -z "$1" ]] && \
                    die "Please provide the path to the private SSH key."
                [[ ! -f "$1" ]]  && die "The private key file does not exist: $1."
                test_key "$1"
                host_priv_key_path="$1"
                ;;
            "-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

    # Check prerequisites
    ensure_devctr
    ensure_build_dir

    # Map the public and private keys to the guest if they are specified.
    [ ! -z "$host_pub_key_path" ] && [ ! -z "$host_priv_key_path" ] &&
        extra_args="--volume $host_pub_key_path:$PUB_KEY_PATH:z \
                    --volume $host_priv_key_path:$PRIV_KEY_PATH:z"

    workdir="$CTR_FC_ROOT_DIR"
    if [ ! -z "$revision" ]; then
      commitish="$revision"
      if ! git cat-file -t "$commitish"; then commitish=origin/"$revision"; fi
      branch_name=tmp-$commitish

      tmp_dir=$(mktemp -d)

      git branch $branch_name $commitish
      git clone -b $branch_name . $tmp_dir
      pushd $tmp_dir
      workdir=$tmp_dir
      extra_args="$extra_args --volume $tmp_dir:$tmp_dir:z"
    fi

    # 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.
    run_devctr \
        --privileged \
        --workdir "$workdir" \
        ${extra_args} \
        -- \
        ./tools/release.sh --libc $libc --profile $profile
    ret=$?

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

    if [ ! -z "$revision" ]; then
      popd
      git branch -D $branch_name
      mkdir -p build/"$revision"/examples
      cp $tmp_dir/build/cargo_target/$(uname -m)-unknown-linux-$libc/$profile/* build/"$revision"
      cp $tmp_dir/build/cargo_target/$(uname -m)-unknown-linux-$libc/$profile/examples/* build/"$revision"/examples
      cmd_sh "rm -rf $tmp_dir"
    fi

    return $ret
}

function cmd_make_release {
    ensure_build_dir
    run_devctr \
        --privileged \
        --workdir "$CTR_FC_ROOT_DIR" \
        -- \
        ./tools/release.sh --libc musl --profile release --make-release
    sudo chown -Rc $USER: release*
}

cmd_distclean() {
    # List of folders to remove.
    dirs=("build" "test_results")

    for dir in "${dirs[@]}"; do
        if [ -d "$dir" ]; then
            say "Removing $dir"
            rm -rf "$dir"
        fi
    done

    # Remove devctr if it exists
    if [ $(docker images -q "$DEVCTR_IMAGE" | wc -l) -eq "1" ]; then
        say "Removing $DEVCTR_IMAGE"
        docker rmi -f "$DEVCTR_IMAGE"
    fi
}

cmd_download_ci_artifacts() {
    if [ "$1" = "--force" ]; then
        rm -rf $FC_BUILD_DIR/img
    fi

    ensure_ci_artifacts
}

ensure_ci_artifacts() {
    if ! command -v aws >/dev/null; then
      die "AWS CLI not installed, which is required for downloading artifacts for integration tests."
    fi

    # Fetch all the artifacts so they are local
    say "Fetching CI artifacts from S3"
    FC_VERSION=$(cmd_sh "cd src/firecracker/src; cargo pkgid | cut -d# -f2 | cut -d. -f1-2")
    S3_URL=s3://spec.ccfc.min/firecracker-ci/v$FC_VERSION/$(uname -m)
    ARTIFACTS=$MICROVM_IMAGES_DIR/$(uname -m)
    if [ ! -d "$ARTIFACTS" ]; then
        mkdir -pv $ARTIFACTS
        aws s3 sync --no-sign-request "$S3_URL" "$ARTIFACTS"
        cmd_sh "./tools/setup-ci-artifacts.sh"
    fi
}

apply_linux_61_tweaks() {
    KV=$(uname -r)
    if [[ $KV != 6.1.* ]] || [ $(uname -m) != x86_64 ]; then
        return
    fi
    say "Applying Linux 6.1 boot-time regression mitigations"

    KVM_VENDOR_MOD=$(lsmod |grep -P "^kvm_(amd|intel)" | awk '{print $1}')
    ITLB_MULTIHIT=/sys/devices/system/cpu/vulnerabilities/itlb_multihit
    NX_HUGEPAGES=/sys/module/kvm/parameters/nx_huge_pages

    # If m6a/m6i
    if grep -q "Not affected" $ITLB_MULTIHIT; then
        echo -e "CPU not vulnerable to iTLB multihit, using kvm.nx_huge_pages=never mitigation"
        # we need a lock so another process is not running the same thing and to
        # avoid race conditions.
        lockfile="/tmp/.linux61_tweaks.lock"
        set -C # noclobber
        while true; do
            if echo "$$" > "$lockfile"; then
                echo "Successfully acquired lock"
                if ! grep -q "never" $NX_HUGEPAGES; then
                    echo "Reloading KVM modules with nx_huge_pages=never"
                    sudo modprobe -r $KVM_VENDOR_MOD kvm
                    sudo modprobe kvm nx_huge_pages=never
                    sudo modprobe $KVM_VENDOR_MOD
                fi
                rm "$lockfile"
                break
            else
                sleep 5s
            fi
        done
        tail -v $ITLB_MULTIHIT $NX_HUGEPAGES
    # else (m5d Skylake and CascadeLake)
    else
        echo "CPU vulnerable to iTLB_multihit, checking if favordynmods is enabled"
        mount |grep cgroup |grep -q favordynmods
        if [ $? -ne 0 ]; then
            say_warn "cgroups' favordynmods option not enabled; VM creation performance may be impacted"
        else
            echo "favordynmods is enabled"
        fi
    fi
}


# Modifies the processors CPU governor and P-state configuration (x86_64 only) for consistent performance. This means
# - Disable turbo boost (Intel only) by writing 1 to /sys/devices/system/cpu/intel_pstate/no_turbo
# - Disable turbo boost (AMD only) by writing 0 to /sys/devices/system/cpu/cpufreq/boost
# - Lock the CPUs' P-state to the highest non-turbo one (Intel only) by writing 100 to /sys/devices/system/cpu/intel_pstate/{min,max}_perf_pct
# - Set the cpu frequency governor to performance by writing "performance" to /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor
apply_performance_tweaks() {
  # m6a instances do not support the amd_pstate driver (yet), so nothing we can do there
  if [[ -d /sys/devices/system/cpu/intel_pstate ]]; then
    # Disable turbo boost. Some of our tests are performance tests, and we want minimum variability wrt processor frequency
    # See also https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/processor_state_control.html
    echo 1 |sudo tee /sys/devices/system/cpu/intel_pstate/no_turbo &> /dev/null

    # Save old values to restore later
    MIN_PERF_PCT=$(cat /sys/devices/system/cpu/intel_pstate/min_perf_pct)
    MAX_PERF_PCT=$(cat /sys/devices/system/cpu/intel_pstate/max_perf_pct)

    # Force the CPU to continuously stay in the highest, non-turbo P-state. The P-state will determine the
    # CPU's clock frequency.
    # https://www.kernel.org/doc/html/v4.12/admin-guide/pm/intel_pstate.html
    echo 100 |sudo tee /sys/devices/system/cpu/intel_pstate/min_perf_pct &> /dev/null
    echo 100 |sudo tee /sys/devices/system/cpu/intel_pstate/max_perf_pct &> /dev/null
  elif [[ -f /sys/devices/system/cpu/cpufreq/boost ]]; then
    echo 0 |sudo tee /sys/devices/system/cpu/cpufreq/boost &> /dev/null
  fi

  # The governor is a linux component that can adjust CPU frequency. "performance" tells it to always run CPUs at
  # their maximum safe frequency. It seems to be the default for Amazon Linux, but it doesn't hurt to make this explicit.
  # See also https://wiki.archlinux.org/title/CPU_frequency_scaling
  echo performance | sudo tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor &> /dev/null
}

unapply_performance_tweaks() {
  if [[ -d /sys/devices/system/cpu/intel_pstate ]]; then
    # reenable turbo boost
    echo 0 |sudo tee /sys/devices/system/cpu/intel_pstate/no_turbo &> /dev/null

    # restore p-state limits
    echo $MIN_PERF_PCT |sudo tee /sys/devices/system/cpu/intel_pstate/min_perf_pct &> /dev/null
    echo $MAX_PERF_PCT |sudo tee /sys/devices/system/cpu/intel_pstate/max_perf_pct &> /dev/null
  elif [[ -f /sys/devices/system/cpu/cpufreq/boost ]]; then
    echo 1 | sudo tee /sys/devices/system/cpu/cpufreq/boost &> /dev/null
  fi

  # We do not reset the governor, as keeping track of each CPUs configured governor is not trivial here. On our CI
  # instances, the performance governor is current the default anyway (2023/11/14)
}


# `$0 test` - run integration tests
# Please see `$0 help` for more information.
#
cmd_test() {
    do_ab_test=0
    do_build=1
    do_kvm_check=1
    # Parse any command line args.
    while [ $# -gt 0 ]; do
        case "$1" in
            "-h"|"--help")      { cmd_help; exit 1; } ;;
            "-c"|"--cpuset-cpus")
                shift
                local cpuset_cpus="$1"
                ;;
            "-m"|"--cpuset-mems")
                shift
                local cpuset_mems="$1"
                ;;
            "--performance")
                local performance_tweaks=1;
                ;;
            "--ab")
                do_ab_test=1
                ;;
            "--no-build")
                do_build=0
                ;;
            "--no-kvm-check")
                do_kvm_check=0
                ;;
            "--")               { shift; break;     } ;;
            *)
                die "Unknown argument: $1. Please use --help for help."
            ;;
        esac
        shift
    done

    # Check prerequisites.
    [ $do_kvm_check != 0 ] && ensure_kvm
    ensure_devctr
    ensure_build_dir
    ensure_ci_artifacts
    if [ $do_build != 0 ]; then
      cmd_build --release
      if [ -n "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" ]; then
        cmd_build --release --rev "$BUILDKITE_PULL_REQUEST_BASE_BRANCH"
      fi
    fi

    apply_linux_61_tweaks

    # If we got to here, we've got all we need to continue.
    say "Kernel version: $(uname -r)"
    say "$(sed '/^processor.*: 0$/,/^processor.*: 1$/!d; /^processor.*: 1$/d' /proc/cpuinfo)"
    say "RPM firmware versions: $(rpm -q microcode_ctl amd-ucode-firmware linux-firmware)"

    env |grep -P "^(AWS_EMF_|BUILDKITE|CODECOV_)" > env.list
    if [[ $performance_tweaks -eq 1 ]]; then
      if [[ "$(uname --machine)" == "x86_64" ]]; then
        say "Detected CI and performance tests, tuning CPU frequency scaling and idle states for reduced variability"

        apply_performance_tweaks
      fi

      # It seems that even if the tests using huge pages run sequentially on ag=1 agents, right-sizing the huge pages
      # pool to the total number of huge pages used across all tests results in spurious failures with pool depletion
      # anyway (something else on the host seems to be stealing our huge pages, and we cannot "ear mark" them for
      # Firecracker processes). Thus, just allocate 48GB of them and call it a day.
      say "Setting up huge pages pool"
      num_hugetlbfs_pages=24552

      huge_pages_old=$(cat /sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages)
      huge_pages_new=$(echo $num_hugetlbfs_pages |sudo tee /sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages)
    fi

    if [[ "$huge_pages_new" -ne "$num_hugetlbfs_pages" ]]; then
      die "Failed to allocate $num_hugetlbfs_pages hugetlbfs pages, only got $huge_pages_new"
    fi

    say "Starting test run ..."

    test_script="./tools/test.sh"

    if [ $do_ab_test -eq 1 ]; then
      test_script="./tools/ab_test.py"
    fi

    # 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 \
        --ulimit nofile=4096:4096 \
        --ulimit memlock=-1:-1 \
        --workdir "$CTR_FC_ROOT_DIR" \
        --cpuset-cpus="$cpuset_cpus" \
        --cpuset-mems="$cpuset_mems" \
        --env-file env.list \
        -- \
        $test_script "$@"

    ret=$?

    say "Finished test run ..."

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

    # undo performance tweaks (in case the instance gets recycled for a non-perf test)
    if [[ $performance_tweaks -eq 1 ]]; then
      if [[ "$(uname --machine)" == "x86_64" ]]; then
        unapply_performance_tweaks
      fi

      echo $huge_pages_old |sudo tee /sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages >/dev/null
    fi

    # do not leave behind env.list file
    rm env.list

    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;  } ;;
            "--")               { shift; break;     } ;;
            *)
                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 \
            --ulimit nofile=4096:4096 \
            --ulimit memlock=-1:-1 \
            --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.
        #
        cmd_fix_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)" \
            --ulimit nofile=4096:4096 \
            --ulimit memlock=-1:-1 \
            --device=/dev/kvm:/dev/kvm \
            --workdir "$CTR_FC_ROOT_DIR" \
            --env PS1="$(whoami)@\h:\w\$ " \
            -- \
            bash --norc
        ret=$?
    fi

    return $ret
}

cmd_sh() {
    ensure_build_dir
    run_devctr \
        --privileged \
        --ulimit nofile=4096:4096 \
        --ulimit memlock=-1:-1 \
        --workdir "$CTR_FC_ROOT_DIR" \
        -- \
        bash --norc -c "$*"
}

cmd_sandbox() {
    cmd_build --release
    ensure_ci_artifacts
    cmd_sh "tmux new env PYTEST_ADDOPTS=--pdbcls=IPython.terminal.debugger:TerminalPdb PYTHONPATH=tests IPYTHONDIR=\$PWD/.ipython ipython -i ./tools/sandbox.py $@"
    cmd_fix_perms ".ipython"
}

cmd_sandbox_native() {
    cmd_build --release

    source /etc/os-release
    case $ID$VERSION_ID in
        ubuntu22.04)
            sudo apt install python3-pip python3.11-dev gcc tmux
            ;;
        al2023)
            sudo yum -y install python3.11-pip python3.11-devel gcc tmux
            ;;
    esac
    python3.11 -m venv sandbox
    source sandbox/bin/activate
    pip3.11 install ipython requests requests_unixsocket2 psutil tenacity filelock
    pip3.11 install jsonschema aws_embedded_metrics
    pip3.11 install packaging pytest
    ensure_ci_artifacts
    tmux neww sudo --preserve-env=HOME,PATH,TMUX env PYTHONPATH=tests IPYTHONDIR=\$PWD/.ipython ipython -i ./tools/sandbox.py $@
}

cmd_test_debug() {
    ensure_ci_artifacts
    cmd_sh "tmux new ./tools/test.sh --pdb $@"
}

# 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() {
    cmd_sh "cargo fmt --all -- --config $(tr '\n' ',' <tests/fmt.toml)"
    cmd_sh "cargo sort"
    cmd_sh "cd tests; black . ../tools ../.buildkite"
    cmd_sh "cd tests; isort . ../tools ../.buildkite"
    cmd_sh "mdformat $(git ls-files '*.md' | tr '\n' ' ')"
}

cmd_mkdocs() {
    cmd_sh "cargo doc --workspace --no-deps --document-private-items"
}

cmd_checkstyle() {
    cmd_test --no-build --no-kvm-check -- -n 4 --dist worksteal integration_tests/style || exit 1
    cmd_test --no-build --no-kvm-check -- -n 4 --doctest-modules framework || exit 1
}

# 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_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_vulns () {
    for f in /sys/devices/system/cpu/vulnerabilities/* ; do
        if $(grep -q "Vulnerable" ${f}) ; then
            say_warn "WARNING: `basename $f`: VULNERABLE.";
        fi
    done
}

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

check_EPT() {
    if [ "$(uname --machine)" = "x86_64" ]; then
        (grep -q "Y" /sys/module/kvm_intel/parameters/ept ; [ $? -ne 1 ]) || \
        say_warn "WARNING: EPT DISABLED. Performance will be affected."
    fi
}

check_vm() {
    if [ $(dmesg | grep -c -i "hypervisor detected") -gt 0 ]; then
        say_warn "WARNING: you are running in a virtual machine." \
    "Firecracker is not well tested under nested virtualization."
    fi
}

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
    check_vm
    say "Checking Host Security Configuration."
    say "Please check $PROD_DOC in case of any error."
    check_KSM
    check_swap
    check_EPT
    check_vulns
}

cmd_install() {
    # By default we install release/musl binaries.
    profile="release"
    target="$TARGET_PREFIX""musl"
    install_path="/usr/local/bin"
    binaries=("firecracker" "jailer" "seccompiler-bin" "rebase-snap" "cpu-template-helper")

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

    # Check that the binaries exist first
    for binary in "${binaries[@]}"; do
        bin_path=$( build_bin_path "$target" "$profile" "$binary" )
        if [ ! -f "$bin_path" ]; then
            die "Missing release binary. Needed file: $bin_path\n"\
            "To build the binaries, run:\n\t$0 build --$profile"
        fi
    done

    # Install the binaries
    for binary in "${binaries[@]}"; do
        say "Installing $binary in $install_path"
        install -m 755 -D -t "$install_path" "$( build_bin_path "$target" "$profile" "$binary" )"
    done
}

cmd_build_ci_artifacts() {
    # Check prerequisites
    ensure_devctr

    # We need to run nested Docker here, so run this container as privileged.
    run_devctr \
        --privileged \
        --workdir "$CTR_FC_ROOT_DIR" \
        -- \
        ./resources/rebuild.sh "$@"

    cmd_fix_perms
}


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)        # purposefully ignored
                ;;
            -*)
                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 "$@"