feat: MDS connections use mTLS

Author: nolanleastin

google/auth/compute_engine/_metadata.py

@@ -30,6 +30,8 @@
 from google.auth import metrics
 from google.auth import transport
 from google.auth._exponential_backoff import ExponentialBackoff
+from google.auth.compute_engine import _mtls
+from google.auth.transport import requests
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -42,13 +44,24 @@
     _GCE_METADATA_HOST = os.getenv(
         environment_vars.GCE_METADATA_ROOT, "metadata.google.internal"
     )
-_METADATA_ROOT = "http://{}/computeMetadata/v1/".format(_GCE_METADATA_HOST)
 
-# This is used to ping the metadata server, it avoids the cost of a DNS
-# lookup.
-_METADATA_IP_ROOT = "http://{}".format(
-    os.getenv(environment_vars.GCE_METADATA_IP, "169.254.169.254")
-)
+GCE_MDS_HOSTS = ["metadata.google.internal", "169.254.169.254"]
+
+
+def _get_metadata_root(use_mtls):
+    """Returns the metadata server root URL."""
+    scheme = "https" if use_mtls else "http"
+    return "{}://{}/computeMetadata/v1/".format(scheme, _GCE_METADATA_HOST)
+
+
+def _get_metadata_ip_root(use_mtls):
+    """Returns the metadata server IP root URL."""
+    scheme = "https" if use_mtls else "http"
+    return "{}://{}".format(
+        scheme, os.getenv(environment_vars.GCE_METADATA_IP, "169.254.169.254")
+    )
+
+
 _METADATA_FLAVOR_HEADER = "metadata-flavor"
 _METADATA_FLAVOR_VALUE = "Google"
 _METADATA_HEADERS = {_METADATA_FLAVOR_HEADER: _METADATA_FLAVOR_VALUE}
@@ -102,6 +115,24 @@ def detect_gce_residency_linux():
     return content.startswith(_GOOGLE)
 
 
+def _prepare_request_for_mds(request, use_mtls=False):
+    """Prepares a request for the metadata server.
+
+    This will check if mTLS should be used and return a new request object if so.
+
+    Args:
+        request (google.auth.transport.Request): A callable used to make
+            HTTP requests.
+
+    Returns:
+        google.auth.transport.Request: Request
+            object to use.
+    """
+    if use_mtls:
+        request = requests.Request(_mtls.create_session())
+    return request
+
+
 def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3):
     """Checks to see if the metadata server is available.
 
@@ -115,6 +146,8 @@ def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3):
     Returns:
         bool: True if the metadata server is reachable, False otherwise.
     """
+    use_mtls = _mtls.should_use_mds_mtls()
+    request = _prepare_request_for_mds(request, use_mtls=use_mtls)
     # NOTE: The explicit ``timeout`` is a workaround. The underlying
     #       issue is that resolving an unknown host on some networks will take
     #       20-30 seconds; making this timeout short fixes the issue, but
@@ -129,7 +162,10 @@ def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3):
     for attempt in backoff:
         try:
             response = request(
-                url=_METADATA_IP_ROOT, method="GET", headers=headers, timeout=timeout
+                url=_get_metadata_ip_root(use_mtls),
+                method="GET",
+                headers=headers,
+                timeout=timeout,
             )
 
             metadata_flavor = response.headers.get(_METADATA_FLAVOR_HEADER)
@@ -153,7 +189,7 @@ def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3):
 def get(
     request,
     path,
-    root=_METADATA_ROOT,
+    root=None,
     params=None,
     recursive=False,
     retry_count=5,
@@ -190,6 +226,14 @@ def get(
         google.auth.exceptions.TransportError: if an error occurred while
             retrieving metadata.
     """
+    use_mtls = _mtls.should_use_mds_mtls()
+    # Prepare the request object for mTLS if needed.
+    # This will create a new request object with the mTLS session.
+    request = _prepare_request_for_mds(request, use_mtls=use_mtls)
+
+    if root is None:
+        root = _get_metadata_root(use_mtls)
+
     base_url = urljoin(root, path)
     query_params = {} if params is None else params
 

google/auth/compute_engine/_mtls.py

@@ -0,0 +1,95 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Mutual TLS for Google Compute Engine metadata server."""
+import enum
+import os
+import ssl
+from dataclasses import dataclass
+
+from pathlib import Path
+import requests
+from requests.adapters import HTTPAdapter
+
+from google.auth import environment_vars, exceptions
+
+@dataclass
+class MdsMtlsConfig:
+    ca_cert_path: str = os.path.join(Path.home(), "mtls_mds_certificates", "root.crt") # path to CA certificate
+    client_combined_cert_path: str = os.path.join(Path.home(), "mtls_mds_certificates", "client_creds.key") # path to file containing client certificate and key
+
+class MdsMtlsMode(enum.Enum):
+    """MDS mTLS mode."""
+    STRICT = "strict"
+    NONE = "none"
+    DEFAULT = "default"
+
+def _parse_mds_mode():
+    """Parses the GCE_METADATA_MTLS_MODE environment variable."""
+    mode_str = os.environ.get(environment_vars.GCE_METADATA_MTLS_MODE, "default").lower()
+    try:
+        return MdsMtlsMode(mode_str)
+    except ValueError:
+        raise ValueError(
+            "Invalid value for GCE_METADATA_MTLS_MODE. Must be one of 'strict', 'none', or 'default'."
+        )
+
+
+def _certs_exist(mds_mtls_config: MdsMtlsConfig):
+    """Checks if the mTLS certificates exist."""
+    return os.path.exists(mds_mtls_config.ca_cert_path) and os.path.exists(
+        mds_mtls_config.client_combined_cert_path
+    )
+
+
+class MdsMtlsAdapter(HTTPAdapter):
+    """An HTTP adapter that uses mTLS for the metadata server."""
+
+    def __init__(self, mds_mtls_config: MdsMtlsConfig, *args, **kwargs):
+        self.ssl_context = ssl.create_default_context()
+        self.ssl_context.load_verify_locations(cafile=mds_mtls_config.ca_cert_path)
+        self.ssl_context.load_cert_chain(certfile=mds_mtls_config.client_combined_cert_path)
+        super(MdsMtlsAdapter, self).__init__(*args, **kwargs)
+
+    def init_poolmanager(self, *args, **kwargs):
+        kwargs["ssl_context"] = self.ssl_context
+        return super(MdsMtlsAdapter, self).init_poolmanager(*args, **kwargs)
+
+    def proxy_manager_for(self, *args, **kwargs):
+        kwargs["ssl_context"] = self.ssl_context
+        return super(MdsMtlsAdapter, self).proxy_manager_for(*args, **kwargs)
+
+def create_session(mds_mtls_config: MdsMtlsConfig = MdsMtlsConfig()):
+    """Creates a requests.Session configured for mTLS."""
+    session = requests.Session()
+    adapter = MdsMtlsAdapter(mds_mtls_config)
+    session.mount("https://", adapter)
+    return session
+
+
+def should_use_mds_mtls(mds_mtls_config: MdsMtlsConfig = MdsMtlsConfig()):
+    """Determines if mTLS should be used for the metadata server."""
+    mode = _parse_mds_mode()
+    if mode == MdsMtlsMode.STRICT:
+        if not _certs_exist(mds_mtls_config):
+            raise exceptions.MutualTLSChannelError(
+                "mTLS certificates not found in strict mode."
+            )
+        return True
+    elif mode == MdsMtlsMode.NONE:
+        return False
+    else:  # Default mode
+        return _certs_exist(mds_mtls_config)

google/auth/environment_vars.py

@@ -60,6 +60,12 @@
 """Environment variable providing an alternate ip:port to be used for ip-only
 GCE metadata requests."""
 
+GCE_METADATA_MTLS_MODE = "GCE_METADATA_MTLS_MODE"
+"""Environment variable controlling the mTLS behavior for GCE metadata requests.
+
+Can be one of "strict", "none", or "default".
+"""
+
 GOOGLE_API_USE_CLIENT_CERTIFICATE = "GOOGLE_API_USE_CLIENT_CERTIFICATE"
 """Environment variable controlling whether to use client certificate or not.