#!/bin/bash

# Copyright (c) 2011 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.

# This utility writes a firmware image to a tegra based board.

# Include common CrOS utilities.
. "/usr/lib/crosutils/common.sh"

get_default_board

# Command line options
DEFINE_string board "${DEFAULT_BOARD}" "The name of the board to flash."
DEFINE_string variant "" "Board variant."
DEFINE_string firmware "" "Path to firmware image to write."
DEFINE_string flasher "" "Path to flasher image to use."
DEFINE_string script "" "Path to flasher script to use."
DEFINE_string bct "" "Path to bct image to write."
DEFINE_string dtb "" "Path to dtb (device tree blob) file to use."
DEFINE_string map "" "Path to System.map file to read for load address."
DEFINE_string config "" "Directory for temporary configuration files."
DEFINE_boolean sign $FLAGS_FALSE "Sign and append a BCT to the firmware image."
DEFINE_boolean dry_run $FLAGS_FALSE "Do everything except calling nvflash."

# Parse command line.
FLAGS "$@" || exit 1
eval set -- "${FLAGS_ARGV}"

# Die on errors.
set -e

CROS_LOG_PREFIX="cros_write_firmware"

# No arguments are expected; error if any are passed.
if [ $# -ne 0 ]; then
  error "Unexpted arguments: $*"
  exit 1
fi

if [ -z "${FLAGS_board}" ]; then
  error "--board required."
  exit 1
fi

get_board_and_variant $FLAGS_board $FLAGS_variant

# Display the U-Boot version string.
function show_u_boot_info() {
  local firmware=$1
  local header=$(strings "${firmware}" | grep "U-Boot" | head -1 )

  info "Using ${header}"
}

# Given two values in decimal format return the first value rounded up to a
# multiple of the second value.  The second value must be a power of two.
function round() {
  awk "BEGIN { print and(($1 + $2 - 1), compl($2 - 1)) }"
}

# Given a value in decimal format return that value converted to hexidecimal
# and formatted as 0xXXXXXXXX.
function to_hex() {
  awk "BEGIN { printf \"0x%x\", $1 }"
}

# Given flasher and dtb, return an estimated size that should be sufficiently
# large for the final flasher image (which is the two combined)
function estimate_final_flasher_size() {
  local flasher_size=$(stat -c"%s" "$1")
  local dtb_size=$(stat -c"%s" "$2")
  round $((${flasher_size} + ${dtb_size})) 65536
}

# Create the flashing script.  This script will be loaded and run by the flasher
# build of U-Boot.  Most of the script is constant, we just need to append the
# sizes and address of the payload.
function make_flashing_script() {
  local firmware=$1
  local script=$2
  local firmware_size=$(stat -c"%s" "${firmware}")
  local image_address=$((${text_base} + ${script_offset} + ${script_size}))

  echo "setenv address       $(to_hex ${image_address})"
  echo "setenv firmware_size $(to_hex ${firmware_size})"
  echo "setenv length        $(to_hex $(round ${firmware_size} 4096))"

  cat "${script}"
}

# Parse the System.map file generated by the U-Boot build to get the entry
# point address.
function get_text_base() {
  local map=$1

  grep -m1 -E "^[0-9a-fA-F]{8} T _start$" ${map} | cut -d " " -f 1
}

###############################################################################

# Generate default values for the firmware, flasher, BCT, map and flash config
# file paths from the board and variant information provided.
firmware=$(find_u_boot_component "image.bin" "${FLAGS_firmware}")
flasher=$(find_u_boot_component "u-boot.bin" "${FLAGS_flasher}")
script=$(find_u_boot_component "flasher.script" "${FLAGS_script}")
bct=$(find_u_boot_component "bct/board.bct" "${FLAGS_bct}")
dtb=$(find_u_boot_component "dtb/${FLAGS_board/_/-}.dtb" "${FLAGS_dtb}")
map=$(find_u_boot_component "System.map" "${FLAGS_map}")

# Default value that hard-coded in u-boot-flasher.bin; use $(()) to convert
# hexadecimal into decimal values
script_offset=$((0x40000))
script_size=$((0x1000))

# Extract the load and entry point for U-Boot.
text_base="0x$(get_text_base "${map}")"

# Verify that all of the files and tools we will need are present.
check_for_file "firmware" " " "${firmware}"
check_for_file "flasher" "  " "${flasher}"
check_for_file "script" "   " "${script}"
check_for_file "BCT" "      " "${bct}"
check_for_file "map" "      " "${map}"
check_for_file "DTB" "      " "${dtb}"

check_for_tool "nvflash" "nvflash"
check_for_tool "dtget" "dtc"
check_for_tool "dtput" "dtc"

show_u_boot_info "${firmware}"

# Estimate final script offset and size. This is only an estimation because the
# dtb size may (or may not) be increased after we write bootcmd into it. So
# later we will have to check whether the flasher size is not larger than our
# estimation.
script_offset="$(estimate_final_flasher_size "${flasher}" "${dtb}")"
script_size=$(round $(stat -c"%s" ${script}) 4096)

info "Flasher script offset: $(to_hex ${script_offset})"
info "Flasher script size:   $(to_hex ${script_size})"

# If the user has not specified a config directory then create one on the fly
# and remove it when we are done.
if [ -z "${FLAGS_config}" ]; then
  config=$(mktemp -d)

  trap "rm -rf ${config}" EXIT
else
  mkdir -p "${FLAGS_config}"

  config=${FLAGS_config}
fi

# Optionally sign the firmware.  This is usefull if you want to write a bare
# U-Boot to flash and boot with that instead of a fully build ChromiumOS
# firmware image.
if [ $FLAGS_sign -eq $FLAGS_TRUE ]; then
  signed_firmware=${config}/signed.bin

  cros_sign_bootstub \
    --bct "${bct}" \
    --bootstub "${firmware}" \
    --output "${signed_firmware}" \
    --config "${config}" \
    --text_base "${text_base}"
else
  signed_firmware=${firmware}
fi

# make a copy of dtb
cp ${dtb} ${config}/u-boot.dtb
dtb="${config}/u-boot.dtb"

# write bootcmd to the dtb
#
# FIXME "%[^^]" (matching anything except '^' character) is a
# hack of putting a bootcmd string that:
#   * contains whitespace character
#   * but does not contain '^' character (so the choice of '^'
#     is somewhat arbitrarily)
# input the dtb file with a '\0' termination.
#
# We cannot use "%s", which stops at whitespace, and we cannot
# use "%${#bootcmd}c", which does not add '\0' to the scanned
# string. So our only option left is "%[^^]", which does not
# stop at whitespace and adds '\0' to the end.
#
bootcmd="source $(to_hex $((${text_base} + ${script_offset})))"
dtput -t s -f "%[^^]" ${dtb} /config/bootcmd "${bootcmd}"
info "Flasher bootcmd: '$(dtget -t s ${dtb} /config/bootcmd)'"

# create flasher
cat ${flasher} ${dtb} > ${config}/u-boot-flasher.bin
flasher="${config}/u-boot-flasher.bin"

# Create the flashing script and turn it into a U-Boot uImage.
make_flashing_script "${signed_firmware}" "${script}" > ${config}/script

mkimage -A arm -T script -d ${config}/script ${config}/script.uimg > /dev/null

# check whether our estimated script offset and size makes sense
final_script_offset=$(stat -c"%s" "${flasher}")
if [ ${final_script_offset} -gt ${script_offset} ]; then
  die "script offset > estimation: $(to_hex ${final_script_offset})"
fi
final_script_size=$(stat -c"%s" "${config}/script.uimg")
if [ ${final_script_size} -gt ${script_size} ]; then
  die "script size > estimation: $(to_hex ${final_script_size})"
fi

# Construct RAM image of flasher + script + firmware.  The firmware should
# already be signed by cbootimage.  The RAM image has a fixed 256K partition
# for the flasher, followed by the script and payload to be written.
b1=$((${script_offset} / 512))
b2=$(((${script_offset} + ${script_size}) / 512))
dd if="${flasher}"            of="${config}/image"          &> /dev/null
dd if="${config}/script.uimg" of="${config}/image" seek=$b1 &> /dev/null
dd if="${signed_firmware}"    of="${config}/image" seek=$b2 &> /dev/null

if [ $FLAGS_dry_run -eq $FLAGS_FALSE ]; then
  # Write the RAM image image to the board.  The flasher will write the payload
  # into the boot device using the flashing script.
  sudo nvflash \
    --bct "${bct}" \
    --setbct \
    --bl "${config}/image" \
    --go \
    --setentry "${text_base}" "${text_base}" ||
    (echo && die "Failed to download U-Boot flasher + payload using nvflash.")

  # TODO(robotboy): Make a U-Boot tool that can compute the same checksum on
  # the host and display it here.  Unfortunately U-Boot uses a slightly
  # different crc32 than zlib, so it's not possible to use zlib's crc32 to do
  # this.  It has to be done by building a tool that uses U-Boots version.
  info "The serial terminal or LCD panel should show before and after CRCs of"
  info "the image written to the boot flash.  If these values match each"
  info "other your flash was probably successful."
fi