From d1a35a63c2a2ed32cb80fbebf098aa67160ea276 Mon Sep 17 00:00:00 2001
From: Chad Smith <chad.smith@canonical.com>
Date: Fri, 9 Feb 2018 09:53:42 -0700
Subject: EC2: Fix get_instance_id called against cached datasource pickle.

Fix an issue in EC2 where the datasource.identity had not been
initialized before being used when restoring datasource from pickle.
This is exposed in upgrade and reboot path.

LP: #1748354
---
 cloudinit/sources/DataSourceEc2.py          |  6 +++
 tests/unittests/test_datasource/test_ec2.py | 59 ++++++++++++++++++++++++++++-
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/cloudinit/sources/DataSourceEc2.py b/cloudinit/sources/DataSourceEc2.py
index e14553b3..21e9ef84 100644
--- a/cloudinit/sources/DataSourceEc2.py
+++ b/cloudinit/sources/DataSourceEc2.py
@@ -147,6 +147,12 @@ class DataSourceEc2(sources.DataSource):
     def get_instance_id(self):
         if self.cloud_platform == Platforms.AWS:
             # Prefer the ID from the instance identity document, but fall back
+            if not getattr(self, 'identity', None):
+                # If re-using cached datasource, it's get_data run didn't
+                # setup self.identity. So we need to do that now.
+                api_version = self.get_metadata_api_version()
+                self.identity = ec2.get_instance_identity(
+                    api_version, self.metadata_address).get('document', {})
             return self.identity.get(
                 'instanceId', self.metadata['instance-id'])
         else:
diff --git a/tests/unittests/test_datasource/test_ec2.py b/tests/unittests/test_datasource/test_ec2.py
index 0f7267bb..dff8b1ec 100644
--- a/tests/unittests/test_datasource/test_ec2.py
+++ b/tests/unittests/test_datasource/test_ec2.py
@@ -2,6 +2,7 @@
 
 import copy
 import httpretty
+import json
 import mock
 
 from cloudinit import helpers
@@ -9,6 +10,29 @@ from cloudinit.sources import DataSourceEc2 as ec2
 from cloudinit.tests import helpers as test_helpers
 
 
+DYNAMIC_METADATA = {
+    "instance-identity": {
+        "document": json.dumps({
+            "devpayProductCodes": None,
+            "marketplaceProductCodes": ["1abc2defghijklm3nopqrs4tu"],
+            "availabilityZone": "us-west-2b",
+            "privateIp": "10.158.112.84",
+            "version": "2017-09-30",
+            "instanceId": "my-identity-id",
+            "billingProducts": None,
+            "instanceType": "t2.micro",
+            "accountId": "123456789012",
+            "imageId": "ami-5fb8c835",
+            "pendingTime": "2016-11-19T16:32:11Z",
+            "architecture": "x86_64",
+            "kernelId": None,
+            "ramdiskId": None,
+            "region": "us-west-2"
+        })
+    }
+}
+
+
 # collected from api version 2016-09-02/ with
 # python3 -c 'import json
 # from cloudinit.ec2_utils import get_instance_metadata as gm
@@ -85,7 +109,7 @@ DEFAULT_METADATA = {
     "public-keys": {"brickies": ["ssh-rsa AAAAB3Nz....w== brickies"]},
     "reservation-id": "r-01efbc9996bac1bd6",
     "security-groups": "my-wide-open",
-    "services": {"domain": "amazonaws.com", "partition": "aws"}
+    "services": {"domain": "amazonaws.com", "partition": "aws"},
 }
 
 
@@ -340,6 +364,39 @@ class TestEc2(test_helpers.HttprettyTestCase):
              'type': 'physical'}]}
         self.assertEqual(expected, ds.network_config)
 
+    @httpretty.activate
+    def test_ec2_get_instance_id_refreshes_identity_on_upgrade(self):
+        """get_instance-id gets DataSourceEc2Local.identity if not present.
+
+        This handles an upgrade case where the old pickled datasource didn't
+        set up self.identity, but 'systemctl cloud-init init' runs
+        get_instance_id which traces on missing self.identity. lp:1748354.
+        """
+        self.datasource = ec2.DataSourceEc2Local
+        ds = self._setup_ds(
+            platform_data=self.valid_platform_data,
+            sys_cfg={'datasource': {'Ec2': {'strict_id': False}}},
+            md=DEFAULT_METADATA)
+        # Mock 404s on all versions except latest
+        all_versions = (
+            [ds.min_metadata_version] + ds.extended_metadata_versions)
+        for ver in all_versions[:-1]:
+            register_mock_metaserver(
+                'http://169.254.169.254/{0}/meta-data/instance-id'.format(ver),
+                None)
+        ds.metadata_address = 'http://169.254.169.254'
+        register_mock_metaserver(
+            '{0}/{1}/meta-data/'.format(ds.metadata_address, all_versions[-1]),
+            DEFAULT_METADATA)
+        # Register dynamic/instance-identity document which we now read.
+        register_mock_metaserver(
+            '{0}/{1}/dynamic/'.format(ds.metadata_address, all_versions[-1]),
+            DYNAMIC_METADATA)
+        ds._cloud_platform = ec2.Platforms.AWS
+        # Setup cached metadata on the Datasource
+        ds.metadata = DEFAULT_METADATA
+        self.assertEqual('my-identity-id', ds.get_instance_id())
+
     @httpretty.activate
     @mock.patch('cloudinit.net.dhcp.maybe_perform_dhcp_discovery')
     def test_valid_platform_with_strict_true(self, m_dhcp):
-- 
cgit v1.2.3