[357] timeout on idle TCP connections in TCPServer

for now with a fixed size

src/lib/asiodns/tcp_server.cc

@@ -72,6 +72,21 @@ TCPServer::TCPServer(io_service& io_service, int fd, int af,
     }
 }
 
+namespace {
+    // Called by the timeout_ deadline timer if the client takes too long.
+    // If not aborted, cancels the given socket
+    // (in which case TCPServer::operator() will be called to continue,
+    // with an 'aborted' error code
+    void do_timeout(asio::ip::tcp::socket& socket,
+                    const asio::error_code& error)
+    {
+        if (error != asio::error::operation_aborted) {
+            socket.cancel();
+        } else {
+        }
+    }
+}
+
 void
 TCPServer::operator()(asio::error_code ec, size_t length) {
     /// Because the coroutine reentry block is implemented as
@@ -114,6 +129,11 @@ TCPServer::operator()(asio::error_code ec, size_t length) {
         /// asynchronous read call.
         data_.reset(new char[MAX_LENGTH]);
 
+        timeout_.reset(new asio::deadline_timer(io_));
+        timeout_->expires_from_now(boost::posix_time::milliseconds(1000));
+        timeout_->async_wait(boost::bind(&do_timeout, boost::ref(*socket_),
+                             asio::placeholders::error));
+
         /// Read the message, in two parts.  First, the message length:
         CORO_YIELD async_read(*socket_, asio::buffer(data_.get(),
                               TCP_MESSAGE_LENGTHSIZE), *this);
@@ -209,6 +229,9 @@ TCPServer::operator()(asio::error_code ec, size_t length) {
         // will simply exit at that time).
         CORO_YIELD async_write(*socket_, bufs, *this);
 
+        // All done, cancel the timeout timer
+        timeout_->cancel();
+
         // TODO: should we keep the connection open for a short time
         // to see if new requests come in?
         socket_->close();

src/lib/asiodns/tcp_server.h

@@ -34,7 +34,7 @@ namespace asiodns {
 /// \brief A TCP-specific \c DNSServer object.
 ///
 /// This class inherits from both \c DNSServer and from \c coroutine,
-/// defined in coroutine.h. 
+/// defined in coroutine.h.
 class TCPServer : public virtual DNSServer, public virtual coroutine {
 public:
     /// \brief Constructor
@@ -122,6 +122,12 @@ private:
 
     boost::shared_ptr<isc::asiolink::IOEndpoint> peer_;
     boost::shared_ptr<isc::asiolink::IOSocket> iosock_;
+
+    // Timer used to timeout on tcp connections
+    // This is a shared pointer because we need to have something
+    // that outlives the operator() call and is copyable (for CORO_FORK)
+    // even though it is only set after fork
+    boost::shared_ptr<asio::deadline_timer> timeout_;
 };
 
 } // namespace asiodns

src/lib/asiodns/tests/dns_server_unittest.cc

@@ -258,7 +258,8 @@ class TCPClient : public SimpleClient {
     // this includes connect, send message and recevice message
     static const unsigned int SERVER_TIME_OUT = 2;
     TCPClient(asio::io_service& service, const ip::tcp::endpoint& server)
-        : SimpleClient(service, SERVER_TIME_OUT)
+        : SimpleClient(service, SERVER_TIME_OUT),
+          send_data_(true), send_data_len_(true)
     {
         server_ = server;
         socket_.reset(new ip::tcp::socket(service));
@@ -280,13 +281,23 @@ class TCPClient : public SimpleClient {
                                 std::string(received_data_ + 2));
     }
 
+    // if set to false, does not actually send data length
+    void setSendDataLen(bool send_data_len) {
+        send_data_len_ = send_data_len;
+    }
+
+    // if set to false, does not actually send data
+    void setSendData(bool send_data) {
+        send_data_ = send_data;
+    }
+
     private:
     void stopWaitingforResponse() {
         socket_->close();
     }
 
     void connectHandler(const asio::error_code& error) {
-        if (!error) {
+        if (!error && send_data_len_) {
             data_to_send_len_ = htons(data_to_send_len_);
             socket_->async_send(buffer(&data_to_send_len_, 2),
                                 boost::bind(&TCPClient::sendMessageBodyHandler,
@@ -297,7 +308,7 @@ class TCPClient : public SimpleClient {
     void sendMessageBodyHandler(const asio::error_code& error,
                                 size_t send_bytes)
     {
-        if (!error && send_bytes == 2) {
+        if (!error && send_bytes == 2 && send_data_) {
             socket_->async_send(buffer(data_to_send_.c_str(),
                                        data_to_send_.size() + 1),
                     boost::bind(&TCPClient::finishSendHandler, this, _1, _2));
@@ -316,6 +327,8 @@ class TCPClient : public SimpleClient {
     ip::tcp::endpoint server_;
     std::string data_to_send_;
     uint16_t data_to_send_len_;
+    bool send_data_;
+    bool send_data_len_;
 };
 
 // \brief provide the context which including two clients and
@@ -565,6 +578,21 @@ TYPED_TEST(DNSServerTest, stopTCPServerAfterOneQuery) {
     EXPECT_TRUE(this->serverStopSucceed());
 }
 
+TYPED_TEST(DNSServerTest, TCPTimeoutOnLen) {
+    this->tcp_client_->setSendDataLen(false);
+    this->testStopServerByStopper(this->tcp_server_, this->tcp_client_,
+                                  this->tcp_client_);
+    EXPECT_EQ("", this->tcp_client_->getReceivedData());
+    EXPECT_FALSE(this->serverStopSucceed());
+}
+
+TYPED_TEST(DNSServerTest, TCPTimeout) {
+    this->tcp_client_->setSendData(false);
+    this->testStopServerByStopper(this->tcp_server_, this->tcp_client_,
+                                  this->tcp_client_);
+    EXPECT_EQ("", this->tcp_client_->getReceivedData());
+    EXPECT_FALSE(this->serverStopSucceed());
+}
 
 // Test whether tcp server stopped successfully before server start to serve
 TYPED_TEST(DNSServerTest, stopTCPServerBeforeItStartServing) {