fix(spanner): avoid bytearray memory copies and handle premature stream closure

Author: sagnghos

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

@@ -93,13 +93,19 @@ def login(
         try:
             call = self._stub.Login(req_iterator, timeout=timeout)
 
+            def _safe_next(stage: str) -> login_pb2.LoginResponse:
+                try:
+                    return next(call)
+                except StopIteration:
+                    raise ValueError(f"Server closed stream prematurely during {stage}")
+
             # Step 1: Handshake Request
             handshake_req = login_pb2.LoginRequest(
                 username=username,
                 handshake_request=authentication_pb2.PasswordAuthenticationHandshakeRequest(),
             )
             req_iterator.send(handshake_req)
-            handshake_resp = next(call)
+            handshake_resp = _safe_next("handshake")
 
             if not handshake_resp.HasField("handshake_response"):
                 raise ValueError("Failed to receive handshake response from server")
@@ -119,15 +125,15 @@ def login(
             # Step 2: Initial OPAQUE Request
             initial_req = authenticator.initial_request()
             req_iterator.send(initial_req)
-            initial_resp = next(call)
+            initial_resp = _safe_next("initial OPAQUE exchange")
 
             # Step 3: Final OPAQUE Request
             final_req = authenticator.final_request(initial_resp)
             req_iterator.send(final_req)
             req_iterator.close()
 
             # Final Response with AccessToken
-            final_resp = next(call)
+            final_resp = _safe_next("final OPAQUE exchange")
             if not final_resp.HasField("access_token"):
                 raise ValueError(
                     "Server failed to return an access token in final response"

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

@@ -437,7 +437,7 @@ def initial_request(self) -> login_pb2.LoginRequest:
             raise ValueError("Authenticator already used or password not available")
 
         try:
-            blinded_message, blind_scalar = blind(bytes(self._password))
+            blinded_message, blind_scalar = blind(self._password)
             self._blind = bytearray(blind_scalar)
 
             self._client_nonce = nonce()
@@ -497,7 +497,7 @@ def final_request(
             )
 
         try:
-            oprf = finalize(bytes(self._blind), evaluated_message)
+            oprf = finalize(self._blind, evaluated_message)
             stretched_oprf = stretch(oprf, self.hash_parameters)
             randomized_password = extract(concat(oprf, stretched_oprf))
 
@@ -529,12 +529,8 @@ def final_request(
                 server_public_key,
             )
 
-            dh1 = diffie_hellman(
-                bytes(self._client_private_keyshare), server_public_keyshare
-            )
-            dh2 = diffie_hellman(
-                bytes(self._client_private_keyshare), server_public_key
-            )
+            dh1 = diffie_hellman(self._client_private_keyshare, server_public_keyshare)
+            dh2 = diffie_hellman(self._client_private_keyshare, server_public_key)
             dh3 = diffie_hellman(client_private_key, server_public_keyshare)
 
             input_key_material = concat(dh1, dh2, dh3)

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

@@ -277,6 +277,23 @@ def test_login_exception_cancels_call(self):
 
             mock_call.cancel.assert_called_once()
 
+    def test_login_premature_stream_closure(self):
+        mock_call = mock.MagicMock()
+        mock_call.__next__.side_effect = StopIteration
+
+        with mock.patch(
+            "google.cloud.spanner_v1.omni.proto.login_pb2_grpc.LoginServiceStub"
+        ) as mock_stub_cls:
+            mock_stub = mock_stub_cls.return_value
+            mock_stub.Login.return_value = mock_call
+
+            client = LoginClient(self.mock_channel)
+            with self.assertRaises(ValueError) as cm:
+                client.login("user", "pass")
+            self.assertIn(
+                "Server closed stream prematurely during handshake", str(cm.exception)
+            )
+
 
 if __name__ == "__main__":
     unittest.main()