ubuntu-drivers: call db_x_loadtemplatefile to accept NVIDIA EULA

Emit a script allowing cloud-init to set linux/nvidia/latelink
debconf selection to true. This avoids having to call
debconf-set-selections and allows cloud-init to pre-confgure
linux-restricted-modules to link NVIDIA drivers to the running kernel.

Cloud-init loads this debconf template and sets the value to true in the
debconf database by sourcing debconf's /usr/share/debconf/confmodule and
uses db_x_loadtemplatefile to register cloud-init's setting for
linux/nvidia/latelink.

LP: #1840080

3 files changed, 130 insertions, 36 deletions

diff --git a/cloudinit/config/cc_ubuntu_drivers.py b/cloudinit/config/cc_ubuntu_drivers.py
index 4da34ee0..297451d6 100644
--- a/cloudinit/config/cc_ubuntu_drivers.py
+++ b/cloudinit/config/cc_ubuntu_drivers.py
@@ -2,13 +2,14 @@
 
 """Ubuntu Drivers: Interact with third party drivers in Ubuntu."""
 
+import os
 from textwrap import dedent
 
-from cloudinit.config import cc_apt_configure
 from cloudinit.config.schema import (
     get_schema_doc, validate_cloudconfig_schema)
 from cloudinit import log as logging
 from cloudinit.settings import PER_INSTANCE
+from cloudinit import temp_utils
 from cloudinit import type_utils
 from cloudinit import util
 
@@ -65,6 +66,33 @@ OLD_UBUNTU_DRIVERS_STDERR_NEEDLE = (
 __doc__ = get_schema_doc(schema)  # Supplement python help()
 
 
+# Use a debconf template to configure a global debconf variable
+# (linux/nvidia/latelink) setting this to "true" allows the
+# 'linux-restricted-modules' deb to accept the NVIDIA EULA and the package
+# will automatically link the drivers to the running kernel.
+
+# EOL_XENIAL: can then drop this script and use python3-debconf which is only
+# available in Bionic and later. Can't use python3-debconf currently as it
+# isn't in Xenial and doesn't yet support X_LOADTEMPLATEFILE debconf command.
+
+NVIDIA_DEBCONF_CONTENT = """\
+Template: linux/nvidia/latelink
+Type: boolean
+Default: true
+Description: Late-link NVIDIA kernel modules?
+ Enable this to link the NVIDIA kernel modules in cloud-init and
+ make them available for use.
+"""
+
+NVIDIA_DRIVER_LATELINK_DEBCONF_SCRIPT = """\
+#!/bin/sh
+# Allow cloud-init to trigger EULA acceptance via registering a debconf
+# template to set linux/nvidia/latelink true
+. /usr/share/debconf/confmodule
+db_x_loadtemplatefile "$1" cloud-init
+"""
+
+
 def install_drivers(cfg, pkg_install_func):
     if not isinstance(cfg, dict):
         raise TypeError(
@@ -90,17 +118,27 @@ def install_drivers(cfg, pkg_install_func):
     if version_cfg:
         driver_arg += ':{}'.format(version_cfg)
 
-    LOG.debug("Installing NVIDIA drivers (%s=%s, version=%s)",
+    LOG.debug("Installing and activating NVIDIA drivers (%s=%s, version=%s)",
               cfgpath, nv_acc, version_cfg if version_cfg else 'latest')
 
-    # Setting NVIDIA latelink confirms acceptance of EULA for the package
-    # linux-restricted-modules
-    # Reference code defining debconf variable is here
-    # https://git.launchpad.net/~ubuntu-kernel/ubuntu/+source/
-    # linux-restricted-modules/+git/eoan/tree/debian/templates/
-    # nvidia.templates.in
-    selections = b'linux-restricted-modules linux/nvidia/latelink boolean true'
-    cc_apt_configure.debconf_set_selections(selections)
+    # Register and set debconf selection linux/nvidia/latelink = true
+    tdir = temp_utils.mkdtemp(needs_exe=True)
+    debconf_file = os.path.join(tdir, 'nvidia.template')
+    debconf_script = os.path.join(tdir, 'nvidia-debconf.sh')
+    try:
+        util.write_file(debconf_file, NVIDIA_DEBCONF_CONTENT)
+        util.write_file(
+            debconf_script,
+            util.encode_text(NVIDIA_DRIVER_LATELINK_DEBCONF_SCRIPT),
+            mode=0o755)
+        util.subp([debconf_script, debconf_file])
+    except Exception as e:
+        util.logexc(
+            LOG, "Failed to register NVIDIA debconf template: %s", str(e))
+        raise
+    finally:
+        if os.path.isdir(tdir):
+            util.del_dir(tdir)
 
     try:
         util.subp(['ubuntu-drivers', 'install', '--gpgpu', driver_arg])
diff --git a/cloudinit/config/tests/test_ubuntu_drivers.py b/cloudinit/config/tests/test_ubuntu_drivers.py
index 6a763bd9..46952692 100644
--- a/cloudinit/config/tests/test_ubuntu_drivers.py
+++ b/cloudinit/config/tests/test_ubuntu_drivers.py
@@ -1,6 +1,7 @@
 # This file is part of cloud-init. See LICENSE file for license information.
 
 import copy
+import os
 
 from cloudinit.tests.helpers import CiTestCase, skipUnlessJsonSchema, mock
 from cloudinit.config.schema import (
@@ -9,11 +10,27 @@ from cloudinit.config import cc_ubuntu_drivers as drivers
 from cloudinit.util import ProcessExecutionError
 
 MPATH = "cloudinit.config.cc_ubuntu_drivers."
+M_TMP_PATH = MPATH + "temp_utils.mkdtemp"
 OLD_UBUNTU_DRIVERS_ERROR_STDERR = (
     "ubuntu-drivers: error: argument <command>: invalid choice: 'install' "
     "(choose from 'list', 'autoinstall', 'devices', 'debug')\n")
 
 
+class AnyTempScriptAndDebconfFile(object):
+
+    def __init__(self, tmp_dir, debconf_file):
+        self.tmp_dir = tmp_dir
+        self.debconf_file = debconf_file
+
+    def __eq__(self, cmd):
+        if not len(cmd) == 2:
+            return False
+        script, debconf_file = cmd
+        if bool(script.startswith(self.tmp_dir) and script.endswith('.sh')):
+            return debconf_file == self.debconf_file
+        return False
+
+
 class TestUbuntuDrivers(CiTestCase):
     cfg_accepted = {'drivers': {'nvidia': {'license-accepted': True}}}
     install_gpgpu = ['ubuntu-drivers', 'install', '--gpgpu', 'nvidia']
@@ -28,20 +45,23 @@ class TestUbuntuDrivers(CiTestCase):
                 {'drivers': {'nvidia': {'license-accepted': "TRUE"}}},
                 schema=drivers.schema, strict=True)
 
-    @mock.patch(MPATH + "cc_apt_configure.debconf_set_selections")
+    @mock.patch(M_TMP_PATH)
     @mock.patch(MPATH + "util.subp", return_value=('', ''))
     @mock.patch(MPATH + "util.which", return_value=False)
     def _assert_happy_path_taken(
-            self, config, m_which, m_subp, m_debconf_set_selections):
+            self, config, m_which, m_subp, m_tmp):
         """Positive path test through handle. Package should be installed."""
+        tdir = self.tmp_dir()
+        debconf_file = os.path.join(tdir, 'nvidia.template')
+        m_tmp.return_value = tdir
         myCloud = mock.MagicMock()
         drivers.handle('ubuntu_drivers', config, myCloud, None, None)
         self.assertEqual([mock.call(['ubuntu-drivers-common'])],
                          myCloud.distro.install_packages.call_args_list)
-        self.assertEqual([mock.call(self.install_gpgpu)],
-                         m_subp.call_args_list)
-        m_debconf_set_selections.assert_called_with(
-            b'linux-restricted-modules linux/nvidia/latelink boolean true')
+        self.assertEqual(
+            [mock.call(AnyTempScriptAndDebconfFile(tdir, debconf_file)),
+             mock.call(self.install_gpgpu)],
+            m_subp.call_args_list)
 
     def test_handle_does_package_install(self):
         self._assert_happy_path_taken(self.cfg_accepted)
@@ -52,20 +72,33 @@ class TestUbuntuDrivers(CiTestCase):
             new_config['drivers']['nvidia']['license-accepted'] = true_value
             self._assert_happy_path_taken(new_config)
 
-    @mock.patch(MPATH + "cc_apt_configure.debconf_set_selections")
-    @mock.patch(MPATH + "util.subp", side_effect=ProcessExecutionError(
-        stdout='No drivers found for installation.\n', exit_code=1))
+    @mock.patch(M_TMP_PATH)
+    @mock.patch(MPATH + "util.subp")
     @mock.patch(MPATH + "util.which", return_value=False)
-    def test_handle_raises_error_if_no_drivers_found(self, m_which, m_subp, _):
+    def test_handle_raises_error_if_no_drivers_found(
+            self, m_which, m_subp, m_tmp):
         """If ubuntu-drivers doesn't install any drivers, raise an error."""
+        tdir = self.tmp_dir()
+        debconf_file = os.path.join(tdir, 'nvidia.template')
+        m_tmp.return_value = tdir
         myCloud = mock.MagicMock()
+
+        def fake_subp(cmd):
+            if cmd[0].startswith(tdir):
+                return
+            raise ProcessExecutionError(
+                stdout='No drivers found for installation.\n', exit_code=1)
+        m_subp.side_effect = fake_subp
+
         with self.assertRaises(Exception):
             drivers.handle(
                 'ubuntu_drivers', self.cfg_accepted, myCloud, None, None)
         self.assertEqual([mock.call(['ubuntu-drivers-common'])],
                          myCloud.distro.install_packages.call_args_list)
-        self.assertEqual([mock.call(self.install_gpgpu)],
-                         m_subp.call_args_list)
+        self.assertEqual(
+            [mock.call(AnyTempScriptAndDebconfFile(tdir, debconf_file)),
+             mock.call(self.install_gpgpu)],
+            m_subp.call_args_list)
         self.assertIn('ubuntu-drivers found no drivers for installation',
                       self.logs.getvalue())
 
@@ -113,19 +146,25 @@ class TestUbuntuDrivers(CiTestCase):
                       myLog.debug.call_args_list[0][0][0])
         self.assertEqual(0, m_install_drivers.call_count)
 
-    @mock.patch(MPATH + "cc_apt_configure.debconf_set_selections")
+    @mock.patch(M_TMP_PATH)
     @mock.patch(MPATH + "util.subp", return_value=('', ''))
     @mock.patch(MPATH + "util.which", return_value=True)
-    def test_install_drivers_no_install_if_present(self, m_which, m_subp, _):
+    def test_install_drivers_no_install_if_present(
+            self, m_which, m_subp, m_tmp):
         """If 'ubuntu-drivers' is present, no package install should occur."""
+        tdir = self.tmp_dir()
+        debconf_file = os.path.join(tdir, 'nvidia.template')
+        m_tmp.return_value = tdir
         pkg_install = mock.MagicMock()
         drivers.install_drivers(self.cfg_accepted['drivers'],
                                 pkg_install_func=pkg_install)
         self.assertEqual(0, pkg_install.call_count)
         self.assertEqual([mock.call('ubuntu-drivers')],
                          m_which.call_args_list)
-        self.assertEqual([mock.call(self.install_gpgpu)],
-                         m_subp.call_args_list)
+        self.assertEqual(
+            [mock.call(AnyTempScriptAndDebconfFile(tdir, debconf_file)),
+             mock.call(self.install_gpgpu)],
+            m_subp.call_args_list)
 
     def test_install_drivers_rejects_invalid_config(self):
         """install_drivers should raise TypeError if not given a config dict"""
@@ -134,21 +173,33 @@ class TestUbuntuDrivers(CiTestCase):
             drivers.install_drivers("mystring", pkg_install_func=pkg_install)
         self.assertEqual(0, pkg_install.call_count)
 
-    @mock.patch(MPATH + "cc_apt_configure.debconf_set_selections")
-    @mock.patch(MPATH + "util.subp", side_effect=ProcessExecutionError(
-        stderr=OLD_UBUNTU_DRIVERS_ERROR_STDERR, exit_code=2))
+    @mock.patch(M_TMP_PATH)
+    @mock.patch(MPATH + "util.subp")
     @mock.patch(MPATH + "util.which", return_value=False)
     def test_install_drivers_handles_old_ubuntu_drivers_gracefully(
-            self, m_which, m_subp, _):
+            self, m_which, m_subp, m_tmp):
         """Older ubuntu-drivers versions should emit message and raise error"""
+        tdir = self.tmp_dir()
+        debconf_file = os.path.join(tdir, 'nvidia.template')
+        m_tmp.return_value = tdir
         myCloud = mock.MagicMock()
+
+        def fake_subp(cmd):
+            if cmd[0].startswith(tdir):
+                return
+            raise ProcessExecutionError(
+                stderr=OLD_UBUNTU_DRIVERS_ERROR_STDERR, exit_code=2)
+        m_subp.side_effect = fake_subp
+
         with self.assertRaises(Exception):
             drivers.handle(
                 'ubuntu_drivers', self.cfg_accepted, myCloud, None, None)
         self.assertEqual([mock.call(['ubuntu-drivers-common'])],
                          myCloud.distro.install_packages.call_args_list)
-        self.assertEqual([mock.call(self.install_gpgpu)],
-                         m_subp.call_args_list)
+        self.assertEqual(
+            [mock.call(AnyTempScriptAndDebconfFile(tdir, debconf_file)),
+             mock.call(self.install_gpgpu)],
+            m_subp.call_args_list)
         self.assertIn('WARNING: the available version of ubuntu-drivers is'
                       ' too old to perform requested driver installation',
                       self.logs.getvalue())
@@ -160,17 +211,21 @@ class TestUbuntuDriversWithVersion(TestUbuntuDrivers):
         'drivers': {'nvidia': {'license-accepted': True, 'version': '123'}}}
     install_gpgpu = ['ubuntu-drivers', 'install', '--gpgpu', 'nvidia:123']
 
-    @mock.patch(MPATH + "cc_apt_configure.debconf_set_selections")
+    @mock.patch(M_TMP_PATH)
     @mock.patch(MPATH + "util.subp", return_value=('', ''))
     @mock.patch(MPATH + "util.which", return_value=False)
-    def test_version_none_uses_latest(self, m_which, m_subp, _):
+    def test_version_none_uses_latest(self, m_which, m_subp, m_tmp):
+        tdir = self.tmp_dir()
+        debconf_file = os.path.join(tdir, 'nvidia.template')
+        m_tmp.return_value = tdir
         myCloud = mock.MagicMock()
         version_none_cfg = {
             'drivers': {'nvidia': {'license-accepted': True, 'version': None}}}
         drivers.handle(
             'ubuntu_drivers', version_none_cfg, myCloud, None, None)
         self.assertEqual(
-            [mock.call(['ubuntu-drivers', 'install', '--gpgpu', 'nvidia'])],
+            [mock.call(AnyTempScriptAndDebconfFile(tdir, debconf_file)),
+             mock.call(['ubuntu-drivers', 'install', '--gpgpu', 'nvidia'])],
             m_subp.call_args_list)
 
     def test_specifying_a_version_doesnt_override_license_acceptance(self):
diff --git a/cloudinit/tests/helpers.py b/cloudinit/tests/helpers.py
index f41180fd..23fddd07 100644
--- a/cloudinit/tests/helpers.py
+++ b/cloudinit/tests/helpers.py
@@ -198,7 +198,8 @@ class CiTestCase(TestCase):
                 prefix="ci-%s." % self.__class__.__name__)
         else:
             tmpd = tempfile.mkdtemp(dir=dir)
-        self.addCleanup(functools.partial(shutil.rmtree, tmpd))
+        self.addCleanup(
+            functools.partial(shutil.rmtree, tmpd, ignore_errors=True))
         return tmpd
 
     def tmp_path(self, path, dir=None):