Fix +tests for GCE datasource not handling per-instance SSH keys

3 files changed, 96 insertions, 46 deletions

diff --git a/ChangeLog b/ChangeLog
index c5ad7c60..bab40552 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -34,6 +34,7 @@
    [Steve Langasek]
  - Add functionality to fixate the uid of a newly added user.
  - Don't overwrite the hostname if the user has changed it after we set it.
+ - GCE datasource does not handle instance ssh keys (LP: 1403617)
 0.7.6:
  - open 0.7.6
  - Enable vendordata on CloudSigma datasource (LP: #1303986)
diff --git a/cloudinit/sources/DataSourceGCE.py b/cloudinit/sources/DataSourceGCE.py
index 608c07f1..f4ed915d 100644
--- a/cloudinit/sources/DataSourceGCE.py
+++ b/cloudinit/sources/DataSourceGCE.py
@@ -30,6 +30,31 @@ BUILTIN_DS_CONFIG = {
 REQUIRED_FIELDS = ('instance-id', 'availability-zone', 'local-hostname')
 
 
+class GoogleMetadataFetcher(object):
+    headers = {'X-Google-Metadata-Request': True}
+
+    def __init__(self, metadata_address):
+        self.metadata_address = metadata_address
+
+    def get_value(self, path, is_text):
+        value = None
+        try:
+            resp = url_helper.readurl(url=self.metadata_address + path,
+                                      headers=self.headers)
+        except url_helper.UrlError as exc:
+            msg = "url %s raised exception %s"
+            LOG.debug(msg, path, exc)
+        else:
+            if resp.code == 200:
+                if is_text:
+                    value = util.decode_binary(resp.contents)
+                else:
+                    value = resp.contents
+            else:
+                LOG.debug("url %s returned code %s", path, resp.code)
+        return value
+
+
 class DataSourceGCE(sources.DataSource):
     def __init__(self, sys_cfg, distro, paths):
         sources.DataSource.__init__(self, sys_cfg, distro, paths)
@@ -50,17 +75,15 @@ class DataSourceGCE(sources.DataSource):
             return public_key
 
     def get_data(self):
-        # GCE metadata server requires a custom header since v1
-        headers = {'X-Google-Metadata-Request': True}
-
         # url_map: (our-key, path, required, is_text)
         url_map = [
-            ('instance-id', 'instance/id', True, True),
-            ('availability-zone', 'instance/zone', True, True),
-            ('local-hostname', 'instance/hostname', True, True),
-            ('public-keys', 'project/attributes/sshKeys', False, True),
-            ('user-data', 'instance/attributes/user-data', False, False),
-            ('user-data-encoding', 'instance/attributes/user-data-encoding',
+            ('instance-id', ('instance/id',), True, True),
+            ('availability-zone', ('instance/zone',), True, True),
+            ('local-hostname', ('instance/hostname',), True, True),
+            ('public-keys', ('project/attributes/sshKeys',
+                             'instance/attributes/sshKeys'), False, True),
+            ('user-data', ('instance/attributes/user-data',), False, False),
+            ('user-data-encoding', ('instance/attributes/user-data-encoding',),
              False, True),
         ]
 
@@ -69,40 +92,25 @@ class DataSourceGCE(sources.DataSource):
             LOG.debug("%s is not resolvable", self.metadata_address)
             return False
 
+        metadata_fetcher = GoogleMetadataFetcher(self.metadata_address)
         # iterate over url_map keys to get metadata items
-        found = False
-        for (mkey, path, required, is_text) in url_map:
-            try:
-                resp = url_helper.readurl(url=self.metadata_address + path,
-                                          headers=headers)
-                if resp.code == 200:
-                    found = True
-                    if is_text:
-                        self.metadata[mkey] = util.decode_binary(resp.contents)
-                    else:
-                        self.metadata[mkey] = resp.contents
+        running_on_gce = False
+        for (mkey, paths, required, is_text) in url_map:
+            value = None
+            for path in paths:
+                new_value = metadata_fetcher.get_value(path, is_text)
+                if new_value is not None:
+                    value = new_value
+            if value:
+                running_on_gce = True
+            if required and value is None:
+                msg = "required key %s returned nothing. not GCE"
+                if not running_on_gce:
+                    LOG.debug(msg, mkey)
                 else:
-                    if required:
-                        msg = "required url %s returned code %s. not GCE"
-                        if not found:
-                            LOG.debug(msg, path, resp.code)
-                        else:
-                            LOG.warn(msg, path, resp.code)
-                        return False
-                    else:
-                        self.metadata[mkey] = None
-            except url_helper.UrlError as e:
-                if required:
-                    msg = "required url %s raised exception %s. not GCE"
-                    if not found:
-                        LOG.debug(msg, path, e)
-                    else:
-                        LOG.warn(msg, path, e)
-                    return False
-                msg = "Failed to get %s metadata item: %s."
-                LOG.debug(msg, path, e)
-
-                self.metadata[mkey] = None
+                    LOG.warn(msg, mkey)
+                return False
+            self.metadata[mkey] = value
 
         if self.metadata['public-keys']:
             lines = self.metadata['public-keys'].splitlines()
@@ -116,7 +124,7 @@ class DataSourceGCE(sources.DataSource):
             else:
                 LOG.warn('unknown user-data-encoding: %s, ignoring', encoding)
 
-        return found
+        return running_on_gce
 
     @property
     def launch_index(self):
diff --git a/tests/unittests/test_datasource/test_gce.py b/tests/unittests/test_datasource/test_gce.py
index 4280abc4..1fb100f7 100644
--- a/tests/unittests/test_datasource/test_gce.py
+++ b/tests/unittests/test_datasource/test_gce.py
@@ -113,10 +113,6 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase):
         self.assertEqual(GCE_META.get('instance/attributes/user-data'),
                          self.ds.get_userdata_raw())
 
-        # we expect a list of public ssh keys with user names stripped
-        self.assertEqual(['ssh-rsa AA2..+aRD0fyVw== root@server'],
-                         self.ds.get_public_ssh_keys())
-
     # test partial metadata (missing user-data in particular)
     @httpretty.activate
     def test_metadata_partial(self):
@@ -141,3 +137,48 @@ class TestDataSourceGCE(test_helpers.HttprettyTestCase):
         decoded = b64decode(
             GCE_META_ENCODING.get('instance/attributes/user-data'))
         self.assertEqual(decoded, self.ds.get_userdata_raw())
+
+    @httpretty.activate
+    def test_missing_required_keys_return_false(self):
+        for required_key in ['instance/id', 'instance/zone',
+                             'instance/hostname']:
+            meta = GCE_META_PARTIAL.copy()
+            del meta[required_key]
+            httpretty.register_uri(httpretty.GET, MD_URL_RE,
+                                   body=_new_request_callback(meta))
+            self.assertEqual(False, self.ds.get_data())
+            httpretty.reset()
+
+    @httpretty.activate
+    def test_project_level_ssh_keys_are_used(self):
+        httpretty.register_uri(httpretty.GET, MD_URL_RE,
+                               body=_new_request_callback())
+        self.ds.get_data()
+
+        # we expect a list of public ssh keys with user names stripped
+        self.assertEqual(['ssh-rsa AA2..+aRD0fyVw== root@server'],
+                         self.ds.get_public_ssh_keys())
+
+    @httpretty.activate
+    def test_instance_level_ssh_keys_are_used(self):
+        key_content = 'ssh-rsa JustAUser root@server'
+        meta = GCE_META.copy()
+        meta['instance/attributes/sshKeys'] = 'user:{0}'.format(key_content)
+
+        httpretty.register_uri(httpretty.GET, MD_URL_RE,
+                               body=_new_request_callback(meta))
+        self.ds.get_data()
+
+        self.assertIn(key_content, self.ds.get_public_ssh_keys())
+
+    @httpretty.activate
+    def test_instance_level_keys_replace_project_level_keys(self):
+        key_content = 'ssh-rsa JustAUser root@server'
+        meta = GCE_META.copy()
+        meta['instance/attributes/sshKeys'] = 'user:{0}'.format(key_content)
+
+        httpretty.register_uri(httpretty.GET, MD_URL_RE,
+                               body=_new_request_callback(meta))
+        self.ds.get_data()
+
+        self.assertEqual([key_content], self.ds.get_public_ssh_keys())