# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#
# Common vm functions for use in crosutils.

DEFAULT_PRIVATE_KEY="${GCLIENT_ROOT}/src/scripts/mod_for_test_scripts/\
ssh_keys/testing_rsa"

DEFINE_string kvm_pid "" \
  "Use this pid file.  If it exists and is set, use the vm specified by pid."
DEFINE_boolean copy ${FLAGS_FALSE} "Copy the image file before starting the VM."
DEFINE_string mem_path "" "VM memory image to save or restore."
DEFINE_boolean no_graphics ${FLAGS_FALSE} "Runs the KVM instance silently."
DEFINE_boolean persist "${FLAGS_FALSE}" "Persist vm."
DEFINE_boolean scsi ${FLAGS_FALSE} "Loads disk as a virtio-scsi-disk. "\
"This option is used for testing Google Compute Engine-compatible images."
DEFINE_boolean snapshot ${FLAGS_FALSE} "Don't commit changes to image."
DEFINE_integer ssh_port 9222 "Port to tunnel ssh traffic over."
DEFINE_string ssh_private_key "${DEFAULT_PRIVATE_KEY}" \
    "Path to the private key to use to ssh into test image as the root user."
DEFINE_string vnc "" "VNC Server to display to instead of SDL "\
"(e.g. pass ':1' to listen on 0.0.0.0:5901)."
DEFINE_string usb_devices "" \
    "Usb devices for passthrough. Specified in a comma-separated list
     where each item is of the form <vendor_id>:<product_id>
     (eg. --usb_devices=1050:0211,0409:005a)"
DEFINE_boolean moblab ${FLAGS_FALSE} "Setup environment for moblab"
DEFINE_string qemu_binary \
  "${DEFAULT_CHROOT_DIR}"/usr/bin/qemu-system-x86_64 \
  "The qemu binary to be used. Defaults to qemu shipped with the SDK."
# TODO(pwang): replace SandyBridge with Haswell-noTSX once all builders
# running VMTest are switched to GCE.
DEFINE_string kvm_cpu "SandyBridge,-invpcid,-tsc-deadline" \
    "Configures guest CPU model and features. The default should be the
     least restrictive option supported by all of Chrome OS infrastructure.
     Other locally valid choices are for instance 'qemu64' (least features)
     or 'host' (maximum features)."
DEFINE_string kvm_m 8G "Configure guest RAM"
DEFINE_integer kvm_smp 0 \
  "Configure guest CPU counts (0 means auto-detect cpu count)"
DEFINE_string serial "" \
  "Configure qemu serial output. By default, serial is output to a file"

KVM_PID_FILE=/tmp/kvm.$$.pid
LIVE_VM_IMAGE=

# Pick which qemu/kvm binary will be used. Must be called before any function
# that needs ${KVM_BINARY}, and *after* the command line has been parsed by the
# calling script. Otherwise the default value can not be overriden by the user.
set_kvm() {
  # The value of the flag is only valid after the command line has been parsed.
  KVM_BINARY="${FLAGS_qemu_binary}"
  if [[ ! -x "${KVM_BINARY}" ]]; then
    if ! KVM_BINARY=$(which qemu-system-x86_64 2> /dev/null); then
      die "no QEMU binary found"
    fi
  fi
  info "QEMU binary: ${KVM_BINARY}"

  # Make sure it's a recent enough version.
  # The version string typically looks like this:
  #"QEMU emulator version 2.5.0, Copyright (c) 2003-2008 Fabrice Bellard"
  # but in Debian/ubuntu distributions, some QEMU binaries have a
  # space and a package version rather the comma just after the version number:
  #"QEMU emulator version 2.0.0 (Debian 2.0.0+dfsg-2ubuntu1.22), Copyright[...]"
  local ver version
  version=$("${KVM_BINARY}" --version)
  ver="${version#QEMU emulator version }"
  info "QEMU version: ${version}"
  case ${ver} in
  [23].[0-9]*) ;;
  *) die "Old/unknown/unsupported version of QEMU: ${ver}" ;;
  esac
}

get_pid() {
  sudo cat "${KVM_PID_FILE}"
}

get_host_files_prefix() {
  echo "${KVM_PID_FILE%.pid}"
}

# Configure paths to KVM pipes. Must not be called until after KVM_PID_FILE
# has been updated. (See, e.g., start_kvm.)
set_kvm_pipes() {
  local base="$(get_host_files_prefix)"
  KVM_PIPE_PREFIX="${base}.monitor"
  KVM_PIPE_IN="${KVM_PIPE_PREFIX}.in"  # to KVM
  KVM_PIPE_OUT="${KVM_PIPE_PREFIX}.out"  # from KVM
  KVM_SERIAL_FILE="${base}.serial"
}

# General purpose blocking kill on a pid.
# This function sends a specified kill signal [0-9] to a pid and waits for it
# die up to a given timeout.  It exponentially backs off it's timeout starting
# at 1 second.
# $1 the process id.
# $2 signal to send (-#).
# $3 max timeout in seconds.
# Returns 0 on success.
blocking_kill() {
  local timeout=1
  sudo kill -$2 $1
  while ps -p $1 > /dev/null && [ ${timeout} -le $3 ]; do
    sleep ${timeout}
    timeout=$((timeout*2))
  done
  ! ps -p ${1} > /dev/null
}

# Send a command to the KVM monitor. The caller is responsible for
# escaping the command, so that it survives sudo sh -c "$arg".
# Additionally, |set_kvm_pipes| must have been called before this
# function.
send_monitor_command() {
  local command="${1}"
  sudo sh -c "echo ${1} > ${KVM_PIPE_IN}"
}

# Send a command to the KVM monitor, and wait for KVM to issue another
# prompt. The caller is responsible for escaping the command, so that
# it survives sudo sh -c "$arg".  Additionally, |set_kvm_pipes| must
# have been called before this function.
send_monitor_command_and_wait() {
  local command="${1}"
  sudo sh -c "echo ${1} > ${KVM_PIPE_IN}"
  # Wait for the command prompt. Note that we send an empty command
  # before waiting, because the monitor's command prompt doesn't
  # include a newline. (And grep waits for a newline.)
  sudo sh -c "echo > ${KVM_PIPE_IN}"
  sudo grep -F -q "(qemu)" "${KVM_PIPE_OUT}"
}

# Return a command which will read stdin, and write a (compressed)
# bytestream to stdout, for the compression format implied by
# |filename|.
get_compressor() {
  local filename="${1}"
  local extra_flag="${2:-}"
  case "${filename}" in
    *.gz)
      compressor="pigz -c ${extra_flag}"
      ;;
    *.bz2)
      compressor="pbzip2 -c ${extra_flag}"
      ;;
    *)
      compressor="cat"
      ;;
  esac
  echo "${compressor}"
}

# Return a command which will read stdin, and write a (decompressed)
# bytestream to stdout, for the compression format implied by
# |filename|.
get_decompressor() {
  get_compressor "${1}" "-d"
}

# Disable ksmd which causes soft kernel lockups on kernel 3.13.
# For details b/71727384.
disable_ksmd() {
  local run_path="/sys/kernel/mm/ksm/run"
  local ksm=$(<"${run_path}")
  if [[ "${ksm}" == "1" ]]; then
    info "Disabling currently enabled ksmd to work around kernel lockups."
    # This is for interactive use. If a calling script hangs here please set
    # this environment in advance when you are able to sudo.
    echo 0 | sudo tee "${run_path}" > /dev/null
  fi
}

# $1: Path to the virtual image to start.
# $2: Name of the board to virtualize.
start_kvm() {
  local vm_image="$1"
  local board="$2"
  local extra_args=( "${@:3}" )

  disable_ksmd
  set_kvm

  # Append 'check' option to warn if requested CPU is not fully supported.
  local kvm_cpu="${FLAGS_kvm_cpu},check"

  # Override default pid file.
  local start_vm=0
  [ -n "${FLAGS_kvm_pid}" ] && KVM_PID_FILE=${FLAGS_kvm_pid}
  if [ -f "${KVM_PID_FILE}" ]; then
    local pid=$(get_pid)
    # Check if the process exists.
    if ps -p ${pid} > /dev/null ; then
      echo "Using a pre-created KVM instance specified by ${FLAGS_kvm_pid}." >&2
      start_vm=1
    else
      # Let's be safe in case they specified a file that isn't a pid file.
      echo "File ${KVM_PID_FILE} exists but specified pid doesn't." >&2
    fi
  fi

  # No kvm specified by pid file found, start a new one.
  if [ ${start_vm} -eq 0 ]; then
    echo "Starting a KVM instance" >&2
    local kvm_flag=""
    local nographics=""
    local usesnapshot=""
    if [ ${FLAGS_no_graphics} -eq ${FLAGS_TRUE} ]; then
      nographics="-display none"
    fi
    if [ -n "${FLAGS_vnc}" ]; then
      nographics="-vnc ${FLAGS_vnc}"
    fi

    if [ ${FLAGS_snapshot} -eq ${FLAGS_TRUE} ]; then
      snapshot="-snapshot"
    fi

    # When using the regular qemu system binary, force KVM.
    case "${KVM_BINARY}" in
      */qemu-system-x86_64)
        kvm_flag="-enable-kvm"
        ;;
    esac

    if [ ${FLAGS_copy} -eq ${FLAGS_TRUE} ]; then
      local our_copy=$(mktemp "${vm_image}.copy.XXXXXXXXXX")
      if cp "${vm_image}" "${our_copy}"; then
          info "Copied ${vm_image} to ${our_copy}."
          vm_image="${our_copy}"
      else
          die "Copy failed. Aborting."
      fi
    fi

    local kvm_smp=${FLAGS_kvm_smp}
    if [[ ${kvm_smp} -eq 0 ]]; then
      local kvm_smp_default=8
      info "Set kvm_smp to min(${kvm_smp_default}, ${NUM_JOBS})."
      kvm_smp=$((kvm_smp_default > NUM_JOBS ? NUM_JOBS : kvm_smp_default))
    fi

    local net_option="-device virtio-net,netdev=eth0"
    local net_user="-netdev user,id=eth0,net=10.0.2.0/27"
    net_user+=",hostfwd=tcp:127.0.0.1:${FLAGS_ssh_port}-:22"

    local incoming=""
    local incoming_option=""
    if [ -n "${FLAGS_mem_path}" ]; then
      local decompressor=$(get_decompressor "${FLAGS_mem_path}")
      incoming="-incoming"
      incoming_option="exec: ${decompressor} ${FLAGS_mem_path}"
    fi

    local usb_passthrough=""
    if [ -n "${FLAGS_usb_devices}" ]; then
      local bus_id
      local usb_devices=(${FLAGS_usb_devices//,/ })
      for bus_id in "${usb_devices[@]}"; do
        local device=(${bus_id//:/ })
        if [ ${#device[@]} -ne 2 ]; then
          continue
        fi
        usb_passthrough+=" -device usb-host,vendorid=$((0x${device[0]}))"
        usb_passthrough+=",productid=$((0x${device[1]}))"
      done

      if [ -n "${usb_passthrough}" ]; then
        usb_passthrough="-usb ${usb_passthrough}"
      fi
    fi

    local base="$(get_host_files_prefix)"
    local moblab_env=""
    if [ ${FLAGS_moblab} -eq ${FLAGS_TRUE} ]; then
      # Increase moblab memory size.
      moblab_env="-m 4G"

      # Add hostforwarding for important moblab pages to the SLIRP connection.
      MOB_MONITOR_PORT=$(( FLAGS_ssh_port + 1 ))
      AFE_PORT=$(( FLAGS_ssh_port + 2 ))
      DEVSERVER_PORT=$(( FLAGS_ssh_port + 3 ))

      net_user+=",hostfwd=tcp:127.0.0.1:${MOB_MONITOR_PORT}-:9991"
      net_user+=",hostfwd=tcp:127.0.0.1:${AFE_PORT}-:80"
      net_user+=",hostfwd=tcp:127.0.0.1:${DEVSERVER_PORT}-:8080"

      info "Mob* Monitor: 127.0.0.1:${MOB_MONITOR_PORT}"
      info "Autotest: 127.0.0.1:${AFE_PORT}"
      info "Devserver: 127.0.0.1:${DEVSERVER_PORT}"
    fi

    set_kvm_pipes
    for pipe in "${KVM_PIPE_IN}" "${KVM_PIPE_OUT}"; do
      sudo rm -f "${pipe}"  # assumed safe because, the PID is not running
      sudo mknod "${pipe}" p
      sudo chmod 600 "${pipe}"
    done

    local serial=""
    if [[ -n "${FLAGS_serial}" ]]; then
      serial="${FLAGS_serial}"
    else
      sudo touch "${KVM_SERIAL_FILE}"
      sudo chmod a+r "${KVM_SERIAL_FILE}"
      serial="file:${KVM_SERIAL_FILE}"
    fi

    local drive
    drive="-drive file=${vm_image},index=0,media=disk,cache=unsafe"
    if [ ${FLAGS_scsi} -eq ${FLAGS_TRUE} ]; then
      drive=$(echo "-drive if=none,id=hd,file=${vm_image},cache=unsafe"\
          "-device virtio-scsi-pci,id=scsi "\
          "-device scsi-hd,drive=hd")
    fi

    # Note: the goofiness around the expansion of |incoming_option| is
    # to ensure that it is quoted if set, but _not_ quoted if
    # unset. (QEMU chokes on empty arguments).
    local cmd=(
      "${KVM_BINARY}" ${kvm_flag}
      -m ${FLAGS_kvm_m}
      -smp ${kvm_smp}
      -cpu ${kvm_cpu}
      -vga virtio
      -pidfile "${KVM_PID_FILE}"
      -chardev pipe,id=control_pipe,path="${KVM_PIPE_PREFIX}"
      -serial "${serial}"
      -mon chardev=control_pipe
      -daemonize
      ${net_option}
      ${nographics}
      ${snapshot}
      ${net_user}
      ${incoming} ${incoming_option:+"$incoming_option"}
      ${usb_passthrough}
      ${moblab_env}
      ${drive}
      "${extra_args[@]}"
    )
    info "Launching: ${cmd[*]}"
    sudo "${cmd[@]}"

    info "KVM started with pid stored in ${KVM_PID_FILE}"
    if [[ -z "${FLAGS_serial}" ]]; then
      info "Serial output, if available, can be found in ${KVM_SERIAL_FILE}"
    fi
    LIVE_VM_IMAGE="${vm_image}"
  fi
}

# Checks to see if we can access the target virtual machine with ssh.
ssh_ping() {
  # TODO(sosa): Remove outside chroot use once all callers work inside chroot.
  local cmd
  if [ $INSIDE_CHROOT -ne 1 ]; then
    cmd="${GCLIENT_ROOT}/src/scripts/ssh_test.sh"
  else
    cmd=/usr/lib/crosutils/ssh_test.sh
  fi
  "${cmd}" \
    --ssh_port=${FLAGS_ssh_port} \
    --private_key=${FLAGS_ssh_private_key} \
    --remote=127.0.0.1 >&2
}
# Tries to ssh into live image $1 times.  After first failure, a try involves
# shutting down and restarting kvm.
retry_until_ssh() {
  local can_ssh_into=1
  local max_retries=3
  local retries=0
  ssh_ping && can_ssh_into=0

  while [ ${can_ssh_into} -eq 1 ] && [ ${retries} -lt ${max_retries} ]; do
    echo "Failed to connect to virtual machine, retrying ... " >&2
    stop_kvm || echo "Could not stop kvm.  Retrying anyway." >&2
    start_kvm "${LIVE_VM_IMAGE}"
    ssh_ping && can_ssh_into=0
    retries=$((retries + 1))
  done
  return ${can_ssh_into}
}

stop_kvm() {
  set_kvm
  if [ "${FLAGS_persist}" -eq "${FLAGS_TRUE}" ]; then
    echo "Persist requested.  Use --ssh_port ${FLAGS_ssh_port} " \
      "--ssh_private_key ${FLAGS_ssh_private_key} " \
      "--kvm_pid ${KVM_PID_FILE} to re-connect to it." >&2
  else
    echo "Stopping the KVM instance" >&2
    set_kvm_pipes
    local pid=$(get_pid)
    if [ -n "${pid}" ]; then
      if ! [ -d "/proc/${pid}" ]; then
        echo "KVM pid ${pid} no longer running." >&2
        return 1
      fi

      if [ -n "${FLAGS_mem_path}" ]; then
        local mem_path="${FLAGS_mem_path}"
        local compressor=$(get_compressor "${mem_path}")
        echo "Saving memory snapshot to ${mem_path}..."
        echo "    freezing VM..."
        send_monitor_command_and_wait "stop"
        echo "    saving memory, piping through ${compressor}..."
        # Create file as current user, so that it will be readable by
        # the current user. (Otherwise, it would be owned by root.)
        touch "${mem_path}"
        send_monitor_command_and_wait \
            "migrate \\\"exec:${compressor} \> ${mem_path}\\\""
        # Flush any disk I/O that is buffered in KVM.
        echo "    flushing disk buffers..."
        send_monitor_command_and_wait "commit all"
        # Quit KVM now, so that we don't modify the filesystem which
        # this memory image depends on.
        echo "    asking KVM to quit..."
        send_monitor_command "quit"
        echo "    done."
      else
        # Initiate the power-off sequence inside the guest. Note that
        # this monitor command does not wait for the guest to power
        # off the system.
        send_monitor_command "system_powerdown"
      fi
      blocking_kill ${pid} 0 16 || blocking_kill ${pid} 9 3
      sudo rm -f "${KVM_PID_FILE}" "${KVM_PIPE_IN}" "${KVM_PIPE_OUT}" \
        "${KVM_SERIAL_FILE}"
    else
      echo "No kvm pid found to stop." >&2
      return 1
    fi
  fi
}