Bash script to attach an EBS volume to an EC2 instance after boot-time

attach_ebs.sh

#!/bin/sh
# attach_ebs.sh - Attach an EBS volume to an EC2 instance.
# Copyright (C) 2020 Peter Willis <[email protected]>
# 
# This script is designed to create and mount a single EBS volume based on its tag:Name
# in order to implement persistent storage. If there is more than one EBS volume
# with the same tag, this script will fail.
# 
# Order of operations:
#   1. Detect EBS volume based on "tag:Name" "$TAG_NAME"
#   2. Wait for it to become available and attach it
#   3. If needed, create GPT partition table
#   4. Try to mount a filesystem; otherwise create a new filesystem
#   5. Mount filesystem

set -e -u -x

[ -n "${FS_TYPE:-}" ]   || FS_TYPE="ext4"
[ -n "${FS_OPTS:-}" ]   || FS_OPTS="errors=remount-ro,nofail,noatime,nodiratime"
[ -n "${FS_LABEL:-}" ]  || FS_LABEL="data-vol"
[ -n "${MK_PARTITION:-}" ] || MK_PARTITION="0"

if [ $# -gt 0 ] ; then
    if [ "$1" = "-h" -o "$1" = "--help" ] ; then
        echo "Usage: $0 [TAG_NAME MOUNT_DEVICE MOUNT_DIR AMI_USER]"
        echo ""
        echo "You can omit the above arguments if there are environment variables of the same name."
        exit 1
    fi
    if [ $# -eq 4 ] ; then
        TAG_NAME="$1"
        MOUNT_DEVICE="$2"
        MOUNT_DIR="$3"
        AMI_USER="$4"
        shift 4
    fi
fi

_shutdown () {
    echo "$0: Error: $1" 1>&2
    echo "$0: Shutting down server in 1 minute." 1>&2
    /sbin/shutdown -h +1
    exit 1
}

# Try 30 times to get the EBS ID of the volume based on its tag:Name
_get_ebsid () {
    local region="$1" tagname="$2"
    for i in `seq 1 30` ; do
        ebsid=$(aws --region "$region" ec2 describe-volumes --filter "Name=tag:Name,Values=$tagname" --query "Volumes[0].{ID:VolumeId}" --output text)
        [ -n "$ebsid" ] && echo "$ebsid" && return
        sleep 1
    done
}

# Get the state of an EBS ID
_ebs_state () {
    local region="$1" ebsid="$2"
    aws --region "$region" ec2 describe-volumes --filter "Name=volume-id,Values=$ebsid" --query "Volumes[0].{STATE:State}" --output text
}

# Wait for an EBS volume to no longer be attached or in-use, then attach it and wait
# for it to be in-use before proceeding.
_attach_ebs () {
    local region="$1" ec2id="$2" tagname="$3" mountdev="$4" state
    local ebsid=$(_get_ebsid "$region" "$tagname")
    local tries=30 sleeptime=30 count=0
    [ -z "$ebsid" ] && _shutdown "Failed to find EBS volume by tag '$TAG_NAME' in 30 seconds; shutting down"
    while [ $count -lt $tries ] ; do
        state="$(_ebs_state "$region" "$ebsid")"
        [ "$state" = "available" ] && break
        echo "$0: Volume $ebsid is in state '$state'; waiting for it to become available ..."
        sleep $sleeptime
        count=$(($count+1))
    done
    [ $count -ge $tries ] && _shutdown "Volume ID '$ebsid' never became available; shutting down"
    if [ ! "$state" = "attached" -a ! "$state" = "in-use" ] ; then
        aws ec2 --region "$region" attach-volume --volume-id "$ebsid" --instance-id "$ec2id" --device "$mountdev"
    fi
    aws ec2 wait volume-in-use \
        --region "$region" \
        --volume-ids "$ebsid" \
        --filters "Name=attachment.instance-id,Values=$ec2id" "Name=attachment.status,Values=attached"
}

_has_fs () {
    blkid "$MOUNT_DEVICE" | grep -i "$FS_TYPE"
}
_mk_part () {
    # If an initial partition doesn't exist, create one.
    [ -n "${PARTED_SCRIPT:-}" ] || PARTED_SCRIPT="mklabel gpt mkpart primary 0% 100%"
    [ -n "${SFDISK_SCRIPT:-}" ] || SFDISK_SCRIPT="label: gpt\n;"
    if ! blkid "$MOUNT_DEVICE" ; then
        # Where supported: GPT partition table and a single large initial partition
        if command -v parted >/dev/null ; then
            parted -a opt --script "$MOUNT_DEVICE" "$PARTED_SCRIPT"
        elif command -v sfdisk >/dev/null ; then
            printf "$SFDISK_SCRIPT\n" | sfdisk "$MOUNT_DEVICE"
        fi
    fi

    # default new partition: append "1" to $MOUNT_DEVICE
    MOUNT_DEVICE="$MOUNT_DEVICE"1
    # wait for partition to show up
    for i in `seq 1 30` ; do
        [ -b "$MOUNT_DEVICE" ] && break
        sleep 1
    done
}
_mk_fs () {
    if ! _has_fs "$MOUNT_DEVICE" ; then
        # Try to just mount it, just in case filesystem detection failed
        if ! mount -t "$FS_TYPE" -o "$FS_OPTS" "$MOUNT_DEVICE" "$MOUNT_DIR" ; then
            # Oh well, we gave it a shot. Format the partition.
            # Both 'ext4' and 'xfs' support a '-L' option for partition label, so might
            # as well tack that on
            mkfs.$FS_TYPE -L "$FS_LABEL" "$MOUNT_DEVICE"
        fi
    fi
}


EC2ID=$(curl -s http://169.254.169.254/latest/meta-data/instance-id)
EC2REGION=$(curl -s 169.254.169.254/latest/meta-data/placement/availability-zone | sed 's/.$//')

if ! _attach_ebs "$EC2REGION" "$EC2ID" "$TAG_NAME" "$MOUNT_DEVICE" ; then
    _shutdown "failed to attach volume"
fi

# Prepend '/dev/' if missing
expr match "$MOUNT_DEVICE" /dev/ >/dev/null || MOUNT_DEVICE="/dev/$MOUNT_DEVICE"

if [ ! -b "$MOUNT_DEVICE" ] ; then
    echo "$0: Error: failed to find block device '$MOUNT_DEVICE'"
    exit 1
fi

[ -d "$MOUNT_DIR" ] || mkdir -p "$MOUNT_DIR"

# Create a filesystem if it doesn't exist yet (new volumes)
if ! _has_fs "$MOUNT_DEVICE" ; then

    # Set MK_PARTITION=1 to create and/or use a partition for the MOUNT_DEVICE
    if [ "${MK_PARTITION:-}" = "1" ] ; then
        _mk_part
    fi

    _mk_fs
fi

mount -t "$FS_TYPE" -o "$FS_OPTS" "$MOUNT_DEVICE" "$MOUNT_DIR"
chown "${AMI_USER}" "$MOUNT_DIR"

# Add new mount to /etc/fstab
grep -q -e "[[:space:]]$MOUNT_DIR[[:space:]]" /etc/fstab || echo "$MOUNT_DEVICE  $MOUNT_DIR  $FS_TYPE  $FS_OPTS  1  2" >> /etc/fstab