xyang-UID2-4369-enable-nitro-logging (#1403)

* xyang-UID2-4369-enable-nitro-logging

Co-authored-by: Release Workflow <unifiedid-admin+release@thetradedesk.com>

Author: clarkxuyang

pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>com.uid2</groupId>
     <artifactId>uid2-operator</artifactId>
-    <version>5.47.75-alpha-168-SNAPSHOT</version>
+    <version>5.47.76-alpha-183-SNAPSHOT</version>
 
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
@@ -22,7 +22,7 @@
         <enclave-aws.version>2.1.0</enclave-aws.version>
         <enclave-azure.version>2.1.0</enclave-azure.version>
         <enclave-gcp.version>2.1.0</enclave-gcp.version>
-        <uid2-shared.version>8.1.10</uid2-shared.version>
+        <uid2-shared.version>8.1.15</uid2-shared.version>
         <image.version>${project.version}</image.version>
         <maven.compiler.source>21</maven.compiler.source>
         <maven.compiler.target>21</maven.compiler.target>

scripts/aws/UID_CloudFormation.template.yml

@@ -189,11 +189,11 @@ Resources:
       SecretString: !Join
         - ''
         - - '{'
-          - '"core_base_url": '
+          - '"core_base_url": "'
           - !If [IsIntegEnvironment, 'https://core-integ.uidapi.com', 'https://core.uidapi.com']
-          - ', "optout_base_url": '
+          - '", "optout_base_url": "'
           - !If [IsIntegEnvironment, 'https://optout-integ.uidapi.com', 'https://optout.uidapi.com']
-          - ', "operator_key": "'
+          - '", "operator_key": "'
           - Ref: APIToken
           - '"'
           - ', "service_instances": 6'

scripts/aws/ec2.py

@@ -11,7 +11,7 @@
 import argparse
 import logging
 from botocore.exceptions import ClientError, NoCredentialsError
-from typing import Dict
+from typing import Dict, List
 import sys
 import time
 import yaml
@@ -105,6 +105,8 @@ def add_defaults(configs: Dict[str, any]) ->  AWSConfidentialComputeConfig:
         try:
             self.configs = add_defaults(json.loads(client.get_secret_value(SecretId=secret_identifier)["SecretString"]))
             self.__validate_aws_specific_config()
+        except json.JSONDecodeError as e:
+            raise OperatorKeyNotFoundError(self.__class__.__name__, f"Can not parse secret {secret_identifier} in {region}")
         except NoCredentialsError as _:
             raise InstanceProfileMissingError(self.__class__.__name__)
         except ClientError as _:
@@ -119,38 +121,80 @@ def __get_max_capacity():
         except Exception as e:
             raise RuntimeError("/etc/nitro_enclaves/allocator.yaml does not have CPU, memory allocated")
 
-    def __setup_vsockproxy(self, log_level: int) -> None:
-        """
-        Sets up the vsock proxy service.
-        """
+    def __setup_vsockproxy(self) -> None:
+        logging.info("Sets up the vSock proxy service")
         thread_count = (multiprocessing.cpu_count() + 1) // 2
         command = [
             "/usr/bin/vsockpx", "-c", "/etc/uid2operator/proxy.yaml",
-            "--workers", str(thread_count), "--log-level", str(log_level), "--daemon"
+            "--workers", str(thread_count), "--daemon"
+        ]
+
+        debug_command = [
+            "/usr/bin/vsockpx", "-c", "/etc/uid2operator/proxy.yaml",
+            "--workers", str(thread_count), "--log-level", "0"
         ]
-        self.run_command(command)
+
+        self.run_service([command, debug_command], "vsock_proxy")
 
     def __run_config_server(self) -> None:
-        """
-        Starts the Flask configuration server.
-        """
+        logging.info("Starts the Flask configuration server")
         os.makedirs("/etc/secret/secret-value", exist_ok=True)
         config_path = "/etc/secret/secret-value/config"
+
+        # Save configs to a file
         with open(config_path, 'w') as config_file:
             json.dump(self.configs, config_file)
+
         os.chdir("/opt/uid2operator/config-server")
         command = ["./bin/flask", "run", "--host", AuxiliaryConfig.LOCALHOST, "--port", AuxiliaryConfig.FLASK_PORT]
-        self.run_command(command, separate_process=True)
+
+        self.run_service([command, command], "flask_config_server", separate_process=True)
 
     def __run_socks_proxy(self) -> None:
+        logging.info("Starts the SOCKS proxy service")
+        command = ["sockd", "-D"]
+
+        # -d specifies debug level
+        debug_command = ["sockd", "-d", "0"]
+
+        self.run_service([command, debug_command], "socks_proxy")
+
+    def run_service(self, command: List[List[str]], log_filename: str, separate_process: bool = False) -> None:
         """
-        Starts the SOCKS proxy service.
+        Runs a service command with logging if debug_mode is enabled.
+
+        :param command: command[0] regular command, command[1] debug mode command
+        :param log_filename: Base name of the log file (e.g., "flask_config_server", "socks_proxy", "vsock_proxy")
+        :param separate_process: Whether to run in a separate process
         """
-        command = ["sockd", "-D"]
-        self.run_command(command)
+        log_file = f"/var/log/{log_filename}.log"
+
+        if self.configs.get("debug_mode") is True:
+            
+            # Remove old log file to start fresh
+            if os.path.exists(log_file):
+                os.remove(log_file)
+
+            # Set up logging
+            logging.basicConfig(
+                filename=log_file,
+                filemode="w",
+                level=logging.DEBUG,
+                format="%(asctime)s %(levelname)s: %(message)s"
+            )
+
+            logging.info(f"Debug mode is on, logging into {log_file}")
+
+            # Run debug mode command
+            with open(log_file, "a") as log:
+                self.run_command(command[1], separate_process=True, stdout=log, stderr=log)
+        else:
+            # Run regular command, possibly daemon
+            self.run_command(command[0], separate_process=separate_process)
 
     def __get_secret_name_from_userdata(self) -> str:
         """Extracts the secret name from EC2 user data."""
+        logging.info("Extracts the secret name from EC2 user data")
         token = self.__get_aws_token()
         response = requests.get(AuxiliaryConfig.get_user_data_url(), headers={"X-aws-ec2-metadata-token": token})
         user_data = response.text
@@ -165,8 +209,7 @@ def __get_secret_name_from_userdata(self) -> str:
 
     def _setup_auxiliaries(self) -> None:
         """Sets up the vsock tunnel, socks proxy and flask server"""
-        log_level = 1 if self.configs["debug_mode"] else 3
-        self.__setup_vsockproxy(log_level)
+        self.__setup_vsockproxy()
         self.__run_config_server()
         self.__run_socks_proxy()
         logging.info("Finished setting up all auxiliaries")
@@ -206,7 +249,7 @@ def __run_nitro_enclave(self):
             "--enclave-name", "uid2operator"
         ]
         if self.configs.get('debug_mode', False):
-            logging.info("Running in debug_mode")
+            logging.info("Running nitro in debug_mode")
             command += ["--debug-mode", "--attach-console"]
         self.run_command(command, separate_process=False)
 
@@ -248,6 +291,7 @@ def __kill_auxiliaries(self) -> None:
     parser = argparse.ArgumentParser(description="Manage EC2-based confidential compute workflows.")
     parser.add_argument("-o", "--operation", choices=["stop", "start"], default="start", help="Operation to perform.")
     args = parser.parse_args()
+
     try:
         ec2 = EC2EntryPoint()
         if args.operation == "stop":

scripts/aws/sockd.conf

@@ -3,10 +3,11 @@ external: ens5
 user.notprivileged: ec2-user
 clientmethod: none
 socksmethod: none
+logoutput: stderr
 
 client pass {
     from: 127.0.0.1/32 to: 127.0.0.1/32
-    log: error # connect disconnect iooperation
+    log: error connect # disconnect iooperation
 }
 
 socks pass {

scripts/confidential_compute.py

@@ -138,13 +138,14 @@ def run_compute(self) -> None:
         pass
 
     @staticmethod
-    def run_command(command, separate_process=False):
+    def run_command(command, separate_process=False, stdout=None, stderr=None):
         logging.info(f"Running command: {' '.join(command)}")
         try:
             if separate_process:
-                subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+                subprocess.Popen(command, stdout=stdout, stderr=stderr)
             else:
-                subprocess.run(command,check=True,text=True)
+                subprocess.run(command, check=True, stdout=stdout, stderr=stderr)
+
         except Exception as e:
             logging.error(f"Failed to run command: {e}", exc_info=True)
             raise RuntimeError (f"Failed to start {' '.join(command)} ")