diff --git a/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java b/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java index cc3e98198..79e5cda19 100644 --- a/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java +++ b/src/main/java/org/whispersystems/textsecuregcm/WhisperServerConfiguration.java @@ -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; diff --git a/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java b/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java index 6f2b964fc..785b5ef4a 100644 --- a/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java +++ b/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java @@ -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; + } + } } \ No newline at end of file diff --git a/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiter.java b/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiter.java index 45b5051c3..ea3d3911c 100644 --- a/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiter.java +++ b/src/main/java/org/whispersystems/textsecuregcm/limits/RateLimiter.java @@ -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);