33import com .google .common .base .Strings ;
44import io .prometheus .client .Histogram ;
55import java .io .IOException ;
6- import java .util .Collections ;
7- import java .util .HashMap ;
8- import java .util .Map ;
6+ import java .lang .reflect .Constructor ;
97import javax .annotation .PostConstruct ;
108import javax .servlet .ServletException ;
119import javax .servlet .http .HttpServlet ;
3331@ Slf4j
3432public abstract class RateLimiterServlet extends HttpServlet {
3533 private static final String KEY_PREFIX_HTTP = "http_" ;
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- }
34+ private static final String ADAPTER_PREFIX = "org.tron.core.services.ratelimiter.adapter." ;
5535
5636 @ Autowired
5737 private RateLimiterContainer container ;
@@ -60,22 +40,42 @@ public abstract class RateLimiterServlet extends HttpServlet {
6040 private void addRateContainer () {
6141 RateLimiterInitialization .HttpRateLimiterItem item = Args .getInstance ()
6242 .getRateLimiterInitialization ().getHttpMap ().get (getClass ().getSimpleName ());
43+ boolean success = false ;
6344 final String name = getClass ().getSimpleName ();
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 );
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 );
7478 }
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
0 commit comments