feat(iast): overhead control engine (OCE) skeleton (#4315)

## Description
Basic structure to implement the features of Overhead control engine (OCE)
OCE control the number of requests to analyze, number of vulnerabilities to report and number of vulnerabilities to analyze. 

## Reviewer Checklist
- [ ] Title is accurate.
- [ ] Description motivates each change.
- [ ] No unnecessary changes were introduced in this PR.
- [ ] PR cannot be broken up into smaller PRs.
- [ ] Avoid breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes unless absolutely necessary.
- [ ] Tests provided or description of manual testing performed is included in the code or PR.
- [ ] Release note has been added for fixes and features, or else `changelog/no-changelog` label added.
- [ ] All relevant GitHub issues are correctly linked.
- [ ] Backports are identified and tagged with Mergifyio.
- [ ] Add to milestone.

Author: avara1986

ddtrace/appsec/iast/__init__.py

@@ -7,7 +7,33 @@
 
 Then, implement the `patch()` function and its wrappers.
 
-When we detect a vulnerability we should report it with `ddtrace.appsec.iast.reporter.report_vulnerability` to add
-this information to the context and `ddtrace.appsec.iast.processor` will send this information to the backend at
-the end of the request
-"""
+In order to have the better performance, the Overhead control engine (OCE) helps us to control the overhead of our
+wrapped functions. We should create a class that inherit from `ddtrace.appsec.iast.taint_sinks._base.VulnerabilityBase`
+and register with `ddtrace.appsec.iast.oce`.
+
+@oce.register
+class MyVulnerability(VulnerabilityBase):
+    vulnerability_type = "MyVulnerability"
+    evidence_type = "kind_of_Vulnerability"
+
+Before that, we should decorate our wrappers with `wrap` method and
+report the vulnerabilities with `report` method. OCE will manage the number of requests, number of vulnerabilities
+to reduce the overhead.
+
+@WeakHash.wrap
+def wrapped_function(wrapped, instance, args, kwargs):
+    # type: (Callable, str, Any, Any, Any) -> Any
+    WeakHash.report(
+        evidence_value=evidence,
+    )
+    return wrapped(*args, **kwargs)
+"""  # noqa: RST201, RST213, RST210
+from ddtrace.appsec.iast.overhead_control_engine import OverheadControl
+
+
+oce = OverheadControl()
+
+
+__all__ = [
+    "oce",
+]

ddtrace/appsec/iast/overhead_control_engine.py

@@ -0,0 +1,87 @@
+"""
+The Overhead control engine (OCE) is an element that by design ensures that the overhead does not go over a maximum
+limit. It will measure operations being executed in a request and it will deactivate detection
+(and therefore reduce the overhead to nearly 0) if a certain threshold is reached.
+"""
+import os
+import threading
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:  # pragma: no cover
+    from typing import Set
+    from typing import Type
+
+MAX_REQUESTS = int(os.environ.get("DD_IAST_MAX_CONCURRENT_REQUESTS", 2))
+MAX_VULNERABILITIES_PER_REQUEST = int(os.environ.get("DD_IAST_VULNERABILITIES_PER_REQUEST", 2))
+
+
+class Operation(object):
+    """Common operation related to Overhead Control Engine (OCE). Every vulnerabilities/taint_sinks should inherit
+    from this class. OCE instance calls these methods to control the overhead produced in each request.
+    """
+
+    _lock = threading.Lock()
+    _vulnerability_quota = MAX_VULNERABILITIES_PER_REQUEST
+
+    @classmethod
+    def reset(cls):
+        cls._vulnerability_quota = MAX_VULNERABILITIES_PER_REQUEST
+
+    @classmethod
+    def acquire_quota(cls):
+        cls._lock.acquire()
+        result = False
+        if cls._vulnerability_quota > 0:
+            cls._vulnerability_quota -= 1
+            result = True
+        cls._lock.release()
+        return result
+
+    @classmethod
+    def has_quota(cls):
+        cls._lock.acquire()
+        result = cls._vulnerability_quota > 0
+        cls._lock.release()
+        return result
+
+
+class OverheadControl(object):
+    _request_quota = MAX_REQUESTS
+    _enabled = False
+    _vulnerabilities = set()  # type: Set[Type[Operation]]
+
+    def acquire_request(self):
+        """Block a request's quota at start of the request.
+
+        TODO: Implement sampling request in this method
+        """
+        if self._request_quota > 0:
+            self._request_quota -= 1
+            self._enabled = True
+
+    def release_request(self):
+        """increment request's quota at end of the request.
+
+        TODO: figure out how to check maximum requests per thread
+        """
+        if self._request_quota < MAX_REQUESTS:
+            self._request_quota += 1
+            self._enabled = False
+        self.vulnerabilities_reset_quota()
+
+    def register(self, klass):
+        # type: (Type[Operation]) -> Type[Operation]
+        """Register vulnerabilities/taint_sinks. This set of elements will restart for each request."""
+        self._vulnerabilities.add(klass)
+        return klass
+
+    @property
+    def request_has_quota(self):
+        # type: () -> bool
+        return self._enabled
+
+    def vulnerabilities_reset_quota(self):
+        # type: () -> None
+        for k in self._vulnerabilities:
+            k.reset()

ddtrace/appsec/iast/processor.py

@@ -3,6 +3,7 @@
 
 import attr
 
+from ddtrace.appsec.iast import oce
 from ddtrace.constants import IAST_CONTEXT_KEY
 from ddtrace.constants import IAST_ENABLED
 from ddtrace.constants import IAST_JSON
@@ -22,7 +23,9 @@
 class AppSecIastSpanProcessor(SpanProcessor):
     def on_span_start(self, span):
         # type: (Span) -> None
-        pass
+        if span.span_type != SpanTypes.WEB:
+            return
+        oce.acquire_request()
 
     def on_span_finish(self, span):
         # type: (Span) -> None
@@ -35,9 +38,12 @@ def on_span_finish(self, span):
         """
         if span.span_type != SpanTypes.WEB:
             return
+
         span.set_metric(IAST_ENABLED, 1.0)
 
         data = _context.get_item(IAST_CONTEXT_KEY, span=span)
 
         if data:
             span.set_tag_str(IAST_JSON, json.dumps(attr.asdict(data)))
+
+        oce.release_request()

ddtrace/appsec/iast/reporter.py

@@ -1,19 +1,7 @@
 from typing import Set
-from typing import TYPE_CHECKING
 
 import attr
 
-from ddtrace.appsec.iast.stacktrace import get_info_frame
-from ddtrace.constants import IAST_CONTEXT_KEY
-from ddtrace.internal import _context
-
-
-if TYPE_CHECKING:  # pragma: no cover
-
-    from typing import Text
-
-    from ddtrace.span import Span
-
 
 @attr.s(eq=False)
 class Evidence(object):
@@ -45,31 +33,3 @@ class Source(object):
 class IastSpanReporter(object):
     sources = attr.ib(type=Set[Source], factory=set)
     vulnerabilities = attr.ib(type=Set[Vulnerability], factory=set)
-
-
-def report_vulnerability(span, vulnerability_type, evidence_type, evidence_value=""):
-    # type: (Span, Text, Text, Text) -> None
-    """Build a IastSpanReporter instance to report it in the `AppSecIastSpanProcessor` as a string JSON"""
-    report = _context.get_item(IAST_CONTEXT_KEY, span=span)
-    file_name, line_number = get_info_frame()
-    if report:
-        report.vulnerabilities.add(
-            Vulnerability(
-                type=vulnerability_type,
-                evidence=Evidence(type=evidence_type, value=evidence_value),
-                location=Location(path=file_name, line=line_number),
-            )
-        )
-
-    else:
-        report = IastSpanReporter(
-            vulnerabilities={
-                Vulnerability(
-                    type=vulnerability_type,
-                    evidence=Evidence(type=evidence_type, value=evidence_value),
-                    location=Location(path=file_name, line=line_number),
-                )
-            }
-        )
-
-    _context.set_item(IAST_CONTEXT_KEY, report, span=span)

ddtrace/appsec/iast/taint_sinks/_base.py

@@ -1,28 +1,86 @@
 from typing import TYPE_CHECKING
 
 from ddtrace import tracer
+from ddtrace.appsec.iast import oce
+from ddtrace.appsec.iast.overhead_control_engine import Operation
+from ddtrace.appsec.iast.reporter import Evidence
+from ddtrace.appsec.iast.reporter import IastSpanReporter
+from ddtrace.appsec.iast.reporter import Location
+from ddtrace.appsec.iast.reporter import Vulnerability
+from ddtrace.appsec.iast.stacktrace import get_info_frame
+from ddtrace.constants import IAST_CONTEXT_KEY
+from ddtrace.internal import _context
 from ddtrace.internal.logger import get_logger
+from ddtrace.vendor.wrapt import wrap_function_wrapper
 
 
 if TYPE_CHECKING:  # pragma: no cover
     from typing import Any
     from typing import Callable
+    from typing import Text
 
 log = get_logger(__name__)
 
 
-def inject_span(func):
-    # type: (Callable) -> Callable
-    """Get the current root Span and attach it to the wrapped function. We need the span to report the vulnerability
-    and update the context with the report information.
-    """
+class VulnerabilityBase(Operation):
+    vulnerability_type = ""
+    evidence_type = ""
 
-    def wrapper(wrapped, instance, args, kwargs):
-        # type: (Callable, Any, Any, Any) -> Any
-        span = tracer.current_root_span()
-        if span:
-            return func(wrapped, span, instance, args, kwargs)
-        log.warning("No root span in the current execution. Skipping IAST Taint sink.")
-        return wrapped(*args, **kwargs)
+    @classmethod
+    def wrap(cls, func):
+        # type: (Callable) -> Callable
+        def wrapper(wrapped, instance, args, kwargs):
+            # type: (Callable, Any, Any, Any) -> Any
+            """Get the current root Span and attach it to the wrapped function. We need the span to report the vulnerability
+            and update the context with the report information.
+            """
+            if oce.request_has_quota and cls.has_quota():
+                return func(wrapped, instance, args, kwargs)
+            else:
+                log.debug("IAST: no vulnerability quota to analyze more sink points")
+            return wrapped(*args, **kwargs)
 
-    return wrapper
+        return wrapper
+
+    @classmethod
+    def report(cls, evidence_value=""):
+        # type: (Text) -> None
+        """Build a IastSpanReporter instance to report it in the `AppSecIastSpanProcessor` as a string JSON
+
+        TODO: check deduplications if DD_IAST_DEDUPLICATION_ENABLED is true
+        """
+        if cls.acquire_quota():
+            span = tracer.current_root_span()
+            if not span:
+                log.debug("No root span in the current execution. Skipping IAST taint sink.")
+                return None
+
+            report = _context.get_item(IAST_CONTEXT_KEY, span=span)
+            file_name, line_number = get_info_frame()
+            if report:
+                report.vulnerabilities.add(
+                    Vulnerability(
+                        type=cls.vulnerability_type,
+                        evidence=Evidence(type=cls.evidence_type, value=evidence_value),
+                        location=Location(path=file_name, line=line_number),
+                    )
+                )
+
+            else:
+                report = IastSpanReporter(
+                    vulnerabilities={
+                        Vulnerability(
+                            type=cls.vulnerability_type,
+                            evidence=Evidence(type=cls.evidence_type, value=evidence_value),
+                            location=Location(path=file_name, line=line_number),
+                        )
+                    }
+                )
+            _context.set_item(IAST_CONTEXT_KEY, report, span=span)
+
+
+def _wrap_function_wrapper_exception(module, name, wrapper):
+    try:
+        wrap_function_wrapper(module, name, wrapper)
+    except (ImportError, AttributeError):
+        log.debug("IAST patching. Module %s.%s not exists", module, name)