Fix leaky bucket serialization.

// FREEBIE

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;

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.spaceRemaining       = spaceRemaining;
-    this.lastUpdateTimeMillis = System.currentTimeMillis();
+    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;
+    }
+  }
 }

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);