Skip to content

Commit 1f3931f

Browse files
Little-Peonyclaude
Little-Peony
and
claude
committed
fix(api): add whitelist rate limiter class to prevent arbitrary class loading
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 2de63bb commit 1f3931f

3 files changed

Lines changed: 132 additions & 42 deletions

File tree

‎chainbase/src/main/java/org/tron/core/db/TransactionTrace.java‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -295,13 +295,14 @@ private void resetAccountUsage(AccountCapsule accountCap,
295295
}
296296
long currentSize = accountCap.getWindowSize(ENERGY);
297297
long currentUsage = accountCap.getEnergyUsage();
298-
// Drop the pre consumed frozen energy
298+
// Drop the preconsumed frozen energy
299299
long newArea = currentUsage * currentSize
300300
- (mergedUsage * mergedSize - usage * size);
301301
// If area merging happened during suicide, use the current window size
302302
long newSize = mergedSize == currentSize ? size : currentSize;
303-
// Calc new usage by fixed x-axes
304-
long newUsage = max(0, newArea / newSize, dynamicPropertiesStore.disableJavaLangMath());
303+
// A zero window size means no valid time window exists and thus zero usage
304+
long newUsage = newSize == 0 ? 0
305+
: max(0, newArea / newSize, dynamicPropertiesStore.disableJavaLangMath());
305306
// Reset account usage and window size
306307
accountCap.setEnergyUsage(newUsage);
307308
accountCap.setNewWindowSize(ENERGY, newUsage == 0 ? 0L : newSize);
@@ -312,13 +313,14 @@ private void resetAccountUsageV2(AccountCapsule accountCap,
312313
long currentSize = accountCap.getWindowSize(ENERGY);
313314
long currentSize2 = accountCap.getWindowSizeV2(ENERGY);
314315
long currentUsage = accountCap.getEnergyUsage();
315-
// Drop the pre consumed frozen energy
316+
// Drop the preconsumed frozen energy
316317
long newArea = currentUsage * currentSize - (mergedUsage * mergedSize - usage * size);
317318
// If area merging happened during suicide, use the current window size
318319
long newSize = mergedSize == currentSize ? size : currentSize;
319320
long newSize2 = mergedSize == currentSize ? size2 : currentSize2;
320-
// Calc new usage by fixed x-axes
321-
long newUsage = max(0, newArea / newSize, dynamicPropertiesStore.disableJavaLangMath());
321+
// A zero window size means no valid time window exists and thus zero usage
322+
long newUsage = newSize == 0 ? 0
323+
: max(0, newArea / newSize, dynamicPropertiesStore.disableJavaLangMath());
322324
// Reset account usage and window size
323325
accountCap.setEnergyUsage(newUsage);
324326
accountCap.setNewWindowSizeV2(ENERGY, newUsage == 0 ? 0L : newSize2);

‎framework/src/main/java/org/tron/core/services/http/RateLimiterServlet.java‎

Lines changed: 36 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@
33
import com.google.common.base.Strings;
44
import io.prometheus.client.Histogram;
55
import java.io.IOException;
6-
import java.lang.reflect.Constructor;
6+
import java.util.Collections;
7+
import java.util.HashMap;
8+
import java.util.Map;
79
import javax.annotation.PostConstruct;
810
import javax.servlet.ServletException;
911
import javax.servlet.http.HttpServlet;
@@ -31,7 +33,25 @@
3133
@Slf4j
3234
public abstract class RateLimiterServlet extends HttpServlet {
3335
private static final String KEY_PREFIX_HTTP = "http_";
34-
private static final String ADAPTER_PREFIX = "org.tron.core.services.ratelimiter.adapter.";
36+
37+
// Whitelist of allowed rate limiter adapter class names.
38+
// Keys are derived from Class.getSimpleName() so they stay in sync if classes are renamed.
39+
// Class.forName() is intentionally NOT used; only these pre-approved classes may be loaded.
40+
private static final Map<String, Class<? extends IRateLimiter>> ALLOWED_ADAPTERS;
41+
private static final String DEFAULT_ADAPTER_NAME = DefaultBaseQqsAdapter.class.getSimpleName();
42+
43+
static {
44+
Map<String, Class<? extends IRateLimiter>> m = new HashMap<>();
45+
for (Class<? extends IRateLimiter> c : new Class[]{
46+
GlobalPreemptibleAdapter.class,
47+
QpsRateLimiterAdapter.class,
48+
IPQPSRateLimiterAdapter.class,
49+
DefaultBaseQqsAdapter.class
50+
}) {
51+
m.put(c.getSimpleName(), c);
52+
}
53+
ALLOWED_ADAPTERS = Collections.unmodifiableMap(m);
54+
}
3555

3656
@Autowired
3757
private RateLimiterContainer container;
@@ -40,42 +60,22 @@ public abstract class RateLimiterServlet extends HttpServlet {
4060
private void addRateContainer() {
4161
RateLimiterInitialization.HttpRateLimiterItem item = Args.getInstance()
4262
.getRateLimiterInitialization().getHttpMap().get(getClass().getSimpleName());
43-
boolean success = false;
4463
final String name = getClass().getSimpleName();
45-
if (item != null) {
46-
String cName = "";
47-
String params = "";
48-
Object obj;
49-
try {
50-
cName = item.getStrategy();
51-
params = item.getParams();
52-
// add the specific rate limiter strategy of servlet.
53-
Class<?> c = Class.forName(ADAPTER_PREFIX + cName);
54-
Constructor constructor;
55-
if (c == GlobalPreemptibleAdapter.class || c == QpsRateLimiterAdapter.class
56-
|| c == IPQPSRateLimiterAdapter.class) {
57-
constructor = c.getConstructor(String.class);
58-
obj = constructor.newInstance(params);
59-
container.add(KEY_PREFIX_HTTP, name, (IRateLimiter) obj);
60-
} else {
61-
constructor = c.getConstructor();
62-
obj = constructor.newInstance(QpsStrategy.DEFAULT_QPS_PARAM);
63-
container.add(KEY_PREFIX_HTTP, name, (IRateLimiter) obj);
64-
}
65-
success = true;
66-
} catch (Exception e) {
67-
this.throwTronError(cName, params, name, e);
68-
}
69-
}
70-
if (!success) {
71-
// if the specific rate limiter strategy of servlet is not defined or fail to add,
72-
// then add a default Strategy.
73-
try {
74-
IRateLimiter rateLimiter = new DefaultBaseQqsAdapter(QpsStrategy.DEFAULT_QPS_PARAM);
75-
container.add(KEY_PREFIX_HTTP, name, rateLimiter);
76-
} catch (Exception e) {
77-
this.throwTronError("DefaultBaseQqsAdapter", QpsStrategy.DEFAULT_QPS_PARAM, name, e);
64+
65+
// If no config for this servlet, fall back to the default adapter.
66+
String cName = (item != null) ? item.getStrategy() : DEFAULT_ADAPTER_NAME;
67+
String params = (item != null) ? item.getParams() : QpsStrategy.DEFAULT_QPS_PARAM;
68+
69+
try {
70+
Class<? extends IRateLimiter> c = ALLOWED_ADAPTERS.get(cName);
71+
if (c == null) {
72+
throw new IllegalArgumentException(
73+
"Unknown rate limiter adapter (not in whitelist): " + cName);
7874
}
75+
IRateLimiter rateLimiter = (IRateLimiter) c.getConstructor(String.class).newInstance(params);
76+
container.add(KEY_PREFIX_HTTP, name, rateLimiter);
77+
} catch (Exception e) {
78+
this.throwTronError(cName, params, name, e);
7979
}
8080
}
8181

‎framework/src/test/java/org/tron/core/services/http/RateLimiterServletWhitelistTest.java‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
package org.tron.core.services.http;
2+
3+
import static org.junit.Assert.assertFalse;
4+
import static org.junit.Assert.assertNotNull;
5+
import static org.junit.Assert.assertNull;
6+
import static org.junit.Assert.assertTrue;
7+
8+
import java.lang.reflect.Field;
9+
import java.util.Map;
10+
import org.junit.BeforeClass;
11+
import org.junit.Test;
12+
import org.tron.core.services.ratelimiter.adapter.DefaultBaseQqsAdapter;
13+
import org.tron.core.services.ratelimiter.adapter.GlobalPreemptibleAdapter;
14+
import org.tron.core.services.ratelimiter.adapter.IPQPSRateLimiterAdapter;
15+
import org.tron.core.services.ratelimiter.adapter.IRateLimiter;
16+
import org.tron.core.services.ratelimiter.adapter.QpsRateLimiterAdapter;
17+
18+
/**
19+
* Security test: verifies that RateLimiterServlet uses a strict whitelist
20+
* instead of Class.forName(), preventing arbitrary class loading (RCE)
21+
* via a tampered config file.
22+
*/
23+
public class RateLimiterServletWhitelistTest {
24+
25+
// Derive names from the classes themselves — stays in sync if classes are renamed.
26+
private static final String GLOBAL_PREEMPTIBLE = GlobalPreemptibleAdapter.class.getSimpleName();
27+
private static final String QPS_RATE_LIMITER = QpsRateLimiterAdapter.class.getSimpleName();
28+
private static final String IP_QPS_RATE_LIMITER = IPQPSRateLimiterAdapter.class.getSimpleName();
29+
private static final String DEFAULT_BASE_QPS = DefaultBaseQqsAdapter.class.getSimpleName();
30+
31+
private static Map<String, Class<? extends IRateLimiter>> allowedAdapters;
32+
33+
@SuppressWarnings("unchecked")
34+
@BeforeClass
35+
public static void loadWhitelist() throws Exception {
36+
Field f = RateLimiterServlet.class.getDeclaredField("ALLOWED_ADAPTERS");
37+
f.setAccessible(true);
38+
allowedAdapters = (Map<String, Class<? extends IRateLimiter>>) f.get(null);
39+
}
40+
41+
// Verifies all 4 legitimate adapters are present and map to the correct classes.
42+
@Test
43+
public void testWhitelistContents() {
44+
assertNotNull(allowedAdapters.get(GLOBAL_PREEMPTIBLE));
45+
assertTrue(allowedAdapters.get(GLOBAL_PREEMPTIBLE)
46+
.isAssignableFrom(GlobalPreemptibleAdapter.class));
47+
48+
assertNotNull(allowedAdapters.get(QPS_RATE_LIMITER));
49+
assertTrue(allowedAdapters.get(QPS_RATE_LIMITER)
50+
.isAssignableFrom(QpsRateLimiterAdapter.class));
51+
52+
assertNotNull(allowedAdapters.get(IP_QPS_RATE_LIMITER));
53+
assertTrue(allowedAdapters.get(IP_QPS_RATE_LIMITER)
54+
.isAssignableFrom(IPQPSRateLimiterAdapter.class));
55+
56+
assertNotNull(allowedAdapters.get(DEFAULT_BASE_QPS));
57+
assertTrue(allowedAdapters.get(DEFAULT_BASE_QPS)
58+
.isAssignableFrom(DefaultBaseQqsAdapter.class));
59+
}
60+
61+
// Verifies that arbitrary / malicious class names are rejected by the whitelist.
62+
@Test
63+
public void testInvalidClassNameIsRejected() {
64+
assertNull(allowedAdapters.get("com.evil.MaliciousAdapter"));
65+
assertNull(allowedAdapters.get("../../../../evil.Payload"));
66+
assertNull(allowedAdapters.get(Runtime.class.getName()));
67+
assertNull(allowedAdapters.get(ProcessBuilder.class.getName()));
68+
assertNull(allowedAdapters.get(""));
69+
assertNull(allowedAdapters.get(null));
70+
}
71+
72+
// Verifies the whitelist cannot be modified at runtime (unmodifiable map).
73+
@Test(expected = UnsupportedOperationException.class)
74+
public void testWhitelistIsImmutable() {
75+
allowedAdapters.put("Injected", DefaultBaseQqsAdapter.class);
76+
}
77+
78+
// Verifies structural invariants: exact size and all entries implement IRateLimiter.
79+
@Test
80+
public void testWhitelistStructure() {
81+
assertFalse("Whitelist must not be empty", allowedAdapters.isEmpty());
82+
assertTrue("Whitelist must contain exactly 4 adapters", allowedAdapters.size() == 4);
83+
for (Map.Entry<String, Class<? extends IRateLimiter>> entry : allowedAdapters.entrySet()) {
84+
assertTrue(entry.getKey() + " must implement IRateLimiter",
85+
IRateLimiter.class.isAssignableFrom(entry.getValue()));
86+
}
87+
}
88+
}

0 commit comments

Comments
 (0)