From 8f7c56d7316c9e7a6ae1a7e6ed42826f468aaf08 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 28 Jul 2026 19:55:39 +0000 Subject: [PATCH 1/5] feat(bigtable): authenticate accelerator RPCs with a per-daemon stdin handshake secret --- .../bigtable/data/_accelerator/_daemon.py | 28 ++++- .../data/_async/_accelerator_client.py | 9 +- .../cloud/bigtable/data/_async/client.py | 14 ++- .../data/_sync_autogen/_accelerator_client.py | 7 +- .../bigtable/data/_sync_autogen/client.py | 12 ++- .../unit/data/test_accelerator_daemon.py | 102 ++++++++++++++++++ 6 files changed, 161 insertions(+), 11 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index 4e1b0e4f85f5..c15c4e28298f 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -25,6 +25,7 @@ import json import os +import secrets import shutil import socket import subprocess @@ -107,6 +108,7 @@ def __init__( self._uds_path: str | None = None self._log_path: str | None = None self._proc: subprocess.Popen[bytes] | None = None + self._auth_secret: str | None = None def __enter__(self) -> "AcceleratorDaemon": """Start the daemon on ``with`` entry and return it.""" @@ -139,6 +141,12 @@ def pid(self) -> int: def is_running(self) -> bool: return self._proc is not None and self._proc.poll() is None + @property + def auth_secret(self) -> str: + if self._auth_secret is None: + raise RuntimeError("AcceleratorDaemon has not been started with a secret") + return self._auth_secret + def start(self) -> None: """Spawn the daemon and wait for the UDS to become connectable. @@ -197,6 +205,22 @@ def start(self) -> None: log_file.close() except OSError: pass + + # Mint a 256-bit secret and write it to the daemon's stdin before it + # binds the socket. The daemon reads this line synchronously before + # accepting connections, then validates it on every RPC via metadata. + self._auth_secret = secrets.token_urlsafe(32) + try: + self._proc.stdin.write(f"{self._auth_secret}\n".encode()) # type: ignore[union-attr] + self._proc.stdin.flush() # type: ignore[union-attr] + except (OSError, ValueError) as exc: + self._force_kill() + self._cleanup_tempdir() + raise RuntimeError( + "Failed to send auth secret to accelerator daemon (process died): " + f"{exc}" + ) from exc + try: self._wait_until_ready(self._startup_timeout) except BaseException: @@ -261,9 +285,7 @@ def read_identity(self) -> dict[str, Any]: "cannot verify its identity" ) from exc except (OSError, ValueError) as exc: - raise RuntimeError( - f"Failed to read {_IDENTITY_FILENAME}: {exc}" - ) from exc + raise RuntimeError(f"Failed to read {_IDENTITY_FILENAME}: {exc}") from exc finally: try: os.unlink(path) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_client.py index a343f9ec7176..b1ab1962695f 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_accelerator_client.py @@ -49,8 +49,9 @@ class _AsyncAcceleratorClient: (``_accelerator/_routing.py``) decides which calls reach this object. """ - def __init__(self, uds_path: str): + def __init__(self, uds_path: str, auth_secret: str): self._uds_path = uds_path + self._metadata = (("x-accelerator-token", auth_secret),) self._channel = insecure_channel(f"unix://{uds_path}") self._mutate_row_stub = self._channel.unary_unary( _MUTATE_ROW_METHOD, @@ -71,7 +72,9 @@ def uds_path(self) -> str: async def mutate_row( self, request: MutateRowRequest, *, timeout: float | None = None ) -> MutateRowResponse: - return await self._mutate_row_stub(request, timeout=timeout) + return await self._mutate_row_stub( + request, timeout=timeout, metadata=self._metadata + ) @CrossSync.convert async def read_rows( @@ -84,7 +87,7 @@ async def read_rows( Shape matches what ``_gapic_client.read_rows`` returns so the existing chunk-merging machinery in ``_read_rows.py`` works unchanged. """ - return self._read_rows_stub(request, timeout=timeout) + return self._read_rows_stub(request, timeout=timeout, metadata=self._metadata) @CrossSync.convert async def close(self) -> None: diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index f667bb32d11d..5fe7268cc006 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -1448,7 +1448,19 @@ def _start_accelerator(self) -> None: # the daemon's independent ADC resolution could land on a different # principal. Verify before routing any RPC through it. self._verify_daemon_identity(server) - self._accelerator_client = AcceleratorClientType(server.uds_path) + self._accelerator_client = AcceleratorClientType( + server.uds_path, server.auth_secret + ) + except _AcceleratorUnverified: + server.close() + warnings.warn( + "Accelerator disabled: could not verify that the daemon " + "resolves the same identity as this client; using the native " + "client instead.", + RuntimeWarning, + stacklevel=2, + ) + return except BaseException: server.close() raise diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_client.py index 3b5edd274060..239904956dcb 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_accelerator_client.py @@ -45,8 +45,9 @@ class _AcceleratorClient: (``_accelerator/_routing.py``) decides which calls reach this object. """ - def __init__(self, uds_path: str): + def __init__(self, uds_path: str, auth_secret: str): self._uds_path = uds_path + self._metadata = (("x-accelerator-token", auth_secret),) self._channel = insecure_channel(f"unix://{uds_path}") self._mutate_row_stub = self._channel.unary_unary( _MUTATE_ROW_METHOD, @@ -66,7 +67,7 @@ def uds_path(self) -> str: def mutate_row( self, request: MutateRowRequest, *, timeout: float | None = None ) -> MutateRowResponse: - return self._mutate_row_stub(request, timeout=timeout) + return self._mutate_row_stub(request, timeout=timeout, metadata=self._metadata) def read_rows(self, request: ReadRowsRequest, *, timeout: float | None = None): """Open the server-streaming ReadRows RPC against the daemon. @@ -75,7 +76,7 @@ def read_rows(self, request: ReadRowsRequest, *, timeout: float | None = None): in async, ``for`` in sync) to consume ``ReadRowsResponse`` messages. Shape matches what ``_gapic_client.read_rows`` returns so the existing chunk-merging machinery in ``_read_rows.py`` works unchanged.""" - return self._read_rows_stub(request, timeout=timeout) + return self._read_rows_stub(request, timeout=timeout, metadata=self._metadata) def close(self) -> None: self._channel.close() diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index a1a0b1f4aaf7..e3b925b063b1 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -1141,7 +1141,17 @@ def _start_accelerator(self) -> None: try: server.start() self._verify_daemon_identity(server) - self._accelerator_client = AcceleratorClientType(server.uds_path) + self._accelerator_client = AcceleratorClientType( + server.uds_path, server.auth_secret + ) + except _AcceleratorUnverified: + server.close() + warnings.warn( + "Accelerator disabled: could not verify that the daemon resolves the same identity as this client; using the native client instead.", + RuntimeWarning, + stacklevel=2, + ) + return except BaseException: server.close() raise diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index bace94ab91a9..8d5742c7f0a9 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -102,6 +102,108 @@ def test_missing_file_raises(self, tmp_path): daemon.read_identity() +class TestAuthSecret: + def test_raises_before_start(self, tmp_path): + daemon = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="has not been started"): + _ = daemon.auth_secret + + def test_written_to_stdin_after_popen(self, tmp_path, monkeypatch): + """start() mints a secret and writes 'secret\n' to the daemon's stdin.""" + written = [] + + class FakeStdin: + def write(self, data): + written.append(data) + + def flush(self): + pass + + class FakeProc: + pid = 12345 + stdin = FakeStdin() + + def poll(self): + return None + + def fake_popen(argv, **kwargs): + raise _StopStart() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + daemon = _make_daemon(tmp_path) + with pytest.raises(_StopStart): + daemon.start() + # Popen raises _StopStart before secret is written; check that the + # secret is written when Popen succeeds by using a real pipe below. + + def test_secret_written_to_stdin(self, tmp_path, monkeypatch): + """The secret is urlsafe and ends with a newline on the wire.""" + written = bytearray() + + class FakeStdin: + def write(self, data): + written.extend(data) + + def flush(self): + pass + + class FakeProc: + pid = 12345 + stdin = FakeStdin() + + def poll(self): + return None + + def fake_popen(argv, **kwargs): + return FakeProc() + + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + daemon = _make_daemon(tmp_path) + # _wait_until_ready will fail immediately (process is fake); we just + # need to get past the stdin-write step, so patch _wait_until_ready. + monkeypatch.setattr(daemon, "_wait_until_ready", lambda timeout: None) + daemon.start() + payload = written.decode() + assert payload.endswith("\n"), "secret payload must end with newline" + token = payload.rstrip("\n") + assert token == daemon.auth_secret + # token_urlsafe(32) produces ≥32 chars of base64url characters. + assert len(token) >= 32 + import re + + assert re.fullmatch(r"[A-Za-z0-9_\-]+", token), "expected urlsafe token" + + def test_broken_pipe_kills_and_reraises(self, tmp_path, monkeypatch): + """A broken pipe during secret write force-kills the daemon and raises.""" + + class FakeStdin: + def write(self, data): + raise OSError("broken pipe") + + def flush(self): + pass + + class FakeProc: + pid = 12345 + stdin = FakeStdin() + + def poll(self): + return None + + def fake_popen(argv, **kwargs): + return FakeProc() + + killed = [] + monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) + monkeypatch.setattr( + AcceleratorDaemon, "_force_kill", lambda self: killed.append(True) + ) + daemon = _make_daemon(tmp_path) + with pytest.raises(RuntimeError, match="auth secret"): + daemon.start() + assert killed, "_force_kill should have been called on broken pipe" + + class _StopStart(Exception): """Sentinel to abort AcceleratorDaemon.start() right after Popen so the test never waits on a real socket.""" From 55f3f9807ba96d1c1491159d88395b8297742c62 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 19:41:06 +0000 Subject: [PATCH 2/5] fix(bigtable): clear auth_secret on teardown; fix hollow test and module-level import Change-Id: I533b2c7e1bb78ff725a74db5b00bb701a082b1df --- .../bigtable/data/_accelerator/_daemon.py | 1 + .../unit/data/test_accelerator_daemon.py | 31 +------------------ 2 files changed, 2 insertions(+), 30 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index c15c4e28298f..bbd44d26ca85 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -381,3 +381,4 @@ def _cleanup_tempdir(self) -> None: self._tempdir = None self._uds_path = None self._log_path = None + self._auth_secret = None diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py index 8d5742c7f0a9..e295997e6f50 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_daemon.py @@ -14,6 +14,7 @@ import json import os +import re import unittest.mock import pytest @@ -108,34 +109,6 @@ def test_raises_before_start(self, tmp_path): with pytest.raises(RuntimeError, match="has not been started"): _ = daemon.auth_secret - def test_written_to_stdin_after_popen(self, tmp_path, monkeypatch): - """start() mints a secret and writes 'secret\n' to the daemon's stdin.""" - written = [] - - class FakeStdin: - def write(self, data): - written.append(data) - - def flush(self): - pass - - class FakeProc: - pid = 12345 - stdin = FakeStdin() - - def poll(self): - return None - - def fake_popen(argv, **kwargs): - raise _StopStart() - - monkeypatch.setattr(_daemon.subprocess, "Popen", fake_popen) - daemon = _make_daemon(tmp_path) - with pytest.raises(_StopStart): - daemon.start() - # Popen raises _StopStart before secret is written; check that the - # secret is written when Popen succeeds by using a real pipe below. - def test_secret_written_to_stdin(self, tmp_path, monkeypatch): """The secret is urlsafe and ends with a newline on the wire.""" written = bytearray() @@ -169,8 +142,6 @@ def fake_popen(argv, **kwargs): assert token == daemon.auth_secret # token_urlsafe(32) produces ≥32 chars of base64url characters. assert len(token) >= 32 - import re - assert re.fullmatch(r"[A-Za-z0-9_\-]+", token), "expected urlsafe token" def test_broken_pipe_kills_and_reraises(self, tmp_path, monkeypatch): From 09167bae0e78b9f7d8625350ddd57e2cc7bd70fb Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 20:08:31 +0000 Subject: [PATCH 3/5] fix(bigtable): guard stdin write with explicit None checks Change-Id: I33e67c9868b186353eff3857406b62ffbfa97eb2 --- .../google/cloud/bigtable/data/_accelerator/_daemon.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py index bbd44d26ca85..9958df4b1a04 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_accelerator/_daemon.py @@ -210,9 +210,15 @@ def start(self) -> None: # binds the socket. The daemon reads this line synchronously before # accepting connections, then validates it on every RPC via metadata. self._auth_secret = secrets.token_urlsafe(32) + if self._proc is None or self._proc.stdin is None: + self._force_kill() + self._cleanup_tempdir() + raise RuntimeError( + "Accelerator daemon process or its stdin is not initialized." + ) try: - self._proc.stdin.write(f"{self._auth_secret}\n".encode()) # type: ignore[union-attr] - self._proc.stdin.flush() # type: ignore[union-attr] + self._proc.stdin.write(f"{self._auth_secret}\n".encode()) + self._proc.stdin.flush() except (OSError, ValueError) as exc: self._force_kill() self._cleanup_tempdir() From 4334925623b0b3656bfb9ba99ff26ca1b96183a2 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 20:27:30 +0000 Subject: [PATCH 4/5] fix(bigtable): use _AcceleratorIdentityError in except clause; re-raise for explicit-mode Change-Id: I030357650201f8681a31d3135cbbff726a4aa650 --- .../google/cloud/bigtable/data/_async/client.py | 11 ++--------- .../cloud/bigtable/data/_sync_autogen/client.py | 9 ++------- 2 files changed, 4 insertions(+), 16 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index 5fe7268cc006..48c06e91196a 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -1451,16 +1451,9 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorUnverified: + except _AcceleratorIdentityError: server.close() - warnings.warn( - "Accelerator disabled: could not verify that the daemon " - "resolves the same identity as this client; using the native " - "client instead.", - RuntimeWarning, - stacklevel=2, - ) - return + raise except BaseException: server.close() raise diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index e3b925b063b1..517052189031 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -1144,14 +1144,9 @@ def _start_accelerator(self) -> None: self._accelerator_client = AcceleratorClientType( server.uds_path, server.auth_secret ) - except _AcceleratorUnverified: + except _AcceleratorIdentityError: server.close() - warnings.warn( - "Accelerator disabled: could not verify that the daemon resolves the same identity as this client; using the native client instead.", - RuntimeWarning, - stacklevel=2, - ) - return + raise except BaseException: server.close() raise From f3dd22eb2345e71a58dcf300a410aa40bf3d297e Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Tue, 13 Oct 2026 20:29:28 +0000 Subject: [PATCH 5/5] test(bigtable): verify identity error propagates correctly, not NameError Change-Id: I8d4bf5acf04b72af47e6d65829b41d7d8d800165 --- .../unit/data/test_accelerator_enablement.py | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py index 0381cd5ad827..a9b88e171cff 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_accelerator_enablement.py @@ -23,7 +23,12 @@ import pytest -from google.cloud.bigtable.data._async.client import _DataApiTargetAsync +import unittest.mock + +from google.cloud.bigtable.data._async.client import ( + _AcceleratorIdentityError, + _DataApiTargetAsync, +) class _ConcreteTarget(_DataApiTargetAsync): @@ -125,3 +130,25 @@ def _interrupt(): t = _bare_target(emulator=None, start=_interrupt) with pytest.raises(KeyboardInterrupt): t._maybe_start_accelerator(explicit=False) + + def test_identity_error_auto_warns_and_falls_back(self, monkeypatch): + """Identity mismatch on auto path warns and falls back — no NameError.""" + + def _start(): + raise _AcceleratorIdentityError("principal mismatch") + + t = _bare_target(emulator=None, start=_start) + with pytest.warns(RuntimeWarning, match="Accelerator disabled"): + t._maybe_start_accelerator(explicit=False) + assert t._accelerator_client is None + assert t._accelerator_daemon is None + + def test_identity_error_explicit_raises(self, monkeypatch): + """Identity mismatch on explicit path re-raises _AcceleratorIdentityError, not NameError.""" + + def _start(): + raise _AcceleratorIdentityError("principal mismatch") + + t = _bare_target(emulator=None, start=_start) + with pytest.raises(_AcceleratorIdentityError, match="principal mismatch"): + t._maybe_start_accelerator(explicit=True)