Add support for Gen2 VM resource disks

azurelinuxagent/common/osutil/debian.py

@@ -33,10 +33,10 @@
 from azurelinuxagent.common.osutil.default import DefaultOSUtil
 
 
-class DebianOSUtil(DefaultOSUtil):
+class DebianOSBaseUtil(DefaultOSUtil):
 
     def __init__(self):
-        super(DebianOSUtil, self).__init__()
+        super(DebianOSBaseUtil, self).__init__()
         self.jit_enabled = True
 
     def restart_ssh_service(self):
@@ -61,10 +61,10 @@ def get_dhcp_lease_endpoint(self):
         return self.get_endpoint_from_leases_path('/var/lib/dhcp/dhclient.*.leases')
 
 
-class DebianOS8Util(DebianOSUtil):
+class DebianOSModernUtil(DebianOSBaseUtil):
 
     def __init__(self):
-        super(DebianOS8Util, self).__init__()
+        super(DebianOSModernUtil, self).__init__()
         self.jit_enabled = True
         self.service_name = self.get_service_name()
 

azurelinuxagent/common/osutil/default.py

@@ -97,6 +97,8 @@
 
 BASE_CGROUPS = '/sys/fs/cgroup'
 
+STORAGE_DEVICE_PATH = '/sys/bus/vmbus/devices/'
+GEN2_DEVICE_ID = 'f8b3781a-1e82-4818-a1c3-63d806ec15bb'
 
 class DefaultOSUtil(object):
     def __init__(self):
@@ -1191,6 +1193,67 @@ def get_mount_point(self, mountlist, device):
                     return tokens[2] if len(tokens) > 2 else None
         return None
 
+    @staticmethod
+    def _enumerate_device_id():
+        """
+        Enumerate all storage device IDs.
+
+        Args:
+            None
+
+        Returns:
+            Iterator[Tuple[str, str]]: VmBus and storage devices.
+        """
+
+        if os.path.exists(STORAGE_DEVICE_PATH):
+            for vmbus in os.listdir(STORAGE_DEVICE_PATH):
+                deviceid = fileutil.read_file(os.path.join(STORAGE_DEVICE_PATH, vmbus, "device_id"))
+                guid = deviceid.strip('{}\n')
+                yield vmbus, guid
+
+    @staticmethod
+    def search_for_resource_disk(gen1_device_prefix, gen2_device_id):
+        """
+        Search the filesystem for a device by ID or prefix.
+
+        Args:
+            gen1_device_prefix (str): Gen1 resource disk prefix.
+            gen2_device_id (str): Gen2 resource device ID.
+
+        Returns:
+            str: The found device.
+        """
+
+        device = None
+        # We have to try device IDs for both Gen1 and Gen2 VMs.
+        logger.info('Searching gen1 prefix {0} or gen2 {1}'.format(gen1_device_prefix, gen2_device_id))
+        try:
+            for vmbus, guid in DefaultOSUtil._enumerate_device_id():
+                if guid.startswith(gen1_device_prefix) or guid == gen2_device_id:
+                    for root, dirs, files in os.walk(STORAGE_DEVICE_PATH + vmbus):
+                        root_path_parts = root.split('/')
+                        # For Gen1 VMs we only have to check for the block dir in the
+                        # current device. But for Gen2 VMs all of the disks (sda, sdb,
+                        # sr0) are presented in this device on the same SCSI controller.
+                        # Because of that we need to also read the LUN. It will be:
+                        #   0 - OS disk
+                        #   1 - Resource disk
+                        #   2 - CDROM
+                        if root_path_parts[-1] == 'block' and (
+                                guid != gen2_device_id or
+                                root_path_parts[-2].split(':')[-1] == '1'):
+                            device = dirs[0]
+                            return device
+                        else:
+                            # older distros
+                            for d in dirs:
+                                if ':' in d and "block" == d.split(':')[0]:
+                                    device = d.split(':')[1]
+                                    return device
+        except (OSError, IOError) as exc:
+            logger.warn('Error getting device for {0} or {1}: {2}', gen1_device_prefix, gen2_device_id, ustr(exc))
+        return None
+
     def device_for_ide_port(self, port_id):
         """
         Return device name attached to ide port 'n'.
@@ -1201,27 +1264,13 @@ def device_for_ide_port(self, port_id):
         if port_id > 1:
             g0 = "00000001"
             port_id = port_id - 2
-        device = None
-        path = "/sys/bus/vmbus/devices/"
-        if os.path.exists(path):
-            try:
-                for vmbus in os.listdir(path):
-                    deviceid = fileutil.read_file(os.path.join(path, vmbus, "device_id"))
-                    guid = deviceid.lstrip('{').split('-')
-                    if guid[0] == g0 and guid[1] == "000" + ustr(port_id):
-                        for root, dirs, files in os.walk(path + vmbus):
-                            if root.endswith("/block"):
-                                device = dirs[0]
-                                break
-                            else:
-                                # older distros
-                                for d in dirs:
-                                    if ':' in d and "block" == d.split(':')[0]:
-                                        device = d.split(':')[1]
-                                        break
-                        break
-            except OSError as oe:
-                logger.warn('Could not obtain device for IDE port {0}: {1}', port_id, ustr(oe))
+
+        gen1_device_prefix = '{0}-000{1}'.format(g0, port_id)
+        device = DefaultOSUtil.search_for_resource_disk(
+            gen1_device_prefix=gen1_device_prefix,
+            gen2_device_id=GEN2_DEVICE_ID
+        )
+        logger.info('Found device: {0}'.format(device))
         return device
 
     def set_hostname_record(self, hostname):

azurelinuxagent/common/osutil/factory.py

@@ -22,7 +22,7 @@
 from .arch import ArchUtil
 from .clearlinux import ClearLinuxUtil
 from .coreos import CoreOSUtil
-from .debian import DebianOSUtil, DebianOS8Util
+from .debian import DebianOSBaseUtil, DebianOSModernUtil
 from .freebsd import FreeBSDOSUtil
 from .openbsd import OpenBSDOSUtil
 from .redhat import RedhatOSUtil, Redhat6xOSUtil
@@ -76,7 +76,7 @@ def _get_osutil(distro_name, distro_code_name, distro_version, distro_full_name)
         return AlpineOSUtil()
 
     if distro_name == "kali":
-        return DebianOSUtil()
+        return DebianOSBaseUtil()
 
     if distro_name == "coreos" or distro_code_name == "coreos":
         return CoreOSUtil()
@@ -90,10 +90,10 @@ def _get_osutil(distro_name, distro_code_name, distro_version, distro_full_name)
             return SUSEOSUtil()
 
     if distro_name == "debian":
-        if Version(distro_version) > Version("7"):
-            return DebianOS8Util()
+        if "sid" in distro_version or Version(distro_version) > Version("7"):
+            return DebianOSModernUtil()
         else:
-            return DebianOSUtil()
+            return DebianOSBaseUtil()
 
     if distro_name == "redhat" \
             or distro_name == "centos" \

tests/common/osutil/test_default.py

@@ -926,6 +926,41 @@ def mock_run_command(cmd):
 
         self.assertTrue(len(pid_list) == 0, "the return value is not an empty list: {0}".format(pid_list))
 
+    @patch('os.walk', return_value=[('host3/target3:0:1/3:0:1:0/block', ['sdb'], [])])
+    @patch('azurelinuxagent.common.utils.fileutil.read_file', return_value='{00000000-0001-8899-0000-000000000000}')
+    @patch('os.listdir', return_value=['00000000-0001-8899-0000-000000000000'])
+    @patch('os.path.exists', return_value=True)
+    def test_device_for_ide_port_gen1_success(
+            self,
+            os_path_exists,
+            os_listdir,
+            fileutil_read_file,
+            os_walk):
+        dev = osutil.DefaultOSUtil().device_for_ide_port(1)
+        self.assertEqual(dev, 'sdb', 'The returned device should be the resource disk')
+
+    @patch('os.walk', return_value=[('host0/target0:0:0/0:0:0:1/block', ['sdb'], [])])
+    @patch('azurelinuxagent.common.utils.fileutil.read_file', return_value='{f8b3781a-1e82-4818-a1c3-63d806ec15bb}')
+    @patch('os.listdir', return_value=['f8b3781a-1e82-4818-a1c3-63d806ec15bb'])
+    @patch('os.path.exists', return_value=True)
+    def test_device_for_ide_port_gen2_success(
+            self,
+            os_path_exists,
+            os_listdir,
+            fileutil_read_file,
+            os_walk):
+        dev = osutil.DefaultOSUtil().device_for_ide_port(1)
+        self.assertEqual(dev, 'sdb', 'The returned device should be the resource disk')
+
+    @patch('os.listdir', return_value=['00000000-0000-0000-0000-000000000000'])
+    @patch('os.path.exists', return_value=True)
+    def test_device_for_ide_port_none(
+            self,
+            os_path_exists,
+            os_listdir):
+        dev = osutil.DefaultOSUtil().device_for_ide_port(1)
+        self.assertIsNone(dev, 'None should be returned if no resource disk found')
+
 def osutil_get_dhcp_pid_should_return_a_list_of_pids(test_instance, osutil_instance):
     """
     This is a very basic test for osutil.get_dhcp_pid. It is simply meant to exercise the implementation of that method

tests/common/osutil/test_factory.py

@@ -20,7 +20,7 @@
 from azurelinuxagent.common.osutil.arch import ArchUtil
 from azurelinuxagent.common.osutil.clearlinux import ClearLinuxUtil
 from azurelinuxagent.common.osutil.coreos import CoreOSUtil
-from azurelinuxagent.common.osutil.debian import DebianOSUtil, DebianOS8Util
+from azurelinuxagent.common.osutil.debian import DebianOSBaseUtil, DebianOSModernUtil
 from azurelinuxagent.common.osutil.freebsd import FreeBSDOSUtil
 from azurelinuxagent.common.osutil.openbsd import OpenBSDOSUtil
 from azurelinuxagent.common.osutil.redhat import RedhatOSUtil, Redhat6xOSUtil
@@ -125,7 +125,7 @@ def test_get_osutil_it_should_return_kali(self):
                           distro_code_name="",
                           distro_version="",
                           distro_full_name="")
-        self.assertTrue(type(ret) == DebianOSUtil)
+        self.assertTrue(type(ret) == DebianOSBaseUtil)
         self.assertEquals(ret.get_service_name(), "waagent")
 
     def test_get_osutil_it_should_return_coreos(self):
@@ -163,14 +163,14 @@ def test_get_osutil_it_should_return_debian(self):
                           distro_code_name="",
                           distro_full_name="",
                           distro_version="7")
-        self.assertTrue(type(ret) == DebianOSUtil)
+        self.assertTrue(type(ret) == DebianOSBaseUtil)
         self.assertEquals(ret.get_service_name(), "waagent")
 
         ret = _get_osutil(distro_name="debian",
                           distro_code_name="",
                           distro_full_name="",
                           distro_version="8")
-        self.assertTrue(type(ret) == DebianOS8Util)
+        self.assertTrue(type(ret) == DebianOSModernUtil)
         self.assertEquals(ret.get_service_name(), "walinuxagent")
 
     def test_get_osutil_it_should_return_redhat(self):