change the mds mtls implementation

1. now we do not create a new request. instead, create an mds mtls
   adapter and mount it on the request session.

2. added _validate_gce_mds_configured_environment, which ensures if we
   are using strict, that the host being contacted is default

3. fix unit tests and add new tests

Author: nolanleastin

google/auth/compute_engine/_metadata.py

@@ -24,38 +24,61 @@
 import os
 from urllib.parse import urljoin
 
+import requests
+
 from google.auth import _helpers
 from google.auth import environment_vars
 from google.auth import exceptions
 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__)
 
+_GCE_DEFAULT_MDS_IP = "169.254.169.254"
+_GCE_DEFAULT_HOST = "metadata.google.internal"
+_GCE_DEFAULT_MDS_HOSTS = [_GCE_DEFAULT_HOST, _GCE_DEFAULT_MDS_IP]
+
 # Environment variable GCE_METADATA_HOST is originally named
 # GCE_METADATA_ROOT. For compatibility reasons, here it checks
 # the new variable first; if not set, the system falls back
 # to the old variable.
 _GCE_METADATA_HOST = os.getenv(environment_vars.GCE_METADATA_HOST, None)
 if not _GCE_METADATA_HOST:
     _GCE_METADATA_HOST = os.getenv(
-        environment_vars.GCE_METADATA_ROOT, "metadata.google.internal"
+        environment_vars.GCE_METADATA_ROOT, _GCE_DEFAULT_HOST
     )
 
-_GCE_DEFAULT_MDS_IP = "169.254.169.254"
-_GCE_MDS_HOSTS = ["metadata.google.internal", _GCE_DEFAULT_MDS_IP]
 
+def _validate_gce_mds_configured_environment():
+    """Validates the GCE metadata server environment configuration for mTLS.
 
-def _get_metadata_root(use_mtls):
+    Raises:
+        google.auth.exceptions.MutualTLSChannelError: if the environment
+            configuration is invalid for mTLS.
+    """
+    mode = _mtls._parse_mds_mode()
+    if mode == _mtls.MdsMtlsMode.STRICT:
+        if _GCE_METADATA_HOST != _GCE_DEFAULT_HOST:
+            # mTLS is only supported when connecting to the default metadata host.
+            # Raise an exception if we are in strict mode (which requires mTLS)
+            # but the metadata host has been overridden. (which means mTLS will fail)
+            raise exceptions.MutualTLSChannelError(
+                "Mutual TLS is required, but the metadata host has been overridden. "
+                "mTLS is only supported when connecting to the default metadata host."
+            )
+
+
+def _get_metadata_root(use_mtls: bool):
     """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):
+def _get_metadata_ip_root(use_mtls: bool):
     """Returns the metadata server IP root URL."""
     scheme = "https" if use_mtls else "http"
     return "{}://{}".format(
@@ -131,8 +154,14 @@ def _prepare_request_for_mds(request, use_mtls=False):
             If mTLS is enabled, this will be a new request object with mTLS session configured.
             Otherwise, it will be the same as the input request.
     """
-    if use_mtls:
-        request = requests.Request(_mtls.create_session())
+    if not use_mtls:
+        return request
+
+    adapter = _mtls.MdsMtlsAdapter()
+    if not request.session:
+        request.session = requests.Session()
+    for host in _GCE_DEFAULT_MDS_HOSTS:
+        request.session.mount(f"https://{host}/", adapter)
     return request
 
 
@@ -236,6 +265,7 @@ def get(
 
     if root is None:
         root = _get_metadata_root(use_mtls)
+    _validate_gce_mds_configured_environment()
 
     base_url = urljoin(root, path)
     query_params = {} if params is None else params

google/auth/compute_engine/_mtls.py

@@ -21,7 +21,6 @@
 import os
 import ssl
 
-import requests
 from requests.adapters import HTTPAdapter
 
 from google.auth import environment_vars, exceptions
@@ -59,6 +58,13 @@ class MdsMtlsConfig:
     )  # path to file containing client certificate and key
 
 
+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 MdsMtlsMode(enum.Enum):
     """MDS mTLS mode. Used to configure connection behavior when connecting to MDS.
 
@@ -85,17 +91,27 @@ def _parse_mds_mode():
         )
 
 
-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
-    )
+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)
 
 
 class MdsMtlsAdapter(HTTPAdapter):
     """An HTTP adapter that uses mTLS for the metadata server."""
 
-    def __init__(self, mds_mtls_config: MdsMtlsConfig, *args, **kwargs):
+    def __init__(
+        self, mds_mtls_config: MdsMtlsConfig = 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(
@@ -110,26 +126,3 @@ def init_poolmanager(self, *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)

tests/compute_engine/test__metadata.py

@@ -712,12 +712,12 @@ def test__get_metadata_ip_root_no_mtls():
     assert _metadata._get_metadata_ip_root(use_mtls=False) == "http://169.254.169.254"
 
 
-@mock.patch("google.auth.compute_engine._mtls.create_session")
-def test__prepare_request_for_mds_mtls(mock_create_session):
-    request = mock.Mock()
-    new_request = _metadata._prepare_request_for_mds(request, use_mtls=True)
-    mock_create_session.assert_called_once()
-    assert isinstance(new_request, google_auth_requests.Request)
+@mock.patch("google.auth.compute_engine._mtls.MdsMtlsAdapter")
+def test__prepare_request_for_mds_mtls(mock_mds_mtls_adapter):
+    request = google_auth_requests.Request(mock.create_autospec(requests.Session))
+    _metadata._prepare_request_for_mds(request, use_mtls=True)
+    mock_mds_mtls_adapter.assert_called_once()
+    assert request.session.mount.call_count == len(_metadata._GCE_DEFAULT_MDS_HOSTS)
 
 
 def test__prepare_request_for_mds_no_mtls():
@@ -726,53 +726,100 @@ def test__prepare_request_for_mds_no_mtls():
     assert new_request is request
 
 
-@mock.patch("google.auth.compute_engine._mtls.should_use_mds_mtls", return_value=True)
-@mock.patch("google.auth.compute_engine._mtls.create_session")
 @mock.patch("google.auth.metrics.mds_ping", return_value=MDS_PING_METRICS_HEADER_VALUE)
+@mock.patch("google.auth.compute_engine._mtls.MdsMtlsAdapter")
+@mock.patch("google.auth.compute_engine._mtls.should_use_mds_mtls", return_value=True)
+@mock.patch("google.auth.transport.requests.Request")
 def test_ping_mtls(
-    mock_metrics_header_value, mock_create_session, mock_should_use_mtls
+    mock_request, mock_should_use_mtls, mock_mds_mtls_adapter, mock_metrics_header_value
 ):
-    response = mock.create_autospec(requests.Response, instance=True)
-    response.status_code = http_client.OK
+    response = mock.create_autospec(transport.Response, instance=True)
+    response.status = http_client.OK
     response.headers = _metadata._METADATA_HEADERS
-    mock_session = mock.Mock()
-    mock_session.request.return_value = response
-    mock_create_session.return_value = mock_session
+    mock_request.return_value = response
 
-    initial_request = mock.Mock()
-    assert _metadata.ping(initial_request)
+    assert _metadata.ping(mock_request)
 
     mock_should_use_mtls.assert_called_once()
-    mock_create_session.assert_called_once()
-    mock_session.request.assert_called_once_with(
-        "GET",
-        "https://169.254.169.254",
+    mock_mds_mtls_adapter.assert_called_once()
+    mock_request.assert_called_once_with(
+        url="https://169.254.169.254",
+        method="GET",
         headers=MDS_PING_REQUEST_HEADER,
         timeout=_metadata._METADATA_DEFAULT_TIMEOUT,
-        data=None,
     )
 
 
+@mock.patch("google.auth.compute_engine._mtls.MdsMtlsAdapter")
 @mock.patch("google.auth.compute_engine._mtls.should_use_mds_mtls", return_value=True)
-@mock.patch("google.auth.compute_engine._mtls.create_session")
-def test_get_mtls(mock_create_session, mock_should_use_mtls):
-    response = mock.create_autospec(requests.Response, instance=True)
-    response.status_code = http_client.OK
-    response.content = _helpers.to_bytes("{}")
+@mock.patch("google.auth.transport.requests.Request")
+def test_get_mtls(mock_request, mock_should_use_mtls, mock_mds_mtls_adapter):
+    response = mock.create_autospec(transport.Response, instance=True)
+    response.status = http_client.OK
+    response.data = _helpers.to_bytes("{}")
     response.headers = {"content-type": "application/json"}
-    mock_session = mock.Mock()
-    mock_session.request.return_value = response
-    mock_create_session.return_value = mock_session
+    mock_request.return_value = response
 
-    initial_request = mock.Mock()
-    _metadata.get(initial_request, "some/path")
+    _metadata.get(mock_request, "some/path")
 
     mock_should_use_mtls.assert_called_once()
-    mock_create_session.assert_called_once()
-    mock_session.request.assert_called_once_with(
-        "GET",
-        "https://metadata.google.internal/computeMetadata/v1/some/path",
-        data=None,
+    mock_mds_mtls_adapter.assert_called_once()
+    mock_request.assert_called_once_with(
+        url="https://metadata.google.internal/computeMetadata/v1/some/path",
+        method="GET",
         headers=_metadata._METADATA_HEADERS,
         timeout=_metadata._METADATA_DEFAULT_TIMEOUT,
     )
+
+
+@pytest.mark.parametrize(
+    "mds_mode, metadata_host, expect_exception",
+    [
+        (_metadata._mtls.MdsMtlsMode.STRICT, _metadata._GCE_DEFAULT_HOST, False),
+        (_metadata._mtls.MdsMtlsMode.STRICT, "custom.host", True),
+        (_metadata._mtls.MdsMtlsMode.NONE, "custom.host", False),
+        (_metadata._mtls.MdsMtlsMode.DEFAULT, _metadata._GCE_DEFAULT_HOST, False),
+    ],
+)
+@mock.patch("google.auth.compute_engine._mtls._parse_mds_mode")
+def test_validate_gce_mds_configured_environment(
+    mock_parse_mds_mode, mds_mode, metadata_host, expect_exception
+):
+    mock_parse_mds_mode.return_value = mds_mode
+    with mock.patch(
+        "google.auth.compute_engine._metadata._GCE_METADATA_HOST", new=metadata_host
+    ):
+        if expect_exception:
+            with pytest.raises(exceptions.MutualTLSChannelError):
+                _metadata._validate_gce_mds_configured_environment()
+        else:
+            _metadata._validate_gce_mds_configured_environment()
+    mock_parse_mds_mode.assert_called_once()
+
+
+@mock.patch("google.auth.compute_engine._mtls.MdsMtlsAdapter")
+def test__prepare_request_for_mds_mtls_session_exists(mock_mds_mtls_adapter):
+    mock_session = mock.create_autospec(requests.Session)
+    request = google_auth_requests.Request(mock_session)
+    new_request = _metadata._prepare_request_for_mds(request, use_mtls=True)
+
+    mock_mds_mtls_adapter.assert_called_once()
+    assert mock_session.mount.call_count == len(_metadata._GCE_DEFAULT_MDS_HOSTS)
+    assert new_request is request
+
+
+@mock.patch("google.auth.compute_engine._mtls.MdsMtlsAdapter")
+def test__prepare_request_for_mds_mtls_no_session(mock_mds_mtls_adapter):
+    request = google_auth_requests.Request(None)
+    # Explicitly set session to None to avoid a session being created in the Request constructor.
+    request.session = None
+
+    with mock.patch("requests.Session") as mock_session_class:
+        new_request = _metadata._prepare_request_for_mds(request, use_mtls=True)
+
+        mock_session_class.assert_called_once()
+        mock_mds_mtls_adapter.assert_called_once()
+        assert new_request.session.mount.call_count == len(
+            _metadata._GCE_DEFAULT_MDS_HOSTS
+        )
+        assert new_request is request

tests/compute_engine/test__mtls.py

@@ -19,6 +19,7 @@
 
 import mock
 import pytest  # type: ignore
+import requests
 
 from google.auth import environment_vars, exceptions
 from google.auth.compute_engine import _mtls
@@ -127,15 +128,14 @@ def test_mds_mtls_adapter_init(mock_ssl_context, mock_mds_mtls_config):
     )
 
 
-@mock.patch("requests.Session")
-@mock.patch("google.auth.compute_engine._mtls.MdsMtlsAdapter")
-def test_create_session(mock_adapter, mock_session, mock_mds_mtls_config):
-    session_instance = mock_session.return_value
-    session = _mtls.create_session(mock_mds_mtls_config)
-    assert session is session_instance
-    mock_adapter.assert_called_once_with(mock_mds_mtls_config)
-    session_instance.mount.assert_called_once_with(
-        "https://", mock_adapter.return_value
+@mock.patch("ssl.create_default_context")
+@mock.patch("requests.adapters.HTTPAdapter.init_poolmanager")
+def test_mds_mtls_adapter_init_poolmanager(
+    mock_init_poolmanager, mock_ssl_context, mock_mds_mtls_config
+):
+    adapter = _mtls.MdsMtlsAdapter(mock_mds_mtls_config)
+    mock_init_poolmanager.assert_called_with(
+        10, 10, block=False, ssl_context=adapter.ssl_context
     )
 
 
@@ -149,3 +149,23 @@ def test_mds_mtls_adapter_proxy_manager_for(
     mock_proxy_manager_for.assert_called_once_with(
         "test_proxy", ssl_context=adapter.ssl_context
     )
+
+
+@mock.patch("ssl.create_default_context")
+def test_mds_mtls_adapter_session_request(mock_ssl_context, mock_mds_mtls_config):
+    adapter = _mtls.MdsMtlsAdapter(mock_mds_mtls_config)
+    session = requests.Session()
+    session.mount("https://", adapter)
+
+    # Mock the adapter's send method to avoid actual network requests
+    adapter.send = mock.Mock()
+    response = requests.Response()
+    response.status_code = 200
+    adapter.send.return_value = response
+
+    # Make a request
+    response = session.get("https://example.com")
+
+    # Assert that the request was successful
+    assert response.status_code == 200
+    adapter.send.assert_called_once()