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/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,57 @@ 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.
+    """
+    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
+                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 +1138,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 +1176,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 +1185,67 @@ 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):
+    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:
+            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):
     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
@@ -1147,6 +1261,8 @@ def taint_parser_output(func, instance, args, kwargs):
                 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__)

tests/appsec/integrations/fixtures/patch_langchain.py

@@ -0,0 +1,19 @@
+from ddtrace.appsec._iast import enable_iast_propagation
+from ddtrace.appsec._iast._patch_modules import patch_iast
+from ddtrace.contrib.internal.langchain.patch import patch as langchain_patch
+from tests.utils import override_env
+from tests.utils import override_global_config
+
+
+# `pytest` automatically calls this function once when tests are run.
+def pytest_configure():
+    with override_global_config(
+        dict(
+            _iast_enabled=True,
+            _iast_deduplication_enabled=False,
+            _iast_request_sampling=100.0,
+        )
+    ), override_env(dict(_DD_IAST_PATCH_MODULES="tests.appsec.integrations")):
+        patch_iast()
+        enable_iast_propagation()
+        langchain_patch()