Added data transfer limits

README.md

@@ -64,12 +64,15 @@ Once that is done, a new proxy will be available on the port returned. All you h
   - status - the HTTP status code to return for URLs that are blacklisted
  - DELETE /proxy/[port]/blacklist - Clears all URL patterns from the blacklist
  - PUT /proxy/[port]/limit - Limit the bandwidth through the proxy. Takes the following parameters:
-  - downstreamKbps - Sets the downstream kbps
-  - upstreamKbps - Sets the upstream kbps
+  - downstreamKbps - Sets the downstream bandwidth limit in kbps
+  - upstreamKbps - Sets the upstream bandwidth limit kbps
+  - downstreamMaxKB - Specifies how many kilobytes in total the client is allowed to download through the proxy.
+  - upstreamMaxKB - Specifies how many kilobytes in total the client is allowed to upload through the proxy.
   - latency - Add the given latency to each HTTP request
   - enable - (true/false) a boolean that enable bandwidth limiter. By default the limit is disabled, although setting any of the properties above will implicitly enable throttling
   - payloadPercentage - a number ]0, 100] specifying what percentage of data sent is payload. e.g. use this to take into account overhead due to tcp/ip.
   - maxBitsPerSecond - The max bits per seconds you want this instance of StreamManager to respect.
+ - GET /proxy/[port]/limit - Displays the amount of data remaining to be uploaded/downloaded until the limit is reached.
  - POST /proxy/[port]/headers - Set and override HTTP Request headers. For example setting a custom User-Agent.
   - Payload data should be json encoded set of headers (not url-encoded)
  - POST /proxy/[port]/hosts - Overrides normal DNS lookups and remaps the given hosts with the associated IP address

src/main/java/net/lightbody/bmp/proxy/bricks/ProxyResource.java

@@ -12,7 +12,6 @@
 import com.google.sitebricks.http.Get;
 import com.google.sitebricks.http.Post;
 import com.google.sitebricks.http.Put;
-
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintWriter;
@@ -28,9 +27,6 @@
 import javax.script.ScriptEngine;
 import javax.script.ScriptEngineManager;
 import javax.script.ScriptException;
-
-
-
 import net.lightbody.bmp.core.har.Har;
 import net.lightbody.bmp.proxy.ProxyExistsException;
 import net.lightbody.bmp.proxy.ProxyManager;
@@ -40,7 +36,6 @@
 import net.lightbody.bmp.proxy.http.BrowserMobHttpResponse;
 import net.lightbody.bmp.proxy.http.RequestInterceptor;
 import net.lightbody.bmp.proxy.http.ResponseInterceptor;
-
 import org.java_bandwidthlimiter.StreamManager;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -348,6 +343,20 @@ public Reply<?> limit(@Named("port") int port, Request request) {
                 streamManager.enable();
             } catch (NumberFormatException e) { }
         }
+        String upstreamMaxKB = request.param("upstreamMaxKB");
+        if (upstreamMaxKB != null) {
+            try {
+                streamManager.setUpstreamMaxKB(Integer.parseInt(upstreamMaxKB));
+                streamManager.enable();
+            } catch (NumberFormatException e) { }
+        }
+        String downstreamMaxKB = request.param("downstreamMaxKB");
+        if (downstreamMaxKB != null) {
+            try {
+                streamManager.setDownstreamMaxKB(Integer.parseInt(downstreamMaxKB));
+                streamManager.enable();
+            } catch (NumberFormatException e) { }
+        }        
         String latency = request.param("latency");
         if (latency != null) {
             try {
@@ -378,6 +387,16 @@ public Reply<?> limit(@Named("port") int port, Request request) {
         return Reply.saying().ok();
     }
 
+    @Get
+    @At("/:port/limit")
+    public Reply<?> getLimits(@Named("port") int port, Request request) {
+        ProxyServer proxy = proxyManager.get(port);
+        if (proxy == null) {
+            return Reply.saying().notFound();
+        }                
+        return Reply.with(new BandwidthLimitDescriptor(proxy.getStreamManager())).as(Json.class);
+    }
+
     @Put
     @At("/:port/timeout")
     public Reply<?> timeout(@Named("port") int port, Request request) {
@@ -559,4 +578,53 @@ public void setProxyList(Collection<ProxyDescriptor> proxyList) {
             this.proxyList = proxyList;
         }
     }
+
+    public static class BandwidthLimitDescriptor {
+        private long maxUpstreamKB;
+        private long remainingUpstreamKB;
+        private long maxDownstreamKB;
+        private long remainingDownstreamKB;
+
+        public BandwidthLimitDescriptor(){
+        }
+
+        public BandwidthLimitDescriptor(StreamManager manager){
+            this.maxDownstreamKB = manager.getMaxDownstreamKB();
+            this.remainingDownstreamKB = manager.getRemainingDownstreamKB();
+            this.maxUpstreamKB = manager.getMaxUpstreamKB();
+            this.remainingUpstreamKB = manager.getRemainingUpstreamKB();
+        }
+
+        public long getMaxUpstreamKB() {
+            return maxUpstreamKB;
+        }
+
+        public void setMaxUpstreamKB(long maxUpstreamKB) {
+            this.maxUpstreamKB = maxUpstreamKB;
+        }
+
+        public long getRemainingUpstreamKB() {
+            return remainingUpstreamKB;
+        }
+
+        public void setRemainingUpstreamKB(long remainingUpstreamKB) {
+            this.remainingUpstreamKB = remainingUpstreamKB;
+        }
+
+        public long getMaxDownstreamKB() {
+            return maxDownstreamKB;
+        }
+
+        public void setMaxDownstreamKB(long maxDownstreamKB) {
+            this.maxDownstreamKB = maxDownstreamKB;
+        }
+
+        public long getRemainingDownstreamKB() {
+            return remainingDownstreamKB;
+        }
+
+        public void setRemainingDownstreamKB(long remainingDownstreamKB) {
+            this.remainingDownstreamKB = remainingDownstreamKB;
+        }        
+    }
 }

src/main/java/org/java_bandwidthlimiter/BandwidthLimiter.java

@@ -36,6 +36,10 @@ public interface BandwidthLimiter {
     public void setDownstreamKbps(long downstreamKbps);
 
     public void setUpstreamKbps(long upstreamKbps);
+
+    public void setDownstreamMaxKB(long downstreamMaxKB);
+
+    public void setUpstreamMaxKB(long upstreamMaxKB);
 
     public void setLatency(long latency);
 }

src/main/java/org/java_bandwidthlimiter/MaximumTransferExceededException.java

@@ -0,0 +1,22 @@
+package org.java_bandwidthlimiter;
+
+import java.io.IOException;
+
+public class MaximumTransferExceededException extends IOException {    
+	private final boolean isUpstream;
+    private final long limit;
+
+    public boolean isUpstream() {
+        return isUpstream;
+    }
+
+	public long getLimit() {
+		return limit;
+	}
+
+	public MaximumTransferExceededException(long limit, boolean isUpstream) {
+		super("Maximum " + (isUpstream? "upstream" : "downstream") + " transfer allowance of " + limit + " KB exceeded.");
+        this.isUpstream = isUpstream;
+		this.limit = limit;
+	}
+}