0
0
mirror of https://github.com/signalapp/Signal-Server.git synced 2024-09-20 03:52:16 +02:00

Fix leaky bucket serialization.

// FREEBIE
This commit is contained in:
Moxie Marlinspike 2015-01-30 11:48:49 -08:00
parent 75aec0a8d4
commit 0928e4c035
3 changed files with 49 additions and 12 deletions

View File

@ -17,14 +17,13 @@
package org.whispersystems.textsecuregcm;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.whispersystems.textsecuregcm.configuration.RedisConfiguration;
import org.whispersystems.textsecuregcm.configuration.FederationConfiguration;
import org.whispersystems.textsecuregcm.configuration.GraphiteConfiguration;
import org.whispersystems.textsecuregcm.configuration.MemcacheConfiguration;
import org.whispersystems.textsecuregcm.configuration.NexmoConfiguration;
import org.whispersystems.textsecuregcm.configuration.PushConfiguration;
import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration;
import org.whispersystems.textsecuregcm.configuration.RedPhoneConfiguration;
import org.whispersystems.textsecuregcm.configuration.RedisConfiguration;
import org.whispersystems.textsecuregcm.configuration.S3Configuration;
import org.whispersystems.textsecuregcm.configuration.TwilioConfiguration;
import org.whispersystems.textsecuregcm.configuration.WebsocketConfiguration;

View File

@ -17,26 +17,28 @@
package org.whispersystems.textsecuregcm.limits;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
public class LeakyBucket {
@JsonProperty
private final int bucketSize;
@JsonProperty
private final double leakRatePerMillis;
@JsonProperty
private int spaceRemaining;
@JsonProperty
private long lastUpdateTimeMillis;
public LeakyBucket(int bucketSize, double leakRatePerMillis) {
this(bucketSize, leakRatePerMillis, bucketSize, System.currentTimeMillis());
}
private LeakyBucket(int bucketSize, double leakRatePerMillis, int spaceRemaining, long lastUpdateTimeMillis) {
this.bucketSize = bucketSize;
this.leakRatePerMillis = leakRatePerMillis;
this.spaceRemaining = bucketSize;
this.lastUpdateTimeMillis = System.currentTimeMillis();
this.spaceRemaining = spaceRemaining;
this.lastUpdateTimeMillis = lastUpdateTimeMillis;
}
public boolean add(int amount) {
@ -56,4 +58,40 @@ public class LeakyBucket {
return Math.min(this.bucketSize,
(int)Math.floor(this.spaceRemaining + (elapsedTime * this.leakRatePerMillis)));
}
public String serialize(ObjectMapper mapper) throws JsonProcessingException {
return mapper.writeValueAsString(new LeakyBucketEntity(bucketSize, leakRatePerMillis, spaceRemaining, lastUpdateTimeMillis));
}
public static LeakyBucket fromSerialized(ObjectMapper mapper, String serialized) throws IOException {
LeakyBucketEntity entity = mapper.readValue(serialized, LeakyBucketEntity.class);
return new LeakyBucket(entity.bucketSize, entity.leakRatePerMillis,
entity.spaceRemaining, entity.lastUpdateTimeMillis);
}
private static class LeakyBucketEntity {
@JsonProperty
private int bucketSize;
@JsonProperty
private double leakRatePerMillis;
@JsonProperty
private int spaceRemaining;
@JsonProperty
private long lastUpdateTimeMillis;
public LeakyBucketEntity() {}
private LeakyBucketEntity(int bucketSize, double leakRatePerMillis,
int spaceRemaining, long lastUpdateTimeMillis)
{
this.bucketSize = bucketSize;
this.leakRatePerMillis = leakRatePerMillis;
this.spaceRemaining = spaceRemaining;
this.lastUpdateTimeMillis = lastUpdateTimeMillis;
}
}
}

View File

@ -73,7 +73,7 @@ public class RateLimiter {
private void setBucket(String key, LeakyBucket bucket) {
try (Jedis jedis = cacheClient.getResource()) {
String serialized = mapper.writeValueAsString(bucket);
String serialized = bucket.serialize(mapper);
jedis.setex(getBucketName(key), (int) Math.ceil((bucketSize / leakRatePerMillis) / 1000), serialized);
} catch (JsonProcessingException e) {
throw new IllegalArgumentException(e);
@ -85,7 +85,7 @@ public class RateLimiter {
String serialized = jedis.get(getBucketName(key));
if (serialized != null) {
return mapper.readValue(serialized, LeakyBucket.class);
return LeakyBucket.fromSerialized(mapper, serialized);
}
} catch (IOException e) {
logger.warn("Deserialization error", e);