fix(spanner): synchronize token refresh with lock and optimize xor_bytes

Author: sagnghos

packages/google-cloud-spanner/google/cloud/spanner_v1/omni/credentials.py

@@ -22,6 +22,7 @@
 import base64
 import datetime
 import logging
+import threading
 from collections import namedtuple
 from typing import Any, Callable, MutableMapping, Optional, Sequence
 
@@ -261,6 +262,7 @@ def __init__(
 
         self.token: Optional[str] = None
         self.expiry: Optional[datetime.datetime] = None
+        self._lock = threading.Lock()
 
     def init_channel(
         self,
@@ -302,66 +304,81 @@ def create_async_auth_interceptor(
         """Creates async gRPC interceptors that attach the Bearer token."""
         return self.create_async_auth_interceptors()
 
-    def refresh(self, request: Any = None) -> None:
+    def _perform_refresh_token(self, request: Any = None) -> None:
         """Refreshes the access token by performing the OPAQUE login flow with Spanner Omni.
 
         Args:
             request (Any, optional): Unused; part of google.auth.credentials.Credentials interface.
         """
-        login_channel = None
-        try:
-            if self.use_plain_text:
-                login_channel = grpc.insecure_channel(self.target)
-            elif self.ssl_credentials is not None:
-                login_channel = grpc.secure_channel(self.target, self.ssl_credentials)
-            elif self.ca_certificate:
-                with open(self.ca_certificate, "rb") as f:
-                    ca_cert = f.read()
-                if self.client_certificate and self.client_key:
-                    with open(self.client_certificate, "rb") as f:
-                        client_cert = f.read()
-                    with open(self.client_key, "rb") as f:
-                        private_key = f.read()
-                    ssl_creds = grpc.ssl_channel_credentials(
-                        root_certificates=ca_cert,
-                        private_key=private_key,
-                        certificate_chain=client_cert,
+        with self._lock:
+            if self.valid:
+                return
+            login_channel = None
+            try:
+                if self.use_plain_text:
+                    login_channel = grpc.insecure_channel(self.target)
+                elif self.ssl_credentials is not None:
+                    login_channel = grpc.secure_channel(
+                        self.target, self.ssl_credentials
                     )
-                elif self.client_certificate or self.client_key:
-                    raise ValueError(
-                        "Both client_certificate and client_key must be provided for mTLS"
+                elif self.ca_certificate:
+                    with open(self.ca_certificate, "rb") as f:
+                        ca_cert = f.read()
+                    if self.client_certificate and self.client_key:
+                        with open(self.client_certificate, "rb") as f:
+                            client_cert = f.read()
+                        with open(self.client_key, "rb") as f:
+                            private_key = f.read()
+                        ssl_creds = grpc.ssl_channel_credentials(
+                            root_certificates=ca_cert,
+                            private_key=private_key,
+                            certificate_chain=client_cert,
+                        )
+                    elif self.client_certificate or self.client_key:
+                        raise ValueError(
+                            "Both client_certificate and client_key must be provided for mTLS"
+                        )
+                    else:
+                        ssl_creds = grpc.ssl_channel_credentials(
+                            root_certificates=ca_cert
+                        )
+                    login_channel = grpc.secure_channel(self.target, ssl_creds)
+                else:
+                    login_channel = grpc.secure_channel(
+                        self.target, grpc.ssl_channel_credentials()
                     )
+
+                client = LoginClient(login_channel)
+                proto_token = client.login(self.username, self._password)
+
+                token_bytes = proto_token.SerializeToString()
+                self.token = base64.b64encode(token_bytes).decode("ascii")
+
+                if proto_token.HasField("expiration_time"):
+                    seconds = proto_token.expiration_time.seconds
+                    nanos = proto_token.expiration_time.nanos
+                    self.expiry = datetime.datetime.fromtimestamp(
+                        seconds + nanos / 1e9, tz=datetime.timezone.utc
+                    ).replace(tzinfo=None)
                 else:
-                    ssl_creds = grpc.ssl_channel_credentials(root_certificates=ca_cert)
-                login_channel = grpc.secure_channel(self.target, ssl_creds)
-            else:
-                login_channel = grpc.secure_channel(
-                    self.target, grpc.ssl_channel_credentials()
-                )
-
-            client = LoginClient(login_channel)
-            proto_token = client.login(self.username, self._password)
-
-            token_bytes = proto_token.SerializeToString()
-            self.token = base64.b64encode(token_bytes).decode("ascii")
-
-            if proto_token.HasField("expiration_time"):
-                seconds = proto_token.expiration_time.seconds
-                nanos = proto_token.expiration_time.nanos
-                self.expiry = datetime.datetime.fromtimestamp(
-                    seconds + nanos / 1e9, tz=datetime.timezone.utc
-                ).replace(tzinfo=None)
-            else:
-                self.expiry = datetime.datetime.now(datetime.timezone.utc).replace(
-                    tzinfo=None
-                ) + datetime.timedelta(hours=1)
-        except Exception as e:
-            raise google.auth.exceptions.RefreshError(
-                f"Failed to login to Spanner Omni: {e}"
-            ) from e
-        finally:
-            if login_channel is not None:
-                login_channel.close()
+                    self.expiry = datetime.datetime.now(datetime.timezone.utc).replace(
+                        tzinfo=None
+                    ) + datetime.timedelta(hours=1)
+            except Exception as e:
+                raise google.auth.exceptions.RefreshError(
+                    f"Failed to login to Spanner Omni: {e}"
+                ) from e
+            finally:
+                if login_channel is not None:
+                    login_channel.close()
+
+    def refresh(self, request: Any = None) -> None:
+        """Refreshes the access token.
+
+        Args:
+            request (Any, optional): Unused; part of google.auth.credentials.Credentials interface.
+        """
+        self._perform_refresh_token(request)
 
     def apply(
         self, headers: MutableMapping[str, str], token: Optional[str] = None

packages/google-cloud-spanner/google/cloud/spanner_v1/omni/opaque.py

@@ -215,9 +215,7 @@ def xor_bytes(a: bytes, b: bytes) -> bytes:
     """Computes bitwise XOR of two equal-length byte sequences."""
     if len(a) != len(b):
         raise ValueError(f"Byte sequences must have equal length: {len(a)} != {len(b)}")
-    if len(a) == 0:
-        raise ValueError("Byte sequences must not be empty")
-    return bytes(x ^ y for x, y in zip(a, b))
+    return (int.from_bytes(a, "big") ^ int.from_bytes(b, "big")).to_bytes(len(a), "big")
 
 
 def concat(*arrays: bytes) -> bytes:

packages/google-cloud-spanner/tests/unit/omni/test_credentials.py

@@ -120,6 +120,19 @@ def test_refresh_failure_wraps_in_refresh_error(self):
             with self.assertRaises(google.auth.exceptions.RefreshError):
                 creds.refresh()
 
+    def test_refresh_skips_when_valid_inside_lock(self):
+        creds = SpannerOmniCredentials("user", "pass", "localhost:9010")
+        creds.token = "existing_valid_token"
+        creds.expiry = datetime.datetime.now(datetime.timezone.utc).replace(
+            tzinfo=None
+        ) + datetime.timedelta(hours=1)
+
+        with mock.patch(
+            "google.cloud.spanner_v1.omni.credentials.LoginClient"
+        ) as mock_login_client_cls:
+            creds.refresh()
+            mock_login_client_cls.assert_not_called()
+
     def test_apply_and_before_request(self):
         creds = SpannerOmniCredentials("user", "pass", "localhost:9010")
         headers = {}

packages/google-cloud-spanner/tests/unit/omni/test_opaque.py

@@ -60,7 +60,7 @@ def test_xor_bytes(self):
             ),
             (b"abc", b"defghi", True),
             (b"abcdefghi", b"jklmnop", True),
-            (b"", b"", True),
+            (b"", b"", False),
         ]
         for a, b, want_err in tests:
             if want_err: