TPT-3292: support gzipped user-data decoding

linode_metadata/metadata_client.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import base64
+import gzip
 import json
 import logging
 from abc import ABC
@@ -29,6 +30,7 @@
 
 BASE_URL = "http://169.254.169.254/v1"
 DEFAULT_API_TIMEOUT = 10.0
+GZIP_MAGIC = b"\x1f\x8b\x08"
 
 
 class BaseMetadataClient(ABC):
@@ -132,6 +134,21 @@ def _parse_response_body(response: Response, content_type: str) -> Any:
 
         return handler()
 
+    @staticmethod
+    def _decode_user_data(encoded: Union[bytes, str]) -> str:
+        """
+        Decode user-data from base64 and gunzip it when the payload is gzipped.
+
+        :param encoded: Base64-encoded user-data from the Metadata Service.
+        :type encoded: Union[bytes, str]
+        :returns: The decoded user-data string.
+        :rtype: str
+        """
+        raw = base64.b64decode(encoded)
+        if raw.startswith(GZIP_MAGIC):
+            raw = gzip.decompress(raw)
+        return raw.decode("utf-8")
+
     def _get_http_method(
         self, method: str
     ) -> Optional[Callable[..., Union[Response, Awaitable[Response]]]]:
@@ -468,11 +485,14 @@ def refresh_token(self, expiry_seconds: int = 3600) -> MetadataToken:
     def get_user_data(self) -> str:
         """
         Returns the user data configured on your running Linode instance.
+
+        NOTE: The result is automatically decoded from base64 and gunzipped
+        when needed.
         """
         response = self._api_call(
             "GET", "/user-data", content_type="text/plain"
         )
-        return base64.b64decode(response).decode("utf-8")
+        return self._decode_user_data(response)
 
     def get_instance(self) -> InstanceResponse:
         """
@@ -699,11 +719,14 @@ async def refresh_token(self, expiry_seconds: int = 3600) -> MetadataToken:
     async def get_user_data(self) -> str:
         """
         Returns the user data configured on your running Linode instance.
+
+        NOTE: The result is automatically decoded from base64 and gunzipped
+        when needed.
         """
         response = await self._api_call(
             "GET", "/user-data", content_type="text/plain"
         )
-        return base64.b64decode(response).decode("utf-8")
+        return self._decode_user_data(response)
 
     async def get_instance(self) -> InstanceResponse:
         """

test/unit/test_user_data.py

@@ -0,0 +1,19 @@
+import base64
+import gzip
+import io
+
+from linode_metadata.metadata_client import BaseMetadataClient
+
+
+def test_decode_user_data_plain():
+    encoded = base64.b64encode(b"mock-user-data")
+    assert BaseMetadataClient._decode_user_data(encoded) == "mock-user-data"
+
+
+def test_decode_user_data_gzip():
+    buffer = io.BytesIO()
+    with gzip.GzipFile(fileobj=buffer, mode="wb") as gz_file:
+        gz_file.write(b"mock-user-data")
+
+    encoded = base64.b64encode(buffer.getvalue())
+    assert BaseMetadataClient._decode_user_data(encoded) == "mock-user-data"