fix(azure): url_helper: specify User-Agent when using headers_cb with…

… readurl() (#5298)

My previous change #5218 introduced a bug where IMDS request is missing
`Cloud-Init` UserAgent and contains `python-requests` UserAgent instead.

A review showed that no users of `headers_cb` were specifying user-agent and it
seems unexpected that it wouldn't be defined similar to users of `headers`.

This PR updates readurl() to specify User-Agent in `req_args` for users passing in
`headers_cb`.

cloudinit/url_helper.py

@@ -423,19 +423,12 @@ def readurl(
     if retries:
         manual_tries = max(int(retries) + 1, 1)
 
-    def_headers = {
-        "User-Agent": "Cloud-Init/%s" % (version.version_string()),
-    }
-    if headers:
-        def_headers.update(headers)
-    headers = def_headers
-
-    if not headers_cb:
-
-        def _cb(url):
-            return headers
+    user_agent = "Cloud-Init/%s" % (version.version_string())
+    if headers is not None:
+        headers = headers.copy()
+    else:
+        headers = {}
 
-        headers_cb = _cb
     if data:
         req_args["data"] = data
     if sec_between is None:
@@ -447,7 +440,13 @@ def _cb(url):
     # Handle retrying ourselves since the built-in support
     # doesn't handle sleeping between tries...
     for i in count():
-        req_args["headers"] = headers_cb(url)
+        if headers_cb:
+            headers = headers_cb(url)
+
+        if "User-Agent" not in headers:
+            headers["User-Agent"] = user_agent
+
+        req_args["headers"] = headers
         filtered_req_args = {}
         for (k, v) in req_args.items():
             if k == "data":

tests/unittests/test_url_helper.py

@@ -20,6 +20,7 @@
     dual_stack,
     oauth_headers,
     read_file_or_url,
+    readurl,
     wait_for_url,
 )
 from tests.unittests.helpers import CiTestCase, mock, skipIf
@@ -282,6 +283,71 @@ def assert_time(func, max_time=1):
     return out
 
 
+class TestReadUrl:
+    @pytest.mark.parametrize("headers", [{}, {"Metadata": "true"}])
+    def test_headers(self, headers):
+        url = "http://hostname/path"
+        m_response = mock.MagicMock()
+
+        expected_headers = headers.copy()
+        expected_headers["User-Agent"] = "Cloud-Init/%s" % (
+            version.version_string()
+        )
+
+        class FakeSession(requests.Session):
+            @classmethod
+            def request(cls, **kwargs):
+                expected_kwargs = {
+                    "url": url,
+                    "allow_redirects": True,
+                    "method": "GET",
+                    "headers": expected_headers,
+                    "stream": False,
+                }
+
+                assert kwargs == expected_kwargs
+                return m_response
+
+        with mock.patch(
+            M_PATH + "requests.Session", side_effect=[FakeSession()]
+        ):
+            response = readurl(url, headers=headers)
+
+        assert response._response == m_response
+
+    @pytest.mark.parametrize("headers", [{}, {"Metadata": "true"}])
+    def test_headers_cb(self, headers):
+        url = "http://hostname/path"
+        m_response = mock.MagicMock()
+
+        expected_headers = headers.copy()
+        expected_headers["User-Agent"] = "Cloud-Init/%s" % (
+            version.version_string()
+        )
+        headers_cb = lambda _: headers
+
+        class FakeSession(requests.Session):
+            @classmethod
+            def request(cls, **kwargs):
+                expected_kwargs = {
+                    "url": url,
+                    "allow_redirects": True,
+                    "method": "GET",
+                    "headers": expected_headers,
+                    "stream": False,
+                }
+
+                assert kwargs == expected_kwargs
+                return m_response
+
+        with mock.patch(
+            M_PATH + "requests.Session", side_effect=[FakeSession()]
+        ):
+            response = readurl(url, headers_cb=headers_cb)
+
+        assert response._response == m_response
+
+
 event = Event()