|
21 | 21 | import datetime |
22 | 22 | import json |
23 | 23 | import ssl |
| 24 | +import struct |
24 | 25 | from typing import Any, Callable, Literal, Optional |
25 | 26 |
|
26 | 27 | from aiofiles.tempfile import TemporaryDirectory |
|
38 | 39 | from google.cloud.sql.connector.connector import _DEFAULT_UNIVERSE_DOMAIN |
39 | 40 | from google.cloud.sql.connector.utils import generate_keys |
40 | 41 | from google.cloud.sql.connector.utils import write_to_file |
| 42 | +import google.cloud.sql.proto.cloud_sql_metadata_exchange_pb2 as connectorspb |
41 | 43 |
|
42 | 44 |
|
43 | 45 | class FakeCredentials: |
@@ -298,3 +300,65 @@ async def generate_ephemeral(self, request: Any) -> web.Response: |
298 | 300 | } |
299 | 301 | } |
300 | 302 | return web.Response(content_type="application/json", body=json.dumps(response)) |
| 303 | + |
| 304 | + |
| 305 | +def metadata_exchange(sock: ssl.SSLSocket) -> None: |
| 306 | + """ |
| 307 | + Mimics server side metadata exchange behavior in four steps: |
| 308 | +
|
| 309 | + 1. Read a big endian uint32 (4 bytes) from the client. This is the number of |
| 310 | + bytes the message consumes. The length does not include the initial four |
| 311 | + bytes. |
| 312 | +
|
| 313 | + 2. Read the message from the client using the message length and serialize |
| 314 | + it into a MetadataExchangeResponse message. |
| 315 | +
|
| 316 | + The real server implementation will then validate the client has connection |
| 317 | + permissions using the provided OAuth2 token based on the auth type. Here in |
| 318 | + the test implementation, the server does nothing. |
| 319 | +
|
| 320 | + 3. Prepare a response and write the size of the response as a big endian |
| 321 | + uint32 (4 bytes) |
| 322 | +
|
| 323 | + 4. Parse the response to bytes and write those to the client as well. |
| 324 | +
|
| 325 | + Subsequent interactions with the test server use the database protocol. |
| 326 | + """ |
| 327 | + # read metadata message length (4 bytes) |
| 328 | + message_len_buffer_size = struct.Struct("I").size |
| 329 | + message_len_buffer = b"" |
| 330 | + while message_len_buffer_size > 0: |
| 331 | + chunk = sock.recv(message_len_buffer_size) |
| 332 | + if not chunk: |
| 333 | + raise RuntimeError( |
| 334 | + "Connection closed while getting metadata exchange length!" |
| 335 | + ) |
| 336 | + message_len_buffer += chunk |
| 337 | + message_len_buffer_size -= len(chunk) |
| 338 | + |
| 339 | + (message_len,) = struct.unpack(">I", message_len_buffer) |
| 340 | + |
| 341 | + # read metadata exchange message |
| 342 | + buffer = b"" |
| 343 | + while message_len > 0: |
| 344 | + chunk = sock.recv(message_len) |
| 345 | + if not chunk: |
| 346 | + raise RuntimeError("Connection closed while performing metadata exchange!") |
| 347 | + buffer += chunk |
| 348 | + message_len -= len(chunk) |
| 349 | + |
| 350 | + # form metadata exchange request to be received from client |
| 351 | + message = connectorspb.CloudSQLConnectRequest() |
| 352 | + # parse metadata exchange request from buffer |
| 353 | + message.ParseFromString(buffer) |
| 354 | + |
| 355 | + # form metadata exchange response to send to client |
| 356 | + resp = connectorspb.CloudSQLConnectResponse( |
| 357 | + response_code=connectorspb.CloudSQLConnectResponse.OK |
| 358 | + ) |
| 359 | + |
| 360 | + # pack big-endian unsigned integer (4 bytes) |
| 361 | + resp_len = struct.pack(">I", resp.ByteSize()) |
| 362 | + |
| 363 | + # send metadata response length and response message |
| 364 | + sock.sendall(resp_len + resp.SerializeToString()) |
0 commit comments