#!/bin/bash

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

SNAPSHOT_CHUNKSIZE=256  # same as in file.py

dir=$(dirname "$0")
if [ "$1" = "prepare" ] || [ "$1" = "cleanup" ]; then
  # shellcheck disable=SC1090,SC1091
  . "$dir/xen-hotplug-common.sh"
  command=$1
else
  # shellcheck disable=SC1090,SC1091
  . "$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 $SNAPSHOT_CHUNKSIZE"
  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

        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" ] && [ -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
        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: