Skip to content

Commit 4e44c84

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
fix(client): skip mTLS URL rewriting for multi-regional .rep.googleapis.com endpoints and honor GOOGLE_API_USE_MTLS_ENDPOINT in Google Auth sessions
Multi-regional Vertex AI endpoints (`aiplatform.us.rep.googleapis.com` and `aiplatform.eu.rep.googleapis.com`) do not publish `.rep.mtls.googleapis.com` DNS records. When `GOOGLE_API_USE_CLIENT_CERTIFICATE=true` is set (such as on Vertex AI Agent Engine with `AGENT_IDENTITY=true`), rewriting these URLs in `to_mtls_url()` causes requests to fail with `NXDOMAIN` (`Name or service not known`). Additionally, `_request_once`, `_async_request_once`, `download_file`, and `async_download_file` previously checked `session._is_mtls` directly when using `AuthorizedSession` / `AsyncAuthorizedSession`, bypassing `_use_mtls_endpoint()` and ignoring `GOOGLE_API_USE_MTLS_ENDPOINT=never`. - Update `to_mtls_url()` to leave multi-regional `.us.rep.` and `.eu.rep.` endpoints unchanged while continuing to rewrite single-region `.rep.` endpoints. - Gate `to_mtls_url()` in the synchronous and asynchronous Google Auth session paths via `self._use_mtls_endpoint(client_cert_loaded=...)`. PiperOrigin-RevId: 1004213757
1 parent d6af6cd commit 4e44c84

2 files changed

Lines changed: 124 additions & 8 deletions

File tree

‎google/genai/_api_client.py‎

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -230,17 +230,31 @@ def join_url_path(base_url: str, path: str) -> str:
230230
return urlunparse(parsed_base._replace(path=base_path + '/' + path))
231231

232232

233+
_MULTI_REGIONAL_REP_SUFFIXES = tuple(
234+
f'.{loc}.rep.{domain}'
235+
for loc in _MULTI_REGIONAL_LOCATIONS
236+
for domain in ('sandbox.googleapis.com', 'googleapis.com')
237+
)
238+
239+
233240
def to_mtls_url(url: str) -> str:
234241
"""Rewrites a googleapis.com URL to its mTLS endpoint.
235242
236243
For example, `aiplatform.googleapis.com` becomes
237244
`aiplatform.mtls.googleapis.com` and `foo.sandbox.googleapis.com` becomes
238-
`foo.mtls.sandbox.googleapis.com`. URLs on other hosts, and URLs that already
239-
point at an mTLS endpoint, are returned unchanged.
245+
`foo.mtls.sandbox.googleapis.com`. URLs on other hosts, URLs that already
246+
point at an mTLS endpoint, and multi-regional `.rep.googleapis.com` endpoints
247+
(which do not publish `.rep.mtls.googleapis.com` DNS records) are returned
248+
unchanged.
240249
"""
241250
parsed = urlparse(url)
242251
netloc = parsed.netloc
243-
if netloc.endswith(('.mtls.googleapis.com', '.mtls.sandbox.googleapis.com')):
252+
hostname = parsed.hostname or ''
253+
if hostname.endswith(
254+
('.mtls.googleapis.com', '.mtls.sandbox.googleapis.com')
255+
):
256+
return url
257+
if hostname.endswith(_MULTI_REGIONAL_REP_SUFFIXES):
244258
return url
245259
for domain in ('sandbox.googleapis.com', 'googleapis.com'):
246260
if netloc.endswith('.' + domain):
@@ -1600,7 +1614,9 @@ def _request_once(
16001614
self._authorized_session.configure_mtls_channel(
16011615
client_cert_source
16021616
) # type: ignore[no-untyped-call]
1603-
if self._authorized_session._is_mtls:
1617+
if self._use_mtls_endpoint(
1618+
client_cert_loaded=bool(self._authorized_session._is_mtls)
1619+
):
16041620
url = to_mtls_url(url)
16051621
response = self._authorized_session.request( # type: ignore[no-untyped-call]
16061622
method=http_request.method.upper(),
@@ -1677,7 +1693,9 @@ async def _async_request_once(
16771693
await session.configure_mtls_channel( # type: ignore[union-attr]
16781694
client_cert_source
16791695
)
1680-
if session._is_mtls: # type: ignore[union-attr]
1696+
if self._use_mtls_endpoint(
1697+
client_cert_loaded=bool(session._is_mtls) # type: ignore[union-attr]
1698+
):
16811699
url = to_mtls_url(url)
16821700
elif self._aiohttp_use_mtls_endpoint:
16831701
url = to_mtls_url(url)
@@ -1749,7 +1767,9 @@ async def _async_request_once(
17491767
await session.configure_mtls_channel( # type: ignore[union-attr]
17501768
client_cert_source
17511769
)
1752-
if session._is_mtls: # type: ignore[union-attr]
1770+
if self._use_mtls_endpoint(
1771+
client_cert_loaded=bool(session._is_mtls) # type: ignore[union-attr]
1772+
):
17531773
url = to_mtls_url(url)
17541774
elif self._aiohttp_use_mtls_endpoint:
17551775
url = to_mtls_url(url)
@@ -2162,7 +2182,9 @@ def _write_chunks(chunks: Iterator[bytes]) -> None:
21622182
self._authorized_session.configure_mtls_channel(
21632183
client_cert_source
21642184
) # type: ignore[no-untyped-call]
2165-
if self._authorized_session._is_mtls:
2185+
if self._use_mtls_endpoint(
2186+
client_cert_loaded=bool(self._authorized_session._is_mtls)
2187+
):
21662188
url = to_mtls_url(url)
21672189
if destination is not None:
21682190
response = self._authorized_session.request( # type: ignore[no-untyped-call]
@@ -2531,7 +2553,9 @@ async def _write_chunks(chunks: AsyncIterator[bytes]) -> None:
25312553
await session.configure_mtls_channel( # type: ignore[union-attr]
25322554
client_cert_source
25332555
)
2534-
if session._is_mtls: # type: ignore[union-attr]
2556+
if self._use_mtls_endpoint(
2557+
client_cert_loaded=bool(session._is_mtls) # type: ignore[union-attr]
2558+
):
25352559
url = to_mtls_url(url)
25362560
elif self._aiohttp_use_mtls_endpoint:
25372561
url = to_mtls_url(url)

‎google/genai/tests/client/test_client_initialization.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2074,6 +2074,26 @@ async def run():
20742074
"https://us-central1-aiplatform.googleapis.com/v1/models",
20752075
"https://us-central1-aiplatform.mtls.googleapis.com/v1/models",
20762076
),
2077+
(
2078+
"https://secretmanager.us-central1.rep.googleapis.com/v1/secrets",
2079+
"https://secretmanager.us-central1.rep.mtls.googleapis.com/v1/secrets",
2080+
),
2081+
(
2082+
"https://aiplatform.us.rep.googleapis.com/v1/models",
2083+
"https://aiplatform.us.rep.googleapis.com/v1/models",
2084+
),
2085+
(
2086+
"https://aiplatform.eu.rep.googleapis.com/v1/models",
2087+
"https://aiplatform.eu.rep.googleapis.com/v1/models",
2088+
),
2089+
(
2090+
"https://aiplatform.us.rep.sandbox.googleapis.com/v1/models",
2091+
"https://aiplatform.us.rep.sandbox.googleapis.com/v1/models",
2092+
),
2093+
(
2094+
"https://aiplatform.eu.rep.sandbox.googleapis.com/v1/models",
2095+
"https://aiplatform.eu.rep.sandbox.googleapis.com/v1/models",
2096+
),
20772097
(
20782098
"wss://aiplatform.googleapis.com/ws",
20792099
"wss://aiplatform.mtls.googleapis.com/ws",
@@ -2392,3 +2412,75 @@ def test_google_auth_async_used_by_default(
23922412
api_client.has_aiohttp = True
23932413
client = _mtls_test_client()
23942414
assert client._api_client._use_google_auth_async()
2415+
2416+
2417+
@pytest.mark.parametrize("call", ["request", "download"])
2418+
def test_google_auth_sync_respects_mtls_endpoint_env_override(
2419+
mock_client_cert, monkeypatch, call
2420+
):
2421+
monkeypatch.setenv("GOOGLE_API_USE_MTLS_ENDPOINT", "never")
2422+
client = _mtls_test_client()
2423+
base_client = client._api_client
2424+
assert base_client._use_google_auth_sync()
2425+
2426+
mock_response = mock.MagicMock()
2427+
mock_response.status_code = 200
2428+
mock_response.headers = {}
2429+
mock_response.text = "{}"
2430+
mock_response.content = b""
2431+
mock_session = mock.MagicMock()
2432+
mock_session._is_mtls = True
2433+
mock_session.request = mock.Mock(return_value=mock_response)
2434+
base_client._authorized_session = mock_session
2435+
2436+
if call == "download":
2437+
base_client.download_file("files/abc:download")
2438+
else:
2439+
base_client.request("post", "models/gemini:generateContent", {})
2440+
2441+
assert mock_session.request.call_args.kwargs["url"].startswith(
2442+
"https://us-central1-aiplatform.googleapis.com/"
2443+
)
2444+
2445+
2446+
@requires_aiohttp
2447+
@pytest.mark.asyncio
2448+
@pytest.mark.parametrize("call", ["request", "stream", "download"])
2449+
async def test_google_auth_async_respects_mtls_endpoint_env_override(
2450+
mock_client_cert, mock_google_auth_aio, monkeypatch, call
2451+
):
2452+
from google.auth.transport import mtls
2453+
2454+
api_client.has_aiohttp = True
2455+
monkeypatch.setenv("GOOGLE_API_USE_MTLS_ENDPOINT", "never")
2456+
monkeypatch.setattr(mtls, "default_client_cert_source", lambda: None)
2457+
client = _mtls_test_client()
2458+
base_client = client._api_client
2459+
assert base_client._use_google_auth_async()
2460+
2461+
mock_response = mock.create_autospec(aiohttp.ClientResponse, instance=True)
2462+
mock_response.status = 200
2463+
mock_response.headers = {}
2464+
mock_response.text.return_value = "{}"
2465+
mock_response.read.return_value = b""
2466+
mock_session = mock.MagicMock()
2467+
mock_session._is_mtls = True
2468+
mock_session.configure_mtls_channel = mock.AsyncMock()
2469+
mock_session.request = mock.AsyncMock(return_value=mock_response)
2470+
2471+
with mock.patch.object(
2472+
base_client, "_get_aiohttp_session", return_value=mock_session
2473+
):
2474+
if call == "download":
2475+
await base_client.async_download_file("files/abc:download")
2476+
else:
2477+
http_request = base_client._build_request(
2478+
"post", "models/gemini:generateContent", {}, None
2479+
)
2480+
await base_client._async_request_once(
2481+
http_request, stream=call == "stream"
2482+
)
2483+
2484+
assert mock_session.request.call_args.kwargs["url"].startswith(
2485+
"https://us-central1-aiplatform.googleapis.com/"
2486+
)

0 commit comments

Comments
 (0)