From 411578a4984c3c1a164710b9dd121572ea8a3cdd Mon Sep 17 00:00:00 2001
From: Anders Ingemann <anders@ingemann.de>
Date: Wed, 8 Apr 2015 23:15:02 +0200
Subject: [PATCH] *Always* use link_dm_node() when install grub 1.99 Grub
 install fails when in a chrooted environment, it has nothing to do with the
 volume being a loopback volume

---
 bootstrapvz/common/fs/__init__.py            |  9 +++---
 bootstrapvz/common/tasks/grub.py             | 32 ++++++--------------
 bootstrapvz/plugins/prebootstrapped/tasks.py | 11 +++----
 3 files changed, 20 insertions(+), 32 deletions(-)

diff --git a/bootstrapvz/common/fs/__init__.py b/bootstrapvz/common/fs/__init__.py
index 694846e..c393ea2 100644
--- a/bootstrapvz/common/fs/__init__.py
+++ b/bootstrapvz/common/fs/__init__.py
@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 
 
 def get_partitions():
@@ -16,7 +17,8 @@ def get_partitions():
 	return matches
 
 
-def remount(volume, fn):
+@contextmanager
+def unmounted(volume):
 	from bootstrapvz.base.fs.partitionmaps.none import NoPartitions
 
 	p_map = volume.partition_map
@@ -24,9 +26,8 @@ def remount(volume, fn):
 	p_map.root.unmount()
 	if not isinstance(p_map, NoPartitions):
 		p_map.unmap(volume)
-		result = fn()
+		yield
 		p_map.map(volume)
 	else:
-		result = fn()
+		yield
 	p_map.root.mount(destination=root_dir)
-	return result
diff --git a/bootstrapvz/common/tasks/grub.py b/bootstrapvz/common/tasks/grub.py
index b528712..9aabba3 100644
--- a/bootstrapvz/common/tasks/grub.py
+++ b/bootstrapvz/common/tasks/grub.py
@@ -39,26 +39,16 @@ class InstallGrub_1_99(Task):
 
 	@classmethod
 	def run(cls, info):
-
-		from ..fs import remount
 		p_map = info.volume.partition_map
 
-		def link_fn():
+		# GRUB screws up when installing in chrooted environments
+		# so we fake a real harddisk with dmsetup.
+		# Guide here: http://ebroder.net/2009/08/04/installing-grub-onto-a-disk-image/
+		from ..fs import unmounted
+		with unmounted(info.volume):
 			info.volume.link_dm_node()
 			if isinstance(p_map, partitionmaps.none.NoPartitions):
 				p_map.root.device_path = info.volume.device_path
-
-		def unlink_fn():
-			info.volume.unlink_dm_node()
-			if isinstance(p_map, partitionmaps.none.NoPartitions):
-				p_map.root.device_path = info.volume.device_path
-
-		# GRUB cannot deal with installing to loopback devices
-		# so we fake a real harddisk with dmsetup.
-		# Guide here: http://ebroder.net/2009/08/04/installing-grub-onto-a-disk-image/
-		from ..fs.loopbackvolume import LoopbackVolume
-		if isinstance(info.volume, LoopbackVolume):
-			remount(info.volume, link_fn)
 		try:
 			[device_path] = log_check_call(['readlink', '-f', info.volume.device_path])
 			device_map_path = os.path.join(info.root, 'boot/grub/device.map')
@@ -77,13 +67,11 @@ class InstallGrub_1_99(Task):
 			# Install grub
 			log_check_call(['chroot', info.root, 'grub-install', device_path])
 			log_check_call(['chroot', info.root, 'update-grub'])
-		except Exception:
-			if isinstance(info.volume, LoopbackVolume):
-				remount(info.volume, unlink_fn)
-			raise
-
-		if isinstance(info.volume, LoopbackVolume):
-			remount(info.volume, unlink_fn)
+		finally:
+			with unmounted(info.volume):
+				info.volume.unlink_dm_node()
+				if isinstance(p_map, partitionmaps.none.NoPartitions):
+					p_map.root.device_path = info.volume.device_path
 
 
 class InstallGrub_2(Task):
diff --git a/bootstrapvz/plugins/prebootstrapped/tasks.py b/bootstrapvz/plugins/prebootstrapped/tasks.py
index 4711e46..df09782 100644
--- a/bootstrapvz/plugins/prebootstrapped/tasks.py
+++ b/bootstrapvz/plugins/prebootstrapped/tasks.py
@@ -4,7 +4,7 @@ from bootstrapvz.common.tasks import volume
 from bootstrapvz.common.tasks import packages
 from bootstrapvz.providers.virtualbox.tasks import guest_additions
 from bootstrapvz.providers.ec2.tasks import ebs
-from bootstrapvz.common.fs import remount
+from bootstrapvz.common.fs import unmounted
 from shutil import copyfile
 import os.path
 import time
@@ -19,9 +19,9 @@ class Snapshot(Task):
 
 	@classmethod
 	def run(cls, info):
-		def mk_snapshot():
-			return info.volume.snapshot()
-		snapshot = remount(info.volume, mk_snapshot)
+		snapshot = None
+		with unmounted(info.volume):
+			snapshot = info.volume.snapshot()
 		msg = 'A snapshot of the bootstrapped volume was created. ID: ' + snapshot.id
 		log.info(msg)
 
@@ -55,9 +55,8 @@ class CopyImage(Task):
 		loopback_backup_name = 'volume-{id}.{ext}.backup'.format(id=info.run_id, ext=info.volume.extension)
 		destination = os.path.join(info.manifest.bootstrapper['workspace'], loopback_backup_name)
 
-		def mk_snapshot():
+		with unmounted(info.volume):
 			copyfile(info.volume.image_path, destination)
-		remount(info.volume, mk_snapshot)
 		msg = 'A copy of the bootstrapped volume was created. Path: ' + destination
 		log.info(msg)