#!/bin/bash

# Usage: athena-login-snapshot {login-start|login-end|update-start|update-end}

# This script contains the logic to manage LVM snapshots of the root
# volume.

# Requirements:

# * The root volume is a logical volume
# * The volume group of the root volume has 21GB of free space for us
#   to use.

# Login snapshots will have a 10GB copy-on-write store for
# modifications to the snapshot.  This backing store will typically
# only be used if a user installs additional packages in the login
# root.

# This script may choose to reboot the machine in order to clear
# I/O-bound process entries blocking the unmount of /login, though
# that circumstance should be fairly rare.

set -e
exec >>/var/log/athena-reactivate 2>&1

updflag=/var/run/athena-update-in-progress
bootflag=/var/run/athena-reboot-after-update
lockfile=/var/run/athena-snapshot.lock
snapshotsize=10G
event=$1
binddirs="/proc /sys /dev /dev/shm /dev/pts /var/run /var/lock /var/tmp /afs /mit /tmp /media /home"
addgroups="admin lpadmin adm fuse cdrom floppy audio video plugdev scanner dialout"
daemons="$(/usr/sbin/policy-rc.d --daemons)"

rootdev=$(awk '$2 == "/" { dev=$1 } END { print dev }' /proc/mounts)
vgname=$(lvs --noheadings -o vg_name "$rootdev" | awk '{print $1}')
rootlvname=$(lvs --noheadings -o lv_name "$rootdev" | awk '{print $1}')
rootlvpath=/dev/$vgname/$rootlvname
loginlvname=login
loginlvpath=/dev/$vgname/login
uloginlvname=login-update
uloginlvpath=/dev/$vgname/login-update

echo "-----"
echo "** Beginning Athena reactivation ($event) session at $(date)"

finish() {
    echo "** Finishing Athena reactivation ($event) at $(date)"
    echo "-----"
    echo
    exit
}
trap finish EXIT

v() {
    echo "** Running:" "$@"
    echo "** $(date)"
    "$@"
}

(
  flock -x 9
  case $event in
  login-start)
    # Procure a login snapshot and mount it.

    if [ -e "$loginlvpath" ]; then
      # A login snapshot already exists; perhaps the machine rebooted
      # during a login.  Clean it up.
      v lvremove -f "$loginlvpath"
    fi

    if [ -e "$updflag" ]; then
      # An update is in progress.  If we get here, we expect there to
      # be a login-update snapshot created before the update started.
      # (If we already used it up, /etc/nologin should prevent us from
      # getting here until the update ends.)  Rename it to login.
      [ -e "$uloginlvpath" ]
      v lvrename "$vgname" "$uloginlvname" "$loginlvname"
    else
      # No update is in progress.  Create our own snapshot of the root.
      sync
      v lvcreate --snapshot --size "$snapshotsize" --name "$loginlvname" \
        "$rootlvpath"
    fi

    # Mount the login snapshot.
    mkdir -p /login
    v mount "$loginlvpath" /login

    # Enable subtree operations on /media by making it a mount point,
    # then share it.
    v mount --bind /media /media
    v mount --make-shared /media

    # Bind-mount a bunch of stuff from the real root into the chroot.
    for dir in $binddirs; do
      v mount --bind "$dir" "/login$dir"
    done

    # Add the user to a bunch of groups in the chroot.
    for group in $addgroups; do
      v chroot /login gpasswd -a "$USER" "$group"
    done

    # There are some daemons that should be running inside the
    # chroot. For example, changes to CUPS config shouldn't persist
    # between sessions.
    #
    # Start those daemons inside the chroot
    #
    # For extra special bonus points, cupsys was renamed to cups at
    # some point, so we have to try both names and catch the error for
    # the one that doesn't exist
    for daemon in $daemons; do
      v invoke-rc.d $daemon stop || [ $? = 100 ]
      v chroot /login invoke-rc.d $daemon start || [ $? = 100 ]
    done

    v touch /login/ClusterLogin
    v touch /var/run/athena-login

    # Add an schroot.conf entry for the chroot.
    conf=/etc/schroot/schroot.conf
    sed -e '/###ATHENA-BEGIN###/,/###ATHENA-END###/d' $conf > $conf.new
    cat >> $conf.new <<EOF
###ATHENA-BEGIN###
[login]
description=Login root snapshot
location=/login
users=$USER
environment-filter=""
###ATHENA-END###
EOF
    mv $conf.new $conf
    ;;

  login-end)
    # Clean-up the temporary file to indicate the logged in state
    v rm /var/run/athena-login

    # Stop any daemons that were specifically started inside the
    # chroot
    for daemon in $daemons; do
      v chroot /login invoke-rc.d $daemon stop || [ $? = 100 ]
      v invoke-rc.d $daemon start || [ $? = 100 ]
    done

    # Clean up any remaining user processes using the bind mounts.
    # Try sending SIGTERM first to give processes a chance to exit
    # cleanly.
    if [ -n "$USER" -a "$USER" != root ]; then
      for dir in $binddirs; do
        v su -s /bin/sh "$USER" -c "fuser -kmv -TERM /login$dir" || true
      done
      sleep 3
      for dir in $binddirs; do
        v su -s /bin/sh "$USER" -c "fuser -kmv /login$dir" || true
      done
    fi

    # Clean up any processes using the chroot mountpoint.
    v fuser -kvm /login || true
    sleep 2

    # Clean up the bind mounts we made earlier.
    # If any of these fail, the umount of /login will fail below,
    # and we will reboot.
    for dir in $(echo $binddirs|tac -s\ ); do
      v umount "/login$dir" || true
    done

    # Unmount /media, which we bind-mounted to itself earlier so it
    # could be shared and then bind-mounted.
    v umount /media || true

    # Attempt to unmount /login.
    if ! v umount /login; then
      # There may be an unkillable process in I/O wait keeping the
      # mountpoint busy.  We need to reboot the machine.
      if [ -e "$updflag" ]; then
        # ... but we don't want to reboot during an update.  Schedule it
        # for the end of the update.
        if [ ! -e /etc/nologin ]; then
          echo "An update and reboot is in progress, please try again later." \
            > /etc/nologin.update
          v ln /etc/nologin.update /etc/nologin
        fi
        v touch "$bootflag"
      else
        # We can just reboot now.
        echo "** Rebooting because of umount /login failure"
        v reboot
      fi
    fi
    v lvremove -f "$loginlvpath"

    if [ -e "$updflag" -a ! -e "$uloginlvpath" ]; then
      # An update is in progress and we just used up its snapshot.  We
      # must block further logins until the update completes.
      if [ ! -e /etc/nologin ]; then
        echo "An update is in progress, please try again later." \
          > /etc/nologin.update
        ln /etc/nologin.update /etc/nologin
      fi
    fi
    ;;

  update-start)
    # Before starting the update, create a root snapshot for use by
    # the next login.  We give this snapshot a different name in case
    # there is already a login in process.
    if [ -e "$uloginlvpath" ]; then
      # It already exists; perhaps the machine rebooted during an
      # update.  Clean it up.
      v lvremove -f "$uloginlvpath"
    fi
    sync
    v lvcreate --snapshot --size "$snapshotsize" --name "$uloginlvname" \
      "$rootlvpath"

    # Touch the flag file signifying an update in progress.
    v touch "$updflag"
    ;;

  update-end)
    if [ -e "$uloginlvpath" ]; then
      # It appears our login snapshot was never used.  Clean it up.
      v lvremove -f "$uloginlvpath"
    fi

    if [ -e /etc/nologin.update ]; then
      # Our login snapshot was used and that login ended before we
      # did, causing further logins to block.  Now that the update has
      # ended, we can unblock logins.
      v rm -f /etc/nologin.update /etc/nologin
    fi

    if [ -e "$bootflag" ]; then
      # We need to reboot in order to unmount /login.
      echo "Rebooting in order to unmount /login."
      v reboot 
    fi

    # Remove the flag file signifying an update in progress.
    v rm -f "$updflag"
    ;;
  esac
) 9> $lockfile
