From b1091b3f552a7badb7bf90faa68cfcc699845c94 Mon Sep 17 00:00:00 2001
From: Jesse Wilson <jwilson@squareup.com>
Date: Sun, 2 Nov 2014 10:52:52 -0500
Subject: [PATCH] Don't close the SPDY peer too early.

This was occasionally causing tests to fail because the test was
observing the connection close.
---
 .../okhttp/internal/spdy/MockSpdyPeer.java    | 44 +++++++++++++------
 .../internal/spdy/Spdy3ConnectionTest.java    |  3 +-
 2 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/okhttp-tests/src/test/java/com/squareup/okhttp/internal/spdy/MockSpdyPeer.java b/okhttp-tests/src/test/java/com/squareup/okhttp/internal/spdy/MockSpdyPeer.java
index 3dac6d32fe69..630a21881a5e 100644
--- a/okhttp-tests/src/test/java/com/squareup/okhttp/internal/spdy/MockSpdyPeer.java
+++ b/okhttp-tests/src/test/java/com/squareup/okhttp/internal/spdy/MockSpdyPeer.java
@@ -74,7 +74,7 @@ public int frameCount() {
   }
 
   public FrameWriter sendFrame() {
-    outFrames.add(new OutFrame(frameCount++, bytesOut.size(), Integer.MAX_VALUE));
+    outFrames.add(new OutFrame(frameCount++, bytesOut.size(), false));
     return frameWriter;
   }
 
@@ -83,17 +83,27 @@ public FrameWriter sendFrame() {
    * won't be generated naturally.
    */
   public void sendFrame(byte[] frame) throws IOException {
-    outFrames.add(new OutFrame(frameCount++, bytesOut.size(), Integer.MAX_VALUE));
+    outFrames.add(new OutFrame(frameCount++, bytesOut.size(), false));
     bytesOut.write(frame);
   }
 
   /**
-   * Sends a frame, truncated to {@code truncateToLength} bytes. This is only
-   * useful for testing error handling as the truncated frame will be
-   * malformed.
+   * Shortens the last frame from its original length to {@code length}. This
+   * will cause the peer to close the socket as soon as this frame has been
+   * written; otherwise the peer stays open until explicitly closed.
    */
-  public FrameWriter sendTruncatedFrame(int truncateToLength) {
-    outFrames.add(new OutFrame(frameCount++, bytesOut.size(), truncateToLength));
+  public FrameWriter truncateLastFrame(int length) {
+    OutFrame lastFrame = outFrames.remove(outFrames.size() - 1);
+    if (length >= bytesOut.size() - lastFrame.start) throw new IllegalArgumentException();
+
+    // Move everything from bytesOut into a new buffer.
+    Buffer fullBuffer = new Buffer();
+    bytesOut.read(fullBuffer, bytesOut.size());
+
+    // Copy back all but what we're truncating.
+    fullBuffer.read(bytesOut, lastFrame.start + length);
+
+    outFrames.add(new OutFrame(lastFrame.sequence, lastFrame.start, true));
     return frameWriter;
   }
 
@@ -136,18 +146,25 @@ private void readAndWriteFrames() throws IOException {
 
       if (nextOutFrame != null && nextOutFrame.sequence == i) {
         long start = nextOutFrame.start;
-        int truncateToLength = nextOutFrame.truncateToLength;
+        boolean truncated;
         long end;
         if (outFramesIterator.hasNext()) {
           nextOutFrame = outFramesIterator.next();
           end = nextOutFrame.start;
+          truncated = false;
         } else {
           end = outBytes.length;
+          truncated = nextOutFrame.truncated;
         }
 
-        // write a frame
-        int length = (int) Math.min(end - start, truncateToLength);
+        // Write a frame.
+        int length = (int) (end - start);
         out.write(outBytes, (int) start, length);
+
+        // If the last frame was truncated, immediately close the connection.
+        if (truncated) {
+          socket.close();
+        }
       } else {
         // read a frame
         InFrame inFrame = new InFrame(i, reader);
@@ -155,7 +172,6 @@ private void readAndWriteFrames() throws IOException {
         inFrames.add(inFrame);
       }
     }
-    Util.closeQuietly(socket);
   }
 
   public Socket openSocket() throws IOException {
@@ -179,12 +195,12 @@ public Socket openSocket() throws IOException {
   private static class OutFrame {
     private final int sequence;
     private final long start;
-    private final int truncateToLength;
+    private final boolean truncated;
 
-    private OutFrame(int sequence, long start, int truncateToLength) {
+    private OutFrame(int sequence, long start, boolean truncated) {
       this.sequence = sequence;
       this.start = start;
-      this.truncateToLength = truncateToLength;
+      this.truncated = truncated;
     }
   }
 
diff --git a/okhttp-tests/src/test/java/com/squareup/okhttp/internal/spdy/Spdy3ConnectionTest.java b/okhttp-tests/src/test/java/com/squareup/okhttp/internal/spdy/Spdy3ConnectionTest.java
index f8a42ebb294b..640d9dbecb52 100644
--- a/okhttp-tests/src/test/java/com/squareup/okhttp/internal/spdy/Spdy3ConnectionTest.java
+++ b/okhttp-tests/src/test/java/com/squareup/okhttp/internal/spdy/Spdy3ConnectionTest.java
@@ -1109,7 +1109,8 @@ private Buffer data(int byteCount) {
     // write the mocking script
     peer.acceptFrame(); // SYN_STREAM
     peer.sendFrame().synReply(false, 1, headerEntries("a", "android"));
-    peer.sendTruncatedFrame(8 + 100).data(false, 1, data(1024), 1024);
+    peer.sendFrame().data(false, 1, data(1024), 1024);
+    peer.truncateLastFrame(8 + 100);
     peer.play();
 
     // play it back