Skip to content

Commit

Permalink
Performance changes
Browse files Browse the repository at this point in the history
  • Loading branch information
czolg committed Jun 28, 2019
1 parent 9d88178 commit 501c69a
Show file tree
Hide file tree
Showing 5 changed files with 67 additions and 79 deletions.
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
.gradle
.idea/*
.idea
out
build
__pycache__
16 changes: 7 additions & 9 deletions src/main/java/io/appium/mitmproxy/InterceptedMessage.java
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,10 @@ public class InterceptedMessage {
public int responseCode;
public List<String[]> responseHeaders;

private final static ObjectMapper objectMapper = new ObjectMapper();

public InterceptedMessage(ByteBuffer buffer) throws IOException {
buffer = buffer.order(ByteOrder.LITTLE_ENDIAN);
buffer.order(ByteOrder.LITTLE_ENDIAN);

int metadataSize = buffer.getInt();
int request_content_size = buffer.getInt();
Expand All @@ -43,8 +44,7 @@ public InterceptedMessage(ByteBuffer buffer) throws IOException {
responseBody = new byte[response_content_size];
buffer.get(responseBody);

ObjectMapper mapper = new ObjectMapper();
JsonNode metadata = mapper.readTree(metadataBytes);
JsonNode metadata = objectMapper.readTree(metadataBytes);
requestMethod = metadata.get("request").get("method").asText();
requestURL = new URL(metadata.get("request").get("url").asText());
JsonNode headers = metadata.get("request").get("headers");
Expand Down Expand Up @@ -72,13 +72,12 @@ public ByteBuffer serializedResponseToMitmproxy() throws JsonProcessingException

// create JSON for metadata. Which is the responseCode and responseHeaders.
// while we're at it, set the Content-Length header
ObjectMapper mapper = new ObjectMapper();
ObjectNode metadataRoot = mapper.createObjectNode();
ObjectNode metadataRoot = objectMapper.createObjectNode();
metadataRoot.put("status_code", responseCode);

ArrayNode headersNode = mapper.createArrayNode();
ArrayNode headersNode = objectMapper.createArrayNode();
List<ArrayNode> headerNodes = responseHeaders.stream().map((h) -> {
ArrayNode headerPair = mapper.createArrayNode();
ArrayNode headerPair = objectMapper.createArrayNode();
headerPair.add(h[0]);
if (h[0].equals("content-length")) {
headerPair.add(Integer.toString(contentLength));
Expand All @@ -90,11 +89,10 @@ public ByteBuffer serializedResponseToMitmproxy() throws JsonProcessingException
headersNode.addAll(headerNodes);
metadataRoot.set("headers", headersNode);

String metadataJson = mapper.writeValueAsString(metadataRoot);
String metadataJson = objectMapper.writeValueAsString(metadataRoot);
byte[] metadata = metadataJson.getBytes(StandardCharsets.UTF_8);
int metadataLength = metadata.length;


ByteBuffer buffer = ByteBuffer.allocate(8 + metadataLength + contentLength);
buffer.order(ByteOrder.LITTLE_ENDIAN);
buffer.putInt(metadataLength);
Expand Down
71 changes: 29 additions & 42 deletions src/main/java/io/appium/mitmproxy/MitmproxyJava.java
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package io.appium.mitmproxy;

import org.apache.commons.io.IOUtils;
import org.zeroturnaround.exec.ProcessExecutor;
import org.zeroturnaround.exec.ProcessResult;
import org.zeroturnaround.exec.stream.slf4j.Slf4jStream;
Expand All @@ -9,27 +10,36 @@
import java.io.IOException;
import java.io.InputStream;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.URISyntaxException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeoutException;
import java.util.function.Function;

public class MitmproxyJava {

private String mitmproxyPath;
private Function<InterceptedMessage, InterceptedMessage> messageInterceptor;
private int proxyPort;
private MitmproxyServer server;
private Future<ProcessResult> mitmproxyProcess;
public static final int WEBSOCKET_PORT = 8765;

public MitmproxyJava(String mitmproxyPath, Function<InterceptedMessage, InterceptedMessage> messageInterceptor) {
public MitmproxyJava(String mitmproxyPath, Function<InterceptedMessage, InterceptedMessage> messageInterceptor, int proxyPort) {
this.mitmproxyPath = mitmproxyPath;
server = new MitmproxyServer(new InetSocketAddress("localhost", WEBSOCKET_PORT), messageInterceptor);
server.start();
this.messageInterceptor = messageInterceptor;
this.proxyPort = proxyPort;
}

public MitmproxyJava(String mitmproxyPath, Function<InterceptedMessage, InterceptedMessage> messageInterceptor) {
this(mitmproxyPath, messageInterceptor, 8080);
}

public void start() throws IOException, TimeoutException, URISyntaxException {
System.out.println("starting mitmproxy on port 8080");
public void start() throws IOException, TimeoutException {
System.out.println("starting mitmproxy on port " + proxyPort);

server = new MitmproxyServer(new InetSocketAddress("localhost", WEBSOCKET_PORT), messageInterceptor);
server.start();

// python script file is zipped inside our jar. extract it into a temporary file.
String pythonScriptPath = extractPythonScriptToFile();
Expand All @@ -41,35 +51,25 @@ public void start() throws IOException, TimeoutException, URISyntaxException {
.start()
.getFuture();

waitForPortToBeInUse(8080);
System.out.println("mitmproxy started on port 8080");
waitForPortToBeInUse(proxyPort);
System.out.println("mitmproxy started on port " + proxyPort);

}

private String extractPythonScriptToFile() throws URISyntaxException, IOException {
private String extractPythonScriptToFile() throws IOException {
File outfile = File.createTempFile("mitmproxy-python-plugin", ".py");

InputStream instream = getClass().getClassLoader().getResourceAsStream("scripts/proxy.py");
FileOutputStream outstream = new FileOutputStream(outfile);

byte[] buffer = new byte[1024];
try (
InputStream inputStream = getClass().getClassLoader().getResourceAsStream("scripts/proxy.py");
FileOutputStream outputStream = new FileOutputStream(outfile)) {

int length;
/*copying the contents from input stream to
* output stream using read and write methods
*/
while ((length = instream.read(buffer)) > 0){
outstream.write(buffer, 0, length);
IOUtils.copy(inputStream, outputStream);
}

//Closing the input/output file streams
instream.close();
outstream.close();

return outfile.getCanonicalPath();
}

public void stop() throws IOException, InterruptedException {
public void stop() throws InterruptedException {
if (mitmproxyProcess != null) {
mitmproxyProcess.cancel(true);
}
Expand All @@ -79,30 +79,17 @@ public void stop() throws IOException, InterruptedException {

private void waitForPortToBeInUse(int port) throws TimeoutException {
boolean inUse = false;
Socket s = null;
int tries = 0;
int maxTries = 60 * 1000 / 100;

while (!inUse) {
try
{
s = new Socket("localhost", port);
}
catch (IOException e)
{

try (Socket s = new Socket("localhost", port)) {
break;
} catch (IOException e) {
inUse = false;
}
finally
{
if(s != null) {
inUse = true;
try {
s.close();
} catch (Exception e) {
}
break;
}
}

tries++;
if (tries == maxTries) {
throw new TimeoutException("Timed out waiting for mitmproxy to start");
Expand Down
29 changes: 14 additions & 15 deletions src/main/java/io/appium/mitmproxy/MitmproxyServer.java
Original file line number Diff line number Diff line change
Expand Up @@ -31,40 +31,39 @@ public void onClose(WebSocket conn, int code, String reason, boolean remote) {

@Override
public void onMessage(WebSocket conn, String message) {
System.out.println("received message from " + conn.getRemoteSocketAddress() + ": " + message);
System.out.println("received message from " + conn.getRemoteSocketAddress() + ": " + message);
}

@Override
public void onMessage( WebSocket conn, ByteBuffer message ) {
public void onMessage(WebSocket conn, ByteBuffer rawInputMessage) {
InterceptedMessage intercepted = null;
InterceptedMessage modifiedMessage = null;

try {
intercepted = new InterceptedMessage(message);
intercepted = new InterceptedMessage(rawInputMessage);
} catch (IOException e) {
System.out.println("Could not parse message");
System.out.println("Could not parse rawInputMessage");
e.printStackTrace();
}

modifiedMessage = interceptor.apply(intercepted);
InterceptedMessage modifiedMessage = interceptor.apply(intercepted);

// if the supplied interceptor function does not return a message, assume no changes were intended and just
// complete the request
if (modifiedMessage == null) {
modifiedMessage = intercepted;
}

try {
conn.send(modifiedMessage.serializedResponseToMitmproxy());
} catch (JsonProcessingException e) {
System.out.println("Could not encode response to mitmproxy");
e.printStackTrace();
conn.send(rawInputMessage);
} else {
try {
conn.send(modifiedMessage.serializedResponseToMitmproxy());
} catch (JsonProcessingException e) {
System.out.println("Could not encode response to mitmproxy");
e.printStackTrace();
}
}
}

@Override
public void onError(WebSocket conn, Exception ex) {
System.err.println("an error occured on connection " + conn.getRemoteSocketAddress() + ":" + ex);
System.err.println("an error occured on connection " + conn.getRemoteSocketAddress() + ":" + ex);
}

@Override
Expand Down
28 changes: 16 additions & 12 deletions src/test/java/io/appium/mitmproxy/MitmproxyJavaTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -11,28 +11,32 @@
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeoutException;

import static junit.framework.TestCase.assertTrue;

public class MitmproxyJavaTest {

// private static final String MITMDUMP_PATH = "/usr/local/bin/mitmdump";
private static final String MITMDUMP_PATH = "C:\\Python37\\Scripts\\mitmdump";


@Test
public void ConstructorTest() throws URISyntaxException, IOException, InterruptedException {
MitmproxyJava proxy = new MitmproxyJava("/usr/local/bin/mitmdump", (InterceptedMessage m) -> {
public void ConstructorTest() throws InterruptedException, IOException, TimeoutException {
MitmproxyJava proxy = new MitmproxyJava(MITMDUMP_PATH, (InterceptedMessage m) -> {
System.out.println(m.requestURL.toString());
return m;
});
proxy.start();
System.out.println("advanced in test");
proxy.stop();
}

@Test
public void SimpleTest() throws InterruptedException, ExecutionException, TimeoutException, IOException, URISyntaxException, UnirestException {
List<InterceptedMessage> messages = new ArrayList<InterceptedMessage>();
public void SimpleTest() throws InterruptedException, TimeoutException, IOException, URISyntaxException, UnirestException {
List<InterceptedMessage> messages = new ArrayList<>();

MitmproxyJava proxy = new MitmproxyJava("/usr/local/bin/mitmdump", (InterceptedMessage m) -> {
MitmproxyJava proxy = new MitmproxyJava(MITMDUMP_PATH, (InterceptedMessage m) -> {
messages.add(m);
return m;
});
Expand All @@ -51,10 +55,10 @@ public void SimpleTest() throws InterruptedException, ExecutionException, Timeou
}

@Test
public void NullInterceptorReturnTest() throws InterruptedException, ExecutionException, TimeoutException, IOException, URISyntaxException, UnirestException {
List<InterceptedMessage> messages = new ArrayList<InterceptedMessage>();
public void NullInterceptorReturnTest() throws InterruptedException, TimeoutException, IOException, URISyntaxException, UnirestException {
List<InterceptedMessage> messages = new ArrayList<>();

MitmproxyJava proxy = new MitmproxyJava("/usr/local/bin/mitmdump", (InterceptedMessage m) -> {
MitmproxyJava proxy = new MitmproxyJava(MITMDUMP_PATH, (InterceptedMessage m) -> {
messages.add(m);
return null;
});
Expand All @@ -73,10 +77,10 @@ public void NullInterceptorReturnTest() throws InterruptedException, ExecutionEx
}

@Test
public void ResponseModificationTest() throws InterruptedException, ExecutionException, TimeoutException, IOException, URISyntaxException, UnirestException {
List<InterceptedMessage> messages = new ArrayList<InterceptedMessage>();
public void ResponseModificationTest() throws InterruptedException, TimeoutException, IOException, URISyntaxException, UnirestException {
List<InterceptedMessage> messages = new ArrayList<>();

MitmproxyJava proxy = new MitmproxyJava("/usr/local/bin/mitmdump", (InterceptedMessage m) -> {
MitmproxyJava proxy = new MitmproxyJava(MITMDUMP_PATH, (InterceptedMessage m) -> {
messages.add(m);
m.responseBody = "Hi from Test".getBytes(StandardCharsets.UTF_8);
return m;
Expand Down

0 comments on commit 501c69a

Please sign in to comment.