Skip to content

Commit 6893a5d

Browse files
jandro996jandro996
committed
AppSecConfigServiceImpl integration
1 parent c542c6b commit 6893a5d

File tree

2 files changed

+250
-2
lines changed

2 files changed

+250
-2
lines changed

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java‎

Lines changed: 50 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
113113
private String currentRuleVersion;
114114
private List<AppSecModule> modulesToUpdateVersionIn;
115115
private volatile AppSecSCAConfig currentSCAConfig;
116+
private AppSecSCAInstrumentationUpdater scaInstrumentationUpdater;
116117

117118
public AppSecConfigServiceImpl(
118119
Config tracerConfig,
@@ -127,6 +128,33 @@ public AppSecConfigServiceImpl(
127128
}
128129
}
129130

131+
/**
132+
* Sets the Instrumentation instance for SCA hot instrumentation.
133+
* Must be called before {@link #maybeSubscribeConfigPolling()} for SCA to work.
134+
*
135+
* @param instrumentation the Java Instrumentation API instance
136+
*/
137+
public void setInstrumentation(java.lang.instrument.Instrumentation instrumentation) {
138+
if (instrumentation == null) {
139+
log.debug("Instrumentation is null, SCA hot instrumentation will not be available");
140+
return;
141+
}
142+
143+
if (!instrumentation.isRetransformClassesSupported()) {
144+
log.warn(
145+
"SCA requires retransformation support, but it's not available in this JVM. "
146+
+ "SCA vulnerability detection will not work.");
147+
return;
148+
}
149+
150+
try {
151+
this.scaInstrumentationUpdater = new AppSecSCAInstrumentationUpdater(instrumentation);
152+
log.debug("SCA instrumentation updater initialized successfully");
153+
} catch (Exception e) {
154+
log.error("Failed to initialize SCA instrumentation updater", e);
155+
}
156+
}
157+
130158
private void subscribeConfigurationPoller() {
131159
// see also close() method
132160
subscribeAsmFeatures();
@@ -364,7 +392,7 @@ private void subscribeSCA() {
364392
if (newConfig == null) {
365393
log.debug("Received removal for SCA config key: {}", configKey);
366394
currentSCAConfig = null;
367-
// TODO: Trigger retransformation to remove instrumentation when updater exists
395+
triggerSCAInstrumentationUpdate(null);
368396
} else {
369397
log.debug(
370398
"Received SCA config update for key: {} - enabled: {}, targets: {}",
@@ -374,7 +402,7 @@ private void subscribeSCA() {
374402
? newConfig.instrumentationTargets.size()
375403
: 0);
376404
currentSCAConfig = newConfig;
377-
// TODO: Trigger retransformation when AppSecInstrumentationUpdater exists
405+
triggerSCAInstrumentationUpdate(newConfig);
378406
}
379407
});
380408
this.configurationPoller.addCapabilities(CAPABILITY_ASM_SCA_VULNERABILITY_DETECTION);
@@ -395,6 +423,26 @@ private void unsubscribeSCA() {
395423
}
396424
}
397425

426+
/**
427+
* Triggers SCA instrumentation update when configuration changes.
428+
*
429+
* @param newConfig the new SCA configuration, or null to remove instrumentation
430+
*/
431+
private void triggerSCAInstrumentationUpdate(AppSecSCAConfig newConfig) {
432+
if (scaInstrumentationUpdater == null) {
433+
log.debug(
434+
"SCA instrumentation updater not initialized. "
435+
+ "Call setInstrumentation() before subscribing to enable SCA.");
436+
return;
437+
}
438+
439+
try {
440+
scaInstrumentationUpdater.onConfigUpdate(newConfig);
441+
} catch (Exception e) {
442+
log.error("Error updating SCA instrumentation", e);
443+
}
444+
}
445+
398446
private void distributeSubConfigurations(
399447
String key, AppSecModuleConfigurer.Reconfiguration reconfiguration) {
400448
maybeInitializeDefaultConfig();

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecSCAInstrumentationUpdater.java‎

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,200 @@
1+
package com.datadog.appsec.config;
2+
3+
import java.lang.instrument.ClassFileTransformer;
4+
import java.lang.instrument.Instrumentation;
5+
import java.util.ArrayList;
6+
import java.util.HashSet;
7+
import java.util.List;
8+
import java.util.Set;
9+
import java.util.concurrent.locks.Lock;
10+
import java.util.concurrent.locks.ReentrantLock;
11+
import org.slf4j.Logger;
12+
import org.slf4j.LoggerFactory;
13+
14+
/**
15+
* Handles dynamic instrumentation updates for Supply Chain Analysis (SCA) vulnerability detection.
16+
*
17+
* <p>This class receives SCA configuration updates from Remote Config and triggers
18+
* retransformation of classes that match the instrumentation targets.
19+
*
20+
* <p>Thread-safe: Multiple threads can call {@link #onConfigUpdate(AppSecSCAConfig)} concurrently.
21+
*/
22+
public class AppSecSCAInstrumentationUpdater {
23+
24+
private static final Logger log = LoggerFactory.getLogger(AppSecSCAInstrumentationUpdater.class);
25+
26+
private final Instrumentation instrumentation;
27+
private final Lock updateLock = new ReentrantLock();
28+
29+
private volatile AppSecSCAConfig currentConfig;
30+
private ClassFileTransformer currentTransformer;
31+
32+
public AppSecSCAInstrumentationUpdater(Instrumentation instrumentation) {
33+
if (instrumentation == null) {
34+
throw new IllegalArgumentException("Instrumentation cannot be null");
35+
}
36+
if (!instrumentation.isRetransformClassesSupported()) {
37+
throw new IllegalStateException(
38+
"SCA requires retransformation support, but it's not available in this JVM");
39+
}
40+
this.instrumentation = instrumentation;
41+
}
42+
43+
/**
44+
* Called when SCA configuration is updated via Remote Config.
45+
*
46+
* @param newConfig the new SCA configuration, or null if config was removed
47+
*/
48+
public void onConfigUpdate(AppSecSCAConfig newConfig) {
49+
updateLock.lock();
50+
try {
51+
if (newConfig == null) {
52+
log.debug("SCA config removed, disabling instrumentation");
53+
removeInstrumentation();
54+
currentConfig = null;
55+
return;
56+
}
57+
58+
if (!isEnabled(newConfig)) {
59+
log.debug("SCA config disabled, removing instrumentation");
60+
removeInstrumentation();
61+
currentConfig = newConfig;
62+
return;
63+
}
64+
65+
if (newConfig.instrumentationTargets == null || newConfig.instrumentationTargets.isEmpty()) {
66+
log.debug("SCA config has no instrumentation targets");
67+
removeInstrumentation();
68+
currentConfig = newConfig;
69+
return;
70+
}
71+
72+
log.info(
73+
"Applying SCA instrumentation for {} targets",
74+
newConfig.instrumentationTargets.size());
75+
76+
AppSecSCAConfig oldConfig = currentConfig;
77+
currentConfig = newConfig;
78+
79+
applyInstrumentation(oldConfig, newConfig);
80+
} finally {
81+
updateLock.unlock();
82+
}
83+
}
84+
85+
private boolean isEnabled(AppSecSCAConfig config) {
86+
return config.enabled != null && config.enabled;
87+
}
88+
89+
private void applyInstrumentation(AppSecSCAConfig oldConfig, AppSecSCAConfig newConfig) {
90+
// Determine which classes need to be retransformed
91+
Set<String> targetClassNames = extractTargetClassNames(newConfig);
92+
93+
if (targetClassNames.isEmpty()) {
94+
log.debug("No valid target class names found");
95+
return;
96+
}
97+
98+
// Remove old transformer if exists
99+
if (currentTransformer != null) {
100+
log.debug("Removing previous SCA transformer");
101+
instrumentation.removeTransformer(currentTransformer);
102+
currentTransformer = null;
103+
}
104+
105+
// Install new transformer
106+
// TODO: Create AppSecSCATransformer
107+
log.debug("Installing new SCA transformer for targets: {}", targetClassNames);
108+
// currentTransformer = new AppSecSCATransformer(newConfig);
109+
// instrumentation.addTransformer(currentTransformer, true);
110+
111+
// Find loaded classes that match targets
112+
List<Class<?>> classesToRetransform = findLoadedClasses(targetClassNames);
113+
114+
if (classesToRetransform.isEmpty()) {
115+
log.debug("No loaded classes match SCA targets (they may load later)");
116+
return;
117+
}
118+
119+
// Trigger retransformation
120+
log.info("Retransforming {} classes for SCA instrumentation", classesToRetransform.size());
121+
retransformClasses(classesToRetransform);
122+
}
123+
124+
private Set<String> extractTargetClassNames(AppSecSCAConfig config) {
125+
Set<String> classNames = new HashSet<>();
126+
127+
for (AppSecSCAConfig.InstrumentationTarget target : config.instrumentationTargets) {
128+
if (target.className == null || target.className.isEmpty()) {
129+
log.warn("Skipping target with null or empty className");
130+
continue;
131+
}
132+
133+
// Convert internal format (org/foo/Bar) to binary name (org.foo.Bar)
134+
String binaryName = target.className.replace('/', '.');
135+
classNames.add(binaryName);
136+
}
137+
138+
return classNames;
139+
}
140+
141+
private List<Class<?>> findLoadedClasses(Set<String> targetClassNames) {
142+
List<Class<?>> matchedClasses = new ArrayList<>();
143+
144+
Class<?>[] loadedClasses = instrumentation.getAllLoadedClasses();
145+
log.debug("Scanning {} loaded classes for SCA targets", loadedClasses.length);
146+
147+
for (Class<?> clazz : loadedClasses) {
148+
if (targetClassNames.contains(clazz.getName())) {
149+
if (!instrumentation.isModifiableClass(clazz)) {
150+
log.debug("Class {} matches target but is not modifiable", clazz.getName());
151+
continue;
152+
}
153+
matchedClasses.add(clazz);
154+
log.debug("Found loaded class matching SCA target: {}", clazz.getName());
155+
}
156+
}
157+
158+
return matchedClasses;
159+
}
160+
161+
private void retransformClasses(List<Class<?>> classes) {
162+
for (Class<?> clazz : classes) {
163+
try {
164+
log.debug("Retransforming class: {}", clazz.getName());
165+
instrumentation.retransformClasses(clazz);
166+
} catch (Exception e) {
167+
log.error("Failed to retransform class: {}", clazz.getName(), e);
168+
} catch (Throwable t) {
169+
log.error("Throwable during retransformation of class: {}", clazz.getName(), t);
170+
}
171+
}
172+
}
173+
174+
private void removeInstrumentation() {
175+
if (currentTransformer != null) {
176+
log.debug("Removing SCA transformer");
177+
instrumentation.removeTransformer(currentTransformer);
178+
currentTransformer = null;
179+
}
180+
181+
// TODO: Optionally retransform classes to remove instrumentation
182+
// For now, instrumentation stays until JVM restart
183+
}
184+
185+
/**
186+
* Gets the current SCA configuration.
187+
*
188+
* @return the current config, or null if none is active
189+
*/
190+
public AppSecSCAConfig getCurrentConfig() {
191+
return currentConfig;
192+
}
193+
194+
/**
195+
* For testing: checks if a transformer is currently installed.
196+
*/
197+
boolean hasTransformer() {
198+
return currentTransformer != null;
199+
}
200+
}

0 commit comments

Comments
 (0)