Reset host name after bounce has allowed walinuxagent to run successfully.

2 files changed, 99 insertions, 66 deletions

diff --git a/cloudinit/sources/DataSourceAzure.py b/cloudinit/sources/DataSourceAzure.py
index d4211fc4..a19d9ca2 100644
--- a/cloudinit/sources/DataSourceAzure.py
+++ b/cloudinit/sources/DataSourceAzure.py
@@ -17,6 +17,7 @@
 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import base64
+import contextlib
 import crypt
 import fnmatch
 import os
@@ -74,6 +75,28 @@ def set_hostname(hostname, hostname_command='hostname'):
     util.subp([hostname_command, hostname])
 
 
+@contextlib.contextmanager
+def temporary_hostname(temp_hostname, cfg, hostname_command='hostname'):
+    """
+    Set a temporary hostname, restoring the previous hostname on exit.
+
+    Will have the value of the previous hostname when used as a context
+    manager, or None if the hostname was not changed.
+    """
+    policy = cfg['hostname_bounce']['policy']
+    previous_hostname = get_hostname(hostname_command)
+    if (not util.is_true(cfg.get('set_hostname'))
+            or util.is_false(policy)
+            or (previous_hostname == temp_hostname and policy != 'force')):
+        yield None
+        return
+    set_hostname(temp_hostname, hostname_command)
+    try:
+        yield previous_hostname
+    finally:
+        set_hostname(previous_hostname, hostname_command)
+
+
 class DataSourceAzureNet(sources.DataSource):
     def __init__(self, sys_cfg, distro, paths):
         sources.DataSource.__init__(self, sys_cfg, distro, paths)
@@ -162,33 +185,40 @@ class DataSourceAzureNet(sources.DataSource):
         # the directory to be protected.
         write_files(ddir, files, dirmode=0o700)
 
-        # handle the hostname 'publishing'
-        try:
-            handle_set_hostname(mycfg.get('set_hostname'),
-                                self.metadata.get('local-hostname'),
-                                mycfg['hostname_bounce'])
-        except Exception as e:
-            LOG.warn("Failed publishing hostname: %s", e)
-            util.logexc(LOG, "handling set_hostname failed")
-
-        try:
-            invoke_agent(mycfg['agent_command'])
-        except util.ProcessExecutionError:
-            # claim the datasource even if the command failed
-            util.logexc(LOG, "agent command '%s' failed.",
-                        mycfg['agent_command'])
-
-        shcfgxml = os.path.join(ddir, "SharedConfig.xml")
-        wait_for = [shcfgxml]
-
-        fp_files = []
-        for pk in self.cfg.get('_pubkeys', []):
-            bname = str(pk['fingerprint'] + ".crt")
-            fp_files += [os.path.join(ddir, bname)]
+        temp_hostname = self.metadata.get('local-hostname')
+        hostname_command = mycfg['hostname_bounce']['hostname_command']
+        with temporary_hostname(temp_hostname, mycfg,
+                                hostname_command=hostname_command) \
+                as previous_hostname:
+            if (previous_hostname is not None
+                    and util.is_true(mycfg.get('set_hostname'))):
+                cfg = mycfg['hostname_bounce']
+                try:
+                    perform_hostname_bounce(hostname=temp_hostname,
+                                            cfg=cfg,
+                                            prev_hostname=previous_hostname)
+                except Exception as e:
+                    LOG.warn("Failed publishing hostname: %s", e)
+                    util.logexc(LOG, "handling set_hostname failed")
 
-        missing = util.log_time(logfunc=LOG.debug, msg="waiting for files",
-                                func=wait_for_files,
-                                args=(wait_for + fp_files,))
+            try:
+                invoke_agent(mycfg['agent_command'])
+            except util.ProcessExecutionError:
+                # claim the datasource even if the command failed
+                util.logexc(LOG, "agent command '%s' failed.",
+                            mycfg['agent_command'])
+
+            shcfgxml = os.path.join(ddir, "SharedConfig.xml")
+            wait_for = [shcfgxml]
+
+            fp_files = []
+            for pk in self.cfg.get('_pubkeys', []):
+                bname = str(pk['fingerprint'] + ".crt")
+                fp_files += [os.path.join(ddir, bname)]
+
+            missing = util.log_time(logfunc=LOG.debug, msg="waiting for files",
+                                    func=wait_for_files,
+                                    args=(wait_for + fp_files,))
         if len(missing):
             LOG.warn("Did not find files, but going on: %s", missing)
 
@@ -307,48 +337,15 @@ def support_new_ephemeral(cfg):
     return mod_list
 
 
-def handle_set_hostname(enabled, hostname, cfg):
-    if not util.is_true(enabled):
-        return
-
-    if not hostname:
-        LOG.warn("set_hostname was true but no local-hostname")
-        return
-
-    apply_hostname_bounce(hostname=hostname, policy=cfg['policy'],
-                          interface=cfg['interface'],
-                          command=cfg['command'],
-                          hostname_command=cfg['hostname_command'])
-
-
-def perform_hostname_bounce(command, env):
-    shell = not isinstance(command, (list, tuple))
-    # capture=False, see comments in bug 1202758 and bug 1206164.
-    util.log_time(logfunc=LOG.debug, msg="publishing hostname",
-                  get_uptime=True, func=util.subp,
-                  kwargs={'args': command, 'shell': shell, 'capture': False,
-                          'env': env})
-
-
-def apply_hostname_bounce(hostname, policy, interface, command,
-                          hostname_command="hostname"):
+def perform_hostname_bounce(hostname, cfg, prev_hostname):
     # set the hostname to 'hostname' if it is not already set to that.
     # then, if policy is not off, bounce the interface using command
-    prev_hostname = get_hostname()
-
-    set_hostname(hostname, hostname_command)
-
-    msg = ("phostname=%s hostname=%s policy=%s interface=%s" %
-           (prev_hostname, hostname, policy, interface))
-
-    if util.is_false(policy):
-        LOG.debug("pubhname: policy false, skipping [%s]", msg)
-        return
-
-    if prev_hostname == hostname and policy != "force":
-        LOG.debug("pubhname: no change, policy != force. skipping. [%s]", msg)
-        return
+    command = cfg['command']
+    interface = cfg['interface']
+    policy = cfg['policy']
 
+    msg = ("hostname=%s policy=%s interface=%s" %
+           (hostname, policy, interface))
     env = os.environ.copy()
     env['interface'] = interface
     env['hostname'] = hostname
@@ -358,7 +355,12 @@ def apply_hostname_bounce(hostname, policy, interface, command,
         command = BOUNCE_COMMAND
 
     LOG.debug("pubhname: publishing hostname [%s]", msg)
-    perform_hostname_bounce(command, env)
+    shell = not isinstance(command, (list, tuple))
+    # capture=False, see comments in bug 1202758 and bug 1206164.
+    util.log_time(logfunc=LOG.debug, msg="publishing hostname",
+                  get_uptime=True, func=util.subp,
+                  kwargs={'args': command, 'shell': shell, 'capture': False,
+                          'env': env})
 
 
 def crtfile_to_pubkey(fname):
diff --git a/tests/unittests/test_datasource/test_azure.py b/tests/unittests/test_datasource/test_azure.py
index 3adf9bdf..7e789853 100644
--- a/tests/unittests/test_datasource/test_azure.py
+++ b/tests/unittests/test_datasource/test_azure.py
@@ -439,6 +439,11 @@ class TestAzureBounce(TestCase):
         }
         return construct_valid_ovf_env(data=odata)
 
+    def test_disabled_bounce_does_not_change_hostname(self):
+        cfg = {'hostname_bounce': {'policy': 'off'}}
+        self._get_ds(self.get_ovf_env_with_dscfg('test-host', cfg)).get_data()
+        self.assertEqual(0, self.set_hostname.call_count)
+
     @mock.patch('cloudinit.sources.DataSourceAzure.perform_hostname_bounce')
     def test_disabled_bounce_does_not_perform_bounce(
             self, perform_hostname_bounce):
@@ -446,6 +451,13 @@ class TestAzureBounce(TestCase):
         self._get_ds(self.get_ovf_env_with_dscfg('test-host', cfg)).get_data()
         self.assertEqual(0, perform_hostname_bounce.call_count)
 
+    def test_same_hostname_does_not_change_hostname(self):
+        host_name = 'unchanged-host-name'
+        self.get_hostname.return_value = host_name
+        cfg = {'hostname_bounce': {'policy': 'yes'}}
+        self._get_ds(self.get_ovf_env_with_dscfg(host_name, cfg)).get_data()
+        self.assertEqual(0, self.set_hostname.call_count)
+
     @mock.patch('cloudinit.sources.DataSourceAzure.perform_hostname_bounce')
     def test_unchanged_hostname_does_not_perform_bounce(
             self, perform_hostname_bounce):
@@ -480,6 +492,25 @@ class TestAzureBounce(TestCase):
             self.get_ovf_env_with_dscfg(expected_hostname, {})).get_data()
         self.assertEqual(1, perform_hostname_bounce.call_count)
 
+    def test_different_hostnames_sets_hostname_back(self):
+        initial_host_name = 'default-host-name'
+        self.get_hostname.return_value = initial_host_name
+        self._get_ds(
+            self.get_ovf_env_with_dscfg('some-host-name', {})).get_data()
+        self.assertEqual(initial_host_name,
+                         self.set_hostname.call_args_list[-1][0][0])
+
+    @mock.patch('cloudinit.sources.DataSourceAzure.perform_hostname_bounce')
+    def test_failure_in_bounce_still_resets_host_name(
+            self, perform_hostname_bounce):
+        perform_hostname_bounce.side_effect = Exception
+        initial_host_name = 'default-host-name'
+        self.get_hostname.return_value = initial_host_name
+        self._get_ds(
+            self.get_ovf_env_with_dscfg('some-host-name', {})).get_data()
+        self.assertEqual(initial_host_name,
+                         self.set_hostname.call_args_list[-1][0][0])
+
     def test_environment_correct_for_bounce_command(self):
         interface = 'int0'
         hostname = 'my-new-host'