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..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 @@ -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,28 @@ 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) + 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()) + self._proc.stdin.flush() + 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 +291,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) @@ -359,3 +387,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/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..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 @@ -1448,7 +1448,12 @@ 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 _AcceleratorIdentityError: + server.close() + raise 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..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 @@ -1141,7 +1141,12 @@ 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 _AcceleratorIdentityError: + server.close() + raise 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..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 @@ -102,6 +103,78 @@ 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_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 + 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.""" 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)