pyJWT + urllib instrumentations (#43)

Author: sohankshirsagar

.cursor/BUGBOT.md

@@ -0,0 +1,5 @@
+# BUGBOT Notes
+
+## Instrumentation Guidelines
+
+- When adding a new instrumentation, the README must be updated to document the new instrumentation.

README.md

@@ -57,6 +57,8 @@ Tusk Drift currently supports the following packages and versions:
 | psycopg2 | all versions |
 | Redis | `>=4.0.0` |
 | Kinde | `>=2.0.1` |
+| PyJWT | all versions |
+| urllib.request | all versions |
 
 If you're using packages or versions not listed above, please create an issue with the package + version you'd like an instrumentation for.
 

drift/core/content_type_utils.py

@@ -17,7 +17,7 @@
     "application/vnd.api+json": DecodedType.JSON,
     # Plain Text (ALLOWED)
     "text/plain": DecodedType.PLAIN_TEXT,
-    # HTML (BLOCKED)
+    # HTML
     "text/html": DecodedType.HTML,
     "application/xhtml+xml": DecodedType.HTML,
     # CSS (BLOCKED)
@@ -111,9 +111,7 @@
     "application/binary": DecodedType.BINARY,
 }
 
-# Only JSON and plain text are acceptable (matches Node SDK)
-# All other content types will cause trace blocking
-ACCEPTABLE_DECODED_TYPES = {DecodedType.JSON, DecodedType.PLAIN_TEXT}
+ACCEPTABLE_DECODED_TYPES = {DecodedType.JSON, DecodedType.PLAIN_TEXT, DecodedType.HTML}
 
 
 def get_decoded_type(content_type: str | None) -> DecodedType | None:

drift/core/drift_sdk.py

@@ -406,6 +406,16 @@ def _init_auto_instrumentations(self) -> None:
         except ImportError:
             pass
 
+        try:
+            import urllib.request
+
+            from ..instrumentation.urllib import UrllibInstrumentation
+
+            _ = UrllibInstrumentation()
+            logger.debug("urllib instrumentation initialized")
+        except ImportError:
+            pass
+
         # Initialize PostgreSQL instrumentation before Django
         # Instrument BOTH psycopg2 and psycopg if available
         # This allows apps to use either or both
@@ -491,6 +501,15 @@ def _init_auto_instrumentations(self) -> None:
             except Exception as e:
                 logger.debug(f"Socket instrumentation initialization failed: {e}")
 
+            # PyJWT instrumentation for JWT verification bypass
+            try:
+                from ..instrumentation.pyjwt import PyJWTInstrumentation
+
+                _ = PyJWTInstrumentation(mode=self.mode)
+                logger.debug("PyJWT instrumentation registered (REPLAY mode)")
+            except Exception as e:
+                logger.debug(f"PyJWT instrumentation registration failed: {e}")
+
     def create_env_vars_snapshot(self) -> None:
         """Create a span capturing all environment variables.
 

drift/core/mock_utils.py

@@ -183,7 +183,9 @@ def find_mock_response_sync(
         mock_response = sdk.request_mock_sync(mock_request)
 
         if not mock_response or not mock_response.found:
-            logger.debug(f"No matching mock found for {trace_id} with input value: {input_value}")
+            logger.debug(
+                f"No matching mock found for {trace_id} with input value: {input_value}, input schema: {input_schema_merges}, input schema hash: {outbound_span.input_schema_hash}, input value hash: {outbound_span.input_value_hash}"
+            )
             return None
 
         logger.debug(f"Found mock response for {trace_id}")

drift/core/trace_blocking_manager.py

@@ -204,7 +204,7 @@ def should_block_span(span: CleanSpanData) -> bool:
     """Check if a span should be blocked due to size or server error status.
 
     Blocks the trace if:
-    1. The span is a SERVER span with ERROR status (e.g., HTTP >= 300)
+    1. The span is a SERVER span with ERROR status (e.g., HTTP >= 400)
     2. The span exceeds the maximum size limit (1MB)
 
     This matches Node SDK behavior in TdSpanExporter.ts.
@@ -221,7 +221,7 @@ def should_block_span(span: CleanSpanData) -> bool:
     span_name = span.name
     blocking_manager = TraceBlockingManager.get_instance()
 
-    # Check 1: Block SERVER spans with ERROR status (e.g., HTTP >= 300)
+    # Check 1: Block SERVER spans with ERROR status (e.g., HTTP >= 400)
     if span.kind == SpanKind.SERVER and span.status.code == StatusCode.ERROR:
         logger.debug(f"Blocking trace {trace_id} - server span '{span_name}' has error status")
         blocking_manager.block_trace(trace_id, reason="server_error")

drift/instrumentation/django/middleware.py

@@ -381,8 +381,8 @@ def dict_to_schema_merges(merges_dict):
         duration_seconds = duration_ns // 1_000_000_000
         duration_nanos = duration_ns % 1_000_000_000
 
-        # Match Node SDK: >= 300 is considered an error (redirects, client errors, server errors)
-        if status_code >= 300:
+        # Match Node SDK: >= 400 is considered an error
+        if status_code >= 400:
             status = SpanStatus(code=StatusCode.ERROR, message=f"HTTP {status_code}")
         else:
             status = SpanStatus(code=StatusCode.OK, message="")

drift/instrumentation/fastapi/instrumentation.py

@@ -535,8 +535,8 @@ def _finalize_span(
     TuskDrift.get_instance()
 
     status_code = response_data.get("status_code", 200)
-    # Match Node SDK: >= 300 is considered an error (redirects, client errors, server errors)
-    if status_code >= 300:
+    # Match Node SDK: >= 400 is considered an error
+    if status_code >= 400:
         span_info.span.set_status(Status(OTelStatusCode.ERROR, f"HTTP {status_code}"))
     else:
         span_info.span.set_status(Status(OTelStatusCode.OK))

drift/instrumentation/pyjwt/__init__.py

@@ -0,0 +1,5 @@
+"""PyJWT instrumentation for REPLAY mode."""
+
+from .instrumentation import PyJWTInstrumentation
+
+__all__ = ["PyJWTInstrumentation"]

drift/instrumentation/pyjwt/instrumentation.py

@@ -0,0 +1,98 @@
+"""PyJWT instrumentation for REPLAY mode.
+
+Patches PyJWT to disable all verification during test replay:
+1. _merge_options - returns all verification options as False
+2. _verify_signature - no-op (defense in depth)
+3. _validate_claims - no-op (defense in depth)
+
+Only active in REPLAY mode.
+"""
+
+from __future__ import annotations
+
+import logging
+from types import ModuleType
+
+from ...core.types import TuskDriftMode
+from ..base import InstrumentationBase
+
+logger = logging.getLogger(__name__)
+
+
+class PyJWTInstrumentation(InstrumentationBase):
+    """Patches PyJWT to disable verification in REPLAY mode."""
+
+    def __init__(self, mode: TuskDriftMode = TuskDriftMode.DISABLED, enabled: bool = True) -> None:
+        self.mode = mode
+        should_enable = enabled and mode == TuskDriftMode.REPLAY
+
+        super().__init__(
+            name="PyJWTInstrumentation",
+            module_name="jwt",
+            supported_versions="*",
+            enabled=should_enable,
+        )
+
+    def patch(self, module: ModuleType) -> None:
+        if self.mode != TuskDriftMode.REPLAY:
+            return
+
+        self._patch_merge_options()
+        self._patch_signature_verification()
+        self._patch_claim_validation()
+        logger.debug("[PyJWTInstrumentation] All patches applied")
+
+    def _patch_signature_verification(self) -> None:
+        """No-op signature verification."""
+        try:
+            from jwt import api_jws
+
+            def patched_verify_signature(self, *args, **kwargs):
+                logger.debug("[PyJWTInstrumentation] _verify_signature called - skipping verification")
+                return None
+
+            api_jws.PyJWS._verify_signature = patched_verify_signature
+            logger.debug("[PyJWTInstrumentation] Patched PyJWS._verify_signature")
+        except Exception as e:
+            logger.warning(f"[PyJWTInstrumentation] Failed to patch _verify_signature: {e}")
+
+    def _patch_claim_validation(self) -> None:
+        """No-op claim validation."""
+        try:
+            from jwt import api_jwt
+
+            def patched_validate_claims(self, *args, **kwargs):
+                logger.debug("[PyJWTInstrumentation] _validate_claims called - skipping validation")
+                return None
+
+            api_jwt.PyJWT._validate_claims = patched_validate_claims
+            logger.debug("[PyJWTInstrumentation] Patched PyJWT._validate_claims")
+        except Exception as e:
+            logger.warning(f"[PyJWTInstrumentation] Failed to patch _validate_claims: {e}")
+
+    def _patch_merge_options(self) -> None:
+        """Patch _merge_options to always return disabled verification options."""
+        try:
+            from jwt import api_jwt
+
+            disabled_options = {
+                "verify_signature": False,
+                "verify_exp": False,
+                "verify_nbf": False,
+                "verify_iat": False,
+                "verify_aud": False,
+                "verify_iss": False,
+                "verify_sub": False,
+                "verify_jti": False,
+                "require": [],
+                "strict_aud": False,
+            }
+
+            def patched_merge_options(self, options=None):
+                logger.debug("[PyJWTInstrumentation] _merge_options called - returning disabled options")
+                return disabled_options
+
+            api_jwt.PyJWT._merge_options = patched_merge_options
+            logger.debug("[PyJWTInstrumentation] Patched PyJWT._merge_options")
+        except Exception as e:
+            logger.warning(f"[PyJWTInstrumentation] Failed to patch _merge_options: {e}")