diff --git a/README.md b/README.md index d75df2bc7..d84765932 100644 --- a/README.md +++ b/README.md @@ -65,12 +65,15 @@ Once that is done, a new proxy will be available on the port returned. All you h - method - regular expression for matching method., e.g., POST. Emtpy for matching all method. - 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 @@ -119,12 +122,14 @@ system properties will be used to specify the upstream proxy. Command-line Arguments ---------------------- - - -port + - -port \ - Port on which the API listens. Default value is 8080. - -address
- Address to which the API is bound. Default value is 0.0.0.0. - - -proxyPortRange - - - Range of ports reserved for proxies. Only applies if *port* parameter is not supplied in the POST request. Default values are +1 to +500+1. + - -proxyPortRange \-\ + - Range of ports reserved for proxies. Only applies if *port* parameter is not supplied in the POST request. Default values are \+1 to \+500+1. + - -ttl \ + - Proxy will be automatically deleted after a specified time period. Off by default. Embedded Mode ------------- diff --git a/pom.xml b/pom.xml index 5a5de012a..a74efda75 100644 --- a/pom.xml +++ b/pom.xml @@ -39,6 +39,20 @@ scm:git:git@github.com:lightbody/browsermob-proxy.git git@github.com:lightbody/browsermob-proxy.git + + + + guiceyfruit.release + GuiceyFruit Release Repository + http://guiceyfruit.googlecode.com/svn/repo/releases/ + + false + + + true + + + UTF-8 @@ -249,6 +263,12 @@ guice-servlet 3.0 + + + org.guiceyfruit + guiceyfruit-core + 2.0 + net.jcip @@ -294,7 +314,7 @@ net.sf.uadetector uadetector-resources - 2013.10 + 2014.09 org.jboss.arquillian.extension diff --git a/src/main/java/net/lightbody/bmp/proxy/BlacklistEntry.java b/src/main/java/net/lightbody/bmp/proxy/BlacklistEntry.java index 59926b581..c530c8dbd 100644 --- a/src/main/java/net/lightbody/bmp/proxy/BlacklistEntry.java +++ b/src/main/java/net/lightbody/bmp/proxy/BlacklistEntry.java @@ -2,27 +2,64 @@ import java.util.regex.Pattern; -public class BlacklistEntry -{ - private Pattern pattern; - private int responseCode; - private Pattern method; +public class BlacklistEntry { + private final Pattern pattern; + private final int responseCode; + private final Pattern method; - public BlacklistEntry(String pattern, int responseCode, String method) { - this.pattern = Pattern.compile(pattern); - this.responseCode = responseCode; - this.method = Pattern.compile(("".equals(method) || method == null) ? ".*" : method); - } + /** + * Creates a new BlacklistEntry with no HTTP method matching (i.e. all methods will match). + * + * @param pattern URL pattern to blacklist + * @param responseCode response code to return for blacklisted URL + */ + public BlacklistEntry(String pattern, int responseCode) { + this(pattern, responseCode, null); + } + + /** + * Creates a new BlacklistEntry which will match both a URL and an HTTP method + * + * @param pattern URL pattern to blacklist + * @param responseCode response code to return for blacklisted URL + * @param method HTTP method to match (e.g. GET, PUT, PATCH, etc.) + */ + public BlacklistEntry(String pattern, int responseCode, String method) { + this.pattern = Pattern.compile(pattern); + this.responseCode = responseCode; + if (method == null || method.isEmpty()) { + this.method = null; + } else { + this.method = Pattern.compile(method); + } + } + + /** + * Determines if this BlacklistEntry matches the given URL. Attempts to match both the URL and the + * HTTP method. + * + * @param url possibly-blacklisted URL + * @param httpMethod HTTP method this URL is being accessed with + * @return true if the URL matches this BlacklistEntry + */ + public boolean matches(String url, String httpMethod) { + if (method != null) { + return pattern.matcher(url).matches() && method.matcher(httpMethod).matches(); + } else { + return pattern.matcher(url).matches(); + } + } - public Pattern getPattern() { - return this.pattern; - } + public Pattern getPattern() { + return this.pattern; + } + + public int getResponseCode() { + return responseCode; + } - public int getResponseCode() { - return this.responseCode; - } - public Pattern getMethod() { - return this.method; + return method; } -} + +} \ No newline at end of file diff --git a/src/main/java/net/lightbody/bmp/proxy/Main.java b/src/main/java/net/lightbody/bmp/proxy/Main.java index 939e73e25..aca863a76 100644 --- a/src/main/java/net/lightbody/bmp/proxy/Main.java +++ b/src/main/java/net/lightbody/bmp/proxy/Main.java @@ -4,20 +4,6 @@ import com.google.inject.Injector; import com.google.inject.servlet.GuiceServletContextListener; import com.google.sitebricks.SitebricksModule; - -import net.lightbody.bmp.exception.JettyException; -import net.lightbody.bmp.proxy.bricks.ProxyResource; -import net.lightbody.bmp.proxy.guice.ConfigModule; -import net.lightbody.bmp.proxy.guice.JettyModule; -import net.lightbody.bmp.proxy.util.StandardFormatter; - -import org.eclipse.jetty.server.Server; -import org.eclipse.jetty.servlet.ServletContextHandler; -import org.eclipse.jetty.util.log.Log; -import org.slf4j.LoggerFactory; - -import javax.servlet.ServletContextEvent; - import java.io.FileInputStream; import java.io.FileNotFoundException; import java.io.IOException; @@ -28,6 +14,17 @@ import java.util.logging.Level; import java.util.logging.LogManager; import java.util.logging.Logger; +import javax.servlet.ServletContextEvent; +import net.lightbody.bmp.exception.JettyException; +import net.lightbody.bmp.proxy.bricks.ProxyResource; +import net.lightbody.bmp.proxy.guice.ConfigModule; +import net.lightbody.bmp.proxy.guice.JettyModule; +import net.lightbody.bmp.proxy.util.StandardFormatter; +import org.eclipse.jetty.server.Server; +import org.eclipse.jetty.servlet.ServletContextHandler; +import org.eclipse.jetty.util.log.Log; +import org.guiceyfruit.jsr250.Jsr250Module; +import org.slf4j.LoggerFactory; public class Main { private static final String LOGGING_PROPERTIES_FILENAME = "conf/bmp-logging.properties"; @@ -38,7 +35,7 @@ public class Main { public static void main(String[] args) { configureJdkLogging(); - final Injector injector = Guice.createInjector(new ConfigModule(args), new JettyModule(), new SitebricksModule() { + final Injector injector = Guice.createInjector(new ConfigModule(args), new Jsr250Module(), new JettyModule(), new SitebricksModule() { @Override protected void configureSitebricks() { scan(ProxyResource.class.getPackage()); @@ -148,4 +145,4 @@ private static void configureDefaultLogger() { logger.addHandler(handler); } } - \ No newline at end of file + diff --git a/src/main/java/net/lightbody/bmp/proxy/ProxyManager.java b/src/main/java/net/lightbody/bmp/proxy/ProxyManager.java index 636608b00..d61448625 100644 --- a/src/main/java/net/lightbody/bmp/proxy/ProxyManager.java +++ b/src/main/java/net/lightbody/bmp/proxy/ProxyManager.java @@ -3,14 +3,16 @@ import com.google.inject.Inject; import com.google.inject.Provider; import com.google.inject.Singleton; - import com.google.inject.name.Named; - import java.net.InetAddress; import java.net.UnknownHostException; import java.util.Collection; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import javax.annotation.PreDestroy; + +import net.lightbody.bmp.proxy.util.ExpirableMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -22,15 +24,28 @@ public class ProxyManager { private int lastPort; private final int minPort; private final int maxPort; - private Provider proxyServerProvider; - private ConcurrentHashMap proxies = new ConcurrentHashMap(); + private final Provider proxyServerProvider; + private final ConcurrentMap proxies; @Inject - public ProxyManager(Provider proxyServerProvider, @Named("minPort") Integer minPort, @Named("maxPort") Integer maxPort) { + public ProxyManager(Provider proxyServerProvider, @Named("minPort") Integer minPort, @Named("maxPort") Integer maxPort, @Named("ttl") Integer ttl) { this.proxyServerProvider = proxyServerProvider; this.minPort = minPort; this.maxPort = maxPort; this.lastPort = maxPort; + this.proxies = ttl > 0 ? + new ExpirableMap(ttl, new ExpirableMap.OnExpire(){ + @Override + public void run(ProxyServer proxy) { + try { + LOG.debug("Expiring ProxyServer `{}`...", proxy.getPort()); + proxy.stop(); + } catch (Exception ex) { + LOG.warn("Error while stopping an expired proxy", ex); + } + } + }) : + new ConcurrentHashMap(); } public ProxyServer create(Map options, Integer port, String bindAddr) { @@ -115,4 +130,11 @@ public void delete(int port) { ProxyServer proxy = proxies.remove(port); proxy.stop(); } + + @PreDestroy + public void stop(){ + if(proxies instanceof ExpirableMap){ + ((ExpirableMap)proxies).stop(); + } + } } diff --git a/src/main/java/net/lightbody/bmp/proxy/ProxyServer.java b/src/main/java/net/lightbody/bmp/proxy/ProxyServer.java index 3f8c004ac..2520c4181 100644 --- a/src/main/java/net/lightbody/bmp/proxy/ProxyServer.java +++ b/src/main/java/net/lightbody/bmp/proxy/ProxyServer.java @@ -4,11 +4,13 @@ import java.net.NetworkInterface; import java.net.SocketException; import java.net.UnknownHostException; +import java.util.Collection; import java.util.Date; import java.util.List; import java.util.Map; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.regex.Pattern; import net.lightbody.bmp.core.har.Har; import net.lightbody.bmp.core.har.HarEntry; @@ -55,7 +57,7 @@ public class ProxyServer { private StreamManager streamManager; private HarPage currentPage; private BrowserMobProxyHandler handler; - private int pageCount = 1; + private AtomicInteger pageCount = new AtomicInteger(1); private AtomicInteger requestCounter = new AtomicInteger(0); public ProxyServer() { @@ -225,7 +227,7 @@ public boolean checkCondition(long elapsedTimeInMs) { } public Har newHar(String initialPageRef) { - pageCount = 1; + pageCount.set(1); Har oldHar = getHar(); @@ -238,14 +240,14 @@ public Har newHar(String initialPageRef) { public void newPage(String pageRef) { if (pageRef == null) { - pageRef = "Page " + pageCount; + pageRef = "Page " + pageCount.get(); } client.setHarPageRef(pageRef); currentPage = new HarPage(pageRef); client.getHar().getLog().addPage(currentPage); - pageCount++; + pageCount.incrementAndGet(); } public void endPage() { @@ -333,26 +335,71 @@ public void clearRewriteRules() { client.clearRewriteRules(); } + public void blacklistRequests(String pattern, int responseCode) { + client.blacklistRequests(pattern, responseCode, null); + } + public void blacklistRequests(String pattern, int responseCode, String method) { client.blacklistRequests(pattern, responseCode, method); } + /** + * @deprecated use getBlacklistedUrls() + */ + @Deprecated public List getBlacklistedRequests() { return client.getBlacklistedRequests(); } - - public WhitelistEntry getWhitelistRequests() { - return client.getWhitelistRequests(); - } + public Collection getBlacklistedUrls() { + return client.getBlacklistedUrls(); + } + + public boolean isWhitelistEnabled() { + return client.isWhitelistEnabled(); + } + + /** + * @deprecated use getWhitelistUrls() + */ + @Deprecated + public List getWhitelistRequests() { + return client.getWhitelistRequests(); + } + + public Collection getWhitelistUrls() { + return client.getWhitelistUrls(); + } + + public int getWhitelistResponseCode() { + return client.getWhitelistResponseCode(); + } + public void clearBlacklist() { client.clearBlacklist(); } + /** + * Whitelists the specified requests. + *

+ * Note: This method overwrites any existing whitelist. + * + * @param patterns regular expression patterns matching URLs to whitelist + * @param responseCode response code to return for non-whitelisted URLs + */ public void whitelistRequests(String[] patterns, int responseCode) { client.whitelistRequests(patterns, responseCode); } + /** + * Enables an empty whitelist, which will return the specified responseCode for all requests. + * + * @param responseCode HTTP response code to return for all requests + */ + public void enableEmptyWhitelist(int responseCode) { + client.whitelistRequests(new String[0], responseCode); + } + public void clearWhitelist() { client.clearWhitelist(); } diff --git a/src/main/java/net/lightbody/bmp/proxy/Whitelist.java b/src/main/java/net/lightbody/bmp/proxy/Whitelist.java new file mode 100644 index 000000000..0a4a55fca --- /dev/null +++ b/src/main/java/net/lightbody/bmp/proxy/Whitelist.java @@ -0,0 +1,75 @@ +package net.lightbody.bmp.proxy; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.regex.Pattern; + +/** + * A URL whitelist. An empty whitelist is disabled by default. This object is immutable and the list of matching + * patterns is unmodifiable after creation. Enabling, disabling, or modifying the whitelist can be safely and easily + * accomplished by updating the whitelist reference to a new whitelist. + */ +public class Whitelist { + private final Collection patterns; + private final int responseCode; + private final boolean enabled; + + /** + * A disabled Whitelist. + */ + public static final Whitelist WHITELIST_DISABLED = new Whitelist(); + + /** + * Creates an empty, disabled Whitelist. + */ + public Whitelist() { + this.patterns = Collections.emptyList(); + this.responseCode = -1; + this.enabled = false; + } + + /** + * Creates an empty, enabled whitelist with the specified response code. + * + * @param responseCode the response code that the (enabled) Whitelist will return for all URLs. + */ + public Whitelist(int responseCode) { + this.patterns = Collections.emptyList(); + this.responseCode = responseCode; + this.enabled = true; + } + + /** + * Creates an whitelist for the specified patterns, returning the given responseCode when a URL does not match one of the patterns. + * + * @param patterns + * @param responseCode + */ + public Whitelist(String[] patterns, int responseCode) { + List patternList = new ArrayList(patterns.length); + + for (String pattern : patterns) { + patternList.add(Pattern.compile(pattern)); + } + + this.patterns = Collections.unmodifiableList(patternList); + + this.responseCode = responseCode; + + this.enabled = true; + } + + public boolean isEnabled() { + return enabled; + } + + public Collection getPatterns() { + return this.patterns; + } + + public int getResponseCode() { + return this.responseCode; + } +} diff --git a/src/main/java/net/lightbody/bmp/proxy/WhitelistEntry.java b/src/main/java/net/lightbody/bmp/proxy/WhitelistEntry.java deleted file mode 100644 index 4c20a9703..000000000 --- a/src/main/java/net/lightbody/bmp/proxy/WhitelistEntry.java +++ /dev/null @@ -1,25 +0,0 @@ -package net.lightbody.bmp.proxy; - -import java.util.List; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.regex.Pattern; - -public class WhitelistEntry { - private List patterns = new CopyOnWriteArrayList(); - private int responseCode; - - public WhitelistEntry(String[] patterns, int responseCode) { - for (String pattern : patterns) { - this.patterns.add(Pattern.compile(pattern)); - } - this.responseCode = responseCode; - } - - public List getPatterns() { - return this.patterns; - } - - public int getResponseCode() { - return this.responseCode; - } -} diff --git a/src/main/java/net/lightbody/bmp/proxy/bricks/ProxyResource.java b/src/main/java/net/lightbody/bmp/proxy/bricks/ProxyResource.java index 067f2651d..421de29a7 100644 --- a/src/main/java/net/lightbody/bmp/proxy/bricks/ProxyResource.java +++ b/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; @@ -50,7 +45,7 @@ public class ProxyResource { private static final Logger LOG = LoggerFactory.getLogger(ProxyResource.class); - private ProxyManager proxyManager; + private final ProxyManager proxyManager; @Inject public ProxyResource(ProxyManager proxyManager) { @@ -159,7 +154,7 @@ public Reply getBlacklist(@Named("port") int port, Request request) { return Reply.saying().notFound(); } - return Reply.with(proxy.getBlacklistedRequests()).as(Json.class); + return Reply.with(proxy.getBlacklistedUrls()).as(Json.class); } @Put @@ -198,7 +193,7 @@ public Reply getWhitelist(@Named("port") int port, Request request) { return Reply.saying().notFound(); } - return Reply.with(proxy.getWhitelistRequests()).as(Json.class); + return Reply.with(proxy.getWhitelistUrls()).as(Json.class); } @Put @@ -349,6 +344,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 { @@ -379,6 +388,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) { @@ -560,4 +579,53 @@ public void setProxyList(Collection 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; + } + } } diff --git a/src/main/java/net/lightbody/bmp/proxy/guice/ConfigModule.java b/src/main/java/net/lightbody/bmp/proxy/guice/ConfigModule.java index 1c28a6455..2e70c6f38 100644 --- a/src/main/java/net/lightbody/bmp/proxy/guice/ConfigModule.java +++ b/src/main/java/net/lightbody/bmp/proxy/guice/ConfigModule.java @@ -39,6 +39,12 @@ public void configure(Binder binder) { .defaultsTo(8081, 8581) .withValuesSeparatedBy('-'); + ArgumentAcceptingOptionSpec ttlSpec = + parser.accepts("ttl", "Time in seconds until an unused proxy is deleted") + .withOptionalArg() + .ofType(Integer.class) + .defaultsTo(0); + parser.acceptsAll(asList("help", "?"), "This help text"); OptionSet options = parser.parse(args); @@ -77,8 +83,9 @@ public void configure(Binder binder) { binder.bind(Key.get(Integer.class, new NamedImpl("port"))).toInstance(port); binder.bind(Key.get(String.class, new NamedImpl("address"))).toInstance(addressSpec.value(options)); binder.bind(Key.get(Integer.class, new NamedImpl("minPort"))).toInstance(minPort); - binder.bind(Key.get(Integer.class, new NamedImpl("maxPort"))).toInstance(maxPort); - + binder.bind(Key.get(Integer.class, new NamedImpl("maxPort"))).toInstance(maxPort); + binder.bind(Key.get(Integer.class, new NamedImpl("ttl"))).toInstance(ttlSpec.value(options)); + /* * Init User Agent String Parser, update of the UAS datastore will run in background. */ diff --git a/src/main/java/net/lightbody/bmp/proxy/http/BrowserMobHttpClient.java b/src/main/java/net/lightbody/bmp/proxy/http/BrowserMobHttpClient.java index 78f1d45c9..4637ed687 100644 --- a/src/main/java/net/lightbody/bmp/proxy/http/BrowserMobHttpClient.java +++ b/src/main/java/net/lightbody/bmp/proxy/http/BrowserMobHttpClient.java @@ -10,14 +10,14 @@ import java.net.URISyntaxException; import java.nio.charset.Charset; import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; import java.util.Date; -import java.util.HashMap; -import java.util.HashSet; -import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Scanner; import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; @@ -32,7 +32,7 @@ import net.lightbody.bmp.core.har.*; import net.lightbody.bmp.proxy.BlacklistEntry; import net.lightbody.bmp.proxy.Main; -import net.lightbody.bmp.proxy.WhitelistEntry; +import net.lightbody.bmp.proxy.Whitelist; import net.lightbody.bmp.proxy.util.*; import net.sf.uadetector.ReadableUserAgent; import net.sf.uadetector.UserAgentStringParser; @@ -41,7 +41,6 @@ import org.apache.http.Header; import org.apache.http.HeaderElement; import org.apache.http.HttpClientConnection; -import org.apache.http.HttpConnection; import org.apache.http.HttpEntity; import org.apache.http.HttpException; import org.apache.http.HttpHost; @@ -118,100 +117,99 @@ public class BrowserMobHttpClient { private static final Logger LOG = LoggerFactory.getLogger(BrowserMobHttpClient.class); - public static UserAgentStringParser PARSER = UADetectorServiceFactory.getCachingAndUpdatingParser(); + private static volatile UserAgentStringParser parser; private static final int BUFFER = 4096; - private Har har; - private String harPageRef; + private volatile Har har; + private volatile String harPageRef; /** * keep headers */ - private boolean captureHeaders; + private volatile boolean captureHeaders; /** * keep contents */ - private boolean captureContent; + private volatile boolean captureContent; /** * keep binary contents (if captureContent is set to true, default policy is to capture binary contents too) */ - private boolean captureBinaryContent = true; + private volatile boolean captureBinaryContent = true; /** * socket factory dedicated to port 80 (HTTP) */ - private SimulatedSocketFactory socketFactory; + private final SimulatedSocketFactory socketFactory; /** * socket factory dedicated to port 443 (HTTPS) */ - private TrustingSSLSocketFactory sslSocketFactory; + private final TrustingSSLSocketFactory sslSocketFactory; - private PoolingHttpClientConnectionManager httpClientConnMgr; + private final PoolingHttpClientConnectionManager httpClientConnMgr; /** * Builders for httpClient * Each time you change their configuration you should call updateHttpClient() */ - private Builder requestConfigBuilder; - private HttpClientBuilder httpClientBuilder; + private final Builder requestConfigBuilder; + private final HttpClientBuilder httpClientBuilder; /** * The current httpClient which will execute HTTP requests */ - private CloseableHttpClient httpClient; + private volatile CloseableHttpClient httpClient; - private BasicCookieStore cookieStore = new BasicCookieStore(); + private final BasicCookieStore cookieStore = new BasicCookieStore(); /** * List of rejected URL patterns */ - private List blacklistEntries = new CopyOnWriteArrayList(); + private final Collection blacklistEntries = new CopyOnWriteArrayList(); /** * List of accepted URL patterns */ - - private WhitelistEntry whitelistEntry = null; + private volatile Whitelist whitelist = Whitelist.WHITELIST_DISABLED; /** * List of URLs to rewrite */ - private List rewriteRules = new CopyOnWriteArrayList(); + private final List rewriteRules = new CopyOnWriteArrayList(); /** * triggers to process when sending request */ - private List requestInterceptors = new CopyOnWriteArrayList(); + private final List requestInterceptors = new CopyOnWriteArrayList(); /** * triggers to process when receiving response */ - private List responseInterceptors = new CopyOnWriteArrayList(); + private final List responseInterceptors = new CopyOnWriteArrayList(); /** * additional headers sent with request */ - private HashMap additionalHeaders = new LinkedHashMap(); + private final Map additionalHeaders = new ConcurrentHashMap(); /** * request timeout: set to -1 to disable timeout */ - private int requestTimeout = -1; + private volatile int requestTimeout = -1; /** * is it possible to add a new request? */ - private AtomicBoolean allowNewRequests = new AtomicBoolean(true); + private final AtomicBoolean allowNewRequests = new AtomicBoolean(true); /** * DNS lookup handler */ - private BrowserMobHostNameResolver hostNameResolver; + private final BrowserMobHostNameResolver hostNameResolver; /** * does the proxy support gzip compression? (set to false if you go through a browser) @@ -221,9 +219,7 @@ public class BrowserMobHttpClient { /** * set of active requests */ - // not using CopyOnWriteArray because we're WRITE heavy and it is for READ heavy operations - // instead doing it the old fashioned way with a synchronized block - private final Set activeRequests = new HashSet(); + private final Set activeRequests = Collections.newSetFromMap(new ConcurrentHashMap()); /** * credentials used for authentication @@ -253,7 +249,7 @@ public class BrowserMobHttpClient { /** * remaining requests counter */ - private AtomicInteger requestCounter; + private final AtomicInteger requestCounter; /** * Init HTTP client @@ -570,10 +566,8 @@ private BadURIException reportBadURI(String url, String method, URISyntaxExcepti } public void checkTimeout() { - synchronized (activeRequests) { - for (ActiveRequest activeRequest : activeRequests) { - activeRequest.checkTimeout(); - } + for (ActiveRequest activeRequest : activeRequests) { + activeRequest.checkTimeout(); } // Close expired connections @@ -587,7 +581,7 @@ public BrowserMobHttpResponse execute(BrowserMobHttpRequest req) { if (!allowNewRequests.get()) { throw new RuntimeException("No more requests allowed"); } - + try { requestCounter.incrementAndGet(); @@ -626,7 +620,7 @@ private BrowserMobHttpResponse execute(BrowserMobHttpRequest req, int depth) { String userAgent = uaHeaders[0].getValue(); try { // note: this doesn't work for 'Fandango/4.5.1 CFNetwork/548.1.4 Darwin/11.0.0' - ReadableUserAgent uai = PARSER.parse(userAgent); + ReadableUserAgent uai = getUserAgentStringParser().parse(userAgent); String browser = uai.getName(); String version = uai.getVersionNumber().toVersionString(); har.getLog().setBrowser(new HarNameVersion(browser, version)); @@ -656,31 +650,27 @@ private BrowserMobHttpResponse execute(BrowserMobHttpRequest req, int depth) { // handle whitelist and blacklist entries int mockResponseCode = -1; - synchronized (this) { - // guard against concurrent modification of whitelistEntry - if (whitelistEntry != null) { - boolean found = false; - for (Pattern pattern : whitelistEntry.getPatterns()) { - if (pattern.matcher(url).matches()) { - found = true; - break; - } - } - - // url does not match whitelist, set the response code - if (!found) { - mockResponseCode = whitelistEntry.getResponseCode(); + // alias the current whitelist, in case the whitelist is changed while processing this request + Whitelist currentWhitelist = whitelist; + if (currentWhitelist.isEnabled()) { + boolean found = false; + for (Pattern pattern : currentWhitelist.getPatterns()) { + if (pattern.matcher(url).matches()) { + found = true; + break; } } + + // url does not match whitelist, set the response code + if (!found) { + mockResponseCode = currentWhitelist.getResponseCode(); + } } - if (blacklistEntries != null) { - for (BlacklistEntry blacklistEntry : blacklistEntries) { - if (blacklistEntry.getPattern().matcher(url).matches() - && blacklistEntry.getMethod().matcher(method.getMethod()).matches()) { - mockResponseCode = blacklistEntry.getResponseCode(); - break; - } + for (BlacklistEntry blacklistEntry : blacklistEntries) { + if (blacklistEntry.matches(url, method.getMethod())) { + mockResponseCode = blacklistEntry.getResponseCode(); + break; } } @@ -735,10 +725,8 @@ private BrowserMobHttpResponse execute(BrowserMobHttpRequest req, int depth) { BasicHttpContext ctx = new BasicHttpContext(); - ActiveRequest activeRequest = new ActiveRequest(method, ctx, entry.getStartedDateTime()); - synchronized (activeRequests) { - activeRequests.add(activeRequest); - } + ActiveRequest activeRequest = new ActiveRequest(method, entry.getStartedDateTime()); + activeRequests.add(activeRequest); // for dealing with automatic authentication if (authType == AuthType.NTLM) { @@ -854,10 +842,7 @@ public HeaderElement[] getElements() throws ParseException { } } finally { // the request is done, get it out of here - //FIXME: use set backed by ConcurrentHashMap - synchronized (activeRequests) { - activeRequests.remove(activeRequest); - } + activeRequests.remove(activeRequest); if (is != null) { try { @@ -949,7 +934,7 @@ public HeaderElement[] getElements() throws ParseException { List result = new ArrayList(); URLEncodedUtils.parse(result, new Scanner(content), null); - List params = new ArrayList(); + List params = new ArrayList(result.size()); data.setParams(params); for (NameValuePair pair : result) { @@ -1130,18 +1115,20 @@ public void shutdown() { } public void abortActiveRequests() { - allowNewRequests.set(true); + allowNewRequests.set(false); - synchronized (activeRequests) { - for (ActiveRequest activeRequest : activeRequests) { - activeRequest.abort(); - } - activeRequests.clear(); + for (ActiveRequest activeRequest : activeRequests) { + activeRequest.abort(); } + + activeRequests.clear(); } public void setHar(Har har) { this.har = har; + + // eagerly initialize the User Agent String Parser, since it will be needed for the HAR + getUserAgentStringParser(); } public void setHarPageRef(String harPageRef) { @@ -1195,8 +1182,11 @@ public void clearRewriteRules() { rewriteRules.clear(); } - // this method is provided for backwards compatibility before we renamed it to - // blacklistRequests (note the plural) + /** + * this method is provided for backwards compatibility before we renamed it to blacklistRequests (note the plural) + * @deprecated use blacklistRequests(String pattern, int responseCode) + */ + @Deprecated public void blacklistRequest(String pattern, int responseCode, String method) { blacklistRequests(pattern, responseCode, method); } @@ -1205,26 +1195,75 @@ public void blacklistRequests(String pattern, int responseCode, String method) { blacklistEntries.add(new BlacklistEntry(pattern, responseCode, method)); } + /** + * @deprecated Use getBlacklistedUrls() + */ + @Deprecated public List getBlacklistedRequests() { - return blacklistEntries; + List blacklist = new ArrayList(blacklistEntries.size()); + blacklist.addAll(blacklistEntries); + + return blacklist; + } + + public Collection getBlacklistedUrls() { + return blacklistEntries; } public void clearBlacklist() { blacklistEntries.clear(); } - public WhitelistEntry getWhitelistRequests() { - return whitelistEntry; + public boolean isWhitelistEnabled() { + return whitelist.isEnabled(); } - - public synchronized void whitelistRequests(String[] patterns, int responseCode) { - // synchronized to guard against concurrent modification - whitelistEntry = new WhitelistEntry(patterns, responseCode); + + /** + * @deprecated use getWhitelistUrls() + * @return unmodifiable list of whitelisted Patterns + */ + @Deprecated + public List getWhitelistRequests() { + List whitelistPatterns = new ArrayList(whitelist.getPatterns().size()); + whitelistPatterns.addAll(whitelist.getPatterns()); + + return Collections.unmodifiableList(whitelistPatterns); + } + + /** + * Retrieves Patterns of URLs that have been whitelisted. + * + * @return unmodifiable whitelisted URL Patterns + */ + public Collection getWhitelistUrls() { + return whitelist.getPatterns(); + } + + public int getWhitelistResponseCode() { + return whitelist.getResponseCode(); } - public synchronized void clearWhitelist() { - // synchronized to guard against concurrent modification - whitelistEntry = null; + /** + * Whitelist the specified request patterns, returning the specified responseCode for non-whitelisted + * requests. + * + * @param patterns regular expression strings matching URL patterns to whitelist. if empty or null, + * the whitelist will be enabled but will not match any URLs. + * @param responseCode the HTTP response code to return for non-whitelisted requests + */ + public void whitelistRequests(String[] patterns, int responseCode) { + if (patterns == null || patterns.length == 0) { + whitelist = new Whitelist(responseCode); + } else { + whitelist = new Whitelist(patterns, responseCode); + } + } + + /** + * Clears and disables the current whitelist. + */ + public void clearWhitelist() { + whitelist = Whitelist.WHITELIST_DISABLED; } public void addHeader(String name, String value) { @@ -1332,45 +1371,50 @@ public void process(final HttpRequest request, final HttpContext context) throws } class ActiveRequest { - HttpRequestBase request; - BasicHttpContext ctx; - Date start; + private final HttpRequestBase request; + private final Date start; + private final AtomicBoolean aborting = new AtomicBoolean(false); - ActiveRequest(HttpRequestBase request, BasicHttpContext ctx, Date start) { + ActiveRequest(HttpRequestBase request, Date start) { this.request = request; - this.ctx = ctx; this.start = start; } - - void checkTimeout() { + + /** + * Checks the timeout for this request, and aborts if necessary. + * @return true if the request was aborted for exceeding its timeout, otherwise false. + */ + boolean checkTimeout() { + if (aborting.get()) { + return false; + } + if (requestTimeout != -1) { if (request != null && start != null && new Date(System.currentTimeMillis() - requestTimeout).after(start)) { - LOG.info("Aborting request to {} after it failed to complete in {} ms", request.getURI().toString(), requestTimeout); + boolean okayToAbort = aborting.compareAndSet(false, true); + if (okayToAbort) { + LOG.info("Aborting request to {} after it failed to complete in {} ms", request.getURI().toString(), requestTimeout); - abort(); + abort(); + + return true; + } } } + + return false; } public void abort() { request.abort(); - - // try to close the connection? is this necessary? unclear based on preliminary debugging of HttpClient, but - // it doesn't seem to hurt to try - HttpConnection conn = (HttpConnection) ctx.getAttribute("http.connection"); - if (conn != null) { - try { - conn.close(); - } catch (IOException e) { - // this is fine, we're shutting it down anyway - } - } + + // no need to close the connection -- the call to request.abort() releases the connection itself } } private class RewriteRule { - private Pattern match; - private String replace; + private final Pattern match; + private final String replace; private RewriteRule(String match, String replace) { this.match = Pattern.compile(match); @@ -1430,4 +1474,22 @@ public static long copyWithStats(InputStream is, OutputStream os) throws IOExcep return bytesCopied; } + + private static final Object PARSER_INIT_LOCK = new Object(); + + /** + * Retrieve the User Agent String Parser. Create the parser if it has not yet been initialized. + * @return + */ + public static UserAgentStringParser getUserAgentStringParser() { + if (parser == null) { + synchronized (PARSER_INIT_LOCK) { + if (parser == null) { + parser = UADetectorServiceFactory.getCachingAndUpdatingParser(); + } + } + } + + return parser; + } } diff --git a/src/main/java/net/lightbody/bmp/proxy/http/HttpClientInterrupter.java b/src/main/java/net/lightbody/bmp/proxy/http/HttpClientInterrupter.java index a9f23de2b..44e923f8e 100644 --- a/src/main/java/net/lightbody/bmp/proxy/http/HttpClientInterrupter.java +++ b/src/main/java/net/lightbody/bmp/proxy/http/HttpClientInterrupter.java @@ -8,7 +8,7 @@ public class HttpClientInterrupter { private static final Logger LOG = LoggerFactory.getLogger(HttpClientInterrupter.class); - private static Set clients = new CopyOnWriteArraySet(); + private static final Set clients = new CopyOnWriteArraySet(); static { Thread thread = new Thread(new Runnable() { @@ -26,7 +26,7 @@ public void run() { try { Thread.sleep(1000); } catch (InterruptedException e) { - // this is OK + Thread.currentThread().interrupt(); } } } diff --git a/src/main/java/net/lightbody/bmp/proxy/http/SimulatedSocket.java b/src/main/java/net/lightbody/bmp/proxy/http/SimulatedSocket.java index b8ea15066..059caf9e3 100644 --- a/src/main/java/net/lightbody/bmp/proxy/http/SimulatedSocket.java +++ b/src/main/java/net/lightbody/bmp/proxy/http/SimulatedSocket.java @@ -16,7 +16,7 @@ * and get-in-out streams to provide throttling */ public class SimulatedSocket extends Socket { - private StreamManager streamManager; + private final StreamManager streamManager; public SimulatedSocket(StreamManager streamManager) { this.streamManager = streamManager; @@ -68,7 +68,7 @@ private void simulateLatency (Date start, Date end, StreamManager streamManager) try { Thread.sleep(streamManager.getLatency()-connectReal); } catch (InterruptedException e) { - Thread.interrupted(); + Thread.currentThread().interrupt(); } // the end after adding latency end = new Date(); diff --git a/src/main/java/net/lightbody/bmp/proxy/util/ExpirableMap.java b/src/main/java/net/lightbody/bmp/proxy/util/ExpirableMap.java new file mode 100644 index 000000000..f2f70e03b --- /dev/null +++ b/src/main/java/net/lightbody/bmp/proxy/util/ExpirableMap.java @@ -0,0 +1,97 @@ +package net.lightbody.bmp.proxy.util; + +import java.util.Date; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; + +public class ExpirableMap extends ConcurrentHashMap{ + public final static int DEFAULT_CHECK_INTERVAL = 10*60; + public final static int DEFAULT_TTL = 30*60; + private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(); + private final long ttl; + private final Map expires; + private final OnExpire onExpire; + + public ExpirableMap(int ttl, int checkInterval, OnExpire onExpire) { + this.ttl = ttl*1000; + this.onExpire = onExpire; + expires = new HashMap<>(); + scheduler.scheduleWithFixedDelay(new Worker(), checkInterval, checkInterval, TimeUnit.SECONDS); + } + + public ExpirableMap(int ttl, OnExpire onExpire) { + this(ttl, DEFAULT_CHECK_INTERVAL, onExpire); + } + + public ExpirableMap(OnExpire onExpire) { + this(DEFAULT_TTL, DEFAULT_CHECK_INTERVAL, onExpire); + } + + public ExpirableMap() { + this(DEFAULT_TTL, DEFAULT_CHECK_INTERVAL, null); + } + + @Override + public V putIfAbsent(K key, V value) { + synchronized(this){ + expires.put(key, new Date().getTime()+ttl); + return super.putIfAbsent(key, value); + } + } + + @Override + public V put(K key, V value) { + synchronized(this){ + expires.put(key, new Date().getTime()+ttl); + return super.put(key, value); + } + } + + public void stop() { + scheduler.shutdown(); + try { + scheduler.awaitTermination(10, TimeUnit.SECONDS); + } catch (InterruptedException ex) { + scheduler.shutdownNow(); + } + } + + private class Worker implements Runnable{ + + @Override + public void run() { + Map m; + synchronized(ExpirableMap.this){ + m = new HashMap<>(expires); + } + Long now = new Date().getTime(); + for(Entry e : m.entrySet()){ + if(e.getValue() > now){ + continue; + } + synchronized(ExpirableMap.this){ + Long expire = expires.get(e.getKey()); + if(expire == null){ + continue; + } + if(expire <= new Date().getTime()){ + expires.remove(e.getKey()); + V v = ExpirableMap.this.remove(e.getKey()); + if(v != null && onExpire != null){ + onExpire.run(v); + } + } + } + } + } + + } + + public interface OnExpire{ + public abstract void run(V value); + } +} diff --git a/src/main/java/org/java_bandwidthlimiter/BandwidthLimiter.java b/src/main/java/org/java_bandwidthlimiter/BandwidthLimiter.java index cc1c0de78..e1a27c331 100644 --- a/src/main/java/org/java_bandwidthlimiter/BandwidthLimiter.java +++ b/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); } diff --git a/src/main/java/org/java_bandwidthlimiter/MaximumTransferExceededException.java b/src/main/java/org/java_bandwidthlimiter/MaximumTransferExceededException.java new file mode 100644 index 000000000..3b00e68ca --- /dev/null +++ b/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; + } +} diff --git a/src/main/java/org/java_bandwidthlimiter/StreamManager.java b/src/main/java/org/java_bandwidthlimiter/StreamManager.java index 7dbfd4a22..4c1b533de 100644 --- a/src/main/java/org/java_bandwidthlimiter/StreamManager.java +++ b/src/main/java/org/java_bandwidthlimiter/StreamManager.java @@ -32,6 +32,7 @@ import java.io.InputStream; import java.io.OutputStream; import java.util.Random; +import static org.java_bandwidthlimiter.BandwidthLimiter.OneSecond; /** * A class that manages the bandwidth of all the registered streams (both input and output streams). @@ -61,6 +62,8 @@ private class StreamParams { public long remainingBps; public long nextResetTimestamp; public long nextResetSubIntervals; + public long maxBytes; + public long remainingBytes; private long timeToNextReset() { return nextResetTimestamp - System.currentTimeMillis(); @@ -83,11 +86,11 @@ private void adjustBytes(long bytesNumber) { //even calls to setDownstreamKbps and setUpstreamKbps will be forced to honor this upperbound. private long maxBytesPerSecond; - private StreamParams downStream = new StreamParams(); - private StreamParams upStream = new StreamParams(); + private final StreamParams downStream = new StreamParams(); + private final StreamParams upStream = new StreamParams(); private long latency = 0; - private Random randomGenerator = new Random(); + private final Random randomGenerator = new Random(); /** * Create an instance of StreamManager. @@ -101,11 +104,13 @@ private void adjustBytes(long bytesNumber) { * */ public StreamManager(long maxBitsPerSecond) { - setMaxBitsPerSecondThreshold( maxBitsPerSecond ); - setDownstreamKbps(maxBitsPerSecond / 1000); - setUpstreamKbps(maxBitsPerSecond / 1000); - setPayloadPercentage(95); - disable(); + this.maxBytesPerSecond = maxBitsPerSecond/8; + setMaxBps(this.downStream, this.maxBytesPerSecond); + setMaxBps(this.upStream, this.maxBytesPerSecond); + this.actualPayloadPercentage = 0.95; + setMaxBytes(this.downStream, 0); + setMaxBytes(this.upStream, 0); + this.enabled = false; } @@ -154,11 +159,31 @@ public void setUpstreamKbps(long upstreamKbps) { public void setLatency(long latency) { this.latency = latency; } + + /** + * Specifies how many kilobytes in total the client is allowed to download. + * When the limit is used up, MaximumTransferExceededException is thrown + * @param downstreamMaxKB + */ + @Override + public void setDownstreamMaxKB(long downstreamMaxKB) { + setMaxBytes(this.downStream, downstreamMaxKB * 1000); + } + + /** + * Specifies how many kilobytes in total the client is allowed to upload. + * When the limit is used up, MaximumTransferExceededException is thrown + * @param upstreamMaxKB + */ + @Override + public void setUpstreamMaxKB(long upstreamMaxKB) { + setMaxBytes(this.upStream, upstreamMaxKB * 1000); + } public long getLatency() { return latency; } - + /** * To take into account overhead due to underlying protocols (e.g. TCP/IP) * @param payloadPercentage a ] 0 , 100] value. where 100 means that the required @@ -192,7 +217,7 @@ public InputStream registerStream(InputStream in) { * Register an output stream. * A client would then use the returned OutputStream, which will throttle * the one passed as parameter. - * * @param in The OutputStream that will be throttled + * @param out The OutputStream that will be throttled * @return a new throttled OutputStream (wrapping the one given as parameter) * */ @@ -215,6 +240,22 @@ public void setMaxBitsPerSecondThreshold(long maxBitsPerSecond) { setMaxBps(this.downStream, this.downStream.maxBps); setMaxBps(this.upStream, this.upStream.maxBps); } + + public long getMaxUpstreamKB(){ + return this.upStream.maxBytes/1000; + } + + public long getRemainingUpstreamKB(){ + return this.upStream.remainingBytes/1000; + } + + public long getMaxDownstreamKB(){ + return this.downStream.maxBytes/1000; + } + + public long getRemainingDownstreamKB(){ + return this.downStream.remainingBytes/1000; + } private void setMaxBps( StreamParams direction, long maxBps ) { synchronized (direction) { @@ -226,6 +267,13 @@ private void setMaxBps( StreamParams direction, long maxBps ) { direction.reset(); } } + + private void setMaxBytes( StreamParams direction, long maxBytes ) { + synchronized (direction) { + direction.maxBytes = maxBytes; + direction.remainingBytes = maxBytes; + } + } private long timeToNextReset(StreamParams direction) { synchronized (direction) { @@ -255,7 +303,7 @@ private int getAllowedBytesWrite(ManagedOutputStream stream, int bufferLength) { private int getAllowedBytesUnFair(StreamParams direction, int bufferLength) { //this is an unfair allocation of bytes/second because it gives as many //bytes as possible to anyone who ask for them - int allowed = 0; + int allowed; synchronized(direction) { resetCounterIfNecessary(direction); //this stream desires to read up to bufferLength bytes @@ -288,6 +336,11 @@ private int manageRead(ManagedInputStream stream, byte[] b, int off, int len) th if(allowed > 0) { // Read a maximum of "allowed" bytes int bytesRead = stream.doRead(b, off, allowed); + + // check if we exceeded the transfer limit + if(this.downStream.maxBytes > 0 && (this.downStream.remainingBytes -= bytesRead) < 0){ + throw new MaximumTransferExceededException(getMaxDownstreamKB(), false); + } // If less than the "allowed" bytes were read, adjust how many we can still read for this period of time adjustBytes(this.downStream, allowed - bytesRead); @@ -318,7 +371,7 @@ private void manageWrite(ManagedOutputStream stream, byte[] b, int off, int len) assert maxBytesPerSecond > 0; int bytesWritten = 0; - int allowed = 0; + int allowed; // we need a while loop since the write doesn't return a "written bytes" count, // rather it expects that all of them are written // hence we loop here until all of them have been written @@ -326,7 +379,11 @@ private void manageWrite(ManagedOutputStream stream, byte[] b, int off, int len) allowed = getAllowedBytesWrite(stream, len); if(allowed > 0) { stream.doWrite(b, off, allowed); - bytesWritten += allowed; + bytesWritten += allowed; + // check if we exceeded the transfer limit + if(this.upStream.maxBytes > 0 && (this.upStream.remainingBytes -= allowed) < 0){ + throw new MaximumTransferExceededException(this.getMaxUpstreamKB(), true); + } } else { long sleepTime = timeToNextReset(this.upStream); if( sleepTime > 0 ) { @@ -348,7 +405,7 @@ private static void threadSleep(long sleepTime) { try { Thread.sleep(sleepTime); } catch (InterruptedException e) { - Thread.interrupted(); + Thread.currentThread().interrupt(); } } @@ -361,7 +418,7 @@ private class ManagedOutputStream extends OutputStream { long lastActivity; boolean roundUp; //just an helper buffer so we don't allocate it all the time when calling void write(int b) which writes ONE byte! - private byte[] oneByteBuff = new byte[1]; + private final byte[] oneByteBuff = new byte[1]; public ManagedOutputStream(OutputStream stream, StreamManager manager) { assert manager != null; @@ -407,10 +464,12 @@ public void doWrite(byte[] b, int offset, int length) throws IOException { stream.write(b, offset, length); } + @Override public void flush() throws IOException { stream.flush(); } + @Override public void close() throws IOException { stream.close(); } @@ -422,7 +481,7 @@ private class ManagedInputStream extends InputStream { long lastActivity; boolean roundUp; //just an helper buffer so we don't allocate it all the time when calling int read() which reads ONE byte! - private byte[] oneByteBuff = new byte[1]; + private final byte[] oneByteBuff = new byte[1]; public ManagedInputStream(InputStream stream, StreamManager manager) { assert manager != null; @@ -445,17 +504,20 @@ public boolean getRoundUp() { return roundUp; } + @Override public int read() throws IOException { read(oneByteBuff, 0, 1); return oneByteBuff[0]; } + @Override public int read(byte[] b) throws IOException { int length = b.length; int bytesRead = read(b, 0, length); return bytesRead; } + @Override public int read(byte[] b, int off, int len) throws IOException { int readBytes = manager.manageRead(this, b, off, len); if( readBytes > 0 ) { @@ -472,26 +534,32 @@ public int doRead(byte[] b, int offset, int length) throws IOException { return stream.read(b, offset, length); } + @Override public long skip(long n) throws IOException { return stream.skip(n); } + @Override public int available() throws IOException { return stream.available(); } + @Override public void close() throws IOException { stream.close(); } + @Override public void mark(int readLimit) { stream.mark(readLimit); } + @Override public void reset() throws IOException { stream.reset(); } + @Override public boolean markSupported() { return stream.markSupported(); } diff --git a/src/test/java/net/lightbody/bmp/proxy/BlackAndWhiteListTest.java b/src/test/java/net/lightbody/bmp/proxy/BlackAndWhiteListTest.java index 7e58903f5..212a03d3b 100644 --- a/src/test/java/net/lightbody/bmp/proxy/BlackAndWhiteListTest.java +++ b/src/test/java/net/lightbody/bmp/proxy/BlackAndWhiteListTest.java @@ -11,6 +11,7 @@ import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.not; import static org.junit.Assert.assertThat; +import static org.junit.Assert.assertFalse; import static org.junit.Assume.assumeThat; /** @@ -116,6 +117,35 @@ public void testWhitelistCanBeCleared() throws ClientProtocolException, IOExcept assertThat(httpStatusWhenGetting("http://127.0.0.1:8080/a.txt"), is(200)); assertThat(httpStatusWhenGetting("http://127.0.0.1:8080/c.png"), is(200)); } + + @Test + public void testWhitelistCanBeReplaced() throws ClientProtocolException, IOException { + proxy.whitelistRequests(new String[] { ".*\\.txt" }, 404); + + // test that the whitelist is working + assertThat(httpStatusWhenGetting("http://127.0.0.1:8080/a.txt"), is(200)); + assertThat(httpStatusWhenGetting("http://127.0.0.1:8080/c.png"), is(404)); + + proxy.whitelistRequests(new String[] { ".*\\.png" }, 404); + + // check that the new whitelist is working and the old is gone + assertThat(httpStatusWhenGetting("http://127.0.0.1:8080/a.txt"), is(404)); + assertThat(httpStatusWhenGetting("http://127.0.0.1:8080/c.png"), is(200)); + } + + @Test + public void testEmptyWhitelist() throws ClientProtocolException, IOException { + assertThat(httpStatusWhenGetting("http://127.0.0.1:8080/a.txt"), is(200)); + + proxy.enableEmptyWhitelist(404); + + assertThat(httpStatusWhenGetting("http://127.0.0.1:8080/a.txt"), is(404)); + } + + @Test + public void testWhitelistIsDisabledByDefault() { + assertFalse("whitelist should be diabled unless explicitly set", proxy.isWhitelistEnabled()); + } /** * Checks that a proxy blacklist can be cleared successfully. diff --git a/src/test/java/net/lightbody/bmp/proxy/ExpirableMapTest.java b/src/test/java/net/lightbody/bmp/proxy/ExpirableMapTest.java new file mode 100644 index 000000000..bbfd37e8d --- /dev/null +++ b/src/test/java/net/lightbody/bmp/proxy/ExpirableMapTest.java @@ -0,0 +1,57 @@ +package net.lightbody.bmp.proxy; + +import java.util.HashSet; +import java.util.Set; +import net.lightbody.bmp.proxy.util.ExpirableMap; +import org.junit.After; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import org.junit.Before; +import org.junit.Test; + +public class ExpirableMapTest { + private Set strings = new HashSet<>(); + private ExpirableMap m; + + @Before + public void setUp() throws Exception { + m = new ExpirableMap<>(1, 1, new ExpirableMap.OnExpire(){ + @Override + public void run(String s) { + ExpirableMapTest.this.strings.add(s); + } + }); + } + + @After + public void tearDown() throws Exception { + m.stop(); + } + + @Test + public void testKeyExpiration() throws Exception { + + m.put(1, "a"); + m.put(1, "b"); + String s = m.putIfAbsent(2, "c"); + + assertNull(s); + + s = m.putIfAbsent(2, "d"); + + assertEquals("c", s); + + Thread.sleep(2000); + + assertEquals(0, m.size()); + + assertFalse(strings.contains("a")); + + assertTrue(strings.contains("b")); + + assertTrue(strings.contains("c")); + + } +} diff --git a/src/test/java/net/lightbody/bmp/proxy/ProxyServerTest.java b/src/test/java/net/lightbody/bmp/proxy/ProxyServerTest.java index 5b47c6aec..573cda2ff 100644 --- a/src/test/java/net/lightbody/bmp/proxy/ProxyServerTest.java +++ b/src/test/java/net/lightbody/bmp/proxy/ProxyServerTest.java @@ -58,7 +58,7 @@ public static DefaultHttpClient getNewHttpClient(int proxyPort) { return new DefaultHttpClient(ccm, params); } catch (Exception e) { - return new DefaultHttpClient(); + throw new RuntimeException("Unable to get HTTP client", e); } } diff --git a/src/test/java/net/lightbody/bmp/proxy/TimeoutsTest.java b/src/test/java/net/lightbody/bmp/proxy/TimeoutsTest.java new file mode 100644 index 000000000..7f10380ed --- /dev/null +++ b/src/test/java/net/lightbody/bmp/proxy/TimeoutsTest.java @@ -0,0 +1,45 @@ +package net.lightbody.bmp.proxy; + +import java.io.IOException; + +import org.apache.http.client.ClientProtocolException; +import org.apache.http.client.methods.CloseableHttpResponse; +import org.apache.http.client.methods.HttpGet; +import org.apache.http.impl.client.DefaultHttpClient; + +import static org.junit.Assert.assertEquals; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +public class TimeoutsTest { + + private ProxyServer proxyServer; + private DefaultHttpClient client; + + @Before + public void setUp() { + proxyServer = new ProxyServer(0); + proxyServer.start(); + + client = ProxyServerTest.getNewHttpClient(proxyServer.getPort()); + } + + @Test + public void testSmallTimeout() throws IllegalStateException, ClientProtocolException, IOException { + proxyServer.setRequestTimeout(2000); + + HttpGet get = new HttpGet("http://blackhole.webpagetest.org/test"); + + CloseableHttpResponse response = client.execute(get); + + assertEquals("Expected HTTP 502 response due to timeout", 502, response.getStatusLine().getStatusCode()); + } + + @After + public void tearDown() { + proxyServer.stop(); + } + +} diff --git a/src/test/java/org/java_bandwidthlimiter/StreamManagerTest.java b/src/test/java/org/java_bandwidthlimiter/StreamManagerTest.java new file mode 100644 index 000000000..4614dc298 --- /dev/null +++ b/src/test/java/org/java_bandwidthlimiter/StreamManagerTest.java @@ -0,0 +1,87 @@ +package org.java_bandwidthlimiter; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import static org.hamcrest.CoreMatchers.instanceOf; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThat; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import org.junit.Before; +import org.junit.Test; + +public class StreamManagerTest { + private final int BUFFER = 256; + private final long LIMIT_KB = 8; + private byte[] data; + + @Before + public void setUp(){ + data = new byte[10*1024]; + for(int i = 0; i < data.length; i++){ + data[i] = 127; + } + } + + @Test + public void testDownstreamDataLimit(){ + StreamManager sm = new StreamManager( BandwidthLimiter.OneMbps ); + sm.setDownstreamMaxKB(LIMIT_KB); + sm.enable(); + InputStream in = sm.registerStream(new ByteArrayInputStream(data)); + byte[] buffer = new byte[BUFFER]; + int read = 0; + try{ + while(read < data.length){ + read += in.read(buffer); + } + + fail(); + + }catch(IOException ex){ + assertThat(ex, instanceOf(MaximumTransferExceededException.class)); + + MaximumTransferExceededException ex2 = (MaximumTransferExceededException)ex; + + assertFalse(ex2.isUpstream()); + assertEquals(LIMIT_KB, ex2.getLimit()); + assertTrue(read < LIMIT_KB*1000 + BUFFER); + assertTrue(read >= LIMIT_KB*1000 - BUFFER); + } + } + + @Test + public void testUpstreamDataLimit(){ + StreamManager sm = new StreamManager( BandwidthLimiter.OneMbps ); + sm.setUpstreamMaxKB(LIMIT_KB); + sm.enable(); + InputStream in = new ByteArrayInputStream(data); + ByteArrayOutputStream out1 = new ByteArrayOutputStream(data.length); + OutputStream out = sm.registerStream(out1); + byte[] buffer = new byte[BUFFER]; + int read = 0; + try{ + while(read < data.length){ + read += in.read(buffer); + out.write(buffer); + } + + fail(); + + }catch(IOException ex){ + assertThat(ex, instanceOf(MaximumTransferExceededException.class)); + + MaximumTransferExceededException ex2 = (MaximumTransferExceededException)ex; + + assertTrue(ex2.isUpstream()); + assertEquals(LIMIT_KB, ex2.getLimit()); + assertTrue(out1.size() < LIMIT_KB*1000 + BUFFER); + assertTrue(out1.size() >= LIMIT_KB*1000 - BUFFER); + } + } + +}