Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
ProxyManager class, for managing your BrowserMob proxy, manipulating HAR files, etc
package com.rmn.automation.framework.proxy;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.rmn.automation.framework.util.AutomationUtils;
import com.rmn.automation.framework.util.JSONUtils;
import org.apache.http.client.ClientProtocolException;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpDelete;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.client.methods.HttpPut;
import org.apache.http.entity.StringEntity;
import org.apache.http.util.EntityUtils;
import org.browsermob.core.har.Har;
import org.browsermob.core.har.HarLog;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
// This class just serves to manage proxies and hand them out to callers, spinning up new ones as they're needed
public class ProxyManager implements IProxyManager {
private final Set<IAutomationProxy> availableProxies = new HashSet<>();
private Set<IAutomationProxy> usedProxies = new HashSet<>();
private Logger log = LoggerFactory.getLogger("ProxyManager");
private final HttpClient client;
private String baseProxyUrl; // Base proxy URL (e.g. http://myproxy.mycompany.io
private int port; // Base port for the proxy (e.g. 9090)
public ProxyManager(HttpClient client, String baseProxyUrl, int port) {
this.client = client;
this.baseProxyUrl = baseProxyUrl;
this.port = port;
}
/**
* Returns the HTTP proxy URL along with the port to access it
* @return The full URL for the proxy, including port
*/
private String getHttpUrlWithPort() {
return "http://" + baseProxyUrl + ":" + port;
}
/**
* Returns a proxy for use
* @return A Proxy (i.e. an open, dedicated connection to the Proxy server) for immediate use
* @throws ProxyManagerException If the proxy cannot be created for any reason.
*/
public IAutomationProxy getProxy(String testName) throws ProxyManagerException {
// Check to see if there is an existing proxy not in use that we can reuse. If there is, then lets go
// ahead and use that proxy
IAutomationProxy proxy = null;
synchronized (availableProxies) {
if(availableProxies.size() > 0) {
Iterator<IAutomationProxy> iterator = availableProxies.iterator();
proxy = iterator.next();
// Remove this proxy so we can't assign it out until its been returned as unused
iterator.remove();
}
}
// If we got a proxy from our available proxies, go ahead and return it
if(proxy != null) {
// Add it to our used proxies so we still have record of the proxy
log.info("Existing proxy available. Assigning: " + proxy);
} else {
// If there are no available proxies not being used, go ahead and create a new one
proxy = createNewProxy(null);
log.info("Existing proxy not available. Created new proxy: " + proxy);
}
// Add the proxy to the list of used proxies so it won't be reused until its released
usedProxies.add(proxy);
return proxy;
}
/**
* Returns a proxy for the assigned port (allows you to listen on a session that's already started)
* @param port The port you're requesting to listen to
* @return The assigned Proxy you asked for
* @throws ProxyManagerException If the proxy cannot be created for any reason.
*/
public IAutomationProxy getAssignedProxy(String testName, Integer port) throws ProxyManagerException {
return this.createNewProxy(port);
}
/**
* Releases the specified proxy, allowing for it to be reused in future applications
* @param proxy The proxy to release
*/
@Override
public void releaseProxy(IAutomationProxy proxy) {
// Proxy will have different teardown logic depending on whether or not its an assigned or dynamic proxy
if(proxy.isAssignedProxy()) {
this.deleteProxy(proxy);
} else {
if(!usedProxies.contains(proxy)) {
log.error("Proxy could not be found to release: " + proxy);
} else {
log.info("Releasing proxy for reuse: " + proxy);
// Remove the proxy from our list of proxies in use since this proxy is no longer in use
usedProxies.remove(proxy);
// If this proxy was recording HAR content, we need to destroy the proxy as we cannot reuse it
// due to the reason that there is currently no way to stop recording HAR content
if(proxy.isHarEnabled()) {
log.warn(String.format("Proxy %s had HAR recording enabled. Deleting.", proxy.getPort()));
this.deleteProxy(proxy);
} else {
// If HAR recording was NOT enabled, we can add the proxy back to the list of available proxies so we can hand it out again
synchronized (availableProxies) {
availableProxies.add(proxy);
}
}
}
}
}
/**
* Adds the specified headers to the proxy.
* @param proxy The proxy to add headers to
* @param headers The headers you'd like to add
*/
@Override
public void addHeaders(IAutomationProxy proxy, Map<String,String> headers) {
HttpPost request = new HttpPost(getHttpUrlWithPort() + "/proxy/" + proxy.getPort() + "/headers");
request.setHeader("Content-Type", "application/json");
request.setHeader("Accept", "application/json");
ObjectMapper mapper = new ObjectMapper();
String headerString;
try {
headerString = mapper.writeValueAsString(headers);
} catch (IOException e) {
throw new ProxyManagerException("Error converting headers into JSON");
}
request.setEntity(new StringEntity(headerString, StandardCharsets.UTF_8));
try {
client.execute(request, response -> {
// Slight conflict between existing behaviour (return based on status code), and the expectations of ResponseHandler
// (throw HttpResponseException on any undesirable statuses). Going with the former for now.
if (response.getStatusLine().getStatusCode() != 200) {
throw new ProxyManagerException("Error response received");
}
return null; // Dummy return to satisfy interface
});
} catch (ClientProtocolException e) {
// Error in HTTP protocol
log.error("Proxy error adding headers: " + Arrays.toString(e.getStackTrace()));
throw new ProxyManagerException("HTTP error adding headers.", e);
} catch (IOException e) {
// Miscellaneous connection issues, or failure to read from stream
log.error("Proxy error adding headers: " + Arrays.toString(e.getStackTrace()));
throw new ProxyManagerException("Connection error adding headers.", e);
}
}
/**
* Creates a new proxy
* @return
*/
private synchronized IAutomationProxy createNewProxy(final Integer port) throws ProxyManagerException {
HttpPost request = new HttpPost(getHttpUrlWithPort() + "/proxy");
if (port != null) {
// If a port is specified, pass it along in the request so it can be requested
request.setEntity(new StringEntity("port=" + port, StandardCharsets.UTF_8));
}
try {
return client.execute(request, response -> {
String contents = EntityUtils.toString(response.getEntity());
Class<? extends AutomationProxy> proxyClass = (port == null) ? AutomationProxy.class : AssignedAutomationProxy.class;
AutomationProxy proxy = JSONUtils.getMappedJSONObject(proxyClass, contents);
// Set the base URL. The proxy will already be set when the object was de-serialized from the response
proxy.setProxyManager(this);
return proxy;
});
} catch (ClientProtocolException e) {
// Error in HTTP protocol
log.error("Proxy error creating new proxy: " + AutomationUtils.getStackTrace(e));
throw new ProxyManagerException("HTTP error creating new proxy.", e);
} catch (IOException e) {
// Miscellaneous connection issues, or failure to read from stream
log.error("Proxy error creating new proxy: " + AutomationUtils.getStackTrace(e));
throw new ProxyManagerException("Connection error creating new proxy.", e);
}
}
/**
* Deletes the specified proxy from the proxy host
* @param proxy The proxy to delete
* @return Whether or not the deletion was successful
*/
// Make this synchronized so ensure we do not have more than 1 delete request going into the proxy server at any given time
private synchronized boolean deleteProxy(IAutomationProxy proxy) {
HttpDelete request = new HttpDelete(getHttpUrlWithPort() + "/proxy/" + proxy.getPort());
try {
return client.execute(request, response -> {
// Slight conflict between existing behaviour (return based on status code), and the expectations of ResponseHandler
// (throw HttpResponseException on any undesirable statuses). Going with the former for now.
return response.getStatusLine().getStatusCode() == 200;
});
} catch (IOException e) {
// Miscellaneous connection issues, or failure to read from stream
log.error("Proxy error deleting proxy: " + AutomationUtils.getStackTrace(e));
}
return false;
}
/**
* Enables HAR recording for the specified proxy
* @param proxy The proxy you wish to start recording the HAR
*/
@Override
public void enableHar(final IAutomationProxy proxy) throws ProxyManagerException {
HttpPut request = new HttpPut(getHttpUrlWithPort() + String.format("/proxy/%s/har", proxy.getPort()));
try {
// There's no "response consumer" per se. We can only create a generic ResponseHandler that returns nothing of value.
client.execute(request, response -> {
int statusCode = response.getStatusLine().getStatusCode();
if (statusCode != 200 && statusCode != 204) {
String contents = EntityUtils.toString(response.getEntity());
log.error(String.format("Error creating HAR with code %s: %s", statusCode, contents));
} else {
// If HAR recording was enabled, mark it on the proxy so we can turn it off later
proxy.markHarRecordingEnabled(true);
}
return null; // Dummy return to satisfy interface
});
} catch (ClientProtocolException e) {
// Error in HTTP protocol
log.error("Proxy error enabling HAR: " + Arrays.toString(e.getStackTrace()));
throw new ProxyManagerException("HTTP error encountered enabling HAR", e);
} catch (IOException e) {
// Miscellaneous connection issues, or failure to read from stream
log.error("Proxy error enabling HAR: " + Arrays.toString(e.getStackTrace()));
throw new ProxyManagerException("Error encountered enabling HAR", e);
}
}
/**
* Returns the HAR for the specified proxy
* @param proxy The proxy to get the HAR from
* @return The HarLog itself
*/
@Override
public synchronized HarLog getHar(IAutomationProxy proxy) {
HttpGet request = new HttpGet(getHttpUrlWithPort() + String.format("/proxy/%s/har", proxy.getPort()));
try {
return client.execute(request, response -> {
int statusCode = response.getStatusLine().getStatusCode();
String contents = EntityUtils.toString(response.getEntity());
if (statusCode != 200) {
// Slight conflict between existing behaviour (log and return null, as below), and the expectations of ResponseHandler
// (throw HttpResponseException on any undesirable statuses). Going with the former for now.
log.error(String.format("Error getting HAR with code %s: %s", statusCode, contents));
return null;
} else {
Har harLog1 = JSONUtils.getMappedJSONObject(Har.class, contents);
return (harLog1 == null) ? null : harLog1.getLog();
}
});
} catch (IOException e) {
// Miscellaneous connection issues, or failure to read from stream
log.error("Proxy error retrieving HAR: " + Arrays.toString(e.getStackTrace()));
}
return null;
}
/**
* Cleans up all proxies in use
*/
public void cleanUp() {
synchronized(availableProxies) {
if(availableProxies.size() > 0) {
Iterator<IAutomationProxy> iterator = availableProxies.iterator();
while(iterator.hasNext()) {
IAutomationProxy proxy = iterator.next();
this.deleteProxy(proxy);
// Remove this proxy so we can't assign it out until its been returned as unused
iterator.remove();
}
}
}
if(usedProxies.size() > 0) {
for(IAutomationProxy proxy : usedProxies) {
log.warn("Proxy still in use after cleanup called. Proxy: " + proxy);
}
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment