From c3c0c6d538ad1be64f5f4f27d4df59d7ac21e9ba Mon Sep 17 00:00:00 2001
From: "James E. Blair" <jim@acmegating.com>
Date: Mon, 18 Oct 2021 16:30:51 -0700
Subject: [PATCH] Add user-data/custom-data to Azure driver

This is the relatively common feature that allows users to include
cloud_init or similar data.  Azure has two versions of it, user-data
and custom-data.  For Nodepool's purpose they are similar, but a user
may prefer one or the other.

Also adds missing docs for the existing tags attribute.

Change-Id: Ia2f78a827c1909cc527733013167b2f3b5db18a3
---
 doc/source/azure.rst                          | 22 +++++++++
 nodepool/driver/azure/adapter.py              | 47 ++++++++++---------
 nodepool/driver/azure/config.py               | 10 ++++
 nodepool/tests/fixtures/azure.yaml            |  2 +
 nodepool/tests/unit/fake_azure.py             |  3 ++
 nodepool/tests/unit/test_driver_azure.py      |  8 ++++
 .../azure-user-data-0b35b3cb9cca8afd.yaml     |  4 ++
 7 files changed, 75 insertions(+), 21 deletions(-)
 create mode 100644 releasenotes/notes/azure-user-data-0b35b3cb9cca8afd.yaml

diff --git a/doc/source/azure.rst b/doc/source/azure.rst
index dbd3e0b67..e008392b3 100644
--- a/doc/source/azure.rst
+++ b/doc/source/azure.rst
@@ -510,7 +510,29 @@ section of the configuration.
                 list on `azure.microsoft.com`_ for the list of sizes
                 availabile in each region.
 
+          .. attr:: tags
+             :type: dict
+             :default: None
+
+             A dictionary of tags to add to newly created VMs.
+
+          .. attr:: user-data
+             :type: str
+             :default: None
+
+             The `Azure User Data`_ value for newly created VMs.
+
+          .. attr:: custom-data
+             :type: str
+             :default: None
+
+             The `Azure Custom Data`_ value for newly created VMs.
+
 
 .. _`Azure CLI`: https://docs.microsoft.com/en-us/cli/azure/create-an-azure-service-principal-azure-cli?view=azure-cli-latest
 
 .. _azure.microsoft.com: https://azure.microsoft.com/en-us/global-infrastructure/services/?products=virtual-machines
+
+.. _`Azure User Data`: https://docs.microsoft.com/en-us/azure/virtual-machines/user-data
+
+.. _`Azure Custom Data`: https://docs.microsoft.com/en-us/azure/virtual-machines/custom-data
diff --git a/nodepool/driver/azure/adapter.py b/nodepool/driver/azure/adapter.py
index 06da51328..a3fd21175 100644
--- a/nodepool/driver/azure/adapter.py
+++ b/nodepool/driver/azure/adapter.py
@@ -646,30 +646,35 @@ class AzureAdapter(statemachine.Adapter):
             }
             os_profile['adminUsername'] = image.username
             os_profile['linuxConfiguration'] = linux_config
+        if label.custom_data:
+            os_profile['customData'] = label.custom_data
 
+        spec = {
+            'location': self.provider.location,
+            'tags': tags,
+            'properties': {
+                'osProfile': os_profile,
+                'hardwareProfile': {
+                    'vmSize': label.hardware_profile["vm-size"]
+                },
+                'storageProfile': {
+                    'imageReference': image_reference,
+                },
+                'networkProfile': {
+                    'networkInterfaces': [{
+                        'id': nic['id'],
+                        'properties': {
+                            'primary': True,
+                        }
+                    }]
+                },
+            },
+        }
+        if label.user_data:
+            spec['properties']['userData'] = label.user_data
         with self.rate_limiter:
             return self.azul.virtual_machines.create(
-                self.resource_group, hostname, {
-                    'location': self.provider.location,
-                    'tags': tags,
-                    'properties': {
-                        'osProfile': os_profile,
-                        'hardwareProfile': {
-                            'vmSize': label.hardware_profile["vm-size"]
-                        },
-                        'storageProfile': {
-                            'imageReference': image_reference,
-                        },
-                        'networkProfile': {
-                            'networkInterfaces': [{
-                                'id': nic['id'],
-                                'properties': {
-                                    'primary': True,
-                                }
-                            }]
-                        },
-                    },
-                })
+                self.resource_group, hostname, spec)
 
     def _deleteVirtualMachine(self, name):
         for vm in self._listVirtualMachines():
diff --git a/nodepool/driver/azure/config.py b/nodepool/driver/azure/config.py
index 51368d4a9..0229010c3 100644
--- a/nodepool/driver/azure/config.py
+++ b/nodepool/driver/azure/config.py
@@ -16,6 +16,7 @@
 # limitations under the License.
 
 import voluptuous as v
+import base64
 import os
 
 from nodepool.driver import ConfigPool
@@ -143,6 +144,13 @@ class AzureLabel(ConfigValue):
 
         self.hardware_profile = label['hardware-profile']
         self.tags = label.get('tags', {})
+        self.user_data = self._encodeData(label.get('user-data', None))
+        self.custom_data = self._encodeData(label.get('custom-data', None))
+
+    def _encodeData(self, s):
+        if not s:
+            return None
+        return base64.b64encode(s.encode('utf8')).decode('utf8')
 
     @staticmethod
     def getSchema():
@@ -156,6 +164,8 @@ class AzureLabel(ConfigValue):
             'diskimage': str,
             v.Required('hardware-profile'): azure_hardware_profile,
             'tags': dict,
+            'user-data': str,
+            'custom-data': str,
         }
 
 
diff --git a/nodepool/tests/fixtures/azure.yaml b/nodepool/tests/fixtures/azure.yaml
index ac6e0a01b..0ef0c65f5 100644
--- a/nodepool/tests/fixtures/azure.yaml
+++ b/nodepool/tests/fixtures/azure.yaml
@@ -49,3 +49,5 @@ providers:
               department: R&D
               team: DevOps
               systemPurpose: CI
+            user-data: "This is the user data"
+            custom-data: "This is the custom data"
diff --git a/nodepool/tests/unit/fake_azure.py b/nodepool/tests/unit/fake_azure.py
index 193d8f64b..06b166038 100644
--- a/nodepool/tests/unit/fake_azure.py
+++ b/nodepool/tests/unit/fake_azure.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import json
 import time
 import os
@@ -33,6 +34,7 @@ class CRUDManager:
     def __init__(self, cloud):
         self.cloud = cloud
         self.items = []
+        self.requests = []
 
     def list(self, request):
         resp = {'value': self.items}
@@ -139,6 +141,7 @@ class VirtualMachinesCRUD(CRUDManager):
 
     def put(self, request):
         data = json.loads(request.body)
+        self.requests.append(copy.deepcopy(data))
         url = urllib.parse.urlparse(request.path_url)
         name = url.path.split('/')[-1]
         data['id'] = url.path
diff --git a/nodepool/tests/unit/test_driver_azure.py b/nodepool/tests/unit/test_driver_azure.py
index ad847ec08..34cf6a460 100644
--- a/nodepool/tests/unit/test_driver_azure.py
+++ b/nodepool/tests/unit/test_driver_azure.py
@@ -58,6 +58,14 @@ class TestDriverAzure(tests.DBTestCase):
         self.assertEqual(node.attributes,
                          {'key1': 'value1', 'key2': 'value2'})
         self.assertEqual(node.host_keys, ['ssh-rsa FAKEKEY'])
+        self.assertEqual(
+            self.fake_azure.crud['Microsoft.Compute/virtualMachines'].
+            items[0]['properties']['osProfile']['customData'],
+            'VGhpcyBpcyB0aGUgY3VzdG9tIGRhdGE=')  # This is the custom data
+        self.assertEqual(
+            self.fake_azure.crud['Microsoft.Compute/virtualMachines'].
+            requests[0]['properties']['userData'],
+            'VGhpcyBpcyB0aGUgdXNlciBkYXRh')  # This is the user data
 
     def test_azure_diskimage(self):
         configfile = self.setup_config(
diff --git a/releasenotes/notes/azure-user-data-0b35b3cb9cca8afd.yaml b/releasenotes/notes/azure-user-data-0b35b3cb9cca8afd.yaml
new file mode 100644
index 000000000..48fb0b7ec
--- /dev/null
+++ b/releasenotes/notes/azure-user-data-0b35b3cb9cca8afd.yaml
@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    The Azure driver now supports user-data and custom-data.