Merge pull request square#1716 from square/jw/kill-halfway

Add socket policy to disconnect mid-response.
Rainmonth · Jun 21, 2015 · 77fb83f · 77fb83f
2 parents cd3a9c3 + 282b418
commit 77fb83f
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 2 deletions.
diff --git a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java
@@ -79,6 +79,7 @@
 import okio.Timeout;
 
 import static com.squareup.okhttp.mockwebserver.SocketPolicy.DISCONNECT_AT_START;
+import static com.squareup.okhttp.mockwebserver.SocketPolicy.DISCONNECT_DURING_RESPONSE_BODY;
 import static com.squareup.okhttp.mockwebserver.SocketPolicy.FAIL_HANDSHAKE;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
@@ -720,7 +721,7 @@ private void writeHttpResponse(Socket socket, BufferedSink sink, MockResponse re
     Buffer body = response.getBody();
     if (body == null) return;
     sleepIfDelayed(response);
-    throttledTransfer(response, socket, body, sink, Long.MAX_VALUE);
+    throttledTransfer(response, socket, body, sink, body.size());
   }
 
   private void sleepIfDelayed(MockResponse response) {
@@ -747,9 +748,18 @@ private void throttledTransfer(MockResponse throttlePolicy, Socket socket, Buffe
     long bytesPerPeriod = throttlePolicy.getThrottleBytesPerPeriod();
     long periodDelayMs = throttlePolicy.getThrottlePeriod(TimeUnit.MILLISECONDS);
 
+    long halfByteCount = byteCount / 2;
+    boolean disconnectHalfway = throttlePolicy.getSocketPolicy() == DISCONNECT_DURING_RESPONSE_BODY;
+
     while (!socket.isClosed()) {
       for (int b = 0; b < bytesPerPeriod; ) {
-        long toRead = Math.min(Math.min(2048, byteCount), bytesPerPeriod - b);
+        // Ensure we do not read past the allotted bytes in this period.
+        long toRead = Math.min(byteCount, bytesPerPeriod - b);
+        // Ensure we do not read past halfway if the policy will kill the connection.
+        if (disconnectHalfway) {
+          toRead = Math.min(toRead, byteCount - halfByteCount);
+        }
+
         long read = source.read(buffer, toRead);
         if (read == -1) return;
 
@@ -758,6 +768,11 @@ private void throttledTransfer(MockResponse throttlePolicy, Socket socket, Buffe
         b += read;
         byteCount -= read;
 
+        if (disconnectHalfway && byteCount == halfByteCount) {
+          socket.close();
+          return;
+        }
+
         if (byteCount == 0) return;
       }
 

diff --git a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/SocketPolicy.java b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/SocketPolicy.java
@@ -50,6 +50,9 @@ public enum SocketPolicy {
    */
   DISCONNECT_AFTER_REQUEST,
 
+  /** Close connection after writing half of the response body (if present). */
+  DISCONNECT_DURING_RESPONSE_BODY,
+
   /** Don't trust the client during the SSL handshake. */
   FAIL_HANDSHAKE,
 

diff --git a/mockwebserver/src/test/java/com/squareup/okhttp/mockwebserver/MockWebServerTest.java b/mockwebserver/src/test/java/com/squareup/okhttp/mockwebserver/MockWebServerTest.java
@@ -255,6 +255,18 @@ public final class MockWebServerTest {
     assertTrue(String.format("Request + Response: %sms", elapsedMillis), elapsedMillis <= 1100);
   }
 
+  @Test public void disconnectHalfway() throws IOException {
+    server.enqueue(new MockResponse()
+        .setBody("ab")
+        .setSocketPolicy(SocketPolicy.DISCONNECT_DURING_RESPONSE_BODY));
+
+    URLConnection connection = server.getUrl("/").openConnection();
+    assertEquals(2, connection.getHeaderFieldLong("Content-Length", -1));
+    InputStream in = connection.getInputStream();
+    assertEquals('a', in.read());
+    assertEquals(-1, in.read());
+  }
+
   private List<String> headersToList(MockResponse response) {
     Headers headers = response.getHeaders();
     int size = headers.size();