Skip to content

Commit 80ec4b6

Browse files
avara1986avara1986
authored
chore(iast): removing redundant operations (#13410)
Removing redundant operations: - call start_iast_context with span to skip a core.get_span() call - remove redundant calls to asm_config.is_iast_request_enabled and asm_config._iast_enabled - remove @WeakHash.wrap because it's a redundant call to asm_config.is_iast_request_enabled This PR is part of #13347 APPSEC-57161 ## 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)
1 parent a3497e7 commit 80ec4b6

File tree

12 files changed

+43
-97
lines changed

12 files changed

+43
-97
lines changed

‎ddtrace/appsec/_deduplications.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ def _check_deduplication(self):
3232
return asm_config._asm_deduplication_enabled
3333

3434
def __call__(self, *args, **kwargs):
35-
result = None
35+
result = False
3636
if self._check_deduplication():
3737
raw_log_hash = hash("".join([str(arg) for arg in self._extract(args)]))
3838
last_reported_timestamp = self.reported_logs.get(raw_log_hash, M_INF) + self._time_lapse

‎ddtrace/appsec/_iast/_iast_request_context_base.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,10 @@ def _set_span_tag_iast_request_tainted(span):
2828
span.set_tag(IAST_SPAN_TAGS.TELEMETRY_REQUEST_TAINTED, total_objects_tainted)
2929

3030

31-
def start_iast_context():
31+
def start_iast_context(span: Optional["Span"] = None):
3232
if asm_config._iast_enabled:
3333
create_propagation_context()
34-
core.set_item(IAST.REQUEST_CONTEXT_KEY, IASTEnvironment())
34+
core.set_item(IAST.REQUEST_CONTEXT_KEY, IASTEnvironment(span))
3535

3636

3737
def end_iast_context(span: Optional["Span"] = None):
@@ -81,7 +81,7 @@ def _move_iast_data_to_root_span():
8181
def _iast_start_request(span=None, *args, **kwargs):
8282
try:
8383
if asm_config._iast_enabled:
84-
start_iast_context()
84+
start_iast_context(span)
8585
request_iast_enabled = False
8686
if oce.acquire_request(span):
8787
request_iast_enabled = True

‎ddtrace/appsec/_iast/constants.py‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
11
import re
2-
from typing import Any
3-
from typing import Dict
42

53

64
VULN_INSECURE_HASHING_TYPE = "WEAK_HASH"
@@ -18,8 +16,6 @@
1816
VULN_SSRF = "SSRF"
1917
VULN_STACKTRACE_LEAK = "STACKTRACE_LEAK"
2018

21-
VULNERABILITY_TOKEN_TYPE = Dict[int, Dict[str, Any]]
22-
2319
HEADER_NAME_VALUE_SEPARATOR = ": "
2420

2521
MD5_DEF = "md5"

‎ddtrace/appsec/_iast/taint_sinks/_base.py‎

Lines changed: 0 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
import os
22
import sysconfig
3-
from typing import Any
4-
from typing import Callable
53
from typing import Optional
64
from typing import Tuple
75
from typing import Union
@@ -18,7 +16,6 @@
1816
from .._iast_request_context import set_iast_reporter
1917
from .._overhead_control_engine import Operation
2018
from .._stacktrace import get_info_frame
21-
from .._utils import _is_iast_debug_enabled
2219
from ..reporter import Evidence
2320
from ..reporter import IastSpanReporter
2421
from ..reporter import Location
@@ -60,26 +57,6 @@ class VulnerabilityBase(Operation):
6057
vulnerability_type = ""
6158
secure_mark = 0
6259

63-
@classmethod
64-
def wrap(cls, func: Callable) -> Callable:
65-
def wrapper(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
66-
"""Get the current root Span and attach it to the wrapped function. We need the span to report the
67-
vulnerability and update the context with the report information.
68-
"""
69-
if not asm_config.is_iast_request_enabled:
70-
if _is_iast_debug_enabled():
71-
log.debug(
72-
"iast::propagation::context::VulnerabilityBase.wrapper. No request quota or this vulnerability "
73-
"is outside the context"
74-
)
75-
return wrapped(*args, **kwargs)
76-
elif cls.has_quota():
77-
return func(wrapped, instance, args, kwargs)
78-
else:
79-
return wrapped(*args, **kwargs)
80-
81-
return wrapper
82-
8360
@classmethod
8461
@taint_sink_deduplication
8562
def _prepare_report(
@@ -93,13 +70,6 @@ def _prepare_report(
9370
*args,
9471
**kwargs,
9572
) -> bool:
96-
if not asm_config.is_iast_request_enabled:
97-
if _is_iast_debug_enabled():
98-
log.debug(
99-
"iast::propagation::context::VulnerabilityBase._prepare_report. "
100-
"No request quota or this vulnerability is outside the context"
101-
)
102-
return False
10373
if line_number is not None and (line_number == 0 or line_number < -1):
10474
line_number = -1
10575

‎ddtrace/appsec/_iast/taint_sinks/code_injection.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ def unpatch():
4949

5050

5151
def _iast_coi(wrapped, instance, args, kwargs):
52-
if asm_config._iast_enabled and len(args) >= 1:
52+
if len(args) >= 1 and asm_config.is_iast_request_enabled:
5353
_iast_report_code_injection(args[0])
5454

5555
caller_frame = None
@@ -85,7 +85,7 @@ def _iast_report_code_injection(code_string: Text):
8585
reported = False
8686
try:
8787
if asm_config.is_iast_request_enabled:
88-
if isinstance(code_string, IAST.TEXT_TYPES) and CodeInjection.has_quota():
88+
if code_string and isinstance(code_string, IAST.TEXT_TYPES) and CodeInjection.has_quota():
8989
if CodeInjection.is_tainted_pyobject(code_string):
9090
CodeInjection.report(evidence_value=code_string)
9191

‎ddtrace/appsec/_iast/taint_sinks/header_injection.py‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ def unpatch():
101101

102102

103103
def _iast_h(wrapped, instance, args, kwargs):
104-
if asm_config._iast_enabled and args:
104+
if asm_config.is_iast_request_enabled:
105105
_iast_report_header_injection(args)
106106
if hasattr(wrapped, "__func__"):
107107
return wrapped.__func__(instance, *args, **kwargs)
@@ -129,17 +129,16 @@ def _process_header(headers_args):
129129
if header_name_lower == header_to_exclude or header_name_lower.startswith(header_to_exclude):
130130
return
131131

132-
if asm_config.is_iast_request_enabled:
133-
if HeaderInjection.has_quota() and (
134-
HeaderInjection.is_tainted_pyobject(header_name) or HeaderInjection.is_tainted_pyobject(header_value)
135-
):
136-
header_evidence = add_aspect(add_aspect(header_name, HEADER_NAME_VALUE_SEPARATOR), header_value)
137-
HeaderInjection.report(evidence_value=header_evidence)
138-
139-
# Reports Span Metrics
140-
increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, HeaderInjection.vulnerability_type)
141-
# Report Telemetry Metrics
142-
_set_metric_iast_executed_sink(HeaderInjection.vulnerability_type)
132+
if HeaderInjection.has_quota() and (
133+
HeaderInjection.is_tainted_pyobject(header_name) or HeaderInjection.is_tainted_pyobject(header_value)
134+
):
135+
header_evidence = add_aspect(add_aspect(header_name, HEADER_NAME_VALUE_SEPARATOR), header_value)
136+
HeaderInjection.report(evidence_value=header_evidence)
137+
138+
# Reports Span Metrics
139+
increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, HeaderInjection.vulnerability_type)
140+
# Report Telemetry Metrics
141+
_set_metric_iast_executed_sink(HeaderInjection.vulnerability_type)
143142
except Exception as e:
144143
iast_error(f"propagation::sink_point::Error in _iast_report_header_injection. {e}")
145144

‎ddtrace/appsec/_iast/taint_sinks/insecure_cookie.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ def _iast_response_cookies(wrapped, instance, args, kwargs):
150150
cookie_value = kwargs.get("value")
151151

152152
if cookie_value and cookie_key:
153-
if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
153+
if asm_config.is_iast_request_enabled and CookiesVulnerability.has_quota():
154154
report_samesite = False
155155
samesite = kwargs.get("samesite", "")
156156
if samesite:

‎ddtrace/appsec/_iast/taint_sinks/sql_injection.py‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ class SqlInjection(VulnerabilityBase):
2525

2626

2727
def check_and_report_sqli(
28-
args: Tuple[Text], kwargs: Dict[str, Any], integration_name: Text, method: Callable[..., Any]
28+
args: Tuple[Text, ...], kwargs: Dict[str, Any], integration_name: Text, method: Callable[..., Any]
2929
) -> bool:
3030
"""Check for SQL injection vulnerabilities in database operations and report them.
3131
@@ -41,14 +41,8 @@ def check_and_report_sqli(
4141
reported = False
4242
try:
4343
if supported_dbapi_integration(integration_name) and method.__name__ == "execute":
44-
if (
45-
len(args)
46-
and args[0]
47-
and isinstance(args[0], IAST.TEXT_TYPES)
48-
and asm_config.is_iast_request_enabled
49-
and SqlInjection.has_quota()
50-
):
51-
if SqlInjection.is_tainted_pyobject(args[0]):
44+
if len(args) and args[0] and isinstance(args[0], IAST.TEXT_TYPES) and asm_config.is_iast_request_enabled:
45+
if SqlInjection.has_quota() and SqlInjection.is_tainted_pyobject(args[0]):
5246
SqlInjection.report(evidence_value=args[0], dialect=integration_name)
5347
reported = True
5448

‎ddtrace/appsec/_iast/taint_sinks/weak_cipher.py‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -130,12 +130,12 @@ def wrapped_aux_blowfish_function(wrapped, instance, args, kwargs):
130130
return result
131131

132132

133-
@WeakCipher.wrap
134133
def wrapped_rc4_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
135134
if asm_config.is_iast_request_enabled:
136-
WeakCipher.report(
137-
evidence_value="RC4",
138-
)
135+
if WeakCipher.has_quota():
136+
WeakCipher.report(
137+
evidence_value="RC4",
138+
)
139139
# Reports Span Metrics
140140
increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakCipher.vulnerability_type)
141141
# Report Telemetry Metrics
@@ -146,12 +146,12 @@ def wrapped_rc4_function(wrapped: Callable, instance: Any, args: Any, kwargs: An
146146
return wrapped(*args, **kwargs)
147147

148148

149-
@WeakCipher.wrap
150149
def wrapped_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
151150
if asm_config.is_iast_request_enabled:
152151
if hasattr(instance, "_dd_weakcipher_algorithm"):
153-
evidence = instance._dd_weakcipher_algorithm + "_" + str(instance.__class__.__name__)
154-
WeakCipher.report(evidence_value=evidence)
152+
if WeakCipher.has_quota():
153+
evidence = instance._dd_weakcipher_algorithm + "_" + str(instance.__class__.__name__)
154+
WeakCipher.report(evidence_value=evidence)
155155

156156
# Reports Span Metrics
157157
increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakCipher.vulnerability_type)
@@ -163,14 +163,14 @@ def wrapped_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -
163163
return wrapped(*args, **kwargs)
164164

165165

166-
@WeakCipher.wrap
167166
def wrapped_cryptography_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
168167
if asm_config.is_iast_request_enabled:
169168
algorithm_name = instance.algorithm.name.lower()
170169
if algorithm_name in get_weak_cipher_algorithms():
171-
WeakCipher.report(
172-
evidence_value=algorithm_name,
173-
)
170+
if WeakCipher.has_quota():
171+
WeakCipher.report(
172+
evidence_value=algorithm_name,
173+
)
174174

175175
# Reports Span Metrics
176176
increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakCipher.vulnerability_type)

‎ddtrace/appsec/_iast/taint_sinks/weak_hash.py‎

Lines changed: 10 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,10 @@
11
import os
22
import sys
3-
from typing import TYPE_CHECKING # noqa:F401
43
from typing import Any
54
from typing import Callable
65
from typing import Set
7-
from typing import Text # noqa:F401
86

97
from ddtrace.appsec._common_module_patches import try_unwrap
10-
from ddtrace.internal.logger import get_logger
118
from ddtrace.settings.asm import config as asm_config
129

1310
from ..._constants import IAST_SPAN_TAGS
@@ -25,9 +22,6 @@
2522
from ._base import VulnerabilityBase
2623

2724

28-
log = get_logger(__name__)
29-
30-
3125
def get_weak_hash_algorithms() -> Set:
3226
CONFIGURED_WEAK_HASH_ALGORITHMS = None
3327
DD_IAST_WEAK_HASH_ALGORITHMS = os.getenv("DD_IAST_WEAK_HASH_ALGORITHMS")
@@ -65,7 +59,7 @@ def unpatch_iast():
6559
try_unwrap("Crypto.Hash.SHA1", "SHA1Hash.hexdigest")
6660

6761

68-
def get_version() -> Text:
62+
def get_version() -> str:
6963
return ""
7064

7165

@@ -119,35 +113,31 @@ def patch():
119113
_set_metric_iast_instrumented_sink(VULN_INSECURE_HASHING_TYPE, num_instrumented_sinks)
120114

121115

122-
@WeakHash.wrap
123116
def wrapped_digest_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
124117
if asm_config.is_iast_request_enabled:
125118
if WeakHash.has_quota() and instance.name.lower() in get_weak_hash_algorithms():
126119
WeakHash.report(
127120
evidence_value=instance.name,
128121
)
129122

130-
# Reports Span Metrics
131-
increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakHash.vulnerability_type)
132-
# Report Telemetry Metrics
133-
_set_metric_iast_executed_sink(WeakHash.vulnerability_type)
123+
# Reports Span Metrics
124+
increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakHash.vulnerability_type)
125+
# Report Telemetry Metrics
126+
_set_metric_iast_executed_sink(WeakHash.vulnerability_type)
134127

135128
if hasattr(wrapped, "__func__"):
136129
return wrapped.__func__(instance, *args, **kwargs)
137130
return wrapped(*args, **kwargs)
138131

139132

140-
@WeakHash.wrap
141133
def wrapped_md5_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
142134
return wrapped_function(wrapped, MD5_DEF, instance, args, kwargs)
143135

144136

145-
@WeakHash.wrap
146137
def wrapped_sha1_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
147138
return wrapped_function(wrapped, SHA1_DEF, instance, args, kwargs)
148139

149140

150-
@WeakHash.wrap
151141
def wrapped_new_function(wrapped: Callable, instance: Any, args: Any, kwargs: Any) -> Any:
152142
if asm_config.is_iast_request_enabled:
153143
if WeakHash.has_quota() and args[0].lower() in get_weak_hash_algorithms():
@@ -160,11 +150,12 @@ def wrapped_new_function(wrapped: Callable, instance: Any, args: Any, kwargs: An
160150
return wrapped(*args, **kwargs)
161151

162152

163-
def wrapped_function(wrapped: Callable, evidence: Text, instance: Any, args: Any, kwargs: Any) -> Any:
153+
def wrapped_function(wrapped: Callable, evidence: str, instance: Any, args: Any, kwargs: Any) -> Any:
164154
if asm_config.is_iast_request_enabled:
165-
WeakHash.report(
166-
evidence_value=evidence,
167-
)
155+
if WeakHash.has_quota():
156+
WeakHash.report(
157+
evidence_value=evidence,
158+
)
168159
# Reports Span Metrics
169160
increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakHash.vulnerability_type)
170161
# Report Telemetry Metrics

0 commit comments

Comments
 (0)