#21824

Committed 20 Mar 2024 08:05PM CUT coverage: 20.661% (+0.09%) from 20.57%

Build # #21824

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #10211 from IQSS/9356-rate-limiting-command-engine

adding rate limiting for command engine

Run Details

90 of 123 new or added lines in 14 files covered. (73.17%)

1 existing line in 1 file now uncovered.

17074 of 82639 relevant lines covered (20.66%)

0.21 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

91.04

/src/main/java/edu/harvard/iq/dataverse/util/cache/RateLimitUtil.java

package edu.harvard.iq.dataverse.util.cache;

import edu.harvard.iq.dataverse.authorization.users.AuthenticatedUser;
import edu.harvard.iq.dataverse.authorization.users.GuestUser;
import edu.harvard.iq.dataverse.authorization.users.User;
import edu.harvard.iq.dataverse.util.SystemConfig;
import jakarta.json.bind.Jsonb;
import jakarta.json.bind.JsonbBuilder;
import jakarta.json.bind.JsonbException;

import javax.cache.Cache;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.logging.Logger;

import static java.lang.Math.max;
import static java.lang.Math.min;

public class RateLimitUtil {
    private static final Logger logger = Logger.getLogger(RateLimitUtil.class.getCanonicalName());
    static final List<RateLimitSetting> rateLimits = new CopyOnWriteArrayList<>();
    static final Map<String, Integer> rateLimitMap = new ConcurrentHashMap<>();
    public static final int NO_LIMIT = -1;

    static String generateCacheKey(final User user, final String action) {
        return (user != null ? user.getIdentifier() : GuestUser.get().getIdentifier()) +
            (action != null ? ":" + action : "");
    }
    static int getCapacity(SystemConfig systemConfig, User user, String action) {
        if (user != null && user.isSuperuser()) {
            return NO_LIMIT;
        }
        // get the capacity, i.e. calls per hour, from config
        return (user instanceof AuthenticatedUser authUser) ?
                getCapacityByTierAndAction(systemConfig, authUser.getRateLimitTier(), action) :
                getCapacityByTierAndAction(systemConfig, 0, action);
    }
    static boolean rateLimited(final Cache<String, String> rateLimitCache, final String key, int capacityPerHour) {
        if (capacityPerHour == NO_LIMIT) {
            return false;
        }
        long currentTime = System.currentTimeMillis() / 60000L; // convert to minutes
        double tokensPerMinute = (capacityPerHour / 60.0);
        // Get the last time this bucket was added to
        final String keyLastUpdate = String.format("%s:last_update",key);
        long lastUpdate = longFromKey(rateLimitCache, keyLastUpdate);
        long deltaTime = currentTime - lastUpdate;
        // Get the current number of tokens in the bucket
        long tokens = longFromKey(rateLimitCache, key);
        long tokensToAdd = (long) (deltaTime * tokensPerMinute);
        if (tokensToAdd > 0) { // Don't update timestamp if we aren't adding any tokens to the bucket
            tokens = min(capacityPerHour, tokens + tokensToAdd);
            rateLimitCache.put(keyLastUpdate, String.valueOf(currentTime));
        }
        // Update with any added tokens and decrement 1 token for this call if not rate limited (0 tokens)
        rateLimitCache.put(key, String.valueOf(max(0, tokens-1)));
        return tokens < 1;
    }

    static int getCapacityByTierAndAction(SystemConfig systemConfig, Integer tier, String action) {
        if (rateLimits.isEmpty()) {
            init(systemConfig);
        }
        
        if (rateLimitMap.containsKey(getMapKey(tier, action))) {
            return rateLimitMap.get(getMapKey(tier,action));
        } else if (rateLimitMap.containsKey(getMapKey(tier))) {
            return rateLimitMap.get(getMapKey(tier));
        } else {
            return getCapacityByTier(systemConfig, tier);
        }
    }
    static int getCapacityByTier(SystemConfig systemConfig, int tier) {
        int value = NO_LIMIT;
        String csvString = systemConfig.getRateLimitingDefaultCapacityTiers();
        try {
            if (!csvString.isEmpty()) {
                int[] values = Arrays.stream(csvString.split(",")).mapToInt(Integer::parseInt).toArray();
                if (tier < values.length) {
                    value = values[tier];
                }
            }
        } catch (NumberFormatException nfe) {
            logger.warning(nfe.getMessage());
        }
        return value;
    }
    static void init(SystemConfig systemConfig) {
        getRateLimitsFromJson(systemConfig);
        /* Convert the List of Rate Limit Settings containing a list of Actions to a fast lookup Map where the key is:
             for default if no action defined: "{tier}:" and the value is the default limit for the tier
             for each action: "{tier}:{action}" and the value is the limit defined in the setting
        */
        rateLimitMap.clear();
        rateLimits.forEach(r -> {
            r.setDefaultLimit(getCapacityByTier(systemConfig, r.getTier()));
            rateLimitMap.put(getMapKey(r.getTier()), r.getDefaultLimitPerHour());
            r.getActions().forEach(a -> rateLimitMap.put(getMapKey(r.getTier(), a), r.getLimitPerHour()));
        });
    }
    
    @SuppressWarnings("java:S2133") // <- To enable casting to generic in JSON-B we need a class instance, false positive
    static void getRateLimitsFromJson(SystemConfig systemConfig) {
        String setting = systemConfig.getRateLimitsJson();
        rateLimits.clear();
        if (!setting.isEmpty()) {
            try (Jsonb jsonb = JsonbBuilder.create()) {
                rateLimits.addAll(jsonb.fromJson(setting,
                        new ArrayList<RateLimitSetting>() {}.getClass().getGenericSuperclass()));
            } catch (JsonbException e) {
                logger.warning("Unable to parse Rate Limit Json: " + e.getLocalizedMessage() + "   Json:(" + setting + ")");
                rateLimits.add(new RateLimitSetting()); // add a default entry to prevent re-initialization
            // Note: Usually using Exception in a catch block is an antipattern and should be avoided.
            //       As the JSON-B interface does not specify a non-generic type, we have to use this.
            } catch (Exception e) {
                logger.warning("Could not close JSON-B reader");
            }
        }
    }
    static String getMapKey(int tier) {
        return getMapKey(tier, null);
    }
    static String getMapKey(int tier, String action) {
        return tier + ":" + (action != null ? action : "");
    }
    static long longFromKey(Cache<String, String> cache, String key) {
        Object l = cache.get(key);
        return l != null ? Long.parseLong(String.valueOf(l)) : 0L;
    }
}

1	package edu.harvard.iq.dataverse.util.cache;
2
3	import edu.harvard.iq.dataverse.authorization.users.AuthenticatedUser;
4	import edu.harvard.iq.dataverse.authorization.users.GuestUser;
5	import edu.harvard.iq.dataverse.authorization.users.User;
6	import edu.harvard.iq.dataverse.util.SystemConfig;
7	import jakarta.json.bind.Jsonb;
8	import jakarta.json.bind.JsonbBuilder;
9	import jakarta.json.bind.JsonbException;
10
11	import javax.cache.Cache;
12	import java.util.ArrayList;
13	import java.util.Arrays;
14	import java.util.List;
15	import java.util.Map;
16	import java.util.concurrent.ConcurrentHashMap;
17	import java.util.concurrent.CopyOnWriteArrayList;
18	import java.util.logging.Logger;
19
20	import static java.lang.Math.max;
21	import static java.lang.Math.min;
22
NEW 23	public class RateLimitUtil {	×
24	private static final Logger logger = Logger.getLogger(RateLimitUtil.class.getCanonicalName());	1✔
25	static final List<RateLimitSetting> rateLimits = new CopyOnWriteArrayList<>();	1✔
26	static final Map<String, Integer> rateLimitMap = new ConcurrentHashMap<>();	1✔
27	public static final int NO_LIMIT = -1;
28
29	static String generateCacheKey(final User user, final String action) {
30	return (user != null ? user.getIdentifier() : GuestUser.get().getIdentifier()) +	1✔
31	(action != null ? ":" + action : "");	1✔
32	}
33	static int getCapacity(SystemConfig systemConfig, User user, String action) {
34	if (user != null && user.isSuperuser()) {	1✔
35	return NO_LIMIT;	1✔
36	}
37	// get the capacity, i.e. calls per hour, from config
38	return (user instanceof AuthenticatedUser authUser) ?	1✔
39	getCapacityByTierAndAction(systemConfig, authUser.getRateLimitTier(), action) :	1✔
40	getCapacityByTierAndAction(systemConfig, 0, action);	1✔
41	}
42	static boolean rateLimited(final Cache<String, String> rateLimitCache, final String key, int capacityPerHour) {
43	if (capacityPerHour == NO_LIMIT) {	1✔
NEW 44	return false;	×
45	}
46	long currentTime = System.currentTimeMillis() / 60000L; // convert to minutes	1✔
47	double tokensPerMinute = (capacityPerHour / 60.0);	1✔
48	// Get the last time this bucket was added to
49	final String keyLastUpdate = String.format("%s:last_update",key);	1✔
50	long lastUpdate = longFromKey(rateLimitCache, keyLastUpdate);	1✔
51	long deltaTime = currentTime - lastUpdate;	1✔
52	// Get the current number of tokens in the bucket
53	long tokens = longFromKey(rateLimitCache, key);	1✔
54	long tokensToAdd = (long) (deltaTime * tokensPerMinute);	1✔
55	if (tokensToAdd > 0) { // Don't update timestamp if we aren't adding any tokens to the bucket	1✔
56	tokens = min(capacityPerHour, tokens + tokensToAdd);	1✔
57	rateLimitCache.put(keyLastUpdate, String.valueOf(currentTime));	1✔
58	}
59	// Update with any added tokens and decrement 1 token for this call if not rate limited (0 tokens)
60	rateLimitCache.put(key, String.valueOf(max(0, tokens-1)));	1✔
61	return tokens < 1;	1✔
62	}
63
64	static int getCapacityByTierAndAction(SystemConfig systemConfig, Integer tier, String action) {
65	if (rateLimits.isEmpty()) {	1✔
66	init(systemConfig);	1✔
67	}
68
69	if (rateLimitMap.containsKey(getMapKey(tier, action))) {	1✔
70	return rateLimitMap.get(getMapKey(tier,action));	1✔
71	} else if (rateLimitMap.containsKey(getMapKey(tier))) {	1✔
72	return rateLimitMap.get(getMapKey(tier));	1✔
73	} else {
74	return getCapacityByTier(systemConfig, tier);	1✔
75	}
76	}
77	static int getCapacityByTier(SystemConfig systemConfig, int tier) {
78	int value = NO_LIMIT;	1✔
79	String csvString = systemConfig.getRateLimitingDefaultCapacityTiers();	1✔
80	try {
81	if (!csvString.isEmpty()) {	1✔
82	int[] values = Arrays.stream(csvString.split(",")).mapToInt(Integer::parseInt).toArray();	1✔
83	if (tier < values.length) {	1✔
84	value = values[tier];	1✔
85	}
86	}
NEW 87	} catch (NumberFormatException nfe) {	×
NEW 88	logger.warning(nfe.getMessage());	×
89	}	1✔
90	return value;	1✔
91	}
92	static void init(SystemConfig systemConfig) {
93	getRateLimitsFromJson(systemConfig);	1✔
94	/* Convert the List of Rate Limit Settings containing a list of Actions to a fast lookup Map where the key is:
95	for default if no action defined: "{tier}:" and the value is the default limit for the tier
96	for each action: "{tier}:{action}" and the value is the limit defined in the setting
97	*/
98	rateLimitMap.clear();	1✔
99	rateLimits.forEach(r -> {	1✔
100	r.setDefaultLimit(getCapacityByTier(systemConfig, r.getTier()));	1✔
101	rateLimitMap.put(getMapKey(r.getTier()), r.getDefaultLimitPerHour());	1✔
102	r.getActions().forEach(a -> rateLimitMap.put(getMapKey(r.getTier(), a), r.getLimitPerHour()));	1✔
103	});	1✔
104	}	1✔
105
106	@SuppressWarnings("java:S2133") // <- To enable casting to generic in JSON-B we need a class instance, false positive
107	static void getRateLimitsFromJson(SystemConfig systemConfig) {
108	String setting = systemConfig.getRateLimitsJson();	1✔
109	rateLimits.clear();	1✔
110	if (!setting.isEmpty()) {	1✔
111	try (Jsonb jsonb = JsonbBuilder.create()) {	1✔
112	rateLimits.addAll(jsonb.fromJson(setting,	1✔
113	new ArrayList<RateLimitSetting>() {}.getClass().getGenericSuperclass()));	1✔
114	} catch (JsonbException e) {	1✔
115	logger.warning("Unable to parse Rate Limit Json: " + e.getLocalizedMessage() + " Json:(" + setting + ")");	1✔
116	rateLimits.add(new RateLimitSetting()); // add a default entry to prevent re-initialization	1✔
117	// Note: Usually using Exception in a catch block is an antipattern and should be avoided.
118	// As the JSON-B interface does not specify a non-generic type, we have to use this.
NEW 119	} catch (Exception e) {	×
NEW 120	logger.warning("Could not close JSON-B reader");	×
121	}	1✔
122	}
123	}	1✔
124	static String getMapKey(int tier) {
125	return getMapKey(tier, null);	1✔
126	}
127	static String getMapKey(int tier, String action) {
128	return tier + ":" + (action != null ? action : "");	1✔
129	}
130	static long longFromKey(Cache<String, String> cache, String key) {
131	Object l = cache.get(key);	1✔
132	return l != null ? Long.parseLong(String.valueOf(l)) : 0L;	1✔
133	}
134	}

IQSS / dataverse / #21824

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous