From cc079429934b29b176f77b90ff106873bc15f36c Mon Sep 17 00:00:00 2001
From: Xuan Wang <xuanwn@google.com>
Date: Fri, 29 Sep 2023 09:42:51 -0700
Subject: [PATCH] [PSM Interop] Log Google Cloud API debug header (#34334)

Add Google Cloud API debug header and read response.

#### Testing

* Triggered in Prod, no header was logged:
https://source.cloud.google.com/results/invocations/c91c6e81-5e5b-4545-8abf-28594d4493a5
* Triggered in Prod with debug header enabled:
https://source.cloud.google.com/results/invocations/cfda4dd5-e3a9-4519-a4e7-85220f780cbf
* Also tested locally in staging.
* Example log:
```
I0913 01:04:19.688863 139651983066944 traffic_director.py:187] Creating GRPC Health Check "xds-k8s-test-health-check-dev"
I0913 01:04:19.692854 139651983066944 compute.py:538] Creating compute resource:
---
grpcHealthCheck:
  portSpecification: USE_SERVING_PORT
name: xds-k8s-test-health-check-dev
type: GRPC
...

I0913 01:04:19.693941 139651983066944 compute.py:600] Adding debug headers for method: compute.healthChecks.insert
I0913 01:04:20.726086 139651983066944 compute.py:595] Received debug headers:
ADViAh6vFBtRL4YC...
I0913 01:04:20.726562 139651983066944 compute.py:610] Waiting 600 sec for compute operation id: operation-1694567059990-6053323a8a81a-5340529e-bf5d5c03
I0913 01:04:23.383233 139651983066944 xds_k8s_testcase.py:272] --- Finished subTest __main__.BaselineTest.test_traffic_director_grpc_setup.0_create_health_check ---
```
---
 .../framework/infrastructure/gcp/compute.py   | 19 +++++++++++++++++++
 .../infrastructure/traffic_director.py        |  5 ++++-
 .../framework/xds_flags.py                    | 11 +++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py b/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py
index 2fa738d2de5..2a970847e76 100644
--- a/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py
+++ b/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py
@@ -19,6 +19,7 @@ from typing import Any, Dict, List, Optional, Set
 
 from googleapiclient import discovery
 import googleapiclient.errors
+import httplib2
 
 import framework.errors
 from framework.helpers import retryers
@@ -26,6 +27,9 @@ from framework.infrastructure import gcp
 
 logger = logging.getLogger(__name__)
 
+DEBUG_HEADER_IN_RESPONSE = "x-encrypted-debug-headers"
+DEBUG_HEADER_KEY = "X-Return-Encrypted-Headers"
+
 
 class ComputeV1(
     gcp.api.GcpProjectApiResource
@@ -34,6 +38,7 @@ class ComputeV1(
     _WAIT_FOR_BACKEND_SEC = 60 * 10
     _WAIT_FOR_BACKEND_SLEEP_SEC = 4
     _WAIT_FOR_OPERATION_SEC = 60 * 10
+    gfe_debug_header: Optional[str]
 
     @dataclasses.dataclass(frozen=True)
     class GcpResource:
@@ -48,9 +53,11 @@ class ComputeV1(
         self,
         api_manager: gcp.api.GcpApiManager,
         project: str,
+        gfe_debug_header: Optional[str] = None,
         version: str = "v1",
     ):
         super().__init__(api_manager.compute(version), project)
+        self.gfe_debug_header = gfe_debug_header
 
     class HealthCheckProtocol(enum.Enum):
         TCP = enum.auto()
@@ -576,9 +583,21 @@ class ComputeV1(
     def _operation_status_done(operation):
         return "status" in operation and operation["status"] == "DONE"
 
+    @staticmethod
+    def _log_debug_header(resp: httplib2.Response):
+        if DEBUG_HEADER_IN_RESPONSE in resp:
+            logger.info(
+                "Received GCP debug headers: %s",
+                resp[DEBUG_HEADER_IN_RESPONSE],
+            )
+
     def _execute(  # pylint: disable=arguments-differ
         self, request, *, timeout_sec=_WAIT_FOR_OPERATION_SEC
     ):
+        if self.gfe_debug_header:
+            logger.info("Adding debug headers for method: %s", request.methodId)
+            request.headers[DEBUG_HEADER_KEY] = self.gfe_debug_header
+            request.add_response_callback(self._log_debug_header)
         operation = request.execute(num_retries=self._GCP_API_RETRIES)
         logger.debug("Operation %s", operation)
         return self._wait(operation["name"], timeout_sec)
diff --git a/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/traffic_director.py b/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/traffic_director.py
index 6db80d4f865..bc549bbcb30 100644
--- a/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/traffic_director.py
+++ b/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/traffic_director.py
@@ -79,7 +79,10 @@ class TrafficDirectorManager:  # pylint: disable=too-many-public-methods
     ):
         # API
         self.compute = _ComputeV1(
-            gcp_api_manager, project, version=compute_api_version
+            gcp_api_manager,
+            project,
+            version=compute_api_version,
+            gfe_debug_header=xds_flags.GFE_DEBUG_HEADER.value,
         )
 
         # Settings
diff --git a/tools/run_tests/xds_k8s_test_driver/framework/xds_flags.py b/tools/run_tests/xds_k8s_test_driver/framework/xds_flags.py
index fd33407d5e8..b710db04431 100644
--- a/tools/run_tests/xds_k8s_test_driver/framework/xds_flags.py
+++ b/tools/run_tests/xds_k8s_test_driver/framework/xds_flags.py
@@ -173,6 +173,17 @@ SOCKET_DEFAULT_TIMEOUT = flags.DEFINE_float(
     ),
 )
 
+GFE_DEBUG_HEADER = flags.DEFINE_enum(
+    "gfe_debug_header",
+    default=None,
+    enum_values=[
+        "gfe_response_only",
+        "all_response",
+        "request_and_response",
+    ],
+    help="Whether to enable GFE debug headers and what value to use.",
+)
+
 
 def set_socket_default_timeout_from_flag() -> None:
     """A helper to configure default socket timeout from a flag.