diff --git a/lib/cpp/src/thrift/transport/THttpClient.cpp b/lib/cpp/src/thrift/transport/THttpClient.cpp index 3cc83345018..dc85a0e82c8 100644 --- a/lib/cpp/src/thrift/transport/THttpClient.cpp +++ b/lib/cpp/src/thrift/transport/THttpClient.cpp @@ -18,10 +18,11 @@ */ #include -#include +#include #include +#include #include -#include +#include #include #include @@ -38,23 +39,32 @@ THttpClient::THttpClient(std::shared_ptr transport, std::string path, std::shared_ptr config) : THttpTransport(transport, config), - host_(host), + host_(host), path_(path), - onewayResponsePending_(false) { -} + onewayResponsePending_(false), + closeAfterResponse_(false) {} -THttpClient::THttpClient(string host, int port, string path, - std::shared_ptr config) +THttpClient::THttpClient(string host, int port, string path, std::shared_ptr config) : THttpTransport(std::shared_ptr(new TSocket(host, port)), config), host_(host), path_(path), - onewayResponsePending_(false) { -} + onewayResponsePending_(false), + closeAfterResponse_(false) {} THttpClient::~THttpClient() = default; void THttpClient::close() { onewayResponsePending_ = false; + closeAfterResponse_ = false; + readBuffer_.resetBuffer(); + readHeaders_ = true; + chunked_ = false; + chunkedDone_ = false; + chunkSize_ = 0; + contentLength_ = 0; + httpPos_ = 0; + httpBufLen_ = 0; + httpBuf_[0] = '\0'; THttpTransport::close(); } @@ -65,17 +75,28 @@ void THttpClient::parseHeader(char* header) { } char* value = colon + 1; - if (boost::istarts_with(header, "Transfer-Encoding")) { + const string name(header, colon); + if (boost::iequals(name, "Transfer-Encoding")) { if (boost::iends_with(value, "chunked")) { chunked_ = true; } - } else if (boost::istarts_with(header, "Content-Length")) { + } else if (boost::iequals(name, "Content-Length")) { chunked_ = false; contentLength_ = atoi(value); + } else if (boost::iequals(name, "Connection")) { + std::vector options; + boost::split(options, value, boost::is_any_of(",")); + for (const string& option : options) { + if (boost::iequals(boost::trim_copy(option), "close")) { + closeAfterResponse_ = true; + break; + } + } } } bool THttpClient::parseStatusLine(char* status) { + closeAfterResponse_ = false; char* http = status; char* code = strchr(http, ' '); @@ -107,6 +128,13 @@ bool THttpClient::parseStatusLine(char* status) { void THttpClient::flush() { resetConsumedMessageSize(); + uint8_t* buf; + uint32_t len; + writeBuffer_.getBuffer(&buf, &len); + if (len == 0) { + return; + } + if (onewayResponsePending_) { if (transport_->isOpen()) { drainPendingOnewayResponse(); @@ -115,15 +143,14 @@ void THttpClient::flush() { } } + if (closeAfterResponse_) { + close(); + } + if (!transport_->isOpen()) { transport_->open(); } - // Fetch the contents of the write buffer - uint8_t* buf; - uint32_t len; - writeBuffer_.getBuffer(&buf, &len); - // Construct the HTTP header std::ostringstream h; h << "POST " << path_ << " HTTP/1.1" << CRLF << "Host: " << host_ << CRLF diff --git a/lib/cpp/src/thrift/transport/THttpClient.h b/lib/cpp/src/thrift/transport/THttpClient.h index b8c14c4a6e2..96e6e1fa10f 100644 --- a/lib/cpp/src/thrift/transport/THttpClient.h +++ b/lib/cpp/src/thrift/transport/THttpClient.h @@ -65,6 +65,7 @@ class THttpClient : public THttpTransport { std::string host_; std::string path_; bool onewayResponsePending_; + bool closeAfterResponse_; void parseHeader(char* header) override; bool parseStatusLine(char* status) override; diff --git a/lib/cpp/test/OneWayHTTPTest.cpp b/lib/cpp/test/OneWayHTTPTest.cpp index 2c545678d59..643cdb333e0 100644 --- a/lib/cpp/test/OneWayHTTPTest.cpp +++ b/lib/cpp/test/OneWayHTTPTest.cpp @@ -17,22 +17,23 @@ * under the License. */ +#include "gen-cpp/OneWayService.h" #include #include -#include #include -#include +#include +#include +#include #include #include #include #include -#include +#include #include +#include #include #include -#include -#include -#include "gen-cpp/OneWayService.h" +#include BOOST_AUTO_TEST_SUITE(OneWayHTTPTest) @@ -61,6 +62,35 @@ namespace utf = boost::unit_test; // Define this env var to enable some logging (in case you need to debug) #undef ENABLE_STDERR_LOGGING +class TInspectableHttpClient : public THttpClient { +public: + explicit TInspectableHttpClient(std::shared_ptr transport) : THttpClient(transport) {} + + bool closesAfterHeader(const string& header) { + closeAfterResponse_ = false; + std::vector buffer(header.begin(), header.end()); + buffer.push_back('\0'); + parseHeader(buffer.data()); + return closeAfterResponse_; + } + + bool chunksAfterHeader(const string& header) { + chunked_ = false; + std::vector buffer(header.begin(), header.end()); + buffer.push_back('\0'); + parseHeader(buffer.data()); + return chunked_; + } + + uint32_t contentLengthAfterHeader(const string& header) { + contentLength_ = 0; + std::vector buffer(header.begin(), header.end()); + buffer.push_back('\0'); + parseHeader(buffer.data()); + return contentLength_; + } +}; + class OneWayServiceHandler : public onewaytest::OneWayServiceIf { public: OneWayServiceHandler() = default; @@ -126,13 +156,63 @@ class TServerReadyEventHandler : public TServerEventHandler, public Monitor { return nullptr; } bool isListening() const { return isListening_; } - uint64_t acceptedCount() const { return accepted_; } + uint64_t acceptedCount() { + Synchronized sync(*this); + return accepted_; + } private: bool isListening_; uint64_t accepted_; }; +class TClosingHttpServer : public THttpServer { +public: + explicit TClosingHttpServer(std::shared_ptr transport) + : THttpServer(transport, transport->getConfiguration()) {} + + void flush() override { + resetConsumedMessageSize(); + + uint8_t* buf; + uint32_t len; + writeBuffer_.getBuffer(&buf, &len); + + std::ostringstream header; + header << "HTTP/1.1 200 OK\r\n" + << "Content-Type: application/x-thrift\r\n" + << "Transfer-Encoding: chunked\r\n" + << "Connection: keep-alive, close\r\n" + << "Connection: keep-alive\r\n\r\n"; + const string headerText = header.str(); + transport_->write(reinterpret_cast(headerText.data()), + static_cast(headerText.size())); + + if (len > 0) { + std::ostringstream chunkSize; + chunkSize << std::hex << len << "\r\n"; + const string chunkPrefix = chunkSize.str(); + transport_->write(reinterpret_cast(chunkPrefix.data()), + static_cast(chunkPrefix.size())); + transport_->write(buf, len); + transport_->write(reinterpret_cast("\r\n"), 2); + } + transport_->write(reinterpret_cast("0\r\n\r\n"), 5); + transport_->flush(); + + writeBuffer_.resetBuffer(); + readHeaders_ = true; + close(); + } +}; + +class TClosingHttpServerTransportFactory : public apache::thrift::transport::TTransportFactory { +public: + std::shared_ptr getTransport(std::shared_ptr transport) override { + return std::make_shared(transport); + } +}; + class TBlockableBufferedTransport : public TBufferedTransport { public: TBlockableBufferedTransport(std::shared_ptr transport) @@ -286,4 +366,59 @@ BOOST_AUTO_TEST_CASE( JSON_HTTP_OneWayWrapperDoesNotPoisonNextCall ) thread.join(); } +BOOST_AUTO_TEST_CASE(HTTP_ClientReconnectsAfterConnectionClose) { + std::shared_ptr ss = std::make_shared(0); + TThreadedServer server(std::make_shared( + std::make_shared()), + ss, std::make_shared(), + std::make_shared()); + + std::shared_ptr pEventHandler(new TServerReadyEventHandler); + server.setServerEventHandler(pEventHandler); + + RPC0ThreadClass t(server); + boost::thread thread(&RPC0ThreadClass::Run, &t); + + { + Synchronized sync(*(pEventHandler.get())); + while (!pEventHandler->isListening()) { + pEventHandler->wait(); + } + } + + { + std::shared_ptr socket(new TSocket("localhost", ss->getPort())); + socket->setRecvTimeout(10000); + std::shared_ptr httpTransport(new THttpClient(socket, "localhost", "/service")); + std::shared_ptr transport(new TBufferedTransport(httpTransport)); + std::shared_ptr protocol(new TBinaryProtocol(transport)); + onewaytest::OneWayServiceClient client(protocol); + + transport->open(); + client.roundTripRPC(); + BOOST_CHECK_EQUAL(pEventHandler->acceptedCount(), 1U); + BOOST_CHECK_NO_THROW(client.roundTripRPC()); + BOOST_CHECK_EQUAL(pEventHandler->acceptedCount(), 2U); + client.oneWayRPC(); + BOOST_CHECK_NO_THROW(client.roundTripRPC()); + BOOST_CHECK_EQUAL(pEventHandler->acceptedCount(), 4U); + transport->close(); + BOOST_CHECK_EQUAL(pEventHandler->acceptedCount(), 4U); + } + + server.stop(); + thread.join(); +} + +BOOST_AUTO_TEST_CASE(HTTP_ClientRequiresExactConnectionHeaderName) { + TInspectableHttpClient client(std::make_shared()); + + BOOST_CHECK(client.closesAfterHeader("Connection: keep-alive, close")); + BOOST_CHECK(!client.closesAfterHeader("Connection-Timeout: close")); + BOOST_CHECK(client.chunksAfterHeader("Transfer-Encoding: chunked")); + BOOST_CHECK(!client.chunksAfterHeader("Transfer-Encoding-Other: chunked")); + BOOST_CHECK_EQUAL(client.contentLengthAfterHeader("Content-Length: 42"), 42U); + BOOST_CHECK_EQUAL(client.contentLengthAfterHeader("Content-Length-Mismatch: 42"), 0U); +} + BOOST_AUTO_TEST_SUITE_END()