feat(asm): ASM support for GRPC

ddtrace/_trace/trace_handlers.py

@@ -390,6 +390,11 @@ def _on_start_response_pre(request, ctx, flask_config, status_code, headers):
     )
 
 
+def _on_grpc_request(request_message):
+    # TODO
+    pass
+
+
 def _cookies_from_response_headers(response_headers):
     cookies = {}
     for header_tuple in response_headers:
@@ -766,6 +771,7 @@ def listen():
     core.on("botocore.kinesis.GetRecords.post", _on_botocore_kinesis_getrecords_post)
     core.on("redis.async_command.post", _on_redis_command_post)
     core.on("redis.command.post", _on_redis_command_post)
+    core.on("grpc.request", _on_grpc_request)
 
     for context_name in (
         "flask.call",

ddtrace/appsec/_processor.py

@@ -218,7 +218,7 @@ def _update_rules(self, new_rules: Dict[str, Any]) -> bool:
     def on_span_start(self, span: Span) -> None:
         from ddtrace.contrib import trace_utils
 
-        if span.span_type != SpanTypes.WEB:
+        if span.span_type not in (SpanTypes.WEB, SpanTypes.HTTP, SpanTypes.GRPC):
             return
 
         if _asm_request_context.free_context_available():
@@ -278,7 +278,7 @@ def _waf_action(
         be retrieved from the `core`. This can be used when you don't want to store
         the value in the `core` before checking the `WAF`.
         """
-        if span.span_type not in (SpanTypes.WEB, SpanTypes.HTTP):
+        if span.span_type not in (SpanTypes.WEB, SpanTypes.HTTP, SpanTypes.GRPC):
             return None
 
         if core.get_item(WAF_CONTEXT_NAMES.BLOCKED, span=span) or core.get_item(WAF_CONTEXT_NAMES.BLOCKED):
@@ -418,7 +418,7 @@ def _is_needed(self, address: str) -> bool:
 
     def on_span_finish(self, span: Span) -> None:
         try:
-            if span.span_type == SpanTypes.WEB:
+            if span.span_type in (SpanTypes.WEB, SpanTypes.HTTP, SpanTypes.GRPC):
                 # Force to set respond headers at the end
                 headers_res = core.get_item(SPAN_DATA_NAMES.RESPONSE_HEADERS_NO_COOKIES, span=span)
                 if headers_res:
@@ -438,7 +438,7 @@ def on_span_finish(self, span: Span) -> None:
             # release asm context if it was created by the span
             _asm_request_context.unregister(span)
 
-            if span.span_type != SpanTypes.WEB:
+            if span.span_type not in (SpanTypes.WEB, SpanTypes.HTTP, SpanTypes.GRPC):
                 return
 
             to_delete = []

ddtrace/contrib/grpc/patch.py

@@ -2,6 +2,7 @@
 
 from ddtrace import Pin
 from ddtrace import config
+from ddtrace.appsec._iast._utils import _is_iast_enabled
 from ddtrace.internal.schema import schematize_service_name
 from ddtrace.vendor.wrapt import wrap_function_wrapper as _w
 
@@ -159,6 +160,7 @@ def _patch_server():
     Pin().onto(constants.GRPC_PIN_MODULE_SERVER)
 
     _w("grpc", "server", _server_constructor_interceptor)
+    _w("grpc._server", "_receive_message", _server_receive_message)
 
 
 def _patch_aio_server():
@@ -252,6 +254,71 @@ def _server_constructor_interceptor(wrapped, instance, args, kwargs):
     return wrapped(*args, **kwargs)
 
 
+class PatchedCondition(object):
+    def __init__(self, state, call, condition):
+        self.state = state
+        self.call = call
+        self.condition = condition
+
+    def __enter__(self):
+        self.condition.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.state.request and _is_iast_enabled():
+            from ddtrace.internal.constants import HTTP_REQUEST_BLOCKED
+
+            from ...internal import core
+
+            core.dispatch("grpc.request", (self.state.request,))
+            if core.get_item(HTTP_REQUEST_BLOCKED):
+                from grpc._cython import cygrpc
+                from grpc._server import _abort
+
+                self.state.condition = self.condition
+
+                _abort(self.state, self.call, cygrpc.StatusCode.internal, "test123")
+
+        self.condition.__exit__(exc_type, exc_val, exc_tb)
+
+    def acquire(self, blocking=True, timeout=None):
+        return self.condition.acquire(blocking, timeout)
+
+    def release(self):
+        return self.condition.release()
+
+    def wait(self, timeout=None):
+        return self.condition.wait(timeout)
+
+    def notify(self, n=1):
+        return self.condition.notify(n)
+
+    def notify_all(self):
+        return self.condition.notify_all()
+
+    def __str__(self):
+        return str(self.condition)
+
+
+def _server_receive_message(wrapped, instance, args, kwargs):
+    closure = wrapped(*args, **kwargs)
+    state = args[0]
+    call = args[1]
+
+    def wrapped_closure(*call_args, **call_kwargs):
+        if not isinstance(state.condition, PatchedCondition):
+            state.condition = PatchedCondition(state, call, state.condition)
+
+        result = closure(*call_args, **call_kwargs)
+
+        if isinstance(state.condition, PatchedCondition):
+            state.condition = state.condition.condition
+
+        return result
+
+    return wrapped_closure
+
+
 def _aio_server_constructor_interceptor(wrapped, instance, args, kwargs):
     pin = Pin.get_from(GRPC_AIO_PIN_MODULE_SERVER)