Index: squid-3.3.14/src/tunnel.cc
===================================================================
--- squid-3.3.14.orig/src/tunnel.cc
+++ squid-3.3.14/src/tunnel.cc
@@ -45,6 +45,7 @@
 #include "fde.h"
 #include "http.h"
 #include "HttpRequest.h"
+#include "HttpReply.h"
 #include "HttpStateFlags.h"
 #include "ip/QosConfig.h"
 #include "MemBuf.h"
@@ -76,6 +77,12 @@ public:
     static void WriteClientDone(const Comm::ConnectionPointer &, char *buf, size_t len, comm_err_t flag, int xerrno, void *data);
     static void WriteServerDone(const Comm::ConnectionPointer &, char *buf, size_t len, comm_err_t flag, int xerrno, void *data);
 
+    /// Starts reading peer response to our CONNECT request.
+    void readConnectResponse();
+
+    /// Called when we may be done handling a CONNECT exchange with the peer.
+    void connectExchangeCheckpoint();
+
     bool noConnections() const;
     char *url;
     HttpRequest *request;
@@ -86,6 +93,23 @@ public:
         return (server.conn != NULL && server.conn->getPeer() ? server.conn->getPeer()->host : request->GetHost());
     };
 
+    /// Whether we are writing a CONNECT request to a peer.
+    bool waitingForConnectRequest() const { return connectReqWriting; }
+    /// Whether we are reading a CONNECT response from a peer.
+    bool waitingForConnectResponse() const { return connectRespBuf; }
+    /// Whether we are waiting for the CONNECT request/response exchange with the peer.
+    bool waitingForConnectExchange() const { return waitingForConnectRequest() || waitingForConnectResponse(); }
+
+    /// Whether the client sent a CONNECT request to us.
+    bool clientExpectsConnectResponse() const {
+        return !(request != NULL &&
+                 (request->flags.spoofClientIp || request->flags.intercepted));
+    }
+
+    /// Sends "502 Bad Gateway" error response to the client,
+    /// if it is waiting for Squid CONNECT response, closing connections.
+    void informUserOfPeerError(const char *errMsg);
+
     class Connection
     {
 
@@ -105,9 +129,12 @@ public:
         int debugLevelForError(int const xerrno) const;
         void closeIfOpen();
         void dataSent (size_t amount);
+        /// writes 'b' buffer, setting the 'writer' member to 'callback'.
+        void write(const char *b, int size, AsyncCall::Pointer &callback, FREE * free_func);
         int len;
         char *buf;
         int64_t *size_ptr;		/* pointer to size in an ConnStateData for logging */
+        AsyncCall::Pointer writer; ///< pending Comm::Write callback
 
         Comm::ConnectionPointer conn;    ///< The currently connected connection.
 
@@ -121,15 +148,23 @@ public:
 
     Connection client, server;
     int *status_ptr;		/* pointer to status for logging */
+    MemBuf *connectRespBuf; ///< accumulates peer CONNECT response when we need it
+    bool connectReqWriting; ///< whether we are writing a CONNECT request to a peer
+
     void copyRead(Connection &from, IOCB *completion);
 
 private:
     CBDATA_CLASS(TunnelStateData);
-    void copy (size_t len, comm_err_t errcode, int xerrno, Connection &from, Connection &to, IOCB *);
+    bool keepGoingAfterRead(size_t len, comm_err_t errcode, int xerrno, Connection &from, Connection &to);
+    void copy (size_t len, Connection &from, Connection &to, IOCB *);
+    void handleConnectResponse(const size_t chunkSize);
     void readServer(char *buf, size_t len, comm_err_t errcode, int xerrno);
     void readClient(char *buf, size_t len, comm_err_t errcode, int xerrno);
     void writeClientDone(char *buf, size_t len, comm_err_t flag, int xerrno);
     void writeServerDone(char *buf, size_t len, comm_err_t flag, int xerrno);
+
+    static void ReadConnectResponseDone(const Comm::ConnectionPointer &, char *buf, size_t len, comm_err_t errcode, int xerrno, void *data);
+    void readConnectResponseDone(char *buf, size_t len, comm_err_t errcode, int xerrno);
 };
 
 static const char *const conn_established = "HTTP/1.1 200 Connection established\r\n\r\n";
@@ -150,13 +185,14 @@ tunnelServerClosed(const CommCloseCbPara
     TunnelStateData *tunnelState = (TunnelStateData *)params.data;
     debugs(26, 3, HERE << tunnelState->server.conn);
     tunnelState->server.conn = NULL;
+    tunnelState->server.writer = NULL;
 
     if (tunnelState->noConnections()) {
         tunnelStateFree(tunnelState);
         return;
     }
 
-    if (!tunnelState->server.len) {
+    if (!tunnelState->client.writer) {
         tunnelState->client.conn->close();
         return;
     }
@@ -168,13 +204,14 @@ tunnelClientClosed(const CommCloseCbPara
     TunnelStateData *tunnelState = (TunnelStateData *)params.data;
     debugs(26, 3, HERE << tunnelState->client.conn);
     tunnelState->client.conn = NULL;
+    tunnelState->client.writer = NULL;
 
     if (tunnelState->noConnections()) {
         tunnelStateFree(tunnelState);
         return;
     }
 
-    if (!tunnelState->client.len) {
+    if (!tunnelState->server.writer) {
         tunnelState->server.conn->close();
         return;
     }
@@ -265,7 +302,119 @@ TunnelStateData::readServer(char *buf, s
         kb_incr(&(statCounter.server.other.kbytes_in), len);
     }
 
-    copy (len, errcode, xerrno, server, client, WriteClientDone);
+    if (keepGoingAfterRead(len, errcode, xerrno, server, client))
+        copy (len, server, client, WriteClientDone);
+}
+
+/// Called when we read [a part of] CONNECT response from the peer
+void
+TunnelStateData::readConnectResponseDone(char *buf, size_t len, comm_err_t errcode, int xerrno)
+{
+    debugs(26, 3, server.conn << ", read " << len << " bytes, err=" << errcode);
+//    assert(waitingForConnectResponse());
+
+    if (errcode == COMM_ERR_CLOSING)
+        return;
+
+    if (len > 0) {
+        connectRespBuf->appended(len);
+        server.bytesIn(len);
+        kb_incr(&(statCounter.server.all.kbytes_in), len);
+        kb_incr(&(statCounter.server.other.kbytes_in), len);
+    }
+
+    if (keepGoingAfterRead(len, errcode, xerrno, server, client))
+        handleConnectResponse(len);
+}
+
+void
+TunnelStateData::informUserOfPeerError(const char *errMsg)
+{
+    server.len = 0;
+    if (!clientExpectsConnectResponse()) {
+        // closing the connection is the best we can do here
+        debugs(50, 3, server.conn << " closing on error: " << errMsg);
+        server.conn->close();
+        return;
+    }
+    ErrorState *err  = new ErrorState(ERR_CONNECT_FAIL, HTTP_BAD_GATEWAY, request);
+    err->callback = tunnelErrorComplete;
+    err->callback_data = this;
+    *status_ptr = HTTP_BAD_GATEWAY;
+    errorSend(client.conn, err);
+}
+
+/* Read from client side and queue it for writing to the server */
+void
+TunnelStateData::ReadConnectResponseDone(const Comm::ConnectionPointer &, char *buf, size_t len, comm_err_t errcode, int xerrno, void *data)
+{
+    TunnelStateData *tunnelState = (TunnelStateData *)data;
+    assert (cbdataReferenceValid (tunnelState));
+
+    tunnelState->readConnectResponseDone(buf, len, errcode, xerrno);
+}
+
+/// Parses [possibly incomplete] CONNECT response and reacts to it.
+/// If the tunnel is being closed or more response data is needed, returns false.
+/// Otherwise, the caller should handle the remaining read data, if any.
+void
+TunnelStateData::handleConnectResponse(const size_t chunkSize)
+{
+    assert(waitingForConnectResponse());
+
+    // Ideally, client and server should use MemBuf or better, but current code
+    // never accumulates more than one read when shoveling data (XXX) so it does
+    // not need to deal with MemBuf complexity. To keep it simple, we use a
+    // dedicated MemBuf for accumulating CONNECT responses. TODO: When shoveling
+    // is optimized, reuse server.buf for CONNEC response accumulation instead.
+
+    /* mimic the basic parts of HttpStateData::processReplyHeader() */
+    HttpReply rep;  // <<< ???
+    http_status parseErr = HTTP_STATUS_NONE;
+    const bool eof = !chunkSize;
+    const bool parsed = rep.parse(connectRespBuf, eof, &parseErr);
+    if (!parsed) {
+        if (parseErr > 0) { // unrecoverable parsing error
+            informUserOfPeerError("malformed CONNECT response from peer");
+            return;
+        }
+
+        // need more data
+        assert(!eof);
+        assert(!parseErr);
+
+        if (!connectRespBuf->hasSpace()) {
+            informUserOfPeerError("huge CONNECT response from peer");
+            return;
+        }
+
+        // keep reading
+        readConnectResponse();
+        return;
+    }
+
+    // CONNECT response was successfully parsed
+    *status_ptr = rep.sline.status;
+
+    // bail if we did not get an HTTP 200 (Connection Established) response
+    if (rep.sline.status != HTTP_OK) {
+        // if we ever decide to reuse the peer connection, we must extract the error response first
+        informUserOfPeerError("unsupported CONNECT response status code");
+        return;
+    }
+
+    if (rep.hdr_sz < connectRespBuf->contentSize()) {
+        // preserve bytes that the server already sent after the CONNECT response
+        server.len = connectRespBuf->contentSize() - rep.hdr_sz;
+        memcpy(server.buf, connectRespBuf->content()+rep.hdr_sz, server.len);
+    } else {
+        // reset; delay pools were using this field to throttle CONNECT response
+        server.len = 0;
+    }
+
+    delete connectRespBuf;
+    connectRespBuf = NULL;
+    connectExchangeCheckpoint();
 }
 
 void
@@ -308,11 +457,14 @@ TunnelStateData::readClient(char *buf, s
         kb_incr(&(statCounter.client_http.kbytes_in), len);
     }
 
-    copy (len, errcode, xerrno, client, server, WriteServerDone);
+    if (keepGoingAfterRead(len, errcode, xerrno, client, server))
+        copy (len, client, server, WriteServerDone);
 }
 
-void
-TunnelStateData::copy (size_t len, comm_err_t errcode, int xerrno, Connection &from, Connection &to, IOCB *completion)
+/// Updates state after reading from client or server.
+/// Returns whether the caller should use the data just read.
+bool
+TunnelStateData::keepGoingAfterRead (size_t len, comm_err_t errcode, int xerrno, Connection &from, Connection &to)
 {
     debugs(26, 3, HERE << "from={" << from.conn << "}, to={" << to.conn << "}");
 
@@ -348,21 +500,31 @@ TunnelStateData::copy (size_t len, comm_
             to.conn->close();
         }
     } else if (cbdataReferenceValid(this)) {
-        debugs(26, 3, HERE << "Schedule Write");
-        AsyncCall::Pointer call = commCbCall(5,5, "TunnelBlindCopyWriteHandler",
-                                             CommIoCbPtrFun(completion, this));
-        Comm::Write(to.conn, from.buf, len, call, NULL);
+        cbdataInternalUnlock(this); /* ??? */
+        return true;
     }
 
     cbdataInternalUnlock(this);	/* ??? */
+    return false;
 }
 
+void
+TunnelStateData::copy(size_t len, Connection &from, Connection &to, IOCB *completion)
+{
+    debugs(26, 3, HERE << "Schedule Write");
+    AsyncCall::Pointer call = commCbCall(5,5, "TunnelBlindCopyWriteHandler",
+                                         CommIoCbPtrFun(completion, this));
+    to.write(from.buf, len, call, NULL);
+}
+
+
 /* Writes data from the client buffer to the server side */
 void
 TunnelStateData::WriteServerDone(const Comm::ConnectionPointer &, char *buf, size_t len, comm_err_t flag, int xerrno, void *data)
 {
     TunnelStateData *tunnelState = (TunnelStateData *)data;
     assert (cbdataReferenceValid (tunnelState));
+    tunnelState->server.writer = NULL;
 
     tunnelState->writeServerDone(buf, len, flag, xerrno);
 }
@@ -414,6 +576,7 @@ TunnelStateData::WriteClientDone(const C
 {
     TunnelStateData *tunnelState = (TunnelStateData *)data;
     assert (cbdataReferenceValid (tunnelState));
+    tunnelState->client.writer = NULL;
 
     tunnelState->writeClientDone(buf, len, flag, xerrno);
 }
@@ -431,6 +594,13 @@ TunnelStateData::Connection::dataSent(si
 }
 
 void
+TunnelStateData::Connection::write(const char *b, int size, AsyncCall::Pointer &callback, FREE * free_func)
+{
+    writer = callback;
+    Comm::Write(conn, b, size, callback, free_func);
+}
+
+void
 TunnelStateData::writeClientDone(char *buf, size_t len, comm_err_t flag, int xerrno)
 {
     debugs(26, 3, HERE << client.conn << ", " << len << " bytes written, flag=" << flag);
@@ -499,6 +669,17 @@ TunnelStateData::copyRead(Connection &fr
     comm_read(from.conn, from.buf, from.bytesWanted(1, SQUID_TCP_SO_RCVBUF), call);
 }
 
+void
+TunnelStateData::readConnectResponse()
+{
+    assert(waitingForConnectResponse());
+
+    AsyncCall::Pointer call = commCbCall(5,4, "readConnectResponseDone",
+                                         CommIoCbPtrFun(ReadConnectResponseDone, this));
+    comm_read(server.conn, connectRespBuf->space(),
+              server.bytesWanted(1, connectRespBuf->spaceSize()), call);
+}
+
 /**
  * Set the HTTP status for this request and sets the read handlers for client
  * and server side connections.
@@ -523,6 +704,7 @@ tunnelConnectedWriteDone(const Comm::Con
 {
     TunnelStateData *tunnelState = (TunnelStateData *)data;
     debugs(26, 3, HERE << conn << ", flag=" << flag);
+    tunnelState->client.writer = NULL;
 
     if (flag != COMM_OK) {
         *tunnelState->status_ptr = HTTP_INTERNAL_SERVER_ERROR;
@@ -530,9 +712,43 @@ tunnelConnectedWriteDone(const Comm::Con
         return;
     }
 
+    // tunnelState->connectReqWriting = false;
     tunnelStartShoveling(tunnelState);
 }
 
+/// Called when we are done writing CONNECT request to a peer.
+static void
+tunnelConnectReqWriteDone(const Comm::ConnectionPointer &conn, char *buf, size_t size, comm_err_t flag, int xerrno, void *data)
+{
+    TunnelStateData *tunnelState = (TunnelStateData *)data;
+    debugs(26, 3, conn << ", flag=" << flag);
+    tunnelState->server.writer = NULL;
+    assert(tunnelState->waitingForConnectRequest());
+
+    if (flag != COMM_OK) {
+        *tunnelState->status_ptr = HTTP_INTERNAL_SERVER_ERROR;
+        tunnelErrorComplete(conn->fd, data, 0);
+        return;
+    }
+
+    tunnelState->connectReqWriting = false;
+    tunnelState->connectExchangeCheckpoint();
+}
+
+void
+TunnelStateData::connectExchangeCheckpoint()
+{
+    if (waitingForConnectResponse()) {
+        debugs(26, 5, "still reading CONNECT response on " << server.conn);
+    } else if (waitingForConnectRequest()) {
+        debugs(26, 5, "still writing CONNECT request on " << server.conn);
+    } else {
+        assert(!waitingForConnectExchange());
+        debugs(26, 3, "done with CONNECT exchange on " << server.conn);
+        tunnelConnected(server.conn, this);
+    }
+}
+
 /*
  * handle the write completion from a proxy request to an upstream origin
  */
@@ -547,7 +763,7 @@ tunnelConnected(const Comm::ConnectionPo
     else {
         AsyncCall::Pointer call = commCbCall(5,5, "tunnelConnectedWriteDone",
                                              CommIoCbPtrFun(tunnelConnectedWriteDone, tunnelState));
-        Comm::Write(tunnelState->client.conn, conn_established, strlen(conn_established), call, NULL);
+        tunnelState->client.write(conn_established, strlen(conn_established), call, NULL);
     }
 }
 
@@ -731,9 +947,22 @@ tunnelRelayConnectRequest(const Comm::Co
     packerClean(&p);
     mb.append("\r\n", 2);
 
+    debugs(11, 2, "Tunnel Server REQUEST: " << tunnelState->server.conn);
+
     AsyncCall::Pointer writeCall = commCbCall(5,5, "tunnelConnectedWriteDone",
-                                   CommIoCbPtrFun(tunnelConnectedWriteDone, tunnelState));
-    Comm::Write(srv, &mb, writeCall);
+                                   CommIoCbPtrFun(tunnelConnectReqWriteDone, tunnelState));
+
+    tunnelState->server.write(mb.buf, mb.size, writeCall, mb.freeFunc());
+    tunnelState->connectReqWriting = true;
+
+    tunnelState->connectRespBuf = new MemBuf;
+    // SQUID_TCP_SO_RCVBUF: we should not accumulate more than regular I/O buffer
+    // can hold since any CONNECT response leftovers have to fit into server.buf.
+    // 2*SQUID_TCP_SO_RCVBUF: HttpMsg::parse() zero-terminates, which uses space.
+    tunnelState->connectRespBuf->init(SQUID_TCP_SO_RCVBUF, 2*SQUID_TCP_SO_RCVBUF);
+    tunnelState->readConnectResponse();
+
+    assert(tunnelState->waitingForConnectExchange());
 
     AsyncCall::Pointer timeoutCall = commCbCall(5, 4, "tunnelTimeout",
                                      CommTimeoutCbPtrFun(tunnelTimeout, tunnelState));
@@ -782,6 +1011,8 @@ TunnelStateData::operator new (size_t)
 {
     CBDATA_INIT_TYPE(TunnelStateData);
     TunnelStateData *result = cbdataAlloc(TunnelStateData);
+    result->connectReqWriting = false;
+    result->connectRespBuf = NULL;
     return result;
 }
 
@@ -789,6 +1020,7 @@ void
 TunnelStateData::operator delete (void *address)
 {
     TunnelStateData *t = static_cast<TunnelStateData *>(address);
+    delete t->connectRespBuf;
     cbdataFree(t);
 }
 
