#!/usr/bin/env sh
# (C) Martin V\"ath <martin@mvath.de>
# SPDX-License-Identifier: GPL-2.0-only

set -u

BASENAME=${0##*/}

if command -v gettext.sh >/dev/null 2>&1
then
. gettext.sh
TEXTDOMAINDIR='/usr/share/locale'
TEXTDOMAIN='zram-init'
export TEXTDOMAINDIR TEXTDOMAIN
Echo() {
	gettext "$*"
	echo
}
else
Echo() {
	printf '%s\n' "$*"
}
gtx_varopen='${'
gtx_varclose='}'
eval_gettext() {
	gtx_a=$*
	gtx_b=
	while gtx_c=${gtx_a%%"$gtx_varopen"*}
	do	gtx_b=$gtx_b$gtx_c
		[ x"$gtx_c" = x"$gtx_a" ] && break
		gtx_a=${gtx_a#*"$gtx_varopen"}
		gtx_c=${gtx_a%%"$gtx_varclose"*}
		case $gtx_c in
		"$gtx_a"|*[!a-zA-Z0-9_]*)
			gtx_b=$gtx_b$gtx_varopen
			continue;;
		esac
		eval gtx_b=\$gtx_b\$$gtx_c
		gtx_a=${gtx_a#*"$gtx_varclose"}
	done
	printf '%s' "$gtx_b"
	unset gtx_a gtx_b gtx_c
}
fi
xEcho() {
	eval_gettext "$*"
	echo
}

if command -v logger >/dev/null 2>&1
then	EchoLog() {
	Echo "$*"
	logger "$*"
}
xEchoLog() {
	xEcho "$*"
	logger "$*"
}
else	EchoLog() {
	Echo "$*"
}
xEchoLog() {
	xEcho "$*"
}
fi

SysCtl() {
	printf '%s' "$1" >|"$2"
}

SysCtlN() {
	printf '%s\n' "$1" >|"$2"
}

Error() {
	EchoLog "${BASENAME}: " >&2
	EchoLog "$*" >&2
}

xError() {
	xEchoLog "${BASENAME}: " >&2
	xEchoLog "$*" >&2
}

xWarning() {
	xError "Warning: "
	xError "$*"
}

Fatal() {
	Error "$*"
	exit 1
}

xFatal() {
	xError "$*"
	exit 1
}

Push() {
	PushA_=`push.sh 2>/dev/null` || Fatal \
'push.sh from https://github.com/vaeth/push/ (v2.0 or newer) required'
	eval "$PushA_"
	Push "$@"
}

IsNumeric() {
	case ${1:-x} in
	*[!0123456789]*)
		return 1;;
	esac
	:
}

Usage() {
	xEcho 'Usage: ${BASENAME} [options] SIZE|write [DIR]'
	Echo \
'Prepare a zRAM device and use it as swap (resp. mount it under DIR).
SIZE is the maximal size in megabytes.
For umounting/freeing the zRAM device, use SIZE=0.
When using "write" (or anything else starting with "w") an idle writeback
is forced (only makes sense if previously initialized with "-w" or "-W LIMIT").
If DIR is - then only a filesystem is created in /dev/zram$DEV (or the device
is removed if SIZE is 0) but it is not mounted
(options "-o", "-c", "-m", and "-T" take no effect in this case, of course).
The latter can be useful e.g. for Btrfs if multiple devices should be mounted
jointly afterwards.
The following options are available.
An empty argument means the same as if the option is not specified.
-d DEV     Use zRAM device DEV. If not specified, DEV=0 is assumed.
           Make sure to use the matching value for umounting (SIZE=0)!
-D NUM     If modprobe needs to be used, require NUM devices. This is not
           recommended. Rely instead on /etc/modprobe.d/zram.conf with the line
           "options zram num_devices=NUM"
-s NUM     Use up to NUM parallel compression streams for the device
-S MAX     Use maximally MAX megabytes of uncompressed memory for the device
-b DEV     Use DEV as backing device
-I         If combined with "-b DEV", store incompressible pages to
           backing device
-w         If combined with "-b DEV", enable idle writeback to backing device
-W LIMIT   As "-w" but additionally set writeback_limit to LIMIT * 4kB.
-a ALGO    Set compression algorithm to ALGO (e.g. zstd, lz4, or lzo)
-c OWNER   If specified, chown to OWNER (resp. OWNER:GROUP) after mounting.
           If not specified, the default owner (root:root) is not changed.
-m MODE    If specified, chmod DIR to MODE after mounting.
-o OPTS    If specified, mount DIR with "-o OPTS".
-p PRIO    Use priority PRIO for the swap device.
           If not specified, PRIO=16383 is assumed.
           Use PRIO=- if you want to keep the default priority (-1).
-t TYPE    Use a filesystem of type TYPE if DIR is specified.
           TYPE can be either ext2, ext4, btrfs or xfs
           If not specified, TYPE=ext4 is assumed.
-B BSIZE   Override default blocksize (ext2, ext4) (BSIZE=1024|2048|4096)
-i IRATIO  If specified override default bytes/inodes in fs (ext2, ext4)
-N INODES  If specified override inode count (ext2, ext4)
-L LABEL   Specify the label LABEL for the new filesystem
-U UUID    Specify the uuid UUID for the new filesystem
-T         If specified, do not use the discard (TRIM) feature of ext4/swap.
           Use this option with linux-3.14 or earlier or when you want a slight
           speed increase at the cost of possibly wasting a lot of memory
-l         Do not use zramctl even if available
-k         Do no attempt to umount/free a previously used zRAM under this
           device

If you have push.sh in $PATH, you can also use accumulatively:
-K ARG     Pass ARG to the respective mkswap or mkfs.* call
-M ARG     Pass ARG to the respective swapon/mount call
-2 ARG     Pass ARG to the tune2fs call (ignored unless for ext2 or ext4)
-Z ARG     Pass ARG to the zramctl call

Call with LANG=C to disable translations'
	exit ${1:-1}
}

PATH=${PATH-}${PATH:+:}/usr/sbin:/sbin
dev=
num=
streams=
algo=
prio=
fstype=
opts=
mode=
owgr=
blocksize=
iratio=
inodes=
label=
uuid=
keep=false
zramctl=:
discard=:
mkfs_opt=
mount_opt=
tune2fs_opt=
zramctl_opt=
mem_limit=
backing_dev=
writeback=false
writeback_limit=
incompressible=false
OPTIND=1
while getopts 's:S:a:b:Ic:d:D:m:o:p:t:B:i:N:L:U:wW:TlkK:M:2:Z:hH' opt
do	case $opt in
	s)	streams=$OPTARG;;
	S)	mem_limit=$OPTARG;;
	a)	algo=$OPTARG;;
	b)	backing_dev=$OPTARG;;
	I)	incompressible=:;;
	c)	owgr=$OPTARG;;
	d)	dev=$OPTARG;;
	D)	num=$OPTARG;;
	m)	mode=$OPTARG;;
	o)	opts=$OPTARG;;
	p)	prio=$OPTARG;;
	t)	fstype=$OPTARG;;
	B)	blocksize=$OPTARG;;
	i)	iratio=$OPTARG;;
	N)	inodes=$OPTARG;;
	L)	label=$OPTARG;;
	U)	uuid=$OPTARG;;
	w)	writeback=:;;
	W)	writeback=:
		writeback_limit=$OPTARG;;
	T)	discard=false;;
	l)	zramctl=false;;
	k)	keep=:;;
	K)	Push mkfs_opt "$OPTARG";;
	M)	Push mount_opt "$OPTARG";;
	2)	Push tune2fs_opt "$OPTARG";;
	Z)	Push zramctl_opt "$OPTARG";;
	'?')	exit 1;;
	*)	Usage 0;;
	esac
done
shift $(( $OPTIND - 1 ))
[ $# -eq 1 ] || [ $# -eq 2 ] || Usage

if $zramctl
then	command -v zramctl >/dev/null 2>&1 || zramctl=false
fi

# Defaults for empty/unspecified options:
: "${dev:=0}" "${prio:=16383}" "${fstype:=ext4}"

IsNumeric "$dev" || xFatal 'device ${dev} is not numeric'

devnode=/dev/zram$dev
zclass=/sys/class/zram-control
block=/sys/block/zram$dev

size=${1:-0}

case ${size} in
w*)
	SysCtlN idle "$block/writeback" || \
		xFatal 'failed to idle writeback zram${dev}'
	exit;;
esac

IsNumeric "${size}" || xFatal 'size ${size} is not numeric'
IsNumeric "${mem_limit:-0}" || xFatal 'mem_limit ${mem_limit} is not numeric'
IsNumeric "${writeback_limit:-0}" || \
	xFatal 'writeback_limit ${writeback_limit} is not numeric'

umount=:
if [ ${size} -ne 0 ]
then	umount=false
	size=$1 && [ $size -gt 0 ] || \
		Fatal 'failed to calculate size'
fi

swap=:
if [ $# -eq 2 ]
then	swap=false
	mount=:
	dir=${2%/}
	case $dir in
	'-')
		mount=false;;
	'/'*)
		:;;
	*)
		dir_or_empty=${dir:-(empty)}
		if $umount
		then
			xWarning '${dir_or_empty} is not a directory path. Umounting anyway'
		else
			xFatal '${dir_or_empty} is not a directory path'
		fi;;
	esac
fi

HotAdd() {
	while :
	do	curradd=
		exec 3<&0 && {
			exec <"$zclass/hot_add" && \
				read curradd || curradd=
			exec 0<&3
			exec 3<&-
		}
		IsNumeric "$curradd" || {
			xWarning 'hot_add failed for ${devnode}'
			break
		}
		[ $curradd -lt $dev ] || break
	done
	sleep=0
	while ! test -b "$devnode"
	do	sleep 1
		sleep=$(( $sleep + 1 ))
		[ $sleep -lt 5 ] || xFatal 'failed to create ${devnode}'
	done
}

if ! test -b "$devnode"
then	$umount && exit
	IsNumeric "${num:-0}" || xFatal 'device count ${num} is not numeric'
	modprobe zram ${num:+num_devices=$num} >/dev/null 2>&1
	sleep=0
	while ! test -b "$devnode" && ! test -d "$zclass"
	do	sleep 1
		sleep=$(( $sleep + 1 ))
		[ $sleep -lt 5 ] || \
			xFatal 'cannot create ${devnode}: ${zclass} missing'
	done
	test -b "$devnode" || {
		HotAdd
		keep=:
	}
fi
if ! $zramctl || [ -n "$mem_limit" ] || [ -n "${backing_dev:++}" ]
then	test -d "$block" || xFatal 'cannot find ${block}'
fi

status=0
if ! $keep
then	if $swap
	then	swapoff -- "$devnode" >/dev/null 2>&1 || status=$?
	elif $mount
	then	umount -- "$devnode" >/dev/null 2>&1 || status=$?
	fi
	if [ $status -eq 0 ]
	then	if $zramctl
		then	zramctl --reset -- "$devnode" || status=$?
		else	SysCtl 1 "$block/reset" || status=$?
			if $umount && test -w "$zclass/hot_remove"
			then	SysCtl "$devnode" "$zclass/hot_remove" \
					|| status=$?
			fi
		fi
		[ $status -eq 0 ] || \
			xWarning 'failed to reset zram${dev}'
	fi
fi
$umount && exit $status
test -b "$devnode" || HotAdd


if $zramctl && [ -z "${backing_dev:++}" ]
then	eval "set -- a $zramctl_opt"
	shift
	zramctl --size "${size}MiB" \
		${streams:+--streams "$streams"} \
		${algo:+--algorithm "$algo"} ${1+"$@"} \
		-- "$devnode" || xFatal 'zramctl zram${dev} failed'
else	[ -z "$streams" ] || SysCtl "$streams" "$block/max_comp_streams" \
		|| xWarning \
		'failed to set zram${dev} max_comp_streams to ${streams}'
	[ -z "$algo" ] || SysCtl "$algo" "$block/comp_algorithm" || \
		xWarning 'failed to set zram${dev} comp_algorithm to ${algo}'
	[ -z "${backing_dev:++}" ] || \
		SysCtlN "$backing_dev" "$block/backing_dev" || \
		xWarning 'failed to set zram${dev} backing_dev'
	SysCtl "${size}M" "$block/disksize" || xFatal 'cannot set zram${dev} size'
	! $incompressible || SysCtlN huge "$block/writeback" || \
		xWarning 'failed to set zram${dev} incompressible writeback'
	! $writeback || InitWriteback() {
	SysCtlN all "$block/idle" || {
		xWarning 'failed to set zram${dev} idle writeback'
		return
	}
	[ -n "${writeback_limit:++}" ] || return 0
	SysctlN 1 "$block/writeback_limit_enable" || {
		xWarning 'failed to enable writeback_limit for zram${dev}'
		return
	}
	SysctlN "$writeback_limit" "$block/writeback_limit" || xWarning \
	'failed to set writeback_limit ${writeback_limit} for zram${dev}'
}
	! $writeback || InitWriteback
fi
[ -z "$mem_limit" ] || SysCtl "${mem_limit}M" "$block/mem_limit" || \
	xWarning 'failed to set zram${dev} mem_limit'

eval "set -- a $mkfs_opt"
shift

if $swap
then	mkswap ${label:+-L "$label"} ${uuid:+-U "$uuid"} ${1+"$@"} \
		-- "$devnode" >/dev/null || xFatal 'mkswap ${devnode} failed'
	discardopt=-d
	$discard || discardopt=
	[ x"$prio" != x'-' ] || prio=
	swapon ${prio:+-p "$prio"} $discardopt -- "$devnode" \
		|| xFatal 'swapon ${devnode} failed'
	exit 0
fi

fsopts='-O^huge_file,sparse_super'
case $fstype in
ext2)
	mkfs.ext2 -m0 "$fsopts" ${blocksize:+-b "$blocksize"} \
		${iratio:+-i "$iratio"} ${inodes:+-N "$inodes"} \
		${label:+-L "$label"} ${uuid:+-U "$uuid"} ${1+"$@"} \
		-- "$devnode" >/dev/null \
		|| xFatal 'mkfs.ext2 ${devnode} failed'
	eval "set -- a $tune2fs_opt"
	shift
	tune2fs -c0 -i0 -m0 ${1+"$@"} -- "$devnode" >/dev/null
	$mount || exit 0
	eval "set -- a $mount_opt"
	shift
	mount -t ext2 ${opts:+-o "$opts"} ${1+"$@"} -- "$devnode" "$dir" \
		|| xFatal 'mount ${devnode} failed';;
ext4)
	fsopts=$fsopts',extent,^uninit_bg,dir_nlink,extra_isize,^has_journal'
	mkfs.ext4 -m0 "$fsopts" ${blocksize:+-b "$blocksize"} \
		${iratio:+-i "$iratio"} ${inodes:+-N "$inodes"} \
		${label:+-L "$label"} ${uuid:+-U "$uuid"} ${1+"$@"} \
		-- "$devnode" >/dev/null \
		|| xFatal 'mkfs.ext4 ${devnode} failed'
	eval "set -- a $tune2fs_opt"
	shift
	tune2fs -c0 -i0 -m0 ${1+"$@"} -- "$devnode" >/dev/null
	$mount || exit 0
	! $discard || opts=$opts${opts:+,}'discard'
	eval "set -- a $mount_opt"
	shift
	mount -t ext4 ${opts:+-o "$opts"} ${1+"$@"} -- "$devnode" "$dir" \
		|| xFatal 'mount ${devnode} failed';;
btrfs)
	mkfs.btrfs -d single -m single \
		${label:+-L "$label"} ${uuid:+-U "$uuid"} ${1+"$@"} \
		-- "$devnode" >/dev/null \
		|| xFatal 'mkfs.btrfs ${devnode} failed'
	$mount || exit 0
	eval "set -- a $mount_opt"
	shift
	mount -t btrfs ${opts:+-o "$opts"} ${1+"$@"} -- "$devnode" "$dir" \
		|| xFatal 'mount ${devnode} failed';;
xfs)
	mkfs.xfs ${label:+-L "$label"} ${uuid:+-m uuid="$uuid"} \
		${1+"$@"} -- "$devnode" >/dev/null \
		|| xFatal 'mkfs.xfs ${devnode} failed'
	$mount || exit 0
	eval "set -- a $mount_opt"
	shift
	mount -t xfs ${opts:+-o "$opts"} ${1+"$@"} -- "$devnode" "$dir" \
		|| xFatal 'mount ${devnode} failed';;
*)
	xFatal 'unsupported filesystem ${fstype}';;
esac

# Change owner/mode if requested
[ -z "$owgr" ] || chown -- "$owgr" "$dir"
[ -z "$mode" ] || chmod -- "$mode" "$dir"
