POC - trace HTTP requests in go

httpdbg/__init__.py

@@ -3,6 +3,6 @@
 from httpdbg.records import HTTPRecords
 
 
-__version__ = "1.2.1"
+__version__ = "2.0.0"
 
 __all__ = ["httprecord", "HTTPRecords"]

httpdbg/__main__.py

@@ -8,6 +8,7 @@
 from httpdbg.log import set_env_for_logging
 from httpdbg.server import httpdbg_srv
 from httpdbg.mode_console import run_console
+from httpdbg.mode_go import run_go
 from httpdbg.mode_module import run_module
 from httpdbg.mode_script import run_script
 
@@ -30,11 +31,13 @@ def pyhttpdbg(params, subparams, test_mode=False):
 
     with httpdbg_srv(params.host, params.port) as records:
         records.server = not params.only_client
-        with httprecord(records, params.initiator, server=records.server):
+        with httprecord(records, params.initiator, server=records.server, go=params.go):
             if params.module:
                 run_module(subparams)
             elif params.script:
                 run_script(subparams)
+            elif params.go:
+                run_go(subparams)
             else:
                 run_console(records, test_mode)
 

httpdbg/args.py

@@ -9,7 +9,7 @@
 def read_args(args: List[str]) -> Tuple[argparse.Namespace, List[str]]:
     httpdbg_args = args
     client_args = []
-    for action in ["--console", "--module", "-m", "--script"]:
+    for action in ["--console", "--module", "-m", "--script", "--go"]:
         if action in args:
             httpdbg_args = args[: args.index(action) + 2]
             client_args = args[args.index(action) + 1 :]
@@ -87,4 +87,10 @@ def read_args(args: List[str]) -> Tuple[argparse.Namespace, List[str]]:
         help="run a script (the next args are passed to the script as is)",
     )
 
+    actions.add_argument(
+        "--go",
+        type=str,
+        help="run a go program using the 'go run' command (the next args are passed to the 'go' process as is) (experimental)",
+    )
+
     return parser.parse_args(httpdbg_args), client_args

httpdbg/hooks/all.py

@@ -12,6 +12,7 @@
 from httpdbg.hooks.generic import hook_generic
 from httpdbg.hooks.http import hook_http
 from httpdbg.hooks.httpx import hook_httpx
+from httpdbg.hooks.go import watcher_go
 from httpdbg.hooks.pytest import hook_pytest
 from httpdbg.hooks.requests import hook_requests
 from httpdbg.hooks.socket import hook_socket
@@ -30,24 +31,30 @@ def httprecord(
     client: bool = True,
     server: bool = False,
     ignore: Union[List[Tuple[str, int]], None] = None,
+    go: bool = False,
 ) -> Generator[HTTPRecords, None, None]:
     if records is None:
         records = HTTPRecords(client=client, server=server, ignore=ignore)
 
-    with watcher_external(records, initiators, server):
-        with hook_flask(records):
-            with hook_socket(records):
-                with hook_fastapi(records):
-                    with hook_starlette(records):
-                        with hook_uvicorn(records):
-                            with hook_http(records):
-                                with hook_httpx(records):
-                                    with hook_requests(records):
-                                        with hook_urllib3(records):
-                                            with hook_aiohttp(records):
-                                                with hook_pytest(records):
-                                                    with hook_unittest(records):
-                                                        with hook_generic(
-                                                            records, initiators
-                                                        ):
-                                                            yield records
+    if not go:
+        with watcher_external(records, initiators, server):
+            with hook_flask(records):
+                with hook_socket(records):
+                    with hook_fastapi(records):
+                        with hook_starlette(records):
+                            with hook_uvicorn(records):
+                                with hook_http(records):
+                                    with hook_httpx(records):
+                                        with hook_requests(records):
+                                            with hook_urllib3(records):
+                                                with hook_aiohttp(records):
+                                                    with hook_pytest(records):
+                                                        with hook_unittest(records):
+                                                            with hook_generic(
+                                                                records, initiators
+                                                            ):
+                                                                yield records
+    else:
+        # if we trace the HTTP requests in go process, there is no reason to hook the Python call
+        with watcher_go(records):
+            yield records

httpdbg/hooks/go.py

@@ -0,0 +1,86 @@
+# -*- coding: utf-8 -*-
+from contextlib import contextmanager
+import json
+import os
+import tempfile
+import time
+import threading
+from typing import Generator
+
+from httpdbg.env import HTTPDBG_MULTIPROCESS_DIR
+from httpdbg.initiator import Group
+from httpdbg.initiator import Initiator
+from httpdbg.log import logger
+from httpdbg.records import HTTPRecord
+from httpdbg.records import HTTPRecords
+
+
+@contextmanager
+def watcher_go(
+    records: HTTPRecords,
+) -> Generator[HTTPRecords, None, None]:
+    if HTTPDBG_MULTIPROCESS_DIR not in os.environ:
+        with tempfile.TemporaryDirectory(prefix="httpdbg_") as httpdbg_multiprocess_dir:
+
+            logger().info(f"watcher_go {httpdbg_multiprocess_dir}")
+            os.environ[HTTPDBG_MULTIPROCESS_DIR] = httpdbg_multiprocess_dir
+
+            try:
+                watcher = WatcherGoDirThread(records, httpdbg_multiprocess_dir, 1.0)
+                watcher.start()
+
+                yield records
+            finally:
+                watcher.shutdown()
+    else:
+        yield records
+
+
+class WatcherGoDirThread(threading.Thread):
+    def __init__(self, records: HTTPRecords, directory: str, delay: float) -> None:
+        self.records: HTTPRecords = records
+        self.directory: str = directory
+        self.delay: float = delay
+        self._running: bool = True
+        threading.Thread.__init__(self)
+
+    def get_go_traces(self):
+        traces_filename = os.path.join(self.directory, "httpdbg-go-traces.json")
+        if os.path.exists(traces_filename):
+            with open(traces_filename, "r") as f:
+                traces = json.load(f)
+                for trace in traces:
+                    if trace["trace_id"] not in self.records.requests.keys():
+                        label = trace["initiator"]["code"]
+                        full_label = f'File "{trace["initiator"]["filename"]}", line {trace["initiator"]["lineno"]}, in {trace["initiator"]["func_name"]}'
+                        full_label += f"\n    {trace['initiator']['code']}"
+                        stack = []
+                        for line in trace["initiator"]["stack"]:
+                            stack.append(f"{line['location']}, in {line['func_name']}")
+                            code = line["code"].replace("\t", "    ")
+                            stack.append(f"    {code}\n")
+
+                        initiator = Initiator(label, full_label, stack)
+                        self.records.add_initiator(initiator)
+                        group = Group(label, full_label, False)
+                        self.records.add_group(group)
+                        record = HTTPRecord(
+                            initiator.id,
+                            group.id,
+                        )
+                        record.id = trace["trace_id"]
+                        with open(trace["request_file"], "rb") as reqf:
+                            record.send_data(reqf.read())
+                        with open(trace["response_file"], "rb") as resf:
+                            record.receive_data(resf.read())
+                        self.records.requests[record.id] = record
+
+    def run(self):
+        while self._running:
+            self.get_go_traces()
+            time.sleep(self.delay)
+
+    def shutdown(self):
+        self._running = False
+        self.join(timeout=5)
+        self.get_go_traces()