[Containerapp] Support debug console

src/containerapp/HISTORY.rst

@@ -10,6 +10,7 @@ upcoming
 * 'az containerapp sessionpool create': Add NodeLTS container-type.
 * 'az containerapp env java-component': Support --min-replicas and --max-replicas for Java components
 * 'az containerapp env create': Support `--dapr-connection-string` to set application insight connection string
+* 'az containerapp debug': Open an SSH-like interactive shell within a container app ephemeral container.
 
 0.3.55
 ++++++

src/containerapp/azext_containerapp/_help.py

@@ -2054,3 +2054,12 @@
           az containerapp job registry set -n my-containerapp-job -g MyResourceGroup \\
               --server MyContainerappJobRegistry.azurecr.io --identity system-environment
 """
+
+helps['containerapp debug'] = """
+    type: command
+    short-summary: Open an SSH-like interactive shell within a container app ephemeral container.
+    examples:
+    - name: debug into a particular container app replica and revision
+      text: |
+          az containerapp debug -n MyContainerapp -g MyResourceGroup --revision MyRevision --replica MyReplica --container MyContainer
+"""

src/containerapp/azext_containerapp/_params.py

@@ -440,3 +440,12 @@ def load_arguments(self, _):
         c.argument('logger_name', help="The logger name.")
         c.argument('logger_level', arg_type=get_enum_type(["off", "error", "info", "debug", "trace", "warn"]), help="Set the log level for the specific logger name.")
         c.argument('all', help="The flag to indicate all logger settings.", action="store_true")
+
+    with self.argument_context('containerapp debug') as c:
+        c.argument('container', help="The container name that the ephemeral container will target to.")
+        c.argument('replica',
+                   help="The name of the replica. List replicas with 'az containerapp replica list'. A replica may not exist if there is no traffic to your app.")
+        c.argument('revision',
+                   help="The name of the container app revision. Defaults to the latest revision.")
+        c.argument('name', name_type, id_part=None, help="The name of the Containerapp.")
+        c.argument('resource_group_name', arg_type=resource_group_name_type, id_part=None)

src/containerapp/azext_containerapp/_ssh_utils.py

@@ -0,0 +1,169 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+# pylint: disable=logging-fstring-interpolation
+# pylint: disable=possibly-used-before-assignment
+
+import sys
+import time
+import threading
+
+import websocket
+from azure.cli.command_modules.containerapp._ssh_utils import SSH_PROXY_INFO, SSH_DEFAULT_ENCODING, SSH_PROXY_ERROR, \
+    SSH_PROXY_FORWARD, SSH_CLUSTER_STDOUT, SSH_CLUSTER_STDERR, SSH_TERM_RESIZE_PREFIX, SSH_INPUT_PREFIX
+from azure.cli.command_modules.containerapp._utils import is_platform_windows
+from azure.cli.core.azclierror import CLIInternalError, ValidationError
+from azure.cli.core.commands.client_factory import get_subscription_id
+
+from knack.log import get_logger
+
+from ._clients import ContainerAppClient
+
+# pylint: disable=import-error,ungrouped-imports
+if is_platform_windows():
+    import msvcrt
+    from azure.cli.command_modules.container._vt_helper import (enable_vt_mode, _get_conout_mode,
+                                                                _set_conout_mode, _get_conin_mode, _set_conin_mode)
+
+logger = get_logger(__name__)
+
+
+class WebSocketConnection:
+    def __init__(self, cmd, resource_group_name, name, revision, replica, container):
+        token_response = ContainerAppClient.get_auth_token(cmd, resource_group_name, name)
+        self._token = token_response["properties"]["token"]
+
+        self._logstream_endpoint = self._get_logstream_endpoint(cmd, resource_group_name, name,
+                                                                revision, replica, container)
+        self._url = self._get_url(cmd=cmd, resource_group_name=resource_group_name, name=name, revision=revision,
+                                  replica=replica, container=container)
+
+        self._socket = websocket.WebSocket(enable_multithread=True)
+
+        logger.info("Attempting to connect to %s", self._url)
+        self._socket.connect(self._url, header=[f"Authorization: Bearer {self._token}"])
+
+        self.is_connected = True
+        self._windows_conout_mode = None
+        self._windows_conin_mode = None
+        if is_platform_windows():
+            self._windows_conout_mode = _get_conout_mode()
+            self._windows_conin_mode = _get_conin_mode()
+
+    @classmethod
+    def _get_logstream_endpoint(cls, cmd, resource_group_name, name, revision, replica, container):
+        containers = ContainerAppClient.get_replica(cmd,
+                                                    resource_group_name,
+                                                    name, revision, replica)["properties"]["containers"]
+        container_info = [c for c in containers if c["name"] == container]
+        if not container_info:
+            raise ValidationError(f"No such container: {container}")
+        return container_info[0]["logStreamEndpoint"]
+
+    def _get_url(self, cmd, resource_group_name, name, revision, replica, container):
+        sub = get_subscription_id(cmd.cli_ctx)
+        base_url = self._logstream_endpoint
+        proxy_api_url = base_url[:base_url.index("/subscriptions/")].replace("https://", "")
+
+        return (f"wss://{proxy_api_url}/subscriptions/{sub}/resourceGroups/{resource_group_name}/containerApps/{name}"
+                f"/revisions/{revision}/replicas/{replica}/debug"
+                f"?targetContainer={container}")
+
+    def disconnect(self):
+        logger.warning("Disconnecting...")
+        self.is_connected = False
+        self._socket.close()
+        if self._windows_conout_mode and self._windows_conin_mode:
+            _set_conout_mode(self._windows_conout_mode)
+            _set_conin_mode(self._windows_conin_mode)
+
+    def send(self, *args, **kwargs):
+        return self._socket.send(*args, **kwargs)
+
+    def recv(self, *args, **kwargs):
+        return self._socket.recv(*args, **kwargs)
+
+
+def _decode_and_output_to_terminal(connection: WebSocketConnection, response, encodings):
+    for i, encoding in enumerate(encodings):
+        try:
+            print(response[2:].decode(encoding), end="", flush=True)
+            break
+        except UnicodeDecodeError as e:
+            if i == len(encodings) - 1:  # ran out of encodings to try
+                connection.disconnect()
+                logger.info("Proxy Control Byte: %s", response[0])
+                logger.info("Cluster Control Byte: %s", response[1])
+                logger.info("Hexdump: %s", response[2:].hex())
+                raise CLIInternalError("Failed to decode server data") from e
+            logger.info("Failed to encode with encoding %s", encoding)
+
+
+def read_ssh(connection: WebSocketConnection, response_encodings):
+    # We need to do resize for the whole session
+    _resize_terminal(connection)
+
+    # response_encodings is the ordered list of Unicode encodings to try to decode with before raising an exception
+    while connection.is_connected:
+        response = connection.recv()
+        if not response:
+            connection.disconnect()
+        else:
+            logger.info("Received raw response %s", response.hex())
+            proxy_status = response[0]
+            if proxy_status == SSH_PROXY_INFO:
+                print(f"INFO: {response[1:].decode(SSH_DEFAULT_ENCODING)}")
+            elif proxy_status == SSH_PROXY_ERROR:
+                print(f"ERROR: {response[1:].decode(SSH_DEFAULT_ENCODING)}")
+            elif proxy_status == SSH_PROXY_FORWARD:
+                control_byte = response[1]
+                if control_byte in (SSH_CLUSTER_STDOUT, SSH_CLUSTER_STDERR):
+                    _decode_and_output_to_terminal(connection, response, response_encodings)
+                else:
+                    connection.disconnect()
+                    raise CLIInternalError("Unexpected message received")
+
+
+def _send_stdin(connection: WebSocketConnection, getch_fn):
+    while connection.is_connected:
+        ch = getch_fn()
+        if connection.is_connected:
+            connection.send(b"".join([SSH_INPUT_PREFIX, ch]))
+
+
+def _resize_terminal(connection: WebSocketConnection):
+    from shutil import get_terminal_size
+    size = get_terminal_size()
+    if connection.is_connected:
+        # send twice with different width to make sure the terminal will display username prefix correctly
+        # refer `kubectl debug` command implementation:
+        # https://github.com/kubernetes/kubectl/blob/14f6a11dd84315dc5179ff04156b338def935eaa/pkg/cmd/attach/attach.go#L296
+        connection.send(b"".join([SSH_TERM_RESIZE_PREFIX,
+                                  f'{{"Width": {size.columns + 1}, '
+                                  f'"Height": {size.lines}}}'.encode(SSH_DEFAULT_ENCODING)]))
+        connection.send(b"".join([SSH_TERM_RESIZE_PREFIX,
+                                  f'{{"Width": {size.columns}, '
+                                  f'"Height": {size.lines}}}'.encode(SSH_DEFAULT_ENCODING)]))
+
+
+def _getch_unix():
+    return sys.stdin.read(1).encode(SSH_DEFAULT_ENCODING)
+
+
+def _getch_windows():
+    while not msvcrt.kbhit():
+        time.sleep(0.01)
+    return msvcrt.getch()
+
+
+def get_stdin_writer(connection: WebSocketConnection):
+    if not is_platform_windows():
+        import tty
+        tty.setcbreak(sys.stdin.fileno())  # needed to prevent printing arrow key characters
+        writer = threading.Thread(target=_send_stdin, args=(connection, _getch_unix))
+    else:
+        enable_vt_mode()  # needed for interactive commands (ie vim)
+        writer = threading.Thread(target=_send_stdin, args=(connection, _getch_windows))
+
+    return writer

src/containerapp/azext_containerapp/commands.py

@@ -7,6 +7,8 @@
 # from azure.cli.core.commands import CliCommandType
 # from msrestazure.tools import is_valid_resource_id, parse_resource_id
 from azure.cli.command_modules.containerapp._transformers import (transform_containerapp_output, transform_containerapp_list_output)
+from azure.cli.command_modules.containerapp._validators import validate_ssh
+
 from azext_containerapp._client_factory import ex_handler_factory
 from ._transformers import (transform_sensitive_values,
                             transform_telemetry_data_dog_values,
@@ -24,6 +26,7 @@ def load_command_table(self, args):
         g.custom_command('update', 'update_containerapp', supports_no_wait=True, exception_handler=ex_handler_factory(), table_transformer=transform_containerapp_output, transform=transform_sensitive_values)
         g.custom_command('delete', 'delete_containerapp', supports_no_wait=True, confirmation=True, exception_handler=ex_handler_factory())
         g.custom_command('up', 'containerapp_up', supports_no_wait=False, exception_handler=ex_handler_factory())
+        g.custom_command('debug', 'containerapp_debug', is_preview=True, validator=validate_ssh)
 
     with self.command_group('containerapp replica') as g:
         g.custom_show_command('show', 'get_replica')  # TODO implement the table transformer

src/containerapp/azext_containerapp/custom.py

@@ -4,13 +4,15 @@
 # --------------------------------------------------------------------------------------------
 # pylint: disable=line-too-long, unused-argument, logging-fstring-interpolation, logging-not-lazy, consider-using-f-string, logging-format-interpolation, inconsistent-return-statements, broad-except, bare-except, too-many-statements, too-many-locals, too-many-boolean-expressions, too-many-branches, too-many-nested-blocks, pointless-statement, expression-not-assigned, unbalanced-tuple-unpacking, unsupported-assignment-operation
 
+import threading
 import time
 from urllib.parse import urlparse
 import json
 import requests
 import subprocess
 from concurrent.futures import ThreadPoolExecutor
 
+from azure.cli.command_modules.containerapp._ssh_utils import SSH_BACKUP_ENCODING, SSH_CTRL_C_MSG
 from azure.cli.core import telemetry as telemetry_core
 
 from azure.cli.core.azclierror import (
@@ -112,6 +114,8 @@
     AzureFileProperties as AzureFilePropertiesModel
 )
 
+from ._ssh_utils import (SSH_DEFAULT_ENCODING, WebSocketConnection, read_ssh, get_stdin_writer)
+
 from ._utils import connected_env_check_cert_name_availability, get_oryx_run_image_tags, patchable_check, get_pack_exec_path, is_docker_running, parse_build_env_vars, env_has_managed_identity
 
 from ._constants import (CONTAINER_APPS_RP,
@@ -3224,3 +3228,30 @@ def set_registry_job(cmd, name, resource_group_name, server, username=None, pass
     containerapp_job_registry_set_decorator.construct_payload()
     r = containerapp_job_registry_set_decorator.set()
     return r
+
+
+def containerapp_debug(cmd, resource_group_name, name, container=None, revision=None, replica=None):
+    conn = WebSocketConnection(cmd=cmd, resource_group_name=resource_group_name, name=name, revision=revision,
+                               replica=replica, container=container)
+
+    encodings = [SSH_DEFAULT_ENCODING, SSH_BACKUP_ENCODING]
+    reader = threading.Thread(target=read_ssh, args=(conn, encodings))
+    reader.daemon = True
+    reader.start()
+
+    writer = get_stdin_writer(conn)
+    writer.daemon = True
+    writer.start()
+
+    logger.warning("Use ctrl + D to exit.")
+    while conn.is_connected:
+        if not reader.is_alive() or not writer.is_alive():
+            logger.warning("Reader or Writer for WebSocket is not alive. Closing the connection.")
+            conn.disconnect()
+
+        try:
+            time.sleep(0.1)
+        except KeyboardInterrupt:
+            if conn.is_connected:
+                logger.info("Caught KeyboardInterrupt. Sending ctrl+c to server")
+                conn.send(SSH_CTRL_C_MSG)