#! /bin/sh -m

pid=$$
protected="$$"

var_devices_removed=""
root_devices_removed=""
ramdisks=""
update_method=
logfile="/var/log/isi_disk_firmware.log"

# flag file to tell /etc/rc.d/ scripts that we're shutting down
# to do a disk firmware update.
update_flag="/tmp/isi_disk_firmware_update"

# Flag file to indicate we've already shutdown.
shutdown_flag="/dev/.isi_disk_firmware_reboot.shutdown"
rc_shutdown_flag="/dev/.isi_disk_firmware_reboot.rc_shutdown"
unconfigure_other_imdd_flag="/dev/.isi_disk_firmware_reboot.unconfigure"

# Flags indicate if we've prepared ramdisks or added ramdisks
# to imdd0 and imdd1.
ramdisk_md5=0
ramdisk_md6=0
md5a_added=0
md6a_added=0

cleanup() {
	rm -f -- "$update_flag"
}

die() {
	# Rebuild imdd0 and imdd1
	for device in $root_devices_removed; do
		echo "re-adding device $device to imdd0"
		isi_imdd_config -a imdd0 $device
		status=$?
		if [ $status -ne 0 ] ; then
			echo "failed to add $device to imdd0"
		fi
	done
	for device in $var_devices_removed; do
		echo "re-adding device $device to imdd1"
		isi_imdd_config -a imdd1 $device
		status=$?
		if [ $status -ne 0 ] ; then
			echo "failed to add $device to imdd1"
		fi
	done
	# Remove ramdisks that were added to imdd0 and imdd1
	if [ $md5a_added -eq 1 ] ; then
		echo "removing device /dev/md5a from imdd0"
		isi_imdd_config -x imdd0 /dev/md5a
	fi
	if [ $md6a_added -eq 1 ] ; then
		echo "removing device /dev/md6a from imdd1"
		isi_imdd_config -x imdd1 /dev/md6a
	fi
	# Free ramdisks that were prepared
	if [ $ramdisk_md5 -eq 1 ] ; then
		echo "Detaching ramdisk md5"
		mdconfig -d -u 5
	fi
	if [ $ramdisk_md6 -eq 1 ] ; then
		echo "Detaching ramdisk md6"
		mdconfig -d -u 6
	fi
	echo "$*" 1>&2
	sleep 1
	cleanup
	exit 1
}

trap 'die SIGINT $$' INT 
trap 'die SIGTERM $$' TERM

prepare_ramdisks() {
	# make sure there is enough memory available for ramdisks
	page_size="$(sysctl vm.stats.vm.v_page_size | sed 's/[^0-9]//g')"
	free_pages="$(sysctl vm.stats.vm.v_free_count | sed 's/[^0-9]//g')"
	inactive_pages="$(sysctl vm.stats.vm.v_inactive_count | sed 's/[^0-9]//g')"
	cache_pages="$(sysctl vm.stats.vm.v_cache_count | sed 's/[^0-9]//g')"
	free_bytes="$((($free_pages + $inactive_pages + cache_pages) * $page_size))"

	# only proceed if 10% more than minimum required memory is available
	if [ $free_bytes -lt 1194239693 ] ; then
		die "not enough free memory for ramdisks, 1.1 GB needed"
	fi

	# allocate ramdisks
	mdconfig -a -t malloc -s 1060227 -u 5
	status=$?
	if [ $status -ne 0 ] ; then
		die "couldn't allocate ramdisk md5"
	fi
	ramdisks="$ramdisks 5"
	ramdisk_md5=1
	if [ "$exit_at" = "ramdisk1" ] ; then
		die "die after allocating ramdisk1"
	fi

	mdconfig -a -t malloc -s 1060227 -u 6
	status=$?
	if [ $status -ne 0 ] ; then
		die "couldn't allocate ramdisk md6"
	fi
	ramdisks="$ramdisks 6"
	ramdisk_md6=1
	if [ "$exit_at" = "ramdisk2" ] ; then
		die "die after allocating ramdisk2"
	fi
}

root_devices=""
var_devices=""

pivot_imdd() {
	local root_removed
	local var_removed
	root_removed=""
	var_removed=""
	# figure out which devices imdd is using for / and /var
	root_devices="$(isi_imdd_config -g imdd0 | grep imdd0)"
	root_devices="${root_devices#imdd0}"
	echo "root devices: $root_devices"
	var_devices="$(isi_imdd_config -g imdd1 | grep imdd1)"
	var_devices="${var_devices#imdd1}"
	echo "var devices: $var_devices"

	# copy partition table, disklabel to md5
	echo "building md5a"
	dd if=/dev/imdd0 of=/dev/md5 bs=64k > /dev/null
	status=$?
	if [ $status -ne 0 ] ; then
		die "couldn't build md5a"
	fi
	sync

	echo "adding md5 to imdd0"
	isi_imdd_config -a imdd0 /dev/md5a
	status=$?
	if [ $status -ne 0 ] ; then
		die "couldn't add ramdisk to imdd0"
	fi

	if ! isi_imdd_config -g imdd0 | grep -q md5 ; then
		die "failed to add md5 to imdd0"
	fi
	md5a_added=1
	if [ "$exit_at" = "pivot1" ] ; then
		die "die after adding md5 to imdd0"
	fi

	# copy partition table, disklabel to md6
	echo "building md6a"
	dd if=/dev/imdd1 of=/dev/md6 bs=64k > /dev/null
	status=$?
	if [ $status -ne 0 ] ; then
		die "couldn't build md6a"
	fi
	sync

	echo "adding md6 to imdd1"
	isi_imdd_config -a imdd1 /dev/md6a
	status=$?
	if [ $status -ne 0 ] ; then
		die "couldn't add ramdisk to imdd1"
	fi
	md6a_added=1
	if [ "$exit_at" = "pivot2" ] ; then
		die "die after adding md6 to imdd1"
	fi

	echo "removing disks from imdd"
	isi_imdd_config -g

	for device in $root_devices; do
		echo "removing device $device from imdd0"
		isi_imdd_config -x imdd0 $device
		status=$?
		if [ $status -eq 0 ]; then
			root_removed="$device $root_removed"
		fi
	done
	root_devices_removed=$root_removed
	if [ "$exit_at" = "pivot3" ] ; then
		die "die after removing root devices"
	fi

	for device in $var_devices; do
		echo "removing device $device from imdd1"
		isi_imdd_config -x imdd1 $device
		status=$?
		if [ $status -eq 0 ] ; then
			var_removed="$device $var_removed"
		fi
	done
	var_devices_removed=$var_removed
	if [ "$exit_at" = "pivot4" ] ; then
		die "die after removing var devices"
	fi

	echo "done removing disks from imdd"
	isi_imdd_config -g
}

unconfigure_other_imdd() {
	# unconfigure the imdds that we won't be needing
	for imdd in imdd2 imdd3 imdd4 imdd5 ; do
		echo "unconfiguring imdd $imdd"
		isi_imdd_config -u $imdd
		status=$?
		if [ $status -ne 0 ] ; then
			isi_imdd_config -g
			die "failed to unconfigure imdd $imdd, $status"
		fi
	done

	# some kernels won't let us unconfigure the dongle imdd7.  if
	# it fails, just remove all but one disk from imdd7.  nothing
	# should be writing to it at this point anyway, so it
	# shouldn't hurt us to leave it around.
	if ! isi_imdd_config -u imdd7 2>/dev/null ; then
		imdd7_devices=$(isi_imdd_config -g imdd7 | grep imdd7)
		imdd7_devices="${imdd7_devices#imdd7}"
		imdd7_devices="${imdd7_devices#* }"
		for dev in $imdd7_devices ; do
			echo "removing $dev from imdd7"
			isi_imdd_config -x imdd7 $dev
		done
	fi

	echo "done unconfiguring imdd"
	isi_imdd_config -g
}

restore_imdd() {
	# rebuild imdd0 so that anything we changed gets saved back to disk
	for device in $root_devices; do
		echo "re-adding device $device to imdd0"
		isi_imdd_config -a imdd0 $device
		status=$?
		if [ $status -ne 0 ] ; then
			echo "failed to add $device to imdd0"
		fi
	done

	# rebuild imdd1 so that anything we changed gets saved back to disk
	for device in $var_devices; do
		echo "re-adding device $device to imdd1"
		isi_imdd_config -a imdd1 $device
		status=$?
		if [ $status -ne 0 ] ; then
			echo "failed to add $device to imdd1"
		fi
	done
}

mostly_shutdown() {
	local mountpoints

	echo "shutting down almost everything"

	if [ ! -L "$shutdown_flag" ] ; then
		/sbin/shutdown --no_kload --unreal now 
		status=$?
		if [ $status -ne 0 ] ; then
			die "couldn't shutdown"
		fi
		ln -s console "$shutdown_flag"
	fi

	# we're now committed to going down.
	# this flag affects the behavior of some /etc/rc.d/ scripts.
	touch "$update_flag"
	sync

	# unmount any NFS filesystems before the external interfaces get disabled.
	umount -a -t nfs >/dev/null 2>&1

	if [ ! -L "$rc_shutdown_flag" ] ; then
		echo "running rc.shutdown"
		# rc.shutdown uses stty to change the status char.  we need to
		# avoid that because we're redirecting output to tee, so stty
		# doesn't have a tty to work with.  Edit out the stty line and
		# run the modifed script.
		sed -e 's|\(^stty .*\)|# \1|' /etc/rc.shutdown > /tmp/rc.shutdown
		# isi_lcd_d stop seems to hang if stdin is a tty, so use /dev/null
		/bin/sh /tmp/rc.shutdown < /dev/null
		rm -f -- /tmp/rc.shutdown
		ln -s console "$rc_shutdown_flag"
	fi

	if [ -n "$debug" ] ; then
		echo process list
		ps -axlww
		echo process list complete
	fi

	# kill processes that don't have a shutdown script.
	echo "killing off processes"
	kill_processes

	sync 
	sleep 2

	if [ -n "$debug" ] ; then
		echo process list 2
		ps -axlww
		echo process list 2 complete
	fi

	sync 
	sleep 2

	# umount most filesystems
	echo "unmounting most filesystems"
	mountpoints="$(mount | awk '{print $3}' | sort -r)"
	for mountpoint in $mountpoints ; do
		case "$mountpoint" in
		/) ;;
		/var) ;;
		/dev) ;;
		/etc/memdisk) ;;
		/keystore) ;;
		/tmp/ufp) ;;
		*) 
			echo "unmounting $mountpoint";
			umount_pass="umount: $mountpoint: not a file system root directory"
			umount $mountpoint
			status=$?
			if [ $status -ne 0 ] ; then
				kill_processes
				echo "Trying to unmount $mountpoint again"
				error=$(umount $mountpoint 2>&1 >/dev/null)
				echo $error
				if [ "$error" = "$umount_pass" ] ; then
					continue
				else
					if [ "$error" = "" ] ; then
						continue
					else
						die "failed to unmount $mountpoint"
					fi
				fi
			fi

			;;
		esac
	done
}

kill_processes() {
	local nonsystem
	local nonsystem_prev
	local force_kill
	local kill_count
	local loop_count

	while [ $pid -ne 1 ] ; do
		pid=$(ps -o ppid= -p $pid)
		protected="$pid $protected"
	done

	# Protect rpcbind 
	rpcbind=$(pgrep rpcbind)
	if [ "$rpcbind" != "" ] ; then
		protected="$protected $rpcbind"
	fi

	# Protect isi_km_d
	isi_km_d=$(pgrep isi_km_d)
	if [ "$isi_km_d" != "" ] ; then
		protected="$protected $isi_km_d"
	fi

	# Protect syslogd
	syslog=$(pgrep syslogd)
	if [ "$syslog" != "" ] ; then
		protected="$protected $syslog"
	fi

	# Protect isi_disk_firmware_reboot child processes
	child=$(pgrep -P $$)
	if [ "$child" != "" ] ; then
		protected="$protected $child"
	fi

	# Make sure to protect any other instances
	# of isi_disk_firmware_reboot if they exist
	inst=$(ps -o pid= -o command= | grep -E 'isi_disk_firmware_reboot' | awk '{print $1}')
	if [ "$inst" != "" ] ; then
		protected="$protected $inst"
	fi

	# swapper
	protected="0 $protected"

	loop_count=0
	force_kill=0
	while [ $loop_count -ne 15 ] ; do
		nonsystem=$(ps agx -o pid= -o flags= -o pgid= -o command= | \
                        perl -ane 'my $f = (hex($F[1]) & 0x204)||0; my $g = getpgrp(); \
			print "$F[0] " if (($f == 0) && ($g != $F[2]))' | sort)
		loop_count=`expr $loop_count + 1`
		kill_count=0
		for n in $nonsystem ; do
			kill=1
			for p in $protected ; do
				if [ $n = $p ] ; then
					kill=0
					break
				fi
			done
			if [ $kill -ne 0 ]; then
				[ -n "$debug" ] && echo "kill $n"
				if [ $force_kill -eq 1 ] ; then
					kill -9 "$n" > /dev/null 2>&1
				else
					kill -TERM "$n" > /dev/null 2>&1
				fi
				kill_count=`expr $kill_count + 1`
			fi
		done
		if [ "$nonsystem" = "$nonsystem_prev" ] ; then
			force_kill=1
		fi
		nonsystem_prev=$nonsystem
		if [ $kill_count -eq 0 ] ; then
			break
		fi
	done
}

grease_wd() {
	local pid
	local percent

	# disable drive stall events to keep them from 
	# spewing onto the console when unloading disk heads
	sysctl hw.disk_event.thresh.slowacc_enabled=0

	echo "starting Western Digital greasing procedure at $(date)"

	# run this in the background so that we can display
	# progress on it, but Control-C and Control-Z will
	# only affect this shell script and not isi_wd_grease.
	/usr/bin/isi_hwtools/isi_wd_grease $verbose &
	pid=$!
	sleep 3

	# display progress periodically
	while pgrep isi_wd_grease > /dev/null ; do
		# report progress to stdout
		ps -axww -o command | grep '^isi_wd_grease' | \
			sed -e 's|.*\(bay.*complete\).*|\1|' | \
			sort -n -k2
		echo ""

		# update the front panel
		percent=$(ps -axww -o command | grep '^isi_wd_grease' | \
			sed  -e 's|.* \([0-9][0-9]*\)% complete.*|\1|' | \
			sort -n | head -n 1)

		[ -n "$debug" ] && echo "updating front panel with ${percent}%"
		/usr/bin/isi_hwtools/isi_lcd_display \
			"WD firmware update preparation ${percent}% complete. DO NOT POWER OFF." \
			 2>/dev/null

		sleep 10
	done

	# make sure WD grease succeeds before doing a firmware
	# update.  We can't let the drives cycle power until
	# they've been properly prepared.
	wait $pid
	status=$?
	if [ $status -ne 0 ] ; then
		/usr/bin/isi_hwtools/isi_lcd_display \
			"WD firmware update preparation failed. DO NOT POWER OFF." \
			 2>/dev/null
		die "WD grease workaround failed with status $status at $(date)"
	else
		echo "finished Western Digital greasing procedure at $(date)"
	fi
}

update_firmware() {
	# don't kick out drives when they power cycle after the fw update
	sysctl kern.cam.lost_cache_downs_drive=0

	/usr/bin/isi_hwtools/isi_lcd_display --clear --kill 2>/dev/null

	if /usr/bin/isi_hwtools/isi_wd_grease -cq && [ -z "$no_grease" ] ; then
		grease_wd
	fi

	printf "\nstarting disk firmware updates at $(date)\n"
	/usr/bin/isi_hwtools/isi_lcd_display \
		"Disk firmware updates in progress. DO NOT POWER OFF." 2>/dev/null
	/usr/bin/isi_hwtools/isi_disk_firmware_update $update_method $verbose $debug
	status=$?
	if [ $status -ne 0 ] ; then
		/usr/bin/isi_hwtools/isi_lcd_display \
			"Disk firmware update failed." 2>/dev/null
		die "disk firmware update failed with status $status at $(date)"
	fi
	/usr/bin/isi_hwtools/isi_lcd_display \
		"Finished disk firmware updates." 2>/dev/null
	echo "finished disk firmware updates at $(date)"
}

finish_shutdown() {
	rm -f -- "$update_flag"

	# kill syslog as /etc/rc.d/isi_post_shutdown would have done.
	if [ -e /var/run/syslog.pid ]; then
		echo "killing syslog"
		kill -TERM "$(cat /var/run/syslog.pid)"
	fi

	sync

	# take down the IB interfaces like isi_post_shutdown would have done.
	echo "taking down IB interfaces"
	ifconfig ib0 down
	ifconfig ib1 down
}

update_and_reboot() {
	local bootdisk_type

	echo ""
	echo "Starting disk firmware reboot at $(date)"

	if [ -x /usr/sbin/isi_bootdisk_type ] ; then
		bootdisk_type=$(/usr/sbin/isi_bootdisk_type)
	else
		bootdisk_type="imdd"
	fi

	if [ "${bootdisk_type}" = "imdd" ]; then
		# if we booted from the backend disks, 
		# we must pivot to ramdisks
		prepare_ramdisks
		pivot_imdd
	fi

	mostly_shutdown
	if [ "$exit_at" = "mostly_shutdown" ] ; then
		die "die after mostly_shutdown"
	fi

	if [ "${bootdisk_type}" = "imdd" ]; then
		# disable crash dumps
		dumpon -v off
		# disable stack dumps
		sysctl kern.dump_all_stacks_on_panic=0
		sdumpon -v off
		if [ ! -L $unconfigure_other_imdd_flag ] ; then
			# detach all other imdds from the disks
			unconfigure_other_imdd
			ln -s console $unconfigure_other_imdd_flag
		fi
	fi

	if [ "$exit_at" = "update_firmware" ] ; then
		die "die prior to firmware update"
	fi
	update_firmware

	if [ "${bootdisk_type}" = "imdd" ]; then
		restore_imdd
	fi

	finish_shutdown

	if [ -n "$poweroff" ] ; then
		# GS wants an option to poweroff instead of reboot,
		# so that they can install new fans in some nodes.
		echo "powering off in 5 seconds"
		sleep 1
		sync
		sleep 4
		/sbin/halt_real -p
	else
		# reboot for real now
		echo "rebooting in 5 seconds"
		sleep 1
		sync
		sleep 4
		/sbin/reboot_real
	fi
}

check_ssh() {
	if [ -z "$SSH_CONNECTION" ] ; then
		return 0
	fi

	ssh_addr="$(echo $SSH_CONNECTION | awk '{print $3}')"
	if [ -n "$ssh_addr" ] ; then
		# fail unless the ssh addr is an internal address.
		# the external interfaces will be taken down.
		safe=0
		for addr in $(/usr/bin/isi_hwtools/isi_internal_addrs) ; do
			if [ "$addr" = "$ssh_addr" ] ; then
				safe=1
				break
			fi
		done
		if [ $safe -eq 0 ] ; then
			echo ""
			echo "$0 will take down the external network."
			echo "Please login from the internal network or serial console and try again."
			exit 0
		fi
	fi

	return 0
}

# we unmount a lot of filesystems, and if the parent shell's current
# working dir is on one of those filesystems, the unmount seems to
# make it lose access to everything, even /.  Since this script just
# exits when it fails, we want the parent shell to be someplace safe
# before starting an update.  Since we can't change the parent shell's
# cwd, we just complain about it and refuse to run.
check_cwd() {
	if [ "$PWD" != "/" ] && [ "$PWD" != "/root" ] ; then
		die "$0 must be run from / or /root"
	fi
}

usage() {
	cat <<EOF

Usage: isi_disk_firmware_reboot [flags]

-o = poweroff instead of reboot
-p = parallel disk firmware update
-s = sequential disk firmware update
-v = verbose output

EOF
}

# fail if not root
if [ "$(whoami)" != "root" ] ; then
	die "$0 must be run by root"
fi

# make sure we have the executables we need
[ -x /usr/bin/isi_hwtools/isi_wd_grease ] || die "isi_wd_grease is missing"
[ -x /usr/bin/isi_hwtools/isi_internal_addrs ] || die "isi_internal_addrs is missing"
[ -x /usr/bin/isi_hwtools/isi_disk_firmware_update ] || 
	die "isi_disk_firmware_update is missing"

while getopts "denopsv" option ; do
	case "$option" in
	\?) usage ; exit 1;;
	d) debug="-d" ;;
	e) exit_at=$2; shift ;;
	n) no_grease=1 ;;
	o) poweroff=1 ;;
	p) update_method="--parallel" ;;
	s) update_method="--sequential" ;;
	v) verbose="-v" ;;
	esac
	shift
done

if [ -z "$update_method" ] ; then
	# we're not doing an update, just check if one is needed
	if /usr/bin/isi_hwtools/isi_disk_firmware_update -q ; then
		echo ""
		echo "To update one disk at a time during a reboot, run:"
		printf "\t$0 -s\n"
		echo ""
		echo "To update all disks in parallel during a reboot, run:"
		printf "\t$0 -p\n"
		echo ""
	fi
elif /usr/bin/isi_hwtools/isi_disk_firmware_update ; then
	check_ssh
	check_cwd

	# ignore SIGHUP, so that we continue even if an ssh session drops
	trap '' SIGHUP
	# ignore SIGTTOU, so that we can always write to the tty
	trap '' SIGTTOU

	# go someplace safe
	cd /

	# do the update as a background job so that it's not in the
	# tty's foreground process group.  The shell needs to have job
	# control enabled so that this background job gets put in a
	# background process group, and won't be affected by any
	# signals the tty may send to the foreground process group.
	# isi_lcd_display expects stdin to be a tty for no
	# particularly good reason, so we can't redirect input from
	# /dev/null.
	update_and_reboot | tee -a $logfile 2>&1 &
	wait $!
fi
exit 0
