#! /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

MOUNTER_IMAGE=${1:-}
MOUNTER_PATH=/home/kubernetes/flexvolume_mounter
VOLUME_PLUGIN_DIR=/home/kubernetes/flexvolume

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 [ -z ${MOUNTER_IMAGE} ]; then
  echo "ERROR: No Container Registry mounter image is specified." 
  echo
  usage
fi

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

# Waits for kubelet to restart for 1 minute. 
kubelet_wait() {
  timeout=60
  kubelet_readonly_port=10255
  until [[ $timeout -eq 0 ]]; do
    printf "." 
    if [[ $( curl -s http://localhost:${kubelet_readonly_port}/healthz ) == "ok" ]]; then
      return 0
    fi
    sleep 1
    timeout=$(( timeout-1 ))
  done

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

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 [ -n ${IMAGE_URL:-} ]; then
    docker rmi -f ${IMAGE_URL} &> /dev/null || /bin/true
  fi
  if [ -n ${MOUNTER_DEFAULT_NAME:-} ]; then
    docker rm -f ${MOUNTER_DEFAULT_NAME} &> /dev/null || /bin/true
  fi
}

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. 
generate_chroot_wrapper() {
  if [ !  -d ${MOUNTER_PATH}/flexvolume ]; then
    echo "Failed to set up FlexVolume driver: cannot find directory '/flexvolume' in the mount utility image." 
    exit 1
  fi
  
  for driver_dir in ${MOUNTER_PATH}/flexvolume/*; do
    if [ -d "$driver_dir" ]; then

      filecount=$(ls -1 $driver_dir | wc -l)
      if [ $filecount -gt 1 ]; then
        echo "ERROR: Expected 1 file in the FlexVolume directory but found $filecount." 
        exit 1
      fi
      
      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. 
      driver_path=flexvolume/$( basename $driver_dir )/${driver_file}
      wrapper_dir=${VOLUME_PLUGIN_DIR}/$( basename $driver_dir )
      wrapper_path=${wrapper_dir}/${driver_file}

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

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

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

METADATA=http://metadata.google.internal/computeMetadata/v1
SVC_ACCT_ENDPOINT=$METADATA/instance/service-accounts/default
ACCESS_TOKEN=$(curl -s -H 'Metadata-Flavor: Google' $SVC_ACCT_ENDPOINT/token | cut -d'"' -f 4)
PROJECT_ID=$(curl -s -H 'Metadata-Flavor: Google' $METADATA/project/project-id)
IMAGE_URL=gcr.io/${PROJECT_ID}/${MOUNTER_IMAGE}
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 [ $?  -eq 0 ]; then
  echo
  echo "FlexVolume is ready." 
else
  echo "ERROR: Timed out after 1 minute waiting for kubelet restart." 
fi
