#!/bin/ash
# fatdog-merge-layers.sh - merge one tmpfs layer to another, deleting the original
# Copyright (C) James Budiono 2011, 2012, 2013, 2016, 2018
# License: GNU GPL Version 3 or later
#
# 2018-11-01: Re-factor, remove dead code, add exclusion for both removal and merge
#
# Note: supports aufs only
# Note: (obviously) assumes that both layers are r/w
# Note: Bashism needed for fast string ops. busybox ash with bashism support is enough
#
# Call: $1 merge-from-layer, $2-merge-to-layer
#

### Initial setup
# Note: DON'T use trailing slash here
TMPFS="$1"			# was /initrd/pup_rw
PUPSAVE="$2"		# was /initrd/pup_ro1

# only merge when both layers exist
! [ "$TMPFS" -a "$PUPSAVE" ] && exit 

### Configuration 
LANG=C # speed up bash
DISKFULL=/tmp/diskfull

# no-remove regex - for closed files
NOREMOVE_LIST=".wh..wh.orph|.wh..wh.plnk|.wh..wh.aufs|.wh..wh..opq|\
$TMPFS/var/run|\
$TMPFS/etc/passwd-?$|\
$TMPFS/etc/shadow-?$|\
$TMPFS/etc/group-?$|\
$TMPFS/etc/gshadow-?$|\
.sock$|\
lxqt.conf$"

# No merge regex - for rsync
NOMERGE_LIST="|\
" #/var/db/fontconfig/"
#NOMERGE_LIST="$(echo "$NOMERGE_LIST" | sed 's_|_ --exclude=_g')"
NOMERGE_LIST="${NOMERGE_LIST//|/ --exclude=}" # bashism is faster


### helpers

# print list of all closed files in TMPFS
list_closed_files() {
	# awk is lot faster than bash
	awk -v tmpfs=$TMPFS -v ignore_list="$NOREMOVE_LIST" '
BEGIN {
	# get the list of open files before processing - this replaces lsof
	OPENFILES = "find /proc -type l -wholename \"/proc/*/fd/*\" 2>/dev/null " \
	            "| xargs readlink 2>/dev/null "                               \
	            "| grep -vE \"^socket:|^pipe:|^anon_inode:\" "                \
	            "| sort -u "
	while (OPENFILES | getline) {
		openfiles[tmpfs $0]=1
	}
	close(OPENFILES)

	# now compare this with the list of files in $TMPFS
	ALLFILES = "find \"" tmpfs "\" -not -type d"
	while (ALLFILES | getline) {
		if ($0 ~ ignore_list) continue;
		if (openfiles[$0] != 1) 
			printf("%s\0",$0);
	}
	close(ALLFILES)
}'
}

# check that diropq is working (it sometimes doesn't).
# test method: create a test in pupsave and see if it's visible from above
# $1 = dir base (rootfs)
# $2 = pupsave dir base
test_diropq() {
	local WORKING=0
	local TMPFILE=$(mktemp "$2"/diropqtest.XXXXXXXXX)
	[ -e "$1/${TMPFILE##*/}" ] && WORKING=1
	rm -rf $TMPFILE
	#echo $1 $WORKING
	return $WORKING;
}


### actual snapmerge functions

# these clean stale whiteouts and make aufs self-consistent
# inconsistent aufs leads to crash
clean_whiteouts() {
	# check for new whiteouts - if yes remove files from pupsave
	echo "removing bottom-layer files which have been deleted from top-layer"
	find "$TMPFS" -mount \( -regex '.*/\.wh\.[^/]*' -type f \) | 
	grep -v -E ".wh..wh.orph|.wh..wh.plnk|.wh..wh.aufs" |
	while read -r FILE; do
		#echo $FILE					# $FILE is TMPFS_WHITEOUT
		FULLNAME="${FILE#$TMPFS}"
		#echo $FULLNAME
		BASE="${FULLNAME%/*}"
		#echo $BASE
		LEAF="${FULLNAME##*/}"
		#echo $LEAF
		#echo $BASE/$LEAF
		
		PUPSAVE_FILE="${PUPSAVE}${BASE}/${LEAF:4}"	
		#echo "Deleting $PUPSAVE_FILE"
		rm -rf "$PUPSAVE_FILE"		# delete the file/dir if it's there
		
		# if this is a dir-opaque file, we need to remove the dir in pupsave
		# but only if the dir-opaque is currently working !!
		if [ "$LEAF" = ".wh..wh..opq" ] && test_diropq "${BASE}" "${PUPSAVE}${BASE}"; then
			#echo remove "${PUPSAVE}${BASE}"
			rm -rf "${PUPSAVE}${BASE}"
		fi
	done

	# check for old whiteouts - remove them from pupsave if new files created in tmpfs
	echo "removing bottom-layer whiteouts for newly added top-layer files"
	find "$PUPSAVE" -mount \( -regex '.*/\.wh\.[^/]*' -type f \) | 
	grep -v -E ".wh..wh.orph|.wh..wh.plnk|.wh..wh.aufs|.wh..wh..opq" |
	while read -r FILE; do
		#echo $FILE					# $FILE is PUPSAVE_WHITEOUT
		FULLNAME="${FILE#$PUPSAVE}"
		#echo $FULLNAME
		BASE="${FULLNAME%/*}"
		#echo $BASE
		LEAF="${FULLNAME##*/}"
		#echo $LEAF
		#echo $BASE/$LEAF
		
		TMPFS_FILE="${TMPFS}${BASE}/${LEAF:4}" # bashism - can't find replacement with decent performance
		#echo $TMPFS_FILE

		# delete whiteout only if a new file/dir has been created in the tmpfs layer
		if [ -e "$TMPFS_FILE" -o -L "$TMPFS_FILE" ]; then
			# if TMPFS_FILE is a dir, we need to add diropq when removing its pupsave whiteout
			# this is just in case and it won't work until next reboot anyway
			[ -d "$TMPFS_FILE" ] &&	touch "$TMPFS_FILE/.wh..wh..opq"
			#echo Deleting whiteout $FILE
			rm -f "$FILE"
		fi
	done	
}

# merge using rsync
merge_layers() {
	# by now we should be consistent - so rsync everything
	echo "merging top-layer files to bottom-layer"
	rm -rf $DISKFULL
	if ! rsync $NOMERGE_LIST -aXH -c --force --no-whole-file "$TMPFS"/ "$PUPSAVE"; then
		case $(df | grep "$PUPSAVE" | awk '{print $4}') in 
			0|"") touch $DISKFULL ;;
		esac
	fi
}

# delete merged files that are closed closed files, BUT only if disk is not full
remove_merged_files() {
	if [ ! -f $DISKFULL ] > /dev/null; then
		echo "removing merged top-layer files"
		list_closed_files | xargs -0 rm -rf
	fi	
}

# force re-evalution of all the layers
force_aufs_reval() {
	echo "re-evaluating layers"
	busybox mount -i -t aufs -o remount,udba=reval aufs /
}

# warn diskfull is necessary
warn_disk_full() {
	if [ -f $DISKFULL ]; then
		[ "$DISPLAY" ] && 
		Xdialog --infobox "Your save file is full, please copy important items manually elsewhere." 0 0 10000 ||
		echo "Your save file is full, please copy important items manually elsewhere."
	fi
}


################### main ######################

renice -n -20 $$ > /dev/null # run at highest priority
clean_whiteouts
merge_layers
remove_merged_files
force_aufs_reval
warn_disk_full
