fix(TCPServer): check stopped status after poll; add TCPServerFactory::stop() and allow factory to return nullptr on connection creation request #4892

Author: aleks-f

Net/include/Poco/Net/TCPServerConnectionFactory.h

@@ -35,7 +35,13 @@ class Net_API TCPServerConnectionFactory
 	/// it accepts.
 	///
 	/// Subclasses must override the createConnection()
-	/// method.
+	/// method, which can refuse connections by returning nullptr.
+	/// Some examples when an implementation may refuse new connections:
+	///
+	///  - number of connections exceeded a limit
+	///  - a connection from unwanted client attempted
+	///  - too many connection attempts in a short timespan
+	///  - etc.
 	///
 	/// The TCPServerConnectionFactoryImpl template class
 	/// can be used to automatically instantiate a
@@ -45,23 +51,55 @@ class Net_API TCPServerConnectionFactory
 public:
 	using Ptr = Poco::SharedPtr<TCPServerConnectionFactory>;
 
+	TCPServerConnectionFactory(const TCPServerConnectionFactory&) = delete;
+	TCPServerConnectionFactory& operator = (const TCPServerConnectionFactory&) = delete;
+	TCPServerConnectionFactory(TCPServerConnectionFactory&&) = delete;
+	TCPServerConnectionFactory& operator = (TCPServerConnectionFactory&&) = delete;
+
 	virtual ~TCPServerConnectionFactory();
 		/// Destroys the TCPServerConnectionFactory.
 
 	virtual TCPServerConnection* createConnection(const StreamSocket& socket) = 0;
 		/// Creates an instance of a subclass of TCPServerConnection,
 		/// using the given StreamSocket.
+		/// This function is allowed to return nullptr, in which case an accepted
+		/// socket will be destroyed by the TCPServerDispatcher.
+
+	void stop();
+		/// Stops the factory.
+		/// Normally, this function is called by TCPServerDispatcher
+		/// to indicate that the server is shutting down; the expected
+		/// implementation behavior after this call is to return nullptr
+		/// on all subsequent connection creation attempts.
+
+	bool isStopped() const;
+		/// Returns true if the factory was stopped, false otherwise.
 
 protected:
 	TCPServerConnectionFactory();
 		/// Creates the TCPServerConnectionFactory.
 
 private:
-	TCPServerConnectionFactory(const TCPServerConnectionFactory&);
-	TCPServerConnectionFactory& operator = (const TCPServerConnectionFactory&);
+	std::atomic<bool> _stopped;
 };
 
 
+//
+// inlines
+//
+
+inline void TCPServerConnectionFactory::stop()
+{
+	_stopped = true;
+}
+
+
+inline bool TCPServerConnectionFactory::isStopped() const
+{
+	return _stopped;
+}
+
+
 template <class S>
 class TCPServerConnectionFactoryImpl: public TCPServerConnectionFactory
 	/// This template provides a basic implementation of
@@ -78,7 +116,10 @@ class TCPServerConnectionFactoryImpl: public TCPServerConnectionFactory
 
 	TCPServerConnection* createConnection(const StreamSocket& socket)
 	{
-		return new S(socket);
+		if (!isStopped())
+			return new S(socket);
+
+		return nullptr;
 	}
 };
 

Net/src/TCPServer.cpp

@@ -116,7 +116,6 @@ void TCPServer::stop()
 {
 	if (!_stopped)
 	{
-		_socket.close();
 		_stopped = true;
 		_thread.join();
 		_pDispatcher->stop();
@@ -133,6 +132,8 @@ void TCPServer::run()
 		{
 			if (_socket.poll(timeout, Socket::SELECT_READ))
 			{
+				if (_stopped) return;
+
 				try
 				{
 					StreamSocket ss = _socket.acceptConnection();

Net/src/TCPServerDispatcher.cpp

@@ -112,10 +112,12 @@ void TCPServerDispatcher::run()
 				if (pCNf)
 				{
 					std::unique_ptr<TCPServerConnection> pConnection(_pConnectionFactory->createConnection(pCNf->socket()));
-					poco_check_ptr(pConnection.get());
-					beginConnection();
-					pConnection->start();
-					endConnection();
+					if (pConnection)
+					{
+						beginConnection();
+						pConnection->start();
+						endConnection();
+					}
 				}
 			}
 		}
@@ -173,6 +175,7 @@ void TCPServerDispatcher::enqueue(const StreamSocket& socket)
 void TCPServerDispatcher::stop()
 {
 	FastMutex::ScopedLock lock(_mutex);
+	_pConnectionFactory->stop();
 	_stopped = true;
 	_queue.clear();
 	for (int i = 0; i < _threadPool.allocated(); i++)