chore(iast): refactor import modules hook #13665
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
avara1986
added
changelog/no-changelog
|
|
P403n1x87
reviewed
Jun 13, 2025
Bootstrap import analysisComparison of import times between this PR and base. SummaryThe average import time from this PR is: 273 ± 2 ms. The average import time from base is: 275 ± 2 ms. The import time difference between this PR and base is: -2.01 ± 0.09 ms. Import time breakdownThe following import paths have shrunk:
|
BenchmarksBenchmark execution time: 2025-06-17 10:58:59 Comparing candidate commit 1a516e7 in PR branch Found 0 performance improvements and 0 performance regressions! Performance is the same for 561 metrics, 3 unstable metrics. |
christophe-papazian
approved these changes
Jun 17, 2025
happynancee
pushed a commit
that referenced
this pull request
Jul 7, 2025
This PR fixes several logic issues and removes code that was duplicated
dozens of times.
Move _iast_instrument_starlette_url to _handlers
Since `iast/_patch.py` and `iast/_patch_modules.py` had overlapping
responsibilities and it was unclear what each file was for, the
functions that serve as the entry point for starting IAST have been
moved to main.py.
The new class that handles all the logic for wrapping functions and
modules now lives in patch_modules.
On one hand, we had the function `try_wrap_function_wrapper` defined
twice — once for AppSec and once for IAST:
The **AppSec** version used `ModuleWatchdog`:
```python
# ddtrace/appsec/_common_module_patches.py
def try_wrap_function_wrapper(module_name: str, name: str, wrapper: Callable) -> None:
@ModuleWatchdog.after_module_imported(module_name)
def _(module):
try:
wrap_object(module, name, FunctionWrapper, (wrapper,))
except (ImportError, AttributeError):
log.debug("ASM patching. Module %s.%s does not exist", module_name, name)
```
Whereas the **IAST** version used `wrapt` directly:
```python
# ddtrace/appsec/_iast/_patch.py
def try_wrap_function_wrapper(module: Text, name: Text, wrapper: Callable):
try:
wrap_object(module, name, FunctionWrapper, (wrapper,))
except (ImportError, AttributeError):
iast_instrumentation_wrapt_debug_log(f"Module {module}.{name} not exists")
```
Additionally, the IAST version was sometimes called like this:
```python
try_wrap_function_wrapper(("_%s" % MD5_DEF), "MD5Type.digest", wrapped_md5_function)
```
And other times using `wrapt.when_imported`:
```python
@when_imported("django.shortcuts")
def _(m):
try_wrap_function_wrapper(m, "redirect", _unvalidated_redirect_for_django)
```
These two patterns were repeated dozens of times, and now all of this
logic has been centralized in the class
`ddtrace.appsec._iast._patch_modules.WrapModulesForIAST`.
## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))
## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
alyshawang
pushed a commit
that referenced
this pull request
Jul 25, 2025
This PR fixes several logic issues and removes code that was duplicated
dozens of times.
Move _iast_instrument_starlette_url to _handlers
Since `iast/_patch.py` and `iast/_patch_modules.py` had overlapping
responsibilities and it was unclear what each file was for, the
functions that serve as the entry point for starting IAST have been
moved to main.py.
The new class that handles all the logic for wrapping functions and
modules now lives in patch_modules.
On one hand, we had the function `try_wrap_function_wrapper` defined
twice — once for AppSec and once for IAST:
The **AppSec** version used `ModuleWatchdog`:
```python
# ddtrace/appsec/_common_module_patches.py
def try_wrap_function_wrapper(module_name: str, name: str, wrapper: Callable) -> None:
@ModuleWatchdog.after_module_imported(module_name)
def _(module):
try:
wrap_object(module, name, FunctionWrapper, (wrapper,))
except (ImportError, AttributeError):
log.debug("ASM patching. Module %s.%s does not exist", module_name, name)
```
Whereas the **IAST** version used `wrapt` directly:
```python
# ddtrace/appsec/_iast/_patch.py
def try_wrap_function_wrapper(module: Text, name: Text, wrapper: Callable):
try:
wrap_object(module, name, FunctionWrapper, (wrapper,))
except (ImportError, AttributeError):
iast_instrumentation_wrapt_debug_log(f"Module {module}.{name} not exists")
```
Additionally, the IAST version was sometimes called like this:
```python
try_wrap_function_wrapper(("_%s" % MD5_DEF), "MD5Type.digest", wrapped_md5_function)
```
And other times using `wrapt.when_imported`:
```python
@when_imported("django.shortcuts")
def _(m):
try_wrap_function_wrapper(m, "redirect", _unvalidated_redirect_for_django)
```
These two patterns were repeated dozens of times, and now all of this
logic has been centralized in the class
`ddtrace.appsec._iast._patch_modules.WrapModulesForIAST`.
## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))
## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR fixes several logic issues and removes code that was duplicated dozens of times.
Move _iast_instrument_starlette_url to _handlers
Since
iast/_patch.pyandiast/_patch_modules.pyhad overlapping responsibilities and it was unclear what each file was for, the functions that serve as the entry point for starting IAST have been moved to main.py.The new class that handles all the logic for wrapping functions and modules now lives in patch_modules.
On one hand, we had the function
try_wrap_function_wrapperdefined twice — once for AppSec and once for IAST:The AppSec version used
ModuleWatchdog:Whereas the IAST version used
wraptdirectly:Additionally, the IAST version was sometimes called like this:
And other times using
wrapt.when_imported:These two patterns were repeated dozens of times, and now all of this logic has been centralized in the class
ddtrace.appsec._iast._patch_modules.WrapModulesForIAST.Task: APPSEC-58037
Checklist
Reviewer Checklist