#!/bin/bash

# Copyright 2017 The Kubernetes Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Sets up FlexVolume drivers on GCE COS instances using mounting utilities packaged in a Google
# Container Registry image.
# The user-provided FlexVolume driver(s) must be under /flexvolume of the image filesystem.
# For example, the driver k8s/nfs must be located at /flexvolume/k8s~nfs/nfs .
#
# This script should be used on a clean instance, with no FlexVolume installed.
# Should not be run on instances with an existing full or partial installation.
# Upon failure, the script will clean up the partial installation automatically.
#
# Must be executed under /home/kubernetes/bin with sudo.
# Warning: kubelet will be restarted upon successful execution.

set -o errexit
set -o nounset
set -o pipefail

setvar MOUNTER_IMAGE = ${1:-}
setvar MOUNTER_PATH = '/home/kubernetes/flexvolume_mounter'
setvar VOLUME_PLUGIN_DIR = '/etc/srv/kubernetes/kubelet-plugins/volume/exec'

proc usage {
  echo "usage: $0 imagename[:tag]"
  echo "    imagename  Name of a Container Registry image. By default the latest image is used."
  echo "    :tag       Container Registry image tag."
  exit 1
}

if test -z ${MOUNTER_IMAGE} {
  echo "ERROR: No Container Registry mounter image is specified."
  echo
  usage
}

# Unmounts a mount point lazily. If a mount point does not exist, continue silently,
# and without error.
proc umount_silent {
  umount -l $1 &> /dev/null || /bin/true
}

# Waits for kubelet to restart for 1 minute.
proc kubelet_wait {
  setvar timeout = '60'
  setvar kubelet_readonly_port = '10255'
  while ! {
    printf "."
    if [[ $( curl -s http://localhost:${kubelet_readonly_port}/healthz ) == "ok" ]] {
      return 0
    }
    sleep 1
    setvar timeout = $(( timeout-1 ))
  }

  # Timed out waiting for kubelet to become healthy.
  return 1
}

proc flex_clean {
  echo
  echo "An error has occurred. Cleaning up..."
  echo

  umount_silent ${VOLUME_PLUGIN_DIR}
  rm -rf ${VOLUME_PLUGIN_DIR}
  umount_silent ${MOUNTER_PATH}/var/lib/kubelet
  umount_silent ${MOUNTER_PATH}
  rm -rf ${MOUNTER_PATH}
  
  if test -n ${IMAGE_URL:-} {
    docker rmi -f ${IMAGE_URL} &> /dev/null || /bin/true
  }
  if test -n ${MOUNTER_DEFAULT_NAME:-} {
    docker rm -f ${MOUNTER_DEFAULT_NAME} &> /dev/null || /bin/true
  }
}

trap flex_clean ERR

# Generates a bash script that wraps all calls to the actual driver inside mount utilities
# in the chroot environment. Kubelet sees this script as the FlexVolume driver.
proc generate_chroot_wrapper {
  if test ! -d ${MOUNTER_PATH}/flexvolume {
    echo "Failed to set up FlexVolume driver: cannot find directory '/flexvolume' in the mount utility image."
    exit 1
  }
  
  for driver_dir in ${MOUNTER_PATH}/flexvolume/* {
    if test -d $driver_dir {

      setvar filecount = $(ls -1 $driver_dir | wc -l)
      if test $filecount -gt 1 {
        echo "ERROR: Expected 1 file in the FlexVolume directory but found $filecount."
        exit 1
      }
      
      setvar driver_file = $( ls $driver_dir | head -n 1 )

      # driver_path points to the actual driver inside the mount utility image,
      # relative to image root.
      # wrapper_path is the wrapper script location, which is known to kubelet.
      setvar driver_path = "flexvolume/$( basename $driver_dir )/${driver_file}"
      setvar wrapper_dir = "${VOLUME_PLUGIN_DIR}/$( basename $driver_dir )"
      setvar wrapper_path = "${wrapper_dir}/${driver_file}"

      mkdir -p $wrapper_dir
      cat >$wrapper_path <<< """
#!/bin/bash
chroot ${MOUNTER_PATH} ${driver_path} "'$'@"
"""

      chmod 755 $wrapper_path
      echo "FlexVolume driver installed at ${wrapper_path}"
    }
  }
}

echo
echo "Importing mount utility image from Container Registry..."
echo

setvar METADATA = "http://metadata.google.internal/computeMetadata/v1"
setvar SVC_ACCT_ENDPOINT = "$METADATA/instance/service-accounts/default"
setvar ACCESS_TOKEN = $(curl -s -H 'Metadata-Flavor: Google' $SVC_ACCT_ENDPOINT/token | cut -d'"' -f 4)
setvar PROJECT_ID = $(curl -s -H 'Metadata-Flavor: Google' $METADATA/project/project-id)
setvar IMAGE_URL = "gcr.io/${PROJECT_ID}/${MOUNTER_IMAGE}"
setvar MOUNTER_DEFAULT_NAME = 'flexvolume_mounter'
sudo -u ${SUDO_USER} docker login -u _token -p $ACCESS_TOKEN https://gcr.io > /dev/null
sudo -u ${SUDO_USER} docker run --name=${MOUNTER_DEFAULT_NAME} ${IMAGE_URL}
docker export ${MOUNTER_DEFAULT_NAME} > /tmp/${MOUNTER_DEFAULT_NAME}.tar
docker rm ${MOUNTER_DEFAULT_NAME} > /dev/null
docker rmi ${IMAGE_URL} > /dev/null

echo
echo "Loading mount utilities onto this instance..."
echo

mkdir -p ${MOUNTER_PATH}
tar xf /tmp/${MOUNTER_DEFAULT_NAME}.tar -C ${MOUNTER_PATH}

# Bind the kubelet directory to one under flexvolume_mounter
mkdir -p ${MOUNTER_PATH}/var/lib/kubelet
mount --rbind /var/lib/kubelet/ ${MOUNTER_PATH}/var/lib/kubelet
mount --make-rshared ${MOUNTER_PATH}/var/lib/kubelet

# Remount the flexvolume_mounter environment with /dev enabled.
mount --bind ${MOUNTER_PATH} ${MOUNTER_PATH}
mount -o remount,dev,exec ${MOUNTER_PATH}

echo
echo "Setting up FlexVolume driver..."
echo

mkdir -p ${VOLUME_PLUGIN_DIR}
mount --bind ${VOLUME_PLUGIN_DIR} ${VOLUME_PLUGIN_DIR}
mount -o remount,exec ${VOLUME_PLUGIN_DIR}
generate_chroot_wrapper

echo
echo "Restarting Kubelet..."
echo

systemctl restart kubelet.service
kubelet_wait
if test $? -eq 0 {
  echo
  echo "FlexVolume is ready."
} else {
  echo "ERROR: Timed out after 1 minute waiting for kubelet restart."
}