#!/bin/bash
#BLURB="RK3399 Boot Loader flashing tool"
# Note: the BLURB above is consumed by 'pkgtools' menu within the installed OS,
# when you elect to re-run some of the package setup scripts.
#
# Copyright 2023 Stuart Winter, Donostia, Spain.
# All rights reserved.
#
# Redistribution and use of this script, with or without modification, is
# permitted provided that the following conditions are met:
#
# 1. Redistributions of this script must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
#  THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR IMPLIED
#  WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
#  MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO
#  EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
#  OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
#  WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
#  OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
#  ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
################################################################################
# Script.....: /usr/sbin/bootloader-flash-rk3399
# Called from: /usr/lib/setup/SeTconfig as
#              T_PX/var/lib/pkgtools/setup/setup.12.bootloader-flash-rk3399
#              or from the shell within the installed OS.
# Purpose....: Write U-Boot to SPI flash of supported Hardware Models.
# Version....: 2.03
# Date.......: 18-Nov-2023
# Authors....: Stuart Winter <mozes@slackware.com>
################################################################################
# References and Credits:
# Kamil Trzciński's RockPro64 repository:
# https://github.com/ayufan-rock64
# https://github.com/ayufan-rock64/linux-u-boot/releases
################################################################################
# Change log
# 2.03 - 18-Nov-2023
#  * Offer bail out option before flashing.
# 2.02 - 01-Oct-2023
# * Use '/sbin/slk-hwm-discover' to discover the Hardware Model
# 2.01 - 28-Jul-2023
# * Fixed cosmetic issue when run from 'pkgtool'
# 2.02 - 14-Sep-2023
# * Added XZ/LZMA compressed firmware support
################################################################################

#echo ${BASH_SOURCE[0]##*/}
#setup.12.bootloader-flash-rk3399

[ $( id -u ) -ne 0 ] && { echo "Error: You must be root to use this tool" ; exit 1 ;}

# From within the Installer, this script is called from /usr/lib/setup/SeTconfig
#
# Two arguments are provided:
# 1 -- the target prefix (normally /, but ${T_PX} from the bootdisk)
# 2 -- the name of the root device.
# Location of the root file system:
T_PX=${1}

# We may require slight behaviour changes depending on the context in which
# this script is launched:
[ -f /.installer-version ] && within_installer=Yes

# Temporary space for decompression:
TMPDIR=${T_PX}/var/lib/pkgtools/setup
# Sanity check.  This directory should exist within the OS and the Installer.
[ ! -d "${TMPDIR}" ] && { echo "ERROR: pkgtools temporary directory '${TMPDIR}' does not exist." ; exit 1 ;}
# Append our temporary directory to it.  '/var/lib/pkgtools/setup/tmp' is the standard location
# for temporary space for Slackware package setup scripts.
TMPDIR+=/tmp/bootloader-flash-rk3399

# Location of the firmware:
# $T_PX is where /usr/lib/setup/SeTpartitions stores the mount point for the new OS' root fs:
# On the OS it'll be empty.
firmwaredir=${T_PX}/usr/share/hwm-bw-rk3399/bootloader-firmware/
[ ! -d "${firmwaredir}" ] && { echo "ERROR: Unable to find firmware directory '${firmwaredir}'" ; exit 1 ;}

# Naming convention for SPI versions of U-Boot:
UBOOTSPIBIN=spi-idbloader.img # this is a suffix
platform_detected=0

################ Functions ################################

# Clean the temporary directory, if it was used to house a decompressed firmware file:
function cleanup() {
   rm -rf $TMPDIR
}

# Extract version of U-Boot from the SPI flash.  This assumes
# a Slackware version is present, otherwise we return '[Unknown / Non-Slackware release]'
# Note: The 'Slackware' string is stripped from the output to conserve screen space in the 'dialog' box!
# $1 = SPI flash device
#      e.g. /dev/mtd0
function u_boot_ver_spi() {
   # Assume U-Boot binaries won't be larger than 2MBytes in size.
   # If we don't put an upper limit on the data we read, it takes over a minute
   # to read the entire SPI flash.
   local spi_ver=$( dd status=none bs=4k count=2M iflag=count_bytes if=$1 | strings | sed -E 's/.*U-Boot (.*Slackware$)/U-Boot \1/' | grep -E '^U-Boot.*Slackware$' | sed 's?  Slackware$??' )
   [ -z "${spi_ver}" ] && spi_ver="[Unknown / Non-Slackware release]"
   echo "${spi_ver}" ;}

# Determine the version of U-Boot contained within the Slackware 'a/hwm-bw-rk3399'
# package:
# $1 = file name
# e.g. /usr/share/hwm-bw-rk3399/bootloader-firmware/rk3399_rockpro64-spi-idbloader.img
function u_boot_ver_pkg() {
   local local_fwfile=$1
   # Handle LZMA compressed firmware:
   file ${local_fwfile} | grep -q 'XZ compressed' && furry="xz"
   local pkg_ver=$( ${furry}cat ${local_fwfile} | strings | sed -E 's/.*U-Boot (.*Slackware$)/U-Boot \1/' | grep -E '^U-Boot.*Slackware$' | sed 's?  Slackware$??')
   echo "${pkg_ver}" ;}
#u_boot_ver_spi /dev/mtd0
#u_boot_ver_pkg /usr/share/hwm-bw-rk3399/bootloader-firmware/rk3399_rockpro64-spi-idbloader.img

# Discover the packaged firmware file size:
function fwfilesize_pkg(){
   # Handle LZMA compressed firmware:
   file ${1} | grep -q 'XZ compressed' && {
      # It's compressed:
      xz --robot --list ${1} | grep -E '^file' | awk '{print $5}' ;} || {
      # Uncompressed; report size natively:
      find ${1} -printf "%s" ;}
}

# Create the temporary holding space and ensure that there's sufficient space
# available to hold the decompressed firmware file:
function preallocate_tmpspace() {
  local tmpdir=$1
  local size=$2
  local exitval
  mkdir -p $tmpdir
  fallocate -l $size $tmpdir/testfile-$size > /dev/null 2>&1
  exitval=$?
  # Wipe the temp file but not the directory, as we'll use it.
  rm -f $tmpdir/testfile-$size
  return $exitval
}

# Decompress firmware:
function decompress_pkg_fw() {
   local fw_basename="${1##*/}" # rk3399_rockpro64-spi-idbloader.img.xz
   local out_file
   # Decompress the firmware to a temporary location:
   # XZ-compressed files:
   [ "${fw_basename##*.}" = "xz" ] && {
      out_file=${TMPDIR}/${fw_basename%.xz} # lop off the file extension
      # Decompress to temporary holding space:
      xz -dc ${1} > ${out_file} ;}
   # Return the path + file name:
   echo ${out_file} ;}

# Determine whether the version of U-Boot within the Slackware package
# is identical with that on SPI flash:
# $1=File name of U-Boot firmware within OS
# $2=Flash device name
function cmp_spi_to_pkg() {
   local local_fwfile=$1
   # Handle LZMA compressed firmware:
   file ${local_fwfile} | grep -q 'XZ compressed' && furry="xz"
   # Obtain MD5 sums of both:
   local spi_md5=$( dd status=none bs=4k count=$( fwfilesize_pkg ${local_fwfile} ) iflag=count_bytes if=$2 | md5sum | awk '{print $1}' )
   local pkg_md5=$( ${furry}cat ${local_fwfile} | md5sum | awk '{print $1}' )
   # Return >0 status code if mismatched:
   [ "${spi_md5}" != "${pkg_md5}" ] && return 1
   return 0 ;}

# Flash management:
#
# Note: for Hardware Models that use 'NAND' flash, we need to
# update this function to support both.
# The code is commented for NAND.
# Note: The RK3399 has 'NOR' flash not 'NAND'.
function write_flash() {
   local flash_dev=${1}
   local flash_bin=${2}

# Example for Hardware Models that use 'NAND' flash:
#   echo "Flash device: ${flash_dev}  Operation: ERASING"
#   echo "This will take a few minutes..."
#   flash_erase ${flash_dev} 0 0
#   estat=$?
#   if [ $estat -gt 0 ]; then
#      read -p "An error occurred.  Press ENTER to continue"
#      return $estat
#   fi
#
#   echo
#   echo "Flash device: ${flash_dev}  Operation: WRITING"
#   echo "Image file..: $( basename ${flash_bin} )"
#   echo
#   echo "This can take between 5 to 30 minutes during which"
#   echo "time you will not see any output here."
#   echo
#   echo "Please wait..."

#   echo "This will take approximately two (2) minutes"

   # For Hardware Models with 'NAND', you'd use nandwrite
   # to install the Boot Loader:
   #nandwrite -qp ${flash_dev} < ${flash_bin}

   # flashcp takes care of erase & write directly:
   printf "\nFlash device: ${flash_dev}\n"
   printf "Image file..: ${flash_bin##*/}\n\n"
   printf "This will take approximately two (2) minutes\n\n"
   printf "Please wait...\n\n"
   flashcp -vA ${flash_bin} ${flash_dev}
   estat=$?
   if [ $estat -gt 0 ]; then
      read -p "An error occurred.  Press ENTER to continue"
      return $estat
    else
      return 0
   fi
}

function offerretry() {
   dialog \
      --backtitle "Bootware" \
      --title "INSTALL BOOT LOADER INTO ONBOARD FLASH (SPI FLASH)" --yesno \
"\nA failure was encountered during the flashing process\n\n
Without the Boot Loader being successfully installed to the onboard
flash of this Hardware Model, it may fail to boot.\n\n

\n\nDo you want to retry? (recommendation is 'Yes')\n \

\n" 15 69
   # If they selected no, exit now.
   return $?
}

function err_report_dialog() {
  local exitcode=$1
  local msg_title="$2"
  local msg="$3"
  dialog \
      --backtitle "Error" \
      --title "${msg_title}" \
      --ok-button "OK" \
      --msgbox "${3}" 0 0
  clear
  exit ${exitcode} ;}

function offerready() {
   dialog \
      --backtitle "Bootware" \
      --title "INSTALL BOOT LOADER INTO ONBOARD FLASH (SPI FLASH)" \
      --yes-label "FLASH" \
      --no-label "ABORT" \
      --yesno \
"\nYour machine must remain powered up during the process.\n
\nPress ABORT to abort without flashing.\n\n" 0 0
   return $?
}

function offersuccess() {
   dialog \
      --backtitle "Bootware" \
      --title "INSTALL BOOT LOADER INTO ONBOARD FLASH (SPI FLASH)" --ok-button "OK" \
      --msgbox "\n\nFlashing was successful.\n" 9 60
   clear
   return 0
}

# Ensure that the flash device exists and we can read it.
# An md5 sum of the selected size will be returned.  This is used for Hardware Models
# where we ship a variety of firmware to support physical hardware configuration profiles.
# The MD5sum of the first X bytes is compared to a store of all available firmware files
# in order to determine the hardware profile.
function flash_test () {
   local flash_dev=$1
   local flash_md5_readsize=$2
   local flash_md5
   [ ! -c ${flash_dev} ] && return 3
   flash_md5=$( dd status=none bs=4k count=${flash_md5_readsize} iflag=count_bytes if=${flash_dev} 2>/dev/null | md5sum | awk '{print $1}' )
   estat=$?
   [ ${estat} -gt 0 ] && return ${estat}
   echo "${flash_md5}"
   return 0
}

################ Functions ################################

# Determine the Hardware Model name:
[ -x ${T_PX}/sbin/slk-hwm-discover ] && { export HWM=$( slk-hwm-discover ) ;} || {
  err_report_dialog 0 \
  "Missing OS component" \
  "\n'slk-hwm-discover' is missing from your system. This is a critical component of the Kernel package.\n" ;}

# Set the variables on a per-device basis.
# If you need to add your hardware model here, see the instructions at:
# http://arm.slackware.com/slackwarearm-devel/
#
# If your Hardware Model uses the same SoC, and has a different
# SPI flash binary file, you should have a separate stanza.
# For example: the Pinebook Pro and the RockPro64 share the SoC 'rk3399',
# but diverge at the U-Boot binary.
# The names of these assets are within the Slackware ARM source tree:
# /platform/aarch64/src/uboot.build
#
# TODO:
#
# To support other Hardware Models in the future, we'll break
# this code out in the same manner as the Slackware ARM Linux Kernel Module Loader.
case "$HWM" in
   ###########################################
   # RK3399-based Hardware Models:
   ###########################################
   #
   # Hardware Model-specific settings:
   "Pine64 Pinebook Pro"*) HWMODEL=pinebookpro ;&
   "Pine64 RockPro64"*|"Pine64 Pinebook Pro"*)
      HWMODEL=${HWMODEL:-"rockpro64"}
      platform_detected=1
      SOC=rk3399
      BOARDFLASHDEV=/dev/mtd0
      # On the RockPro64:
      #root@bladswede:~# cat /proc/mtd
      #dev:    size   erasesize  name
      #mtd0: 01000000 00001000 "spi0.0"
      BOARDFLASHNAME=spi0.0
      UBOOTSPIBIN=${SOC}_${HWMODEL}-${UBOOTSPIBIN} ;;
esac

# Sanity check:
if [ $platform_detected -ne 1 ]; then
#   echo "$0 : platform not detected"
#   echo "No boot loader flashing required."
   # Writing a new boot loader to the SPI flash is not a requirement
   # for Slackware ARM / AArch64.
   # Within the Installer we don't output any messages since they are irrelevant and look like an error.
   # Within the OS, we report that the user isn't running this tool on the correct Hardware Model.
   [ -z "${within_installer}" ] && { echo "This boot loader flashing tool does not support this Hardware Model." ;}
   # Exit with 0 status to avoid any breakages within the Installer:
   exit 0
fi

# Ensure the flash device is present and that we can read it:
flash_test ${BOARDFLASHDEV} 100k > /dev/null || { echo "$0 : ERROR - unaddressable flash device ${BOARDFLASHDEV}" ; exit 1 ;}

# If the firmware is compressed, append the file
# extension:
[ -s ${firmwaredir}/${UBOOTSPIBIN}.xz ] && UBOOTSPIBIN+=".xz"

# Ensure that we can find the boot loader:
if [ ! -s ${firmwaredir}/${UBOOTSPIBIN} ]; then
   echo "$0 : ERROR - unable to find boot loader image at"
   echo "${firmwaredir}/${UBOOTSPIBIN}"
   exit 4
fi

#printf "\n\nScanning Boot Loader information, please wait..."
dialog \
   --backtitle "Bootware" \
   --title "Boot Loader Management" \
   --infobox "Scanning Boot Loader information, please wait..." 3 55

# Determine installed and available Boot Loader versions and set appropriate
# messages:
cmp_spi_to_pkg ${firmwaredir}/${UBOOTSPIBIN} $BOARDFLASHDEV && {
  default_flash_opt="--defaultno" # set dialog's default response option to 'No'
  fw_sync_msg="\Zb\Zr\Z2Boot Loader is up to date - flashing \Zunot required\Z1\Zn" ;} || {
  fw_sync_msg="\Zb\Zr\Z1Boot Loader update available - flashing recommended\Z1\Zn" ;}

dialog \
   ${default_flash_opt} \
   --colors \
   --backtitle "Bootware" \
   --title "INSTALL BOOT LOADER INTO ONBOARD FLASH (SPI FLASH)" --yesno \
"

\nHardware Model       : \Zb\Zr\Z8 ${HWM} \Z1\Zn
\nSystem On Chip       : \Zb\Zr\Z8 ${SOC} \Z1\Zn

\nInstalled Boot Loader: \Zb\Zr\Z8 $( u_boot_ver_spi $BOARDFLASHDEV ) \Z1\Zn
\nAvailable Boot Loader: \Zb\Zr\Z8 $( u_boot_ver_pkg ${firmwaredir}/${UBOOTSPIBIN} ) \Z1\Zn

\n\n\n${fw_sync_msg}

\n\nDo you want to flash the available version now?\n

\n" 15 79

# If they selected no, exit now.
[ $? = 1 ] && { clear ; cleanup ; exit 0 ;}

# Path to the firmware we plan to flash to SPI:
fw_flash_bin=${firmwaredir}/${UBOOTSPIBIN}

# If it's compressed, decompress it ready for flashing:
if file ${fw_flash_bin} | grep -q 'XZ compressed'; then
   # Determine the uncompressed file size:
   fw_flash_bin_size=$( fwfilesize_pkg ${fw_flash_bin} )
   # Allocate and ensure sufficient temporary space:
   preallocate_tmpspace $TMPDIR ${fw_flash_bin_size}
   if [ $? -gt 0 ]; then
      dialog \
         --backtitle "Bootware" \
         --title "Error" --ok-button "OK" \
         --msgbox "\n\nInsufficient temporary storage: ${fw_flash_bin_size} bytes required.\n" 9 60
      cleanup > /dev/null
      exit 1
   fi
   # Decompress, storing the uncompressed version's path+file name:
   fw_flash_bin="$( decompress_pkg_fw ${fw_flash_bin} )"
fi

# Final notice and chance to bail out:
offerready || { clear ; cleanup ; exit 0 ;}

# Loop until either flashing was successful or the user abandons it (after being offered to retry):
while true; do
   clear
   write_flash $BOARDFLASHDEV ${fw_flash_bin}
   estat=$?
   if [ $estat -gt 0 ]; then
      offerretry || { clear ; cleanup ; exit 1 ;} # Failed to flash, user didn't want to retry
    else
      # Flashing was successful, and user acknowledged:
      offersuccess ; cleanup
      break;
   fi
done