feat(iast): add support for langchain v0.1.0+

Author: smola

ddtrace/appsec/_iast/_ast/iastpatch.c

@@ -18,8 +18,8 @@ static size_t cached_packages_count = 0;
 
 /* Static Lists */
 static const char* static_allowlist[] = {
-    "jinja2.",     "pygments.", "multipart.", "sqlalchemy.", "python_multipart.", "attrs.",
-    "jsonschema.", "s3fs.",     "mysql.",     "pymysql.",    "markupsafe.",       "werkzeug.utils."
+    "jinja2.", "pygments.", "multipart.", "sqlalchemy.", "python_multipart.", "attrs.", "jsonschema.",
+    "s3fs.",   "mysql.",    "pymysql.",   "markupsafe.", "werkzeug.utils.",   "langchain_core."
 };
 static const size_t static_allowlist_count = sizeof(static_allowlist) / sizeof(static_allowlist[0]);
 

ddtrace/appsec/_iast/_iast_request_context.py

@@ -23,6 +23,7 @@
 
 
 def set_iast_reporter(iast_reporter: IastSpanReporter) -> None:
+    print("set_iast_reporter")
     env = _get_iast_env()
     if env:
         env.iast_reporter = iast_reporter
@@ -33,7 +34,9 @@ def set_iast_reporter(iast_reporter: IastSpanReporter) -> None:
 def get_iast_reporter() -> Optional[IastSpanReporter]:
     env = _get_iast_env()
     if env:
+        print("GOT ENV")
         return env.iast_reporter
+    print("GOT NO ENV")
     return None
 
 

ddtrace/appsec/_iast/_overhead_control_engine.py

@@ -68,6 +68,7 @@ def has_quota(cls) -> bool:
     @classmethod
     def is_not_reported(cls, filename: Text, lineno: int) -> bool:
         if asm_config._iast_deduplication_enabled:
+            print("DAMN DEDUP ENABLED")
             vulnerability_id = (filename, lineno)
             if vulnerability_id in cls._reported_vulnerabilities:
                 return False

ddtrace/appsec/_iast/taint_sinks/_base.py

@@ -189,6 +189,7 @@ def _create_evidence_and_report(
     @classmethod
     def report(cls, evidence_value: TEXT_TYPES = "", dialect: Optional[str] = None) -> None:
         """Build a IastSpanReporter instance to report it in the `AppSecIastSpanProcessor` as a string JSON"""
+        print(f"report {cls.vulnerability_type}")
         if cls.acquire_quota():
             file_name = line_number = function_name = class_name = None
 
@@ -198,6 +199,7 @@ def report(cls, evidence_value: TEXT_TYPES = "", dialect: Optional[str] = None)
             else:
                 file_name, line_number, function_name, class_name = cls._compute_file_line()
                 if file_name is None:
+                    print("NO LOCATION FILE NAME")
                     cls.increment_quota()
                     return
 

ddtrace/appsec/_iast/taint_sinks/command_injection.py

@@ -48,7 +48,6 @@ class CommandInjection(VulnerabilityBase):
 
 def _iast_report_cmdi(shell_args: Union[str, List[str]]) -> None:
     report_cmdi = ""
-
     try:
         if asm_config.is_iast_request_enabled:
             if CommandInjection.has_quota():

ddtrace/contrib/internal/langchain/patch.py

@@ -196,6 +196,9 @@ def traced_llm_generate(langchain, pin, func, instance, args, kwargs):
                 span.set_tag_str("langchain.request.%s.parameters.%s" % (llm_provider, param), str(val))
 
         completions = func(*args, **kwargs)
+
+        _iast_taint_llm_output(prompts, completions)
+
         if _is_openai_llm_instance(instance):
             _tag_openai_token_usage(span, completions.llm_output)
 
@@ -252,6 +255,9 @@ async def traced_llm_agenerate(langchain, pin, func, instance, args, kwargs):
                 span.set_tag_str("langchain.request.%s.parameters.%s" % (llm_provider, param), str(val))
 
         completions = await func(*args, **kwargs)
+
+        _iast_taint_llm_output(prompts, completions)
+
         if _is_openai_llm_instance(instance):
             _tag_openai_token_usage(span, completions.llm_output)
 
@@ -942,6 +948,59 @@ async def traced_base_tool_ainvoke(langchain, pin, func, instance, args, kwargs)
     return tool_output
 
 
+def _iast_taint_llm_output(prompts, completions):
+    """
+    Taints the output of an LLM call if its inputs are tainted.
+
+    Range propagation does not make sense in LLMs. So we get the first source in inputs, if any,
+    and taint the full output with that source.
+    """
+    print(f"_iast_taint_llm_output: prompts={prompts}, completions={completions}")
+    if not asm_config._iast_enabled:
+        return
+    if not isinstance(prompts, (tuple, list)):
+        return
+    if not hasattr(completions, "generations"):
+        return
+    try:
+        generations = completions.generations
+        if not isinstance(generations, list):
+            return
+
+        from ddtrace.appsec._iast._taint_tracking._taint_objects import get_tainted_ranges
+        from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
+
+        source = None
+        for prompt in prompts:
+            if not isinstance(prompt, str):
+                continue
+            tainted_ranges = get_tainted_ranges(prompt)
+            if tainted_ranges:
+                source = tainted_ranges[0].source
+                break
+        if not source:
+            return
+        for gens in generations:
+            for gen in gens:
+                if not hasattr(gen, "text"):
+                    continue
+                text = gen.text
+                if not isinstance(text, str):
+                    continue
+                print("TAINT LLM OUTPUT")
+                new_text = taint_pyobject(
+                    pyobject=text,
+                    source_name=source.name,
+                    source_value=source.value,
+                    source_origin=source.origin,
+                )
+                setattr(gen, "text", new_text)
+    except Exception as e:
+        from ddtrace.appsec._iast._metrics import _set_iast_error_metric
+
+        _set_iast_error_metric("IAST propagation error. langchain _iast_taint_llm_output. {}".format(e))
+
+
 def _patch_embeddings_and_vectorstores():
     """
     Text embedding models override two abstract base methods instead of super calls,
@@ -1081,10 +1140,15 @@ def patch():
     if asm_config._iast_enabled:
         from ddtrace.appsec._iast._metrics import _set_iast_error_metric
 
+        wrap("langchain_core", "prompts.prompt.PromptTemplate.format", iast_propagate_prompt_template_format)
+        wrap("langchain_core", "prompts.prompt.PromptTemplate.aformat", iast_propagate_prompt_template_aformat)
+
         def wrap_output_parser(module, parser):
             # Ensure not double patched
             if not isinstance(deep_getattr(module, "%s.parse" % parser), wrapt.ObjectProxy):
-                wrap(module, "%s.parse" % parser, taint_parser_output)
+                wrap(module, "%s.parse" % parser, iast_propagate_output_parse)
+            if not isinstance(deep_getattr(module, "%s.aparse" % parser), wrapt.ObjectProxy):
+                wrap(module, "%s.aparse" % parser, iast_propagate_output_aparse)
 
         try:
             with_agent_output_parser(wrap_output_parser)
@@ -1114,6 +1178,7 @@ def unpatch():
     unwrap(langchain_core.language_models.llms.BaseLLM, "astream")
     unwrap(langchain_core.tools.BaseTool, "invoke")
     unwrap(langchain_core.tools.BaseTool, "ainvoke")
+
     if langchain_openai:
         unwrap(langchain_openai.OpenAIEmbeddings, "embed_documents")
     if langchain_pinecone:
@@ -1122,16 +1187,70 @@ def unpatch():
     if langchain_community:
         _unpatch_embeddings_and_vectorstores()
 
+    if asm_config._iast_enabled:
+        unwrap(langchain_core.prompts.prompt.PromptTemplate, "format")
+        unwrap(langchain_core.prompts.prompt.PromptTemplate, "aformat")
+
     delattr(langchain, "_datadog_integration")
 
 
-def taint_parser_output(func, instance, args, kwargs):
-    from ddtrace.appsec._iast._metrics import _set_iast_error_metric
-    from ddtrace.appsec._iast._taint_tracking._taint_objects import get_tainted_ranges
-    from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
+def iast_propagate_prompt_template_format(func, instance, args, kwargs):
+    """
+    Propagate taint in PromptTemplate.format, from any input, to the output.
+    """
+    result = func(*args, **kwargs)
+    return _iast_propagate_prompt_template_format_inner(kwargs, result)
+
+
+async def iast_propagate_prompt_template_aformat(func, instance, args, kwargs):
+    """
+    Propagate taint in PromptTemplate.aformat, from any input, to the output.
+    """
+    result = await func(*args, **kwargs)
+    return _iast_propagate_prompt_template_format_inner(kwargs, result)
+
+
+def _iast_propagate_prompt_template_format_inner(kwargs, result):
+    print(f"_iast_propagate_prompt_template_format_inner: {kwargs}, {result}")
+    try:
+        if not asm_config.is_iast_request_enabled:
+            return result
+
+        from ddtrace.appsec._iast._taint_tracking._taint_objects import get_tainted_ranges
+        from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
+
+        source = None
+        for value in kwargs.values():
+            ranges = get_tainted_ranges(value)
+            if ranges:
+                source = ranges[0].source
+                break
+        if source:
+            print("TAINTED TEMPLATE FORMAT")
+            return taint_pyobject(result, source.name, source.value, source.origin)
+    except Exception as e:
+        from ddtrace.appsec._iast._metrics import _set_iast_error_metric
+
+        _set_iast_error_metric("IAST propagation error. langchain iast_propagate_prompt_template_format. {}".format(e))
+    return result
+
 
+def iast_propagate_output_parse(func, instance, args, kwargs):
     result = func(*args, **kwargs)
+    return _iast_propagate_output_parse_inner(args, kwargs, result)
+
+
+async def iast_propagate_output_aparse(func, instance, args, kwargs):
+    result = await func(*args, **kwargs)
+    return _iast_propagate_output_parse_inner(args, kwargs, result)
+
+
+def _iast_propagate_output_parse_inner(args, kwargs, result):
+    print(f"_iast_propagate_output_parse_inner: {args}, {kwargs}, {result}")
     try:
+        from ddtrace.appsec._iast._taint_tracking._taint_objects import get_tainted_ranges
+        from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
+
         try:
             from langchain_core.agents import AgentAction
             from langchain_core.agents import AgentFinish
@@ -1141,12 +1260,17 @@ def taint_parser_output(func, instance, args, kwargs):
         ranges = get_tainted_ranges(args[0])
         if ranges:
             source = ranges[0].source
+            print("WILL TAINT")
             if isinstance(result, AgentAction):
+                print("TAINTED TOOL INPUT")
                 result.tool_input = taint_pyobject(result.tool_input, source.name, source.value, source.origin)
             elif isinstance(result, AgentFinish) and "output" in result.return_values:
+                print("TAINTED OUTPUT")
                 values = result.return_values
                 values["output"] = taint_pyobject(values["output"], source.name, source.value, source.origin)
     except Exception as e:
+        from ddtrace.appsec._iast._metrics import _set_iast_error_metric
+
         _set_iast_error_metric("IAST propagation error. langchain taint_parser_output. {}".format(e))
 
     return result

hatch.toml

@@ -601,6 +601,47 @@ fastapi = ["==0.94.1"]
 python = ["3.8", "3.10", "3.13"]
 fastapi = ["~=0.114.2"]
 
+## ASM appsec_integrations_langchain
+
+[envs.appsec_integrations_langchain]
+template = "appsec_integrations_langchain"
+dependencies = [
+    "pytest",
+    "pytest-asyncio",
+    "pytest-cov",
+    "hypothesis",
+    "langchain{matrix:langchain:}",
+    "langchain-experimental{matrix:langchain-experimental:}",
+]
+
+[envs.appsec_integrations_langchain.env-vars]
+DD_TRACE_AGENT_URL = "http://testagent:9126"
+_DD_IAST_PATCH_MODULES = "benchmarks.,tests.appsec."
+DD_IAST_REQUEST_SAMPLING = "100"
+DD_IAST_DEDUPLICATION_ENABLED = "false"
+
+[envs.appsec_integrations_langchain.scripts]
+test = [
+    "uname -a",
+    "pip freeze",
+    "python -m pytest -vvv {args:tests/appsec/integrations/langchain_tests/}",
+]
+
+[[envs.appsec_integrations_langchain.matrix]]
+python = ["3.9", "3.10", "3.11", "3.12", "3.13"]
+langchain = ["~=0.1"]
+langchain-experimental = ["~=0.1"]
+
+[[envs.appsec_integrations_langchain.matrix]]
+python = ["3.9", "3.10", "3.11", "3.12", "3.13"]
+langchain = ["~=0.2"]
+langchain-experimental = ["~=0.2"]
+
+[[envs.appsec_integrations_langchain.matrix]]
+python = ["3.9", "3.10", "3.11", "3.12", "3.13"]
+langchain = ["~=0.3"]
+langchain-experimental = ["~=0.3"]
+
 ## ASM FastAPI
 
 [envs.appsec_threats_fastapi]

releasenotes/notes/iast-langchain-0.1.0-e437ea90fe66ad31.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Code Security: IAST support for langchain v0.1.0 and above.

tests/appsec/iast/conftest.py

@@ -52,6 +52,7 @@ def _start_iast_context_and_oce(span=None):
     if oce.acquire_request(span):
         start_iast_context()
         request_iast_enabled = True
+
     set_iast_request_enabled(request_iast_enabled)
 
 
@@ -61,13 +62,6 @@ def _end_iast_context_and_oce(span=None):
 
 
 def iast_context(env, request_sampling=100.0, deduplication=False, asm_enabled=False):
-    try:
-        from ddtrace.contrib.internal.langchain.patch import patch as langchain_patch
-        from ddtrace.contrib.internal.langchain.patch import unpatch as langchain_unpatch
-    except Exception:
-        langchain_patch = lambda: True  # noqa: E731
-        langchain_unpatch = lambda: True  # noqa: E731
-
     class MockSpan:
         _trace_id_64bits = 17577308072598193742
 
@@ -87,7 +81,6 @@ class MockSpan:
         cmdi_patch()
         header_injection_patch()
         code_injection_patch()
-        langchain_patch()
         patch_common_modules()
         yield
         unpatch_common_modules()
@@ -97,7 +90,6 @@ class MockSpan:
         cmdi_unpatch()
         header_injection_unpatch()
         code_injection_unpatch()
-        langchain_unpatch()
         _end_iast_context_and_oce()
 
 

tests/appsec/iast/iast_utils.py

@@ -53,6 +53,7 @@ def get_line_and_hash(label: Text, vuln_type: Text, filename=None, fixed_line=No
 def _iast_patched_module_and_patched_source(module_name, new_module_object=False):
     module = importlib.import_module(module_name)
     module_path, patched_source = astpatch_module(module)
+    assert patched_source is not None
     compiled_code = compile(patched_source, module_path, "exec")
     module_changed = types.ModuleType(module_name) if new_module_object else module
     exec(compiled_code, module_changed.__dict__)