#!/bin/sh

# 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

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 /var/run /var/lock /afs /mit /tmp /media /home"
addgroups="cdrom floppy audio video plugdev"

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

(
  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.
      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" ]
      lvrename "$vgname" "$uloginlvname" "$loginlvname"
    else
      # No update is in progress.  Create our own snapshot of the root.
      sync
      lvcreate --snapshot --size "$snapshotsize" --name "$loginlvname" \
        "$rootlvpath"
    fi

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

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

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

    # Prevent daemons from starting inside the chroot.
    (echo "#!/bin/sh"; echo "exit 101") > /login/usr/sbin/policy-rc.d
    chmod 755 /login/usr/sbin/policy-rc.d

    # 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 any processes using the mountpoint.
    fuser -km /login > /dev/null || true
    sleep 2

    # Clean up the bind mounts we made earlier.
    for dir in $binddirs; do
      umount "/login$dir"
    done

    # Attempt to unmount /login.
    if ! 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
          ln /etc/nologin.update /etc/nologin
        fi
        touch "$bootflag"
      else
        # We can just reboot now.
        reboot
      fi
    fi
    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.
      lvremove -f "$uloginlvpath"
    fi
    sync
    lvcreate --snapshot --size "$snapshotsize" --name "$uloginlvname" \
      "$rootlvpath"

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

  update-end)
    if [ -e "$uloginlvpath" ]; then
      # It appears our login snapshot was never used.  Clean it up.
      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.
      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."
      reboot 
    fi

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