Track upstream changes to mockwebserver.

This removes the MockStreamResponse class, replacing it with the
previous MockResponse class. The MockStreamResponse class added
too much to the MockWebServer API for too little benefit: the
BaseResponse class wasn't 'public' but it was required to be public
API for callers to call MockWebServer.enqueue(). This approach is
less object-oriented, but also a lot smaller.

Change-Id: I8496fc3c4dcbfaaf6cd5e97f06d2786b49fa51e1

9 files changed, 643 insertions, 383 deletions

diff --git a/src/main/java/com/google/mockwebserver/BaseMockResponse.java b/src/main/java/com/google/mockwebserver/BaseMockResponse.java
deleted file mode 100644
index a12abf5..0000000
--- a/src/main/java/com/google/mockwebserver/BaseMockResponse.java
+++ /dev/null
@@ -1,178 +0,0 @@
-/*
- * Copyright (C) 2013 Google Inc.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.google.mockwebserver;
-
-import static com.google.mockwebserver.MockWebServer.ASCII;
-
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-import java.util.ArrayList;
-import java.util.Iterator;
-import java.util.List;
-
-/**
- * Base scripted response to be replayed by {@link MockWebServer}.
- */
-abstract class BaseMockResponse<T extends BaseMockResponse<T>> {
-    protected static final String CONTENT_LENGTH = "Content-Length";
-
-    private String status = "HTTP/1.1 200 OK";
-    private List<String> headers = new ArrayList<String>();
-    private int bytesPerSecond = Integer.MAX_VALUE;
-    private SocketPolicy socketPolicy = SocketPolicy.KEEP_OPEN;
-
-    protected BaseMockResponse() {
-    }
-
-    @Override
-    protected Object clone() throws CloneNotSupportedException {
-        final BaseMockResponse<?> result = (BaseMockResponse<?>) super.clone();
-        result.headers = new ArrayList<String>(result.headers);
-        return result;
-    }
-
-    /**
-     * Returns the HTTP response line, such as "HTTP/1.1 200 OK".
-     */
-    public String getStatus() {
-        return status;
-    }
-
-    public T setResponseCode(int code) {
-        this.status = "HTTP/1.1 " + code + " OK";
-        return self();
-    }
-
-    public T setStatus(String status) {
-        this.status = status;
-        return self();
-    }
-
-    /**
-     * Returns the HTTP headers, such as "Content-Length: 0".
-     */
-    public List<String> getHeaders() {
-        return headers;
-    }
-
-    public T clearHeaders() {
-        headers.clear();
-        return self();
-    }
-
-    public T addHeader(String header) {
-        headers.add(header);
-        return self();
-    }
-
-    public T addHeader(String name, Object value) {
-        return addHeader(name + ": " + String.valueOf(value));
-    }
-
-    public T setHeader(String name, Object value) {
-        removeHeader(name);
-        return addHeader(name, value);
-    }
-
-    public T removeHeader(String name) {
-        name += ": ";
-        for (Iterator<String> i = headers.iterator(); i.hasNext();) {
-            String header = i.next();
-            if (name.regionMatches(true, 0, header, 0, name.length())) {
-                i.remove();
-            }
-        }
-        return self();
-    }
-
-    public SocketPolicy getSocketPolicy() {
-        return socketPolicy;
-    }
-
-    public T setSocketPolicy(SocketPolicy socketPolicy) {
-        this.socketPolicy = socketPolicy;
-        return self();
-    }
-
-    public int getBytesPerSecond() {
-        return bytesPerSecond;
-    }
-
-    /**
-     * Set simulated network speed, in bytes per second.
-     */
-    public T setBytesPerSecond(int bytesPerSecond) {
-        this.bytesPerSecond = bytesPerSecond;
-        return self();
-    }
-
-    @Override public String toString() {
-        return status;
-    }
-
-    /**
-     * Write complete response, including all headers and the given body.
-     * Handles applying {@link #setBytesPerSecond(int)} limits.
-     */
-    protected void writeResponse(InputStream body, OutputStream out) throws IOException {
-        out.write((getStatus() + "\r\n").getBytes(ASCII));
-        for (String header : getHeaders()) {
-            out.write((header + "\r\n").getBytes(ASCII));
-        }
-        out.write(("\r\n").getBytes(ASCII));
-        out.flush();
-
-        // Stream data in MTU-sized increments
-        final byte[] buffer = new byte[1452];
-        final long delayMs;
-        if (bytesPerSecond == Integer.MAX_VALUE) {
-            delayMs = 0;
-        } else {
-            delayMs = (1000 * buffer.length) / bytesPerSecond;
-        }
-
-        int read;
-        long sinceDelay = 0;
-        while ((read = body.read(buffer)) != -1) {
-            out.write(buffer, 0, read);
-            out.flush();
-
-            sinceDelay += read;
-            if (sinceDelay >= buffer.length && delayMs > 0) {
-                sinceDelay %= buffer.length;
-                try {
-                    Thread.sleep(delayMs);
-                } catch (InterruptedException e) {
-                    throw new AssertionError();
-                }
-            }
-        }
-    }
-
-    /**
-     * Write complete response. Usually implemented by calling
-     * {@link #writeResponse(InputStream, OutputStream)} with the
-     * implementation-specific body.
-     */
-    public abstract void writeResponse(OutputStream out) throws IOException;
-
-    /**
-     * Return concrete {@code this} to enable builder-style methods.
-     */
-    protected abstract T self();
-}
diff --git a/src/main/java/com/google/mockwebserver/Dispatcher.java b/src/main/java/com/google/mockwebserver/Dispatcher.java
new file mode 100644
index 0000000..0456025
--- /dev/null
+++ b/src/main/java/com/google/mockwebserver/Dispatcher.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright (C) 2012 Google Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.mockwebserver;
+
+/**
+ * Handler for mock server requests.
+ */
+public abstract class Dispatcher {
+    /**
+     * Returns a response to satisfy {@code request}. This method may block (for
+     * instance, to wait on a CountdownLatch).
+     */
+    public abstract MockResponse dispatch(RecordedRequest request) throws InterruptedException;
+
+    /**
+     * Returns the socket policy of the next request.  Default implementation
+     * returns {@link SocketPolicy#KEEP_OPEN}. Mischievous implementations can
+     * return other values to test HTTP edge cases.
+     */
+    public SocketPolicy peekSocketPolicy() {
+        return SocketPolicy.KEEP_OPEN;
+    }
+}
diff --git a/src/main/java/com/google/mockwebserver/MockResponse.java b/src/main/java/com/google/mockwebserver/MockResponse.java
index 8f11996..674a911 100644
--- a/src/main/java/com/google/mockwebserver/MockResponse.java
+++ b/src/main/java/com/google/mockwebserver/MockResponse.java
@@ -17,84 +17,224 @@
 package com.google.mockwebserver;
 
 import static com.google.mockwebserver.MockWebServer.ASCII;
-
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.OutputStream;
+import java.io.InputStream;
 import java.io.UnsupportedEncodingException;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
 
 /**
- * A scripted response to be replayed by {@link MockWebServer}.
+ * A scripted response to be replayed by the mock web server.
  */
-public class MockResponse extends BaseMockResponse<MockResponse> implements Cloneable {
+public final class MockResponse implements Cloneable {
     private static final String CHUNKED_BODY_HEADER = "Transfer-encoding: chunked";
 
+    private String status = "HTTP/1.1 200 OK";
+    private List<String> headers = new ArrayList<String>();
+
+    /** The response body content, or null if {@code bodyStream} is set. */
     private byte[] body;
+    /** The response body content, or null if {@code body} is set. */
+    private InputStream bodyStream;
+
+    private int bytesPerSecond = Integer.MAX_VALUE;
+    private SocketPolicy socketPolicy = SocketPolicy.KEEP_OPEN;
 
+    /**
+     * Creates a new mock response with an empty body.
+     */
     public MockResponse() {
-        this.body = new byte[0];
-        addHeader(CONTENT_LENGTH, 0);
+        setBody(new byte[0]);
     }
 
-    @Override
-    public MockResponse clone() {
+    @Override public MockResponse clone() {
         try {
-            return (MockResponse) super.clone();
+            MockResponse result = (MockResponse) super.clone();
+            result.headers = new ArrayList<String>(result.headers);
+            return result;
         } catch (CloneNotSupportedException e) {
             throw new AssertionError();
         }
     }
 
+    /**
+     * Returns the HTTP response line, such as "HTTP/1.1 200 OK".
+     */
+    public String getStatus() {
+        return status;
+    }
+
+    public MockResponse setResponseCode(int code) {
+        this.status = "HTTP/1.1 " + code + " OK";
+        return this;
+    }
+
+    public MockResponse setStatus(String status) {
+        this.status = status;
+        return this;
+    }
+
+    /**
+     * Returns the HTTP headers, such as "Content-Length: 0".
+     */
+    public List<String> getHeaders() {
+        return headers;
+    }
+
+    /**
+     * Removes all HTTP headers including any "Content-Length" and
+     * "Transfer-encoding" headers that were added by default.
+     */
+    public MockResponse clearHeaders() {
+        headers.clear();
+        return this;
+    }
+
+    /**
+     * Adds {@code header} as an HTTP header. For well-formed HTTP {@code
+     * header} should contain a name followed by a colon and a value.
+     */
+    public MockResponse addHeader(String header) {
+        headers.add(header);
+        return this;
+    }
+
+    /**
+     * Adds a new header with the name and value. This may be used to add
+     * multiple headers with the same name.
+     */
+    public MockResponse addHeader(String name, Object value) {
+        return addHeader(name + ": " + String.valueOf(value));
+    }
+
+    /**
+     * Removes all headers named {@code name}, then adds a new header with the
+     * name and value.
+     */
+    public MockResponse setHeader(String name, Object value) {
+        removeHeader(name);
+        return addHeader(name, value);
+    }
+
+    /**
+     * Removes all headers named {@code name}.
+     */
+    public MockResponse removeHeader(String name) {
+        name += ":";
+        for (Iterator<String> i = headers.iterator(); i.hasNext(); ) {
+            String header = i.next();
+            if (name.regionMatches(true, 0, header, 0, name.length())) {
+                i.remove();
+            }
+        }
+        return this;
+    }
+
+    /**
+     * Returns the raw HTTP payload, or null if this response is streamed.
+     */
+    public byte[] getBody() {
+        return body;
+    }
+
+    /**
+     * Returns an input stream containing the raw HTTP payload.
+     */
+    InputStream getBodyStream() {
+        return bodyStream != null ? bodyStream : new ByteArrayInputStream(body);
+    }
+
     public MockResponse setBody(byte[] body) {
+        setHeader("Content-Length", body.length);
         this.body = body;
-        setHeader(CONTENT_LENGTH, body.length);
+        this.bodyStream = null;
+        return this;
+    }
+
+    public MockResponse setBody(InputStream bodyStream, long bodyLength) {
+        setHeader("Content-Length", bodyLength);
+        this.body = null;
+        this.bodyStream = bodyStream;
         return this;
     }
 
+    /**
+     * Sets the response body to the UTF-8 encoded bytes of {@code body}.
+     */
     public MockResponse setBody(String body) {
         try {
-            return setBody(body.getBytes(ASCII));
+            return setBody(body.getBytes("UTF-8"));
         } catch (UnsupportedEncodingException e) {
             throw new AssertionError();
         }
     }
 
-    public byte[] getBody() {
-        return body;
-    }
+    /**
+     * Sets the response body to {@code body}, chunked every {@code
+     * maxChunkSize} bytes.
+     */
+    public MockResponse setChunkedBody(byte[] body, int maxChunkSize) {
+        removeHeader("Content-Length");
+        headers.add(CHUNKED_BODY_HEADER);
 
-    public MockResponse setChunkedBody(byte[] body, int maxChunkSize) throws IOException {
-        removeHeader(CONTENT_LENGTH);
-        addHeader(CHUNKED_BODY_HEADER);
+        try {
+            ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
+            int pos = 0;
+            while (pos < body.length) {
+                int chunkSize = Math.min(body.length - pos, maxChunkSize);
+                bytesOut.write(Integer.toHexString(chunkSize).getBytes(ASCII));
+                bytesOut.write("\r\n".getBytes(ASCII));
+                bytesOut.write(body, pos, chunkSize);
+                bytesOut.write("\r\n".getBytes(ASCII));
+                pos += chunkSize;
+            }
+            bytesOut.write("0\r\n\r\n".getBytes(ASCII)); // last chunk + empty trailer + crlf
+
+            this.body = bytesOut.toByteArray();
+            return this;
+        } catch (IOException e) {
+            throw new AssertionError(); // In-memory I/O doesn't throw IOExceptions.
+        }
+    }
 
-        ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
-        int pos = 0;
-        while (pos < body.length) {
-            int chunkSize = Math.min(body.length - pos, maxChunkSize);
-            bytesOut.write(Integer.toHexString(chunkSize).getBytes(ASCII));
-            bytesOut.write("\r\n".getBytes(ASCII));
-            bytesOut.write(body, pos, chunkSize);
-            bytesOut.write("\r\n".getBytes(ASCII));
-            pos += chunkSize;
+    /**
+     * Sets the response body to the UTF-8 encoded bytes of {@code body},
+     * chunked every {@code maxChunkSize} bytes.
+     */
+    public MockResponse setChunkedBody(String body, int maxChunkSize) {
+        try {
+            return setChunkedBody(body.getBytes("UTF-8"), maxChunkSize);
+        } catch (UnsupportedEncodingException e) {
+            throw new AssertionError();
         }
-        bytesOut.write("0\r\n\r\n".getBytes(ASCII)); // last chunk + empty trailer + crlf
+    }
+
+    public SocketPolicy getSocketPolicy() {
+        return socketPolicy;
+    }
 
-        this.body = bytesOut.toByteArray();
+    public MockResponse setSocketPolicy(SocketPolicy socketPolicy) {
+        this.socketPolicy = socketPolicy;
         return this;
     }
 
-    public MockResponse setChunkedBody(String body, int maxChunkSize) throws IOException {
-        return setChunkedBody(body.getBytes(ASCII), maxChunkSize);
+    public int getBytesPerSecond() {
+        return bytesPerSecond;
     }
 
-    @Override
-    protected MockResponse self() {
+    /**
+     * Set simulated network speed, in bytes per second. This applies to the
+     * response body only; response headers are not throttled.
+     */
+    public MockResponse setBytesPerSecond(int bytesPerSecond) {
+        this.bytesPerSecond = bytesPerSecond;
         return this;
     }
 
-    @Override
-    public void writeResponse(OutputStream out) throws IOException {
-        super.writeResponse(new ByteArrayInputStream(body), out);
+    @Override public String toString() {
+        return status;
     }
 }
diff --git a/src/main/java/com/google/mockwebserver/MockStreamResponse.java b/src/main/java/com/google/mockwebserver/MockStreamResponse.java
deleted file mode 100644
index 9db1c3c..0000000
--- a/src/main/java/com/google/mockwebserver/MockStreamResponse.java
+++ /dev/null
@@ -1,78 +0,0 @@
-/*
- * Copyright (C) 2013 Google Inc.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.google.mockwebserver;
-
-import java.io.ByteArrayInputStream;
-import java.io.Closeable;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.OutputStream;
-
-/**
- * A scripted response to be replayed by {@link MockWebServer}. This specific
- * variant uses an {@link InputStream} as its data source. Each instance can
- * only be consumed once.
- */
-public class MockStreamResponse extends BaseMockResponse<MockStreamResponse> {
-    private InputStream body;
-
-    public MockStreamResponse() {
-        body = new ByteArrayInputStream(new byte[0]);
-        addHeader(CONTENT_LENGTH, 0);
-    }
-
-    public MockStreamResponse setBody(InputStream body, long bodyLength) {
-        // Release any existing body
-        if (this.body != null) {
-            closeQuietly(this.body);
-        }
-
-        this.body = body;
-        setHeader(CONTENT_LENGTH, bodyLength);
-        return this;
-    }
-
-    @Override
-    public void writeResponse(OutputStream out) throws IOException {
-        if (body == null) {
-            throw new IllegalStateException("Stream already consumed");
-        }
-
-        try {
-            super.writeResponse(body, out);
-        } finally {
-            closeQuietly(body);
-        }
-        body = null;
-    }
-
-    @Override
-    protected MockStreamResponse self() {
-        return this;
-    }
-
-    private static void closeQuietly(Closeable closeable) {
-        if (closeable != null) {
-            try {
-                closeable.close();
-            } catch (RuntimeException rethrown) {
-                throw rethrown;
-            } catch (Exception ignored) {
-            }
-        }
-    }
-}
diff --git a/src/main/java/com/google/mockwebserver/MockWebServer.java b/src/main/java/com/google/mockwebserver/MockWebServer.java
index b774eff..9a89912 100644
--- a/src/main/java/com/google/mockwebserver/MockWebServer.java
+++ b/src/main/java/com/google/mockwebserver/MockWebServer.java
@@ -18,14 +18,12 @@ package com.google.mockwebserver;
 
 import static com.google.mockwebserver.SocketPolicy.DISCONNECT_AT_START;
 import static com.google.mockwebserver.SocketPolicy.FAIL_HANDSHAKE;
-
 import java.io.BufferedInputStream;
 import java.io.BufferedOutputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
-import java.net.HttpURLConnection;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.MalformedURLException;
@@ -38,20 +36,17 @@ import java.net.UnknownHostException;
 import java.security.cert.CertificateException;
 import java.security.cert.X509Certificate;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
-import java.util.Set;
+import java.util.Map;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.LinkedBlockingDeque;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Level;
 import java.util.logging.Logger;
-
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLSocket;
 import javax.net.ssl.SSLSocketFactory;
@@ -69,17 +64,15 @@ public final class MockWebServer {
     private static final Logger logger = Logger.getLogger(MockWebServer.class.getName());
     private final BlockingQueue<RecordedRequest> requestQueue
             = new LinkedBlockingQueue<RecordedRequest>();
-    private final BlockingQueue<BaseMockResponse<?>> responseQueue
-            = new LinkedBlockingDeque<BaseMockResponse<?>>();
-    private final Set<Socket> openClientSockets
-            = Collections.newSetFromMap(new ConcurrentHashMap<Socket, Boolean>());
-    private boolean singleResponse;
+    /** All map values are Boolean.TRUE. (Collections.newSetFromMap isn't available in Froyo) */
+    private final Map<Socket, Boolean> openClientSockets = new ConcurrentHashMap<Socket, Boolean>();
     private final AtomicInteger requestCount = new AtomicInteger();
     private int bodyLimit = Integer.MAX_VALUE;
     private ServerSocket serverSocket;
     private SSLSocketFactory sslSocketFactory;
     private ExecutorService executor;
     private boolean tunnelProxy;
+    private Dispatcher dispatcher = new QueueDispatcher();
 
     private int port = -1;
 
@@ -107,10 +100,14 @@ public final class MockWebServer {
      *
      * @param path the request path, such as "/".
      */
-    public URL getUrl(String path) throws MalformedURLException, UnknownHostException {
-        return sslSocketFactory != null
-                ? new URL("https://" + getHostName() + ":" + getPort() + path)
-                : new URL("http://" + getHostName() + ":" + getPort() + path);
+    public URL getUrl(String path) {
+        try {
+            return sslSocketFactory != null
+                    ? new URL("https://" + getHostName() + ":" + getPort() + path)
+                    : new URL("http://" + getHostName() + ":" + getPort() + path);
+        } catch (MalformedURLException e) {
+            throw new AssertionError(e);
+        }
     }
 
     /**
@@ -159,25 +156,16 @@ public final class MockWebServer {
         return requestCount.get();
     }
 
-    public void enqueue(BaseMockResponse<?> response) {
-        if (response instanceof MockResponse) {
-            responseQueue.add(((MockResponse) response).clone());
-        } else {
-            responseQueue.add(response);
-        }
-    }
-
     /**
-     * By default, this class processes requests coming in by adding them to a
-     * queue and serves responses by removing them from another queue. This mode
-     * is appropriate for correctness testing.
+     * Scripts {@code response} to be returned to a request made in sequence.
+     * The first request is served by the first enqueued response; the second
+     * request by the second enqueued response; and so on.
      *
-     * <p>Serving a single response causes the server to be stateless: requests
-     * are not enqueued, and responses are not dequeued. This mode is appropriate
-     * for benchmarking.
+     * @throws ClassCastException if the default dispatcher has been replaced
+     *     with {@link #setDispatcher(Dispatcher)}.
      */
-    public void setSingleResponse(boolean singleResponse) {
-        this.singleResponse = singleResponse;
+    public void enqueue(MockResponse response) {
+        ((QueueDispatcher) dispatcher).enqueueResponse(response.clone());
     }
 
     /**
@@ -218,7 +206,7 @@ public final class MockWebServer {
                 } catch (Throwable e) {
                     logger.log(Level.WARNING, "MockWebServer server socket close failed", e);
                 }
-                for (Iterator<Socket> s = openClientSockets.iterator(); s.hasNext();) {
+                for (Iterator<Socket> s = openClientSockets.keySet().iterator(); s.hasNext(); ) {
                     try {
                         s.next().close();
                         s.remove();
@@ -234,22 +222,22 @@ public final class MockWebServer {
             }
 
             private void acceptConnections() throws Exception {
-                do {
+                while (true) {
                     Socket socket;
                     try {
                         socket = serverSocket.accept();
-                    } catch (SocketException ignored) {
-                        continue;
+                    } catch (SocketException e) {
+                        return;
                     }
-                    BaseMockResponse<?> peek = responseQueue.peek();
-                    if (peek != null && peek.getSocketPolicy() == DISCONNECT_AT_START) {
-                        responseQueue.take();
+                    final SocketPolicy socketPolicy = dispatcher.peekSocketPolicy();
+                    if (socketPolicy == DISCONNECT_AT_START) {
+                        dispatchBookkeepingRequest(0, socket);
                         socket.close();
                     } else {
-                        openClientSockets.add(socket);
+                        openClientSockets.put(socket, true);
                         serveConnection(socket);
                     }
-                } while (!responseQueue.isEmpty());
+                }
             }
         }));
     }
@@ -279,15 +267,16 @@ public final class MockWebServer {
                     if (tunnelProxy) {
                         createTunnel();
                     }
-                    BaseMockResponse<?> response = responseQueue.peek();
-                    if (response != null && response.getSocketPolicy() == FAIL_HANDSHAKE) {
+                    final SocketPolicy socketPolicy = dispatcher.peekSocketPolicy();
+                    if (socketPolicy == FAIL_HANDSHAKE) {
+                        dispatchBookkeepingRequest(sequenceNumber, raw);
                         processHandshakeFailure(raw, sequenceNumber++);
                         return;
                     }
                     socket = sslSocketFactory.createSocket(
                             raw, raw.getInetAddress().getHostAddress(), raw.getPort(), true);
                     ((SSLSocket) socket).setUseClientMode(false);
-                    openClientSockets.add(socket);
+                    openClientSockets.put(socket, true);
                     openClientSockets.remove(raw);
                 } else {
                     socket = raw;
@@ -296,7 +285,8 @@ public final class MockWebServer {
                 InputStream in = new BufferedInputStream(socket.getInputStream());
                 OutputStream out = new BufferedOutputStream(socket.getOutputStream());
 
-                while (!responseQueue.isEmpty() && processOneRequest(socket, in, out)) {}
+                while (processOneRequest(socket, in, out)) {
+                }
 
                 if (sequenceNumber == 0) {
                     logger.warning("MockWebServer connection didn't make a request");
@@ -305,9 +295,6 @@ public final class MockWebServer {
                 in.close();
                 out.close();
                 socket.close();
-                if (responseQueue.isEmpty()) {
-                    shutdown();
-                }
                 openClientSockets.remove(socket);
             }
 
@@ -317,11 +304,11 @@ public final class MockWebServer {
              */
             private void createTunnel() throws IOException, InterruptedException {
                 while (true) {
-                    BaseMockResponse<?> connect = responseQueue.peek();
+                    final SocketPolicy socketPolicy = dispatcher.peekSocketPolicy();
                     if (!processOneRequest(raw, raw.getInputStream(), raw.getOutputStream())) {
                         throw new IllegalStateException("Tunnel without any CONNECT!");
                     }
-                    if (connect.getSocketPolicy() == SocketPolicy.UPGRADE_TO_SSL_AT_END) {
+                    if (socketPolicy == SocketPolicy.UPGRADE_TO_SSL_AT_END) {
                         return;
                     }
                 }
@@ -337,8 +324,10 @@ public final class MockWebServer {
                 if (request == null) {
                     return false;
                 }
-                BaseMockResponse<?> response = dispatch(request);
-                response.writeResponse(out);
+                requestCount.incrementAndGet();
+                requestQueue.add(request);
+                MockResponse response = dispatcher.dispatch(request);
+                writeResponse(out, response);
                 if (response.getSocketPolicy() == SocketPolicy.DISCONNECT_AT_END) {
                     in.close();
                     out.close();
@@ -355,7 +344,6 @@ public final class MockWebServer {
     }
 
     private void processHandshakeFailure(Socket raw, int sequenceNumber) throws Exception {
-        responseQueue.take();
         X509TrustManager untrusted = new X509TrustManager() {
             @Override public void checkClientTrusted(X509Certificate[] chain, String authType)
                     throws CertificateException {
@@ -379,8 +367,11 @@ public final class MockWebServer {
         } catch (IOException expected) {
         }
         socket.close();
+    }
+
+    private void dispatchBookkeepingRequest(int sequenceNumber, Socket socket) throws InterruptedException {
         requestCount.incrementAndGet();
-        requestQueue.add(new RecordedRequest(null, null, null, -1, null, sequenceNumber, socket));
+        dispatcher.dispatch(new RecordedRequest(null, null, null, -1, null, sequenceNumber, socket));
     }
 
     /**
@@ -394,7 +385,7 @@ public final class MockWebServer {
         } catch (IOException streamIsClosed) {
             return null; // no request because we closed the stream
         }
-        if (request.isEmpty()) {
+        if (request.length() == 0) {
             return null; // no request because the stream is exhausted
         }
 
@@ -402,7 +393,7 @@ public final class MockWebServer {
         int contentLength = -1;
         boolean chunked = false;
         String header;
-        while (!(header = readAsciiUntilCrlf(in)).isEmpty()) {
+        while ((header = readAsciiUntilCrlf(in)).length() != 0) {
             headers.add(header);
             String lowercaseHeader = header.toLowerCase();
             if (contentLength == -1 && lowercaseHeader.startsWith("content-length:")) {
@@ -436,15 +427,11 @@ public final class MockWebServer {
 
         if (request.startsWith("OPTIONS ") || request.startsWith("GET ")
                 || request.startsWith("HEAD ") || request.startsWith("DELETE ")
-                || request .startsWith("TRACE ") || request.startsWith("CONNECT ")) {
+                || request.startsWith("TRACE ") || request.startsWith("CONNECT ")) {
             if (hasBody) {
                 throw new IllegalArgumentException("Request must not have a body: " + request);
             }
-        } else if (request.startsWith("POST ") || request.startsWith("PUT ")) {
-            if (!hasBody) {
-                throw new IllegalArgumentException("Request must have a body: " + request);
-            }
-        } else {
+        } else if (!request.startsWith("POST ") && !request.startsWith("PUT ")) {
             throw new UnsupportedOperationException("Unexpected method: " + request);
         }
 
@@ -452,27 +439,44 @@ public final class MockWebServer {
                 requestBody.numBytesReceived, requestBody.toByteArray(), sequenceNumber, socket);
     }
 
-    /**
-     * Returns a response to satisfy {@code request}.
-     */
-    private BaseMockResponse<?> dispatch(RecordedRequest request) throws InterruptedException {
-        if (responseQueue.isEmpty()) {
-            throw new IllegalStateException("Unexpected request: " + request);
+    private void writeResponse(OutputStream out, MockResponse response) throws IOException {
+        out.write((response.getStatus() + "\r\n").getBytes(ASCII));
+        for (String header : response.getHeaders()) {
+            out.write((header + "\r\n").getBytes(ASCII));
         }
+        out.write(("\r\n").getBytes(ASCII));
+        out.flush();
 
-        // to permit interactive/browser testing, ignore requests for favicons
-        if (request.getRequestLine().equals("GET /favicon.ico HTTP/1.1")) {
-            System.out.println("served " + request.getRequestLine());
-            return new MockResponse()
-                        .setResponseCode(HttpURLConnection.HTTP_NOT_FOUND);
+        final InputStream in = response.getBodyStream();
+        if (in == null) {
+            return;
         }
+        final int bytesPerSecond = response.getBytesPerSecond();
 
-        if (singleResponse) {
-            return responseQueue.peek();
+        // Stream data in MTU-sized increments
+        final byte[] buffer = new byte[1452];
+        final long delayMs;
+        if (bytesPerSecond == Integer.MAX_VALUE) {
+            delayMs = 0;
         } else {
-            requestCount.incrementAndGet();
-            requestQueue.add(request);
-            return responseQueue.take();
+            delayMs = (1000 * buffer.length) / bytesPerSecond;
+        }
+
+        int read;
+        long sinceDelay = 0;
+        while ((read = in.read(buffer)) != -1) {
+            out.write(buffer, 0, read);
+            out.flush();
+
+            sinceDelay += read;
+            if (sinceDelay >= buffer.length && delayMs > 0) {
+                sinceDelay %= buffer.length;
+                try {
+                    Thread.sleep(delayMs);
+                } catch (InterruptedException e) {
+                    throw new AssertionError();
+                }
+            }
         }
     }
 
@@ -513,12 +517,25 @@ public final class MockWebServer {
 
     private void readEmptyLine(InputStream in) throws IOException {
         String line = readAsciiUntilCrlf(in);
-        if (!line.isEmpty()) {
+        if (line.length() != 0) {
             throw new IllegalStateException("Expected empty but was: " + line);
         }
     }
 
     /**
+     * Sets the dispatcher used to match incoming requests to mock responses.
+     * The default dispatcher simply serves a fixed sequence of responses from
+     * a {@link #enqueue(MockResponse) queue}; custom dispatchers can vary the
+     * response based on timing or the content of the request.
+     */
+    public void setDispatcher(Dispatcher dispatcher) {
+        if (dispatcher == null) {
+            throw new NullPointerException();
+        }
+        this.dispatcher = dispatcher;
+    }
+
+    /**
      * An output stream that drops data after bodyLimit bytes.
      */
     private class TruncatingOutputStream extends ByteArrayOutputStream {
diff --git a/src/main/java/com/google/mockwebserver/QueueDispatcher.java b/src/main/java/com/google/mockwebserver/QueueDispatcher.java
new file mode 100644
index 0000000..bc26694
--- /dev/null
+++ b/src/main/java/com/google/mockwebserver/QueueDispatcher.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2012 Google Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.mockwebserver;
+
+import java.net.HttpURLConnection;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+
+/**
+ * Default dispatcher that processes a script of responses.  Populate the script by calling
+ * {@link #enqueueResponse(MockResponse)}.
+ */
+public class QueueDispatcher extends Dispatcher {
+    protected final BlockingQueue<MockResponse> responseQueue
+            = new LinkedBlockingQueue<MockResponse>();
+    private MockResponse failFastResponse;
+
+    @Override public MockResponse dispatch(RecordedRequest request) throws InterruptedException {
+        // to permit interactive/browser testing, ignore requests for favicons
+        final String requestLine = request.getRequestLine();
+        if (requestLine != null && requestLine.equals("GET /favicon.ico HTTP/1.1")) {
+            System.out.println("served " + requestLine);
+            return new MockResponse()
+                    .setResponseCode(HttpURLConnection.HTTP_NOT_FOUND);
+        }
+
+        if (failFastResponse != null && responseQueue.peek() == null) {
+            // Fail fast if there's no response queued up.
+            return failFastResponse;
+        }
+
+        return responseQueue.take();
+    }
+
+    @Override public SocketPolicy peekSocketPolicy() {
+        MockResponse peek = responseQueue.peek();
+        if (peek == null) {
+            return failFastResponse != null
+                    ? failFastResponse.getSocketPolicy()
+                    : SocketPolicy.KEEP_OPEN;
+        }
+        return peek.getSocketPolicy();
+    }
+
+    public void enqueueResponse(MockResponse response) {
+        responseQueue.add(response);
+    }
+
+    public void setFailFast(boolean failFast) {
+        MockResponse failFastResponse = failFast
+                ? new MockResponse().setResponseCode(HttpURLConnection.HTTP_NOT_FOUND)
+                : null;
+        setFailFast(failFastResponse);
+    }
+
+    public void setFailFast(MockResponse failFastResponse) {
+        this.failFastResponse = failFastResponse;
+    }
+}
diff --git a/src/main/java/com/google/mockwebserver/RecordedRequest.java b/src/main/java/com/google/mockwebserver/RecordedRequest.java
index 0728ff3..2864d80 100644
--- a/src/main/java/com/google/mockwebserver/RecordedRequest.java
+++ b/src/main/java/com/google/mockwebserver/RecordedRequest.java
@@ -16,7 +16,9 @@
 
 package com.google.mockwebserver;
 
+import java.io.UnsupportedEncodingException;
 import java.net.Socket;
+import java.util.ArrayList;
 import java.util.List;
 import javax.net.ssl.SSLSocket;
 
@@ -34,7 +36,7 @@ public final class RecordedRequest {
     private final int sequenceNumber;
     private final String sslProtocol;
 
-    RecordedRequest(String requestLine, List<String> headers, List<Integer> chunkSizes,
+    public RecordedRequest(String requestLine, List<String> headers, List<Integer> chunkSizes,
             int bodySize, byte[] body, int sequenceNumber, Socket socket) {
         this.requestLine = requestLine;
         this.headers = headers;
@@ -65,11 +67,50 @@ public final class RecordedRequest {
         return requestLine;
     }
 
+    public String getMethod() {
+        return method;
+    }
+
+    public String getPath() {
+        return path;
+    }
+
+    /**
+     * Returns all headers.
+     */
     public List<String> getHeaders() {
         return headers;
     }
 
     /**
+     * Returns the first header named {@code name}, or null if no such header
+     * exists.
+     */
+    public String getHeader(String name) {
+        name += ":";
+        for (String header : headers) {
+            if (name.regionMatches(true, 0, header, 0, name.length())) {
+                return header.substring(name.length()).trim();
+            }
+        }
+        return null;
+    }
+
+    /**
+     * Returns the headers named {@code name}.
+     */
+    public List<String> getHeaders(String name) {
+        List<String> result = new ArrayList<String>();
+        name += ":";
+        for (String header : headers) {
+            if (name.regionMatches(true, 0, header, 0, name.length())) {
+                result.add(header.substring(name.length()).trim());
+            }
+        }
+        return result;
+    }
+
+    /**
      * Returns the sizes of the chunks of this request's body, or an empty list
      * if the request's body was empty or unchunked.
      */
@@ -93,6 +134,17 @@ public final class RecordedRequest {
     }
 
     /**
+     * Returns the body of this POST request decoded as a UTF-8 string.
+     */
+    public String getUtf8Body() {
+        try {
+            return new String(body, "UTF-8");
+        } catch (UnsupportedEncodingException e) {
+            throw new AssertionError();
+        }
+    }
+
+    /**
      * Returns the index of this request on its HTTP connection. Since a single
      * HTTP connection may serve multiple requests, each request is assigned its
      * own sequence number.
@@ -112,12 +164,4 @@ public final class RecordedRequest {
     @Override public String toString() {
         return requestLine;
     }
-
-    public String getMethod() {
-        return method;
-    }
-
-    public String getPath() {
-        return path;
-    }
-}
+}
\ No newline at end of file
diff --git a/src/test/java/com/google/mockwebserver/CustomDispatcherTest.java b/src/test/java/com/google/mockwebserver/CustomDispatcherTest.java
new file mode 100644
index 0000000..e75b902
--- /dev/null
+++ b/src/test/java/com/google/mockwebserver/CustomDispatcherTest.java
@@ -0,0 +1,100 @@
+/*
+ * Copyright (C) 2012 Google Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.mockwebserver;
+
+import junit.framework.TestCase;
+
+import java.io.IOException;
+import java.net.HttpURLConnection;
+import java.net.URL;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class CustomDispatcherTest extends TestCase {
+
+    private MockWebServer mockWebServer = new MockWebServer();
+
+    @Override
+    public void tearDown() throws Exception {
+        mockWebServer.shutdown();
+    }
+
+    public void testSimpleDispatch() throws Exception {
+        mockWebServer.play();
+        final List<RecordedRequest> requestsMade = new ArrayList<RecordedRequest>();
+        final Dispatcher dispatcher = new Dispatcher() {
+            @Override
+            public MockResponse dispatch(RecordedRequest request) throws InterruptedException {
+                requestsMade.add(request);
+                return new MockResponse();
+            }
+        };
+        assertEquals(0, requestsMade.size());
+        mockWebServer.setDispatcher(dispatcher);
+        final URL url = mockWebServer.getUrl("/");
+        final HttpURLConnection conn = (HttpURLConnection) url.openConnection();
+        conn.getResponseCode(); // Force the connection to hit the "server".
+        // Make sure our dispatcher got the request.
+        assertEquals(1, requestsMade.size());
+    }
+
+    public void testOutOfOrderResponses() throws Exception {
+        AtomicInteger firstResponseCode = new AtomicInteger();
+        AtomicInteger secondResponseCode = new AtomicInteger();
+        mockWebServer.play();
+        final String secondRequest = "/bar";
+        final String firstRequest = "/foo";
+        final CountDownLatch latch = new CountDownLatch(1);
+        final Dispatcher dispatcher = new Dispatcher() {
+            @Override
+            public MockResponse dispatch(RecordedRequest request) throws InterruptedException {
+                if (request.getPath().equals(firstRequest)) {
+                    latch.await();
+                }
+                return new MockResponse();
+            }
+        };
+        mockWebServer.setDispatcher(dispatcher);
+        final Thread startsFirst = buildRequestThread(firstRequest, firstResponseCode);
+        startsFirst.start();
+        final Thread endsFirst = buildRequestThread(secondRequest, secondResponseCode);
+        endsFirst.start();
+        endsFirst.join();
+        assertEquals(0, firstResponseCode.get()); // First response is still waiting.
+        assertEquals(200, secondResponseCode.get()); // Second response is done.
+        latch.countDown();
+        startsFirst.join();
+        assertEquals(200, firstResponseCode.get()); // And now it's done!
+        assertEquals(200, secondResponseCode.get()); // (Still done).
+    }
+
+    private Thread buildRequestThread(final String path, final AtomicInteger responseCode) {
+        return new Thread(new Runnable() {
+            @Override public void run() {
+                final URL url = mockWebServer.getUrl(path);
+                final HttpURLConnection conn;
+                try {
+                    conn = (HttpURLConnection) url.openConnection();
+                    responseCode.set(conn.getResponseCode()); // Force the connection to hit the "server".
+                } catch (IOException e) {
+                }
+            }
+        });
+    }
+
+}
diff --git a/src/test/java/com/google/mockwebserver/MockWebServerTest.java b/src/test/java/com/google/mockwebserver/MockWebServerTest.java
index 353361b..a8d5331 100644
--- a/src/test/java/com/google/mockwebserver/MockWebServerTest.java
+++ b/src/test/java/com/google/mockwebserver/MockWebServerTest.java
@@ -17,6 +17,7 @@
 package com.google.mockwebserver;
 
 import java.io.BufferedReader;
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
@@ -24,21 +25,78 @@ import java.net.HttpURLConnection;
 import java.net.SocketTimeoutException;
 import java.net.URL;
 import java.net.URLConnection;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import junit.framework.TestCase;
 
 public final class MockWebServerTest extends TestCase {
 
     private MockWebServer server = new MockWebServer();
 
-    @Override protected void setUp() throws Exception {
-        super.setUp();
-    }
-
     @Override protected void tearDown() throws Exception {
         server.shutdown();
         super.tearDown();
     }
 
+    public void testRecordedRequestAccessors() {
+        List<String> headers = Arrays.asList(
+                "User-Agent: okhttp",
+                "Cookie: s=square",
+                "Cookie: a=android",
+                "X-Whitespace:  left",
+                "X-Whitespace:right  ",
+                "X-Whitespace:  both  "
+        );
+        List<Integer> chunkSizes = Collections.emptyList();
+        byte[] body = {'A', 'B', 'C'};
+        String requestLine = "GET / HTTP/1.1";
+        RecordedRequest request = new RecordedRequest(
+                requestLine, headers, chunkSizes, body.length, body, 0, null);
+        assertEquals("s=square", request.getHeader("cookie"));
+        assertEquals(Arrays.asList("s=square", "a=android"), request.getHeaders("cookie"));
+        assertEquals("left", request.getHeader("x-whitespace"));
+        assertEquals(Arrays.asList("left", "right", "both"), request.getHeaders("x-whitespace"));
+        assertEquals("ABC", request.getUtf8Body());
+    }
+
+    public void testDefaultMockResponse() {
+        MockResponse response = new MockResponse();
+        assertEquals(Arrays.asList("Content-Length: 0"), response.getHeaders());
+        assertEquals("HTTP/1.1 200 OK", response.getStatus());
+    }
+
+    public void testSetBodyAdjustsHeaders() throws IOException {
+        MockResponse response = new MockResponse().setBody("ABC");
+        assertEquals(Arrays.asList("Content-Length: 3"), response.getHeaders());
+        InputStream in = response.getBodyStream();
+        assertEquals('A', in.read());
+        assertEquals('B', in.read());
+        assertEquals('C', in.read());
+        assertEquals(-1, in.read());
+        assertEquals("HTTP/1.1 200 OK", response.getStatus());
+    }
+
+    public void testMockResponseAddHeader() {
+        MockResponse response = new MockResponse()
+                .clearHeaders()
+                .addHeader("Cookie: s=square")
+                .addHeader("Cookie", "a=android");
+        assertEquals(Arrays.asList("Cookie: s=square", "Cookie: a=android"),
+                response.getHeaders());
+    }
+
+    public void testMockResponseSetHeader() {
+        MockResponse response = new MockResponse()
+                .clearHeaders()
+                .addHeader("Cookie: s=square")
+                .addHeader("Cookie: a=android")
+                .addHeader("Cookies: delicious");
+        response.setHeader("cookie", "r=robot");
+        assertEquals(Arrays.asList("Cookies: delicious", "cookie: r=robot"),
+                response.getHeaders());
+    }
+
     public void testRegularResponse() throws Exception {
         server.enqueue(new MockResponse().setBody("hello world"));
         server.play();
@@ -75,6 +133,29 @@ public final class MockWebServerTest extends TestCase {
         assertEquals("GET /new-path HTTP/1.1", redirect.getRequestLine());
     }
 
+    /**
+     * Test that MockWebServer blocks for a call to enqueue() if a request
+     * is made before a mock response is ready.
+     */
+    public void testDispatchBlocksWaitingForEnqueue() throws Exception {
+        server.play();
+
+        new Thread() {
+            @Override public void run() {
+                try {
+                    Thread.sleep(1000);
+                } catch (InterruptedException ignored) {
+                }
+                server.enqueue(new MockResponse().setBody("enqueued in the background"));
+            }
+        }.start();
+
+        URLConnection connection = server.getUrl("/").openConnection();
+        InputStream in = connection.getInputStream();
+        BufferedReader reader = new BufferedReader(new InputStreamReader(in));
+        assertEquals("enqueued in the background", reader.readLine());
+    }
+
     public void testNonHexadecimalChunkSize() throws Exception {
         server.enqueue(new MockResponse()
                 .setBody("G\r\nxxxxxxxxxxxxxxxx\r\n0\r\n\r\n")
@@ -122,4 +203,30 @@ public final class MockWebServerTest extends TestCase {
         assertEquals(0, server.takeRequest().getSequenceNumber());
         assertEquals(0, server.takeRequest().getSequenceNumber());
     }
+
+    public void testDisconnectAtStart() throws Exception {
+        server.enqueue(new MockResponse().setSocketPolicy(SocketPolicy.DISCONNECT_AT_START));
+        server.enqueue(new MockResponse()); // The jdk's HttpUrlConnection is a bastard.
+        server.enqueue(new MockResponse());
+        server.play();
+        try {
+            server.getUrl("/a").openConnection().getInputStream();
+        } catch (IOException e) {
+            // Expected.
+        }
+        server.getUrl("/b").openConnection().getInputStream(); // Should succeed.
+    }
+
+    public void testStreamingResponseBody() throws Exception {
+        InputStream responseBody = new ByteArrayInputStream("ABC".getBytes("UTF-8"));
+        server.enqueue(new MockResponse().setBody(responseBody, 3));
+        server.play();
+
+        InputStream in = server.getUrl("/").openConnection().getInputStream();
+        assertEquals('A', in.read());
+        assertEquals('B', in.read());
+        assertEquals('C', in.read());
+
+        assertEquals(-1, responseBody.read()); // The body is exhausted.
+    }
 }