#!/bin/bash

# Usage: block-snapshot add|remove img-file cow-file
#
# This creates dm-snapshot device on given arguments

dir=$(dirname "$0")
if [ "$1" = "prepare" ] || [ "$1" = "cleanup" ]; then
  . "$dir/xen-hotplug-common.sh"
  command=$1
else
  . "$dir/block-common.sh"
fi

shopt -s nullglob

if [ -n "$XENBUS_PATH" ]; then
  HOTPLUG_STORE="/var/run/xen-hotplug/${XENBUS_PATH//\//-}"
fi

get_dev() {
  dev=$1

  if [ -L "$dev" ]; then
    dev=$(readlink -f "$dev") || fatal "$dev link does not exist."
  fi

  if [ -f "$dev" ]; then
    file=$dev

    loopdev=$(losetup -j $file | head -1 | cut -d : -f 1)
    if [ -n "$loopdev" ]; then
      # found existing loop to this file
      echo $loopdev
      return
    fi

    # assign new loop device
    loopdev=$(losetup -f 2>/dev/null || find_free_loopback_dev)
    if [ "$loopdev" = '' ]
    then
      release_lock "block"
      fatal 'Failed to find an unused loop device'
    fi

    do_or_die losetup "$loopdev" "$file"
    echo $loopdev
  else
    test -e "$dev" || fatal "$dev does not exist."
    test -b "$dev" || fatal "$dev is not a block device nor file."
    echo "$dev"
  fi
}

get_dm_snapshot_name() {
  local base cow cow2

  base=$1
  cow=$2
  cow2=$3

  name="snapshot-$(stat -c '%D:%i' "$base")-$(stat -c '%D:%i' "$cow")"
  if [ -n "$cow2" ]; then
    name="$name-$(stat -c '%D:%i' "$cow2")"
  fi
  echo "$name"
}

create_dm_snapshot() {
  local base_dev cow_dev base_sz base cow dm_devname

  dm_devname=$1
  base=$2
  cow=$3

  if [ ! -e /dev/mapper/$dm_devname ]; then
    # prepare new snapshot device
    base_dev=$(get_dev $base)
    cow_dev=$(get_dev $cow)
    base_sz=$(blockdev --getsz $base_dev)
    do_or_die dmsetup create $dm_devname --table "0 $base_sz snapshot $base_dev $cow_dev P 256"
  fi

}

create_dm_snapshot_origin() {
  local base_dev base_sz dm_devname base

  dm_devname=$1
  base=$2

  if [ ! -e /dev/mapper/$dm_devname ]; then
    # prepare new snapshot-origin device
    base_dev=$(get_dev $base)
    base_sz=$(blockdev --getsz $base_dev)
    do_or_die dmsetup create $dm_devname --table "0 $base_sz snapshot-origin $base_dev"
  fi
}

t=$(basename $0)
t=${t#block-}

case "$command" in
  add)
    case $t in
      snapshot|origin)
        p=$(xenstore_read_default "$XENBUS_PATH/params" 'MISSING')
        if [ "$p" == "MISSING" ]; then
          fatal "Missing device parameters ($t $XENBUS_PATH/params)"
        fi
        echo $p > "$HOTPLUG_STORE-params"
        echo $t > "$HOTPLUG_STORE-type"
        base=${p%%:*}
        cow=${p#*:}
        cow2=${p##*:}
        if [ "$cow" != "$cow2" ]; then
          cow=${cow%:*}
        else
          cow2=""
        fi

        if [ -L "$base" ]; then
          base=$(readlink -f "$base") || fatal "$base link does not exist."
        fi

        if [ -L "$cow" ]; then
          cow=$(readlink -f "$cow") || fatal "$cow link does not exist."
        fi

        if [ -L "$cow2" ]; then
          cow2=$(readlink -f "$cow2") || fatal "$cow2 link does not exist."
        fi

        # first ensure that snapshot device exists (to write somewhere changes from snapshot-origin)
        dm_devname=$(get_dm_snapshot_name "$base" "$cow")

        claim_lock "block"

        # prepare snapshot device
        create_dm_snapshot $dm_devname "$base" "$cow"

        if [ -n "$cow2" ]; then
          dm_devname_full=$(get_dm_snapshot_name "$base" "$cow" "$cow2")
          create_dm_snapshot "$dm_devname_full" "/dev/mapper/$dm_devname" "$cow2"
          dm_devname="$dm_devname_full"
        fi

        if [ "$t" == "snapshot" ]; then
          #that's all for snapshot, store name of prepared device
          xenstore_write "$XENBUS_PATH/node" "/dev/mapper/$dm_devname"
          echo "/dev/mapper/$dm_devname" > "$HOTPLUG_STORE-node"
          write_dev /dev/mapper/$dm_devname
        elif [ "$t" == "origin" ]; then
          # for origin - prepare snapshot-origin device and store its name
          dm_devname=origin-$(stat -c '%D:%i' "$base")
          create_dm_snapshot_origin $dm_devname "$base"
          xenstore_write "$XENBUS_PATH/node" "/dev/mapper/$dm_devname"
          echo "/dev/mapper/$dm_devname" > "$HOTPLUG_STORE-node"
          write_dev /dev/mapper/$dm_devname
        fi
        # Save domain name for template commit on device remove
        domid=$(xenstore_read "$XENBUS_PATH/frontend-id")
        # Store name of target domain in case of stubdom
        domid=$(xenstore_read_default "/local/domain/$domid/target" "$domid")
        domain=$(xl domname $domid)
        echo $domain > "$HOTPLUG_STORE-domain"

        release_lock "block"
        exit 0
        ;;
    esac
    ;;
  prepare)
    t=$2
    case $t in
      snapshot|origin)
        p=$3
        base=${p%%:*}
        cow=${p#*:}
        cow2=${p##*:}
        if [ "$cow" != "$cow2" ]; then
          cow=${cow%:*}
        else
          cow2=""
        fi

        if [ -L "$base" ]; then
          base=$(readlink -f "$base") || fatal "$base link does not exist."
        fi

        if [ -L "$cow" ]; then
          cow=$(readlink -f "$cow") || fatal "$cow link does not exist."
        fi

        if [ -L "$cow2" ]; then
          cow2=$(readlink -f "$cow2") || fatal "$cow2 link does not exist."
        fi

        # first ensure that snapshot device exists (to write somewhere changes from snapshot-origin)
        dm_devname=$(get_dm_snapshot_name "$base" "$cow")

        claim_lock "block"

        # prepare snapshot device
        create_dm_snapshot $dm_devname "$base" "$cow"

        if [ -n "$cow2" ]; then
          dm_devname_full=$(get_dm_snapshot_name "$base" "$cow" "$cow2")
          create_dm_snapshot "$dm_devname_full" "/dev/mapper/$dm_devname" "$cow2"
          dm_devname="$dm_devname_full"
        fi

        if [ "$t" == "snapshot" ]; then
          #that's all for snapshot, store name of prepared device
          echo "/dev/mapper/$dm_devname"
        elif [ "$t" == "origin" ]; then
          # for origin - prepare snapshot-origin device and store its name
          dm_devname=origin-$(stat -c '%D:%i' "$base")
          create_dm_snapshot_origin $dm_devname "$base"
          echo "/dev/mapper/$dm_devname"
        fi

        release_lock "block"
        exit 0
        ;;
    esac
    ;;
  remove|cleanup)
    if [ "$command" = "cleanup" ]; then
      t=$2
    else
      t=$(cat $HOTPLUG_STORE-type 2>/dev/null || echo 'MISSING')
    fi
    case "$t" in
      snapshot|origin)
        if [ "$command" = "cleanup" ]; then
          node=$3
        else
          node=$(cat "$HOTPLUG_STORE-node" 2> /dev/null)
        fi

        if [ -z "$node" ]; then
          #fatal "No device node to remove"
          #Most likely already removed
          exit 0
        fi

        if [ ! -e "$node" ]; then
          fatal "Device $node does not exists"
        fi

        claim_lock "block"

        use_count=$(dmsetup info $node 2>/dev/null|grep Open|awk '{print $3}')

        if [ -z "$use_count" ]; then
            log info "Device $node already removed"
            release_lock "block"
            exit 0
        fi

        # do not remove snapshot if snapshot origin is still present
        if [ "${node/snapshot/}" != "$node" -a -e "/dev/mapper/origin-$(echo $node|cut -d- -f2)" ]; then
          use_count=1
        fi

        if [ "$use_count" -gt 0 ]; then
          log info "Device $node still in use - not removing"
          release_lock "block"
          exit 0
        fi

        # get list of used (loop) devices
        deps="$(dmsetup deps $node -o blkdevname | cut -d: -f2 | sed -e 's#(\([a-z0-9-]\+\))#/dev/\1#g')"

        # if this is origin
        if [ "${node/origin/}" != "$node" ]; then
          # remove unused snapshots
          for snap in /dev/mapper/snapshot-$(echo $node|cut -d- -f2)-*; do
            use_count=$(dmsetup info $snap|grep Open|awk '{print $3}')
            if [ "$use_count" -eq 0 ]; then
              # unused snapshot - remove it
              deps="$deps $(dmsetup deps $snap -o blkdevname | cut -d: -f2 |\
                   sed -e 's#(\([a-z0-9-]\+\))#/dev/\1#g')"
              log debug "Removing $snap"
              dmsetup remove $snap
            fi
          done
          if [ "$command" = "remove" ]; then
            # Commit template changes
            domain=$(cat "$HOTPLUG_STORE-domain")
            if [ "$domain" ]; then
              # Dont stop on errors
              if [ -r /var/lib/qubes/qubes-test.xml -a \
                  "${domain#test-}" != "$domain" ]; then
                export QUBES_XML_PATH=/var/lib/qubes/qubes-test.xml
              fi
              /usr/bin/qvm-template-commit --offline-mode "$domain" || true
            fi
          fi
        fi

        if [ -e $node ]; then
          log debug "Removing $node"
          dmsetup remove $node
        fi

        # try to free unused devices
        for dev in $deps; do
          if [ -b "$dev" ]; then
            log debug "Removing $dev"
            case $dev in
              /dev/loop*)
                losetup -d $dev 2> /dev/null || true
                ;;
              /dev/dm-*)
                dmsetup remove $dev 2> /dev/null || true
                ;;
            esac
          fi
        done

        if [ -n "$HOTPLUG_STORE" ]; then
          rm $HOTPLUG_STORE-*
        fi
        release_lock "block"
        
        exit 0
        ;;
    esac
    ;;
esac

# vim:sw=2:et: