Revert "fix server HTTP keep-alive"

Author: CheSema

src/Server/HTTP/HTTPServerRequest.cpp

@@ -56,15 +56,11 @@ HTTPServerRequest::HTTPServerRequest(HTTPContextPtr context, HTTPServerResponse
     /// and retry with exactly the same (incomplete) set of rows.
     /// That's why we have to check body size if it's provided.
     if (getChunkedTransferEncoding())
-    {
         stream = std::make_unique<HTTPChunkedReadBuffer>(std::move(in), HTTP_MAX_CHUNK_SIZE);
-        stream_is_bounded = true;
-    }
     else if (hasContentLength())
     {
         size_t content_length = getContentLength();
         stream = std::make_unique<LimitReadBuffer>(std::move(in), LimitReadBuffer::Settings{.read_no_less = content_length, .read_no_more = content_length, .expect_eof = true});
-        stream_is_bounded = true;
     }
     else if (getMethod() != HTTPRequest::HTTP_GET && getMethod() != HTTPRequest::HTTP_HEAD && getMethod() != HTTPRequest::HTTP_DELETE)
     {

src/Server/HTTP/HTTPServerRequest.h

@@ -47,14 +47,6 @@ class HTTPServerRequest : public HTTPRequest
     X509Certificate peerCertificate() const;
 #endif
 
-    bool canKeepAlive() const
-    {
-        if (stream && stream_is_bounded)
-            return stream->eof();
-
-        return false;
-    }
-
 private:
     /// Limits for basic sanity checks when reading a header
     enum Limits
@@ -73,7 +65,6 @@ class HTTPServerRequest : public HTTPRequest
     Poco::Net::SocketAddress client_address;
     Poco::Net::SocketAddress server_address;
 
-    bool stream_is_bounded = false;
     bool secure;
 
     void readRequest(ReadBuffer & in);

src/Server/HTTP/HTTPServerResponse.cpp

@@ -10,6 +10,7 @@
 #include <Poco/Net/HTTPHeaderStream.h>
 #include <Poco/Net/HTTPStream.h>
 #include <IO/NullWriteBuffer.h>
+#include <sstream>
 
 
 namespace DB
@@ -39,6 +40,9 @@ std::shared_ptr<WriteBuffer> HTTPServerResponse::send()
         // the connection would be poisoned.
         // Next request over that connection reads previously unreaded message as a HTTP status line
 
+        // Send header
+        Poco::Net::HTTPHeaderOutputStream hs(session);
+        write(hs);
         // make sure that nothing is sent to the client if it was HTTP_HEAD request
         stream = std::make_shared<NullWriteBuffer>(write_event);
 
@@ -49,36 +53,75 @@ std::shared_ptr<WriteBuffer> HTTPServerResponse::send()
         // but if we do, then it is safer to close the connection at the end
         setKeepAlive(false);
 
+        // Send header
+        Poco::Net::HTTPHeaderOutputStream hs(session);
+        write(hs);
         stream = std::make_shared<AutoFinalizedWriteBuffer<WriteBufferFromPocoSocket>>(session.socket(), write_event);
     }
     else if (getChunkedTransferEncoding())
     {
+        // Send header
+        Poco::Net::HTTPHeaderOutputStream hs(session);
+        write(hs);
         stream = std::make_shared<AutoFinalizedWriteBuffer<HTTPWriteBufferChunked>>(session.socket(), write_event);
     }
     else if (hasContentLength())
     {
+        // Send header
+        Poco::Net::HTTPHeaderOutputStream hs(session);
+        write(hs);
         stream = std::make_shared<AutoFinalizedWriteBuffer<HTTPWriteBufferFixedLength>>(session.socket(), getContentLength(), write_event);
     }
     else
     {
         setKeepAlive(false);
-
+        // Send header
+        Poco::Net::HTTPHeaderOutputStream hs(session);
+        write(hs);
         stream = std::make_shared<AutoFinalizedWriteBuffer<WriteBufferFromPocoSocket>>(session.socket(), write_event);
     }
 
-    Poco::Net::HTTPHeaderOutputStream hs(session);
-    beginWrite(hs);
-    hs << "\r\n";
-    hs.flush();
-
+    send_started = true;
     return stream;
 }
 
-/// Only this method is called inside WriteBufferFromHTTPServerResponse
-void HTTPServerResponse::beginWrite(std::ostream & ostr)
+std::pair<std::shared_ptr<WriteBuffer>, std::shared_ptr<WriteBuffer>> HTTPServerResponse::beginSend()
 {
-    allowKeepAliveIFFRequestIsFullyRead();
+    poco_assert(!stream);
+    poco_assert(!header_stream);
+
+    /// NOTE: Code is not exception safe.
 
+    if ((request && request->getMethod() == HTTPRequest::HTTP_HEAD) || getStatus() < 200 || getStatus() == HTTPResponse::HTTP_NO_CONTENT
+        || getStatus() == HTTPResponse::HTTP_NOT_MODIFIED)
+    {
+        throw Poco::Exception("HTTPServerResponse::beginSend is invalid for HEAD request");
+    }
+
+    if (hasContentLength())
+    {
+        throw Poco::Exception("HTTPServerResponse::beginSend is invalid for response with Content-Length header");
+    }
+
+    // Write header to buffer
+    std::stringstream header; //STYLE_CHECK_ALLOW_STD_STRING_STREAM
+    beginWrite(header);
+    // Send header
+    auto str = header.str();
+    header_stream = std::make_shared<AutoFinalizedWriteBuffer<WriteBufferFromPocoSocket>>(session.socket(), write_event, str.size());
+    header_stream->write(str.data(), str.size());
+
+    if (getChunkedTransferEncoding())
+        stream = std::make_shared<AutoFinalizedWriteBuffer<HTTPWriteBufferChunked>>(session.socket(), write_event);
+    else
+        stream = std::make_shared<AutoFinalizedWriteBuffer<WriteBufferFromPocoSocket>>(session.socket(), write_event);
+
+    send_started = true;
+    return std::make_pair(header_stream, stream);
+}
+
+void HTTPServerResponse::beginWrite(std::ostream & ostr) const
+{
     HTTPResponse::beginWrite(ostr);
     send_started = true;
 }
@@ -87,11 +130,9 @@ void HTTPServerResponse::sendBuffer(const void * buffer, std::size_t length)
 {
     setContentLength(static_cast<int>(length));
     setChunkedTransferEncoding(false);
-
     // Send header
     Poco::Net::HTTPHeaderOutputStream hs(session);
-    beginWrite(hs);
-    hs << "\r\n";
+    write(hs);
     hs.flush();
 
     if (request && request->getMethod() != HTTPRequest::HTTP_HEAD)
@@ -125,17 +166,8 @@ void HTTPServerResponse::redirect(const std::string & uri, HTTPStatus status)
 
     // Send header
     Poco::Net::HTTPHeaderOutputStream hs(session);
-    beginWrite(hs);
-    hs << "\r\n";
+    write(hs);
     hs.flush();
 }
 
-void HTTPServerResponse::allowKeepAliveIFFRequestIsFullyRead()
-{
-    /// Connection can only be reused if we've fully read the previous request and all its POST data.
-    /// Otherwise we'd misinterpret the leftover data as part of the next request's header.
-    /// HTTPServerRequest::canKeepAlive() checks that request stream is bounded and is fully read.
-    if (!request || !request->canKeepAlive())
-        setKeepAlive(false);
-}
 }

src/Server/HTTP/HTTPServerResponse.h

@@ -234,10 +234,17 @@ class HTTPServerResponse : public HTTPResponse
     /// or redirect() has been called.
     std::shared_ptr<WriteBuffer> send();
 
-    /// Dangerous, it is not a virtual method in HTTPResponse but it is redefined here
+    /// Sends the response headers to the client
+    /// but do not finish headers with \r\n,
+    /// allowing to continue sending additional header fields.
+    ///
+    /// Must not be called after send(), sendFile(), sendBuffer()
+    /// or redirect() has been called.
+    std::pair<std::shared_ptr<WriteBuffer>, std::shared_ptr<WriteBuffer>> beginSend();
+
     /// Override to correctly mark that the data send had been started for
     /// zero-copy response (i.e. replicated fetches).
-    void beginWrite(std::ostream & ostr);
+    void beginWrite(std::ostream & ostr) const;
 
     /// Sends the response header to the client, followed
     /// by the contents of the given buffer.
@@ -277,8 +284,6 @@ class HTTPServerResponse : public HTTPResponse
 
     const Poco::Net::HTTPServerSession & getSession() const { return session; }
 
-    void allowKeepAliveIFFRequestIsFullyRead();
-
 private:
     Poco::Net::HTTPServerSession & session;
     HTTPServerRequest * request = nullptr;

src/Server/HTTPHandler.cpp

@@ -617,6 +617,31 @@ void HTTPHandler::processQuery(
         {},
         handle_exception_in_output_format,
         query_finish_callback);
+
+    /// HTTPChunkedReadBuffer doesn't consume the final \r\n0\r\n\r\n if the caller reads exactly all bytes,
+    /// without checking for eof() (i.e. trying to read past the end) after that.
+    /// If those leftovers are then seen by the next request's HTTPServerRequest::readRequest,
+    /// it would in theory produce the error:
+    /// Invalid HTTP version string: /?query=I,
+    /// because that extra 0 field shifts all fields by one, and uri ends up interpreted as version.
+    /// There are a few options how to prevent that:
+    /// 1. Officially require that IInputFormat must drain the input buffer, make sure all implementations do that.
+    /// 2. Make the HTTP handler drain the request's POST read buffer after the request.
+    /// 3. Make HTTPChunkedReadBuffer eagerly drain the final \r\n0\r\n\r\n before returning the final bytes of data,
+    //     but that seems very inconvenient to implement because of how zero-copying is done in HTTPChunkedReadBuffer::nextImpl()
+    /// You can find an option (2) implemented below.
+    if (in_post_maybe_compressed->available())
+    {
+        String remaining;
+        readStringUntilEOF(remaining, *in_post_maybe_compressed);
+        auto hex_string = hexString(remaining.data(), remaining.size());
+#if defined(DEBUG_OR_SANITIZER_BUILD)
+        throw Exception(ErrorCodes::LOGICAL_ERROR,
+            "There is still some data in the buffer: {}", hex_string);
+#else
+        LOG_WARNING(log, "There is still some data in the buffer: {}", hex_string);
+#endif
+    }
 }
 
 bool HTTPHandler::trySendExceptionToClient(