Skip to content

Commit a451470

Browse files
smolaavara1986
andauthored
chore(iast): fix for dict add inplace (#13372)
[APPSEC-57605](https://datadoghq.atlassian.net/browse/APPSEC-57605) ## 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) [APPSEC-57605]: https://datadoghq.atlassian.net/browse/APPSEC-57605?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ --------- Co-authored-by: Alberto Vara <[email protected]>
1 parent 80ec4b6 commit a451470

File tree

3 files changed

+708
-381
lines changed

3 files changed

+708
-381
lines changed

benchmarks/bm/iast_fixtures/str_methods.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,36 @@ def do_operator_add_inplace_params(a, b):
5555
return a
5656

5757

58+
def do_operator_add_inplace_dict_key(a, key, b):
59+
a[key] += b
60+
return a[key]
61+
62+
63+
def _get_dictionary(dictionary):
64+
return dictionary
65+
66+
67+
def do_operator_add_inplace_dict_key_from_function(a, key, b):
68+
_get_dictionary(a)[key] += b
69+
return a[key]
70+
71+
72+
class MyClassWithDict(object):
73+
data = {}
74+
75+
def __init__(self, my_cool_dict):
76+
self.data = my_cool_dict
77+
78+
def get_data(self):
79+
return self.data
80+
81+
82+
def do_operator_add_inplace_dict_key_from_class(a, key, b):
83+
my_cool_data = MyClassWithDict(a)
84+
my_cool_data.get_data()[key] += b
85+
return a[key]
86+
87+
5888
def do_operator_add_inplace_3_params(a, b, c):
5989
a += b
6090
a += c

ddtrace/appsec/_iast/_ast/visitor.py

Lines changed: 23 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
from ..constants import DEFAULT_WEAK_RANDOMNESS_FUNCTIONS
2020

2121

22-
PY3 = sys.version_info[0] >= 3
2322
PY39_PLUS = sys.version_info >= (3, 9, 0)
2423

2524
_PREFIX = IAST.PATCH_ADDED_SYMBOL_PREFIX
@@ -230,7 +229,7 @@ def _merge_dicts(*args_functions: Set[str]) -> Set[str]:
230229

231230
@staticmethod
232231
def _is_string_node(node: Any) -> bool:
233-
if PY3 and (isinstance(node, ast.Constant) and isinstance(node.value, IAST.TEXT_TYPES)):
232+
if isinstance(node, ast.Constant) and isinstance(node.value, IAST.TEXT_TYPES):
234233
return True
235234

236235
return False
@@ -652,46 +651,39 @@ def visit_AugAssign(self, augassign_node: ast.AugAssign) -> Any:
652651
"""
653652
Replace an inplace add or multiply (+= / *=)
654653
"""
655-
if isinstance(augassign_node.target, ast.Subscript):
656-
# Can't augassign to function call, ignore this node
657-
augassign_node.target.avoid_convert = True # type: ignore[attr-defined]
658-
self.generic_visit(augassign_node)
659-
return augassign_node
660-
661654
self.generic_visit(augassign_node)
655+
662656
if augassign_node.op.__class__ == ast.Add:
663657
# Optimization: ignore augassigns where the right side term is an integer since
664658
# they can't apply to strings
665659
if self._is_numeric_node(augassign_node.value):
666660
return augassign_node
667661

668662
replacement_func = self._aspect_operators["INPLACE_ADD"]
669-
else:
670-
return augassign_node
671-
672-
# We must change the ctx of the target and value to Load() for using
673-
# them as function arguments while keeping it as Store() for the
674-
# Assign.targets, thus the manual copy
675663

676-
func_arg1 = copy.deepcopy(augassign_node.target)
677-
func_arg1.ctx = ast.Load()
678-
func_arg2 = copy.deepcopy(augassign_node.value)
679-
func_arg2.ctx = ast.Load() # type: ignore[attr-defined]
664+
# Regular inplace add for non-subscript targets
665+
func_arg1 = copy.deepcopy(augassign_node.target)
666+
func_arg1.ctx = ast.Load()
667+
func_arg2 = copy.deepcopy(augassign_node.value)
668+
if hasattr(func_arg2, "ctx"):
669+
func_arg2.ctx = ast.Load()
670+
671+
call_node = self._call_node(
672+
augassign_node,
673+
func=self._attr_node(augassign_node, replacement_func),
674+
args=[func_arg1, func_arg2],
675+
)
680676

681-
call_node = self._call_node(
682-
augassign_node,
683-
func=self._attr_node(augassign_node, replacement_func),
684-
args=[func_arg1, func_arg2],
685-
)
677+
self.ast_modified = True
678+
return self._node(
679+
ast.Assign,
680+
augassign_node,
681+
targets=[augassign_node.target],
682+
value=call_node,
683+
type_comment=None,
684+
)
686685

687-
self.ast_modified = True
688-
return self._node(
689-
ast.Assign,
690-
augassign_node,
691-
targets=[augassign_node.target],
692-
value=call_node,
693-
type_comment=None,
694-
)
686+
return augassign_node
695687

696688
def visit_FormattedValue(self, fmt_value_node: ast.FormattedValue) -> Any:
697689
"""

0 commit comments

Comments
 (0)