#!/bin/bash

VE_PARALLEL=0

VZCTL=/usr/sbin/vzctl
VZLIST=/usr/sbin/vzlist
VZTACTL=/usr/sbin/vztactl

VZCONF=/etc/vz/vz.conf
NETCLASSES=/etc/vz/conf/networks_classes
SYSFS_PERM=/sys/fs/cgroup/ve/ve.default_sysfs_permissions
SUSPEND_CONF_DIR=/etc/vz/vzreboot

get_parallel()
{
	[ -n "${VE_PARALLEL}" -a "${VE_PARALLEL}" != "0" ] && return
	VE_PARALLEL=`awk '
BEGIN { num=0; }
$1 == "processor" { num++; }
END { print num * 4; }' /proc/cpuinfo`
}

setup_ve_sysfs()
{
	[ ! -f "$SYSFS_PERM"  ] && return 0

PERM="/ rx
block -
class rx
class/block rx
class/net rx
class/tty rx
class/mem rx
class/misc rx
devices rx
devices/system rx
devices/system/cpu rx
devices/system/cpu/kernel_max r
devices/system/cpu/possible r
devices/virtual rx
devices/virtual/misc rx
devices/virtual/misc/tun rx
devices/virtual/misc/tun/dev r
devices/virtual/misc/tun/uevent rw
devices/virtual/net rx
devices/virtual/mem rx
devices/virtual/tty rx
devices/virtual/tty0 rx
devices/virtual/tty1 rx
devices/virtual/tty2 rx
devices/virtual/tty3 rx
devices/virtual/tty4 rx
devices/virtual/tty5 rx
devices/virtual/tty6 rx
devices/virtual/tty7 rx
devices/virtual/tty8 rx
devices/virtual/tty9 rx
devices/virtual/tty10 rx
devices/virtual/tty11 rx
devices/virtual/tty12 rx
dev rx
dev/block rx
dev/char rx
fs rx
fs/cgroup rw
kernel rx
kernel/fscaps r
kernel/profiling r
kernel/uevent_seqnum r
module/ rx
devices/virtual/mem rx
devices/virtual/mem/full rx
devices/virtual/mem/full/dev r
devices/virtual/mem/full/uevent rw
devices/virtual/mem/null rx
devices/virtual/mem/null/dev r
devices/virtual/mem/null/uevent rw
devices/virtual/mem/random rx
devices/virtual/mem/random/dev r
devices/virtual/mem/random/uevent rw
devices/virtual/mem/urandom rx
devices/virtual/mem/urandom/dev r
devices/virtual/mem/urandom/uevent rw
devices/virtual/mem/zero rx
devices/virtual/mem/zero/dev r
devices/virtual/mem/zero/uevent rw
devices/virtual/tty/ptmx rx
devices/virtual/tty/ptmx/dev r
devices/virtual/tty/ptmx/uevent rw"

	local ifs=$IFS
	IFS="
"
	for i in $PERM; do
        	echo "$i" > $SYSFS_PERM 2>/dev/null
	done
	IFS=$ifs

	echo "devices/system/cpu/ rx" > $SYSFS_PERM
	for i in `ls -1d /sys/devices/system/cpu/cpu[0-9]*`; do
		echo "$i" | sed s,/sys/,,g | sed 's/$/ rx/' > $SYSFS_PERM
	done
}

get_physical_interfaces()
{
	local link target
	for link in $(find /sys/class/net -type l); do
		target=$(readlink -f "${link}")
		# Only physical
		[[ "${target}" == "/sys/devices/virtual"* ]] && continue
		basename "${link}"
	done
}

setup_venet()
{
	[ ! -e /sys/class/net/venet0 ] && ip l a name venet0 type venet 
	ip l s dev venet0 up
	ip -6 a a fe80::1/128 dev venet0 2>/dev/null
	sysctl -w net.ipv4.conf.venet0.send_redirects=0
	sysctl -w net.ipv4.conf.venet0.rp_filter=0
}

setup_net()
{
	local hwcsum_set=1 tso_set=1 val
	local rx=on tx=on sg=on tso=on
	local ethtool_args= iface_params

	setup_venet

	# We do not want to export the whole config file.
	eval $(. ${VZCONF} && echo USE_NET_TSO="${USE_NET_TSO}" \
		USE_NET_HWCSUM="${USE_NET_HWCSUM}")

	if [ "${USE_NET_HWCSUM}" = "yes" ]; then
		ethtool_args="${ethtool_args} rx on tx on sg on"
	elif [ "${USE_NET_HWCSUM}" = "no" ]; then
		ethtool_args="${ethtool_args} rx off tx off sg off"
	else
		hwcsum_set=0
	fi

	if [ "${USE_NET_TSO}" = "yes" ]; then
		ethtool_args="${ethtool_args} tso on"
	elif [ "${USE_NET_TSO}" = "no" ]; then
		ethtool_args="${ethtool_args} tso off"
	else
		tso_set=0
	fi

	if [[ ${hwcsum_set} -eq 0 || ${tso_set} -eq 0 ]]; then
		# Get minimal common subset of capabilities.
		for iface in $(get_physical_interfaces); do
			iface_params=$(ethtool -k "${iface}" | awk '
				/^rx-checksumming:/{print "rx", $2};
				/^tx-checksumming:/{print "tx", $2};
				/^scatter-gather:/{print "sg", $2};
				/^tcp-segmentation-offload:/{print "tso", $2}')
			for param in rx tx sg tso; do
				val=$(echo "${iface_params}" | awk "/${param}/{print \$2}")
				[ "${val}" = "off" ] && eval ${param}=off
			done
		done
		[ ${hwcsum_set} -eq 0 ] && ethtool_args="${ethtool_args} rx ${rx} tx ${tx} sg ${sg}"
		[ ${tso_set} -eq 0 ] && ethtool_args="${ethtool_args} tso ${tso}"
	fi
	[ -n "${ethtool_args}" ] && ethtool -K venet0 ${ethtool_args} &>/dev/null
	return 0
}

setup_tc()
{
	$VZTACTL class_load -s $NETCLASSES 2>&1
}

mount_vz()
{
	local m ret

	[ -f /etc/fstab ] || return
	eval $(. ${VZCONF} && echo VZMOUNTS="${VZMOUNTS}")

	for m in ${VZMOUNTS}; do
		if ! grep -wq -E "${m}.*noauto" /etc/fstab; then
			continue
		fi
		if grep -wq "${m}" /proc/mounts; then
			continue
		fi
		echo -n $"Checking " OpenVZ " partition $m file system..."
		fsck -V -T -a ${m} 2>&1
		ret=$?
		# The exit code returned by fsck is the sum of:
		# 0    - No errors
		# 1    - File system errors corrected
		if [ $ret -ne 0 -a $ret -ne 1 ]; then
			return 1
		fi

		echo "Mount " OpenVZ " partition: $m"
		mount ${m}
		if [ $? -ne 0 ]; then
			return 1
		fi
	done
}

cleanup_stale_reg()
{
	[ -x /usr/sbin/shaman ] && $VZCTL cleanup-stail-registration || :
}

wait_pids()
{
	for pid in $@; do
		wait $pid
	done
}

start_env()
{
	local pid pids iter

	target=`readlink /etc/systemd/system/default.target`
	if [ -n "$target" ]; then
		target=`basename $target`
	else
		target=multi-user.target
	fi
	if /usr/bin/systemctl -q --no-legend list-jobs | grep -qw "$target"; then 
		return
	fi

	get_parallel

	iter=0
	for id in `prlctl list -Ha -o id --vmtype ct`; do
		if ! prlctl list -i $id | grep "^Autostart: on"; then
			continue
		fi
		prlctl start $id >/dev/null 2>&1 &
		pid=$!
		eval VE_${pid}=${id}
		pids="$pids $pid"
		let iter++
		if [ $iter -ge $VE_PARALLEL ]; then
			wait_pids $pids
			pids=
			iter=0
		fi
	done

	wait_pids $pids
	return 0
}

start()
{
	if ! grep -q '^[[:space:]]*SERVER_UUID=' $VZCONF; then
		[ -f /etc/vz/dispatcher.xml ] && uuid=`xmllint --xpath '/ParallelsPreferences/ServerIdentification/ServerUuid/text()' /etc/vz/dispatcher.xml | tr -d {}`
		[ -z "$uuid" ] && uuid=$(/usr/bin/uuidgen)
		[ -n "$uuid" ] && echo -e "\nSERVER_UUID=$uuid" >> $VZCONF
	fi

	setup_ve_sysfs
	setup_net
	setup_tc
	mount_vz
	cleanup_stale_reg
	start_env
}

stop()
{
	local cmd pid pids iter velist veid
	local action=suspend

	eval $(. ${VZCONF} && echo VE_STOP_MODE="${VE_STOP_MODE}" \
		VE_PARALLEL="${VE_PARALLEL}")

	[ "$VE_STOP_MODE" = "stop" ] && action=stop
	mkdir -p ${SUSPEND_CONF_DIR}
		
	get_parallel

	velist=`$VZLIST -H -o ctid`
	for veid in $velist; do
		# Run command
		cmd=$action
		if [ -z "$VE_STOP_MODE" ]; then
			a=`$VZLIST -H -o autostop $veid`
			[ "$a" = "stop" -o "$a" = "suspend" ] && cmd=$a
		fi
		$VZCTL --skiplock $cmd $veid &
		pid=$!
		pids="$pids $pid"
		let iter++
		if [ $iter -ge $VE_PARALLEL ]; then
			wait_pids $pids
			pids=
			iter=0
		fi
	done

	wait_pids $pids

	# forcible stop the rest
	velist=`$VZLIST -o ctid`
	for veid in $velist; do
		$VZCTL --skiplock stop $veid --force > /dev/null 2>&1
	done

	# Unmount
	velist=`$VZLIST -ao ctid,status | awk '$2=="mounted" {print $1}'`
	for veid in $velist; do
		$VZCTL --skiplock umount $veid > /dev/null 2>&1
	done
}

shaperrestart()
{
	setup_tc
	/usr/libexec/libvzctl/scripts/vz-setrate shaperon
	exit $?
}

case "$1" in
  mount)
	mount_vz
	;;
  start)
	start
	;;
  stop)
        stop
        ;;
  shaperrestart)
	shaperrestart
	;;
  *)
        echo "Usage: $0 start|stop"
        exit 1
esac
