From 0d494fb49ef4858e099243edbb5940816623dd81 Mon Sep 17 00:00:00 2001
From: Rick Wright <rickw@google.com>
Date: Thu, 29 Jan 2015 15:58:46 -0800
Subject: [PATCH] Disable resize on disks larger than 2TB

Change-Id: I9764fe2a06cf47e8c0daf38df41c288c280bd6f7
---
 bootstrapvz/providers/gce/__init__.py         |  5 ++
 .../local-premount/gce-disable-growroot       | 52 +++++++++++++++++++
 bootstrapvz/providers/gce/tasks/__init__.py   |  3 ++
 bootstrapvz/providers/gce/tasks/initd.py      | 24 ++++++++-
 4 files changed, 83 insertions(+), 1 deletion(-)
 create mode 100755 bootstrapvz/providers/gce/assets/initramfs-tools/scripts/local-premount/gce-disable-growroot

diff --git a/bootstrapvz/providers/gce/__init__.py b/bootstrapvz/providers/gce/__init__.py
index 88477a8..1b49d1c 100644
--- a/bootstrapvz/providers/gce/__init__.py
+++ b/bootstrapvz/providers/gce/__init__.py
@@ -9,6 +9,7 @@ import tasks.packages
 from bootstrapvz.common.tasks import apt
 from bootstrapvz.common.tasks import loopback
 from bootstrapvz.common.tasks import initd
+from bootstrapvz.common.tasks import kernel
 from bootstrapvz.common.tasks import ssh
 from bootstrapvz.common.tasks import volume
 
@@ -55,6 +56,10 @@ def resolve_tasks(taskset, manifest):
 	if manifest.volume['partitions']['type'] != 'none':
 		taskset.add(initd.AdjustExpandRootScript)
 
+	if manifest.volume['partitions']['type'] != 'mbr':
+		taskset.update([tasks.initd.AddGrowRootDisable,
+		                kernel.UpdateInitramfs])
+
 	if 'gcs_destination' in manifest.image:
 		taskset.add(tasks.image.UploadImage)
 		if 'gce_project' in manifest.image:
diff --git a/bootstrapvz/providers/gce/assets/initramfs-tools/scripts/local-premount/gce-disable-growroot b/bootstrapvz/providers/gce/assets/initramfs-tools/scripts/local-premount/gce-disable-growroot
new file mode 100755
index 0000000..0e7d7d9
--- /dev/null
+++ b/bootstrapvz/providers/gce/assets/initramfs-tools/scripts/local-premount/gce-disable-growroot
@@ -0,0 +1,52 @@
+# Selectively disable growroot               -*- shell-script -*-
+set -e
+
+message() { echo "DISABLE-GROWROOT:" "$@" ; }
+error_exit() { message "$@"; exit 1; }
+
+. /scripts/functions
+
+# initramfs-tools exports the following variables, used below:
+# $ROOT - Generally "/dev/disk/by-uuid/<id>" which is a link to /dev/sda1
+# $ROOTFLAGS - Generally empty
+# $ROOTFSTYPE - Generally empty
+# $rootmnt - Set to "/root"
+
+# Follow link to get real root location
+if [ ! -L "${ROOT}" ]; then
+  real_root=${ROOT}
+else
+  real_root=$(readlink -f "${ROOT}")
+fi
+
+# Remove partition number to get disk
+disk=$(echo ${real_root} | sed 's/[0-9]*$//')
+
+# Determine number of 512-byte sectors in 2TB
+two_tb=$((2*(1024**4)))
+max_sectors=$((${two_tb}/512))
+
+# Determine number of sectors on disk
+geometry=$(sfdisk ${disk} --show-pt-geometry)
+cyl=$(echo $geometry | cut -d " " -f 2)
+heads=$(echo $geometry | cut -d " " -f 4)
+secs=$(echo $geometry | cut -d " " -f 6)
+sectors=$((${cyl}*${heads}*${secs}))
+
+# If disk is >2TB, disable growroot
+if [ "$sectors" -gt "$max_sectors" ]; then
+  message "Disk size >2TB - Not expanding root partition"
+  # Temporarily mount filesystem
+  if [ -z "${ROOTFSTYPE}" ]; then
+    fstype=$(get_fstype "${ROOT}")
+  else
+    fstype=${ROOTFSTYPE}
+  fi
+  mount -w ${fstype:+-t ${fstype} }${ROOTFLAGS} ${ROOT} ${rootmnt} ||
+      error_exit "failed to mount ${ROOT}."
+  # Disable growroot
+  touch "${rootmnt}/etc/growroot-disabled"
+  # Unmount filesystem
+  umount "${rootmnt}" || error_exit "failed to umount ${rootmnt}";
+fi
+
diff --git a/bootstrapvz/providers/gce/tasks/__init__.py b/bootstrapvz/providers/gce/tasks/__init__.py
index e69de29..494ec79 100644
--- a/bootstrapvz/providers/gce/tasks/__init__.py
+++ b/bootstrapvz/providers/gce/tasks/__init__.py
@@ -0,0 +1,3 @@
+import os.path
+
+assets = os.path.normpath(os.path.join(os.path.dirname(__file__), '../assets'))
diff --git a/bootstrapvz/providers/gce/tasks/initd.py b/bootstrapvz/providers/gce/tasks/initd.py
index c826bfe..f2b20bc 100644
--- a/bootstrapvz/providers/gce/tasks/initd.py
+++ b/bootstrapvz/providers/gce/tasks/initd.py
@@ -1,6 +1,9 @@
 from bootstrapvz.base import Task
 from bootstrapvz.common import phases
 from bootstrapvz.common.tasks import initd
+from bootstrapvz.common.tasks import kernel
+from . import assets
+import os.path
 
 
 class AdjustExpandRootDev(Task):
@@ -10,7 +13,26 @@ class AdjustExpandRootDev(Task):
 
 	@classmethod
 	def run(cls, info):
-		import os.path
 		from bootstrapvz.common.tools import sed_i
 		script = os.path.join(info.root, 'etc/init.d/expand-root')
 		sed_i(script, '/dev/xvda', '/dev/sda')
+
+
+class AddGrowRootDisable(Task):
+	description = 'Add script to selectively disable growroot'
+	phase = phases.system_modification
+	successors = [kernel.UpdateInitramfs]
+
+	@classmethod
+	def run(cls, info):
+		import stat
+		rwxr_xr_x = (stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR |
+		             stat.S_IRGRP                | stat.S_IXGRP |
+		             stat.S_IROTH                | stat.S_IXOTH)
+		from shutil import copy
+		script_src = os.path.join(assets,
+                            'initramfs-tools/scripts/local-premount/gce-disable-growroot')
+		script_dst = os.path.join(info.root,
+                            'etc/initramfs-tools/scripts/local-premount/gce-disable-growroot')
+		copy(script_src, script_dst)
+		os.chmod(script_dst, rwxr_xr_x)