Created
April 10, 2018 19:57
Star
You must be signed in to star a gist
ProxyManager class, for managing your BrowserMob proxy, manipulating HAR files, etc
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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