Merge pull request #22183 from hrydgard/net-thread-tracking

Net: reap finished connection threads, and don't echo unsendable close codes

Common/Net/HTTPServer.cpp

@@ -44,14 +44,37 @@
 
 
 void NewThreadExecutor::Run(std::function<void()> func) {
-	threads_.push_back(std::thread(func));
+	// Every connection gets a thread, and we only ever joined them at shutdown - so a server that
+	// had served N connections was still holding N joinable std::threads. Reap the finished ones.
+	Prune();
+
+	auto done = std::make_shared<std::atomic<bool>>(false);
+	Worker worker;
+	worker.done = done;
+	worker.thread = std::thread([func, done]() {
+		func();
+		done->store(true, std::memory_order_release);
+	});
+	workers_.push_back(std::move(worker));
+}
+
+void NewThreadExecutor::Prune() {
+	for (size_t i = 0; i < workers_.size(); ) {
+		if (workers_[i].done->load(std::memory_order_acquire)) {
+			// Set right at the end of the thread body, so this join returns essentially at once.
+			workers_[i].thread.join();
+			workers_.erase(workers_.begin() + i);
+		} else {
+			++i;
+		}
+	}
 }
 
 NewThreadExecutor::~NewThreadExecutor() {
 	// If Run was ever called...
-	for (auto &thread : threads_)
-		thread.join();
-	threads_.clear();
+	for (auto &worker : workers_)
+		worker.thread.join();
+	workers_.clear();
 }
 
 namespace http {

Common/Net/HTTPServer.h

@@ -1,8 +1,11 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
 #include <map>
+#include <memory>
 #include <thread>
+#include <vector>
 
 #include "Common/Net/HTTPHeaders.h"
 #include "Common/Net/Resolve.h"
@@ -13,7 +16,15 @@ public:
 	void Run(std::function<void()> func);
 
 private:
-	std::vector<std::thread> threads_;
+	// Reap threads that have finished. Only called from the thread that calls Run().
+	void Prune();
+
+	struct Worker {
+		std::thread thread;
+		// Set by the worker as its last act, read by whoever calls Run() next.
+		std::shared_ptr<std::atomic<bool>> done;
+	};
+	std::vector<Worker> workers_;
 };
 
 namespace net {

Common/Net/WebsocketServer.cpp

@@ -493,6 +493,32 @@ bool WebSocketServer::ReadPending() {
 	return true;
 }
 
+// Which close codes we're allowed to put on the wire, per RFC 6455 7.4.1. 1004, 1005, 1006 and
+// 1015 are reserved for local use only, and anything below 1000 is undefined.
+static bool IsCloseCodeSendable(uint16_t code) {
+	if (code >= 3000 && code <= 4999) {
+		// Registered and private-use ranges.
+		return true;
+	}
+	switch ((WebSocketClose)code) {
+	case WebSocketClose::NORMAL:
+	case WebSocketClose::GOING_AWAY:
+	case WebSocketClose::PROTOCOL_ERROR:
+	case WebSocketClose::UNSUPPORTED_DATA:
+	case WebSocketClose::INVALID_DATA:
+	case WebSocketClose::POLICY_VIOLATION:
+	case WebSocketClose::MESSAGE_TOO_LONG:
+	case WebSocketClose::MISSING_EXTENSION:
+	case WebSocketClose::INTERNAL_ERROR:
+	case WebSocketClose::SERVICE_RESTART:
+	case WebSocketClose::TRY_AGAIN_LATER:
+	case WebSocketClose::BAD_GATEWAY:
+		return true;
+	default:
+		return false;
+	}
+}
+
 bool WebSocketServer::ReadControlFrame(int opcode, size_t sz) {
 	std::vector<uint8_t> payload;
 	payload.resize(sz);
@@ -522,8 +548,10 @@ bool WebSocketServer::ReadControlFrame(int opcode, size_t sz) {
 	} else if (opcode == (int)Opcode::CLOSE) {
 		if (payload.size() >= 2) {
 			uint16_t reason = (payload[0] << 8) | payload[1];
-			// Send back a close right away.
-			Close(WebSocketClose(reason));
+			// Send back a close right away - but not their code verbatim. NO_STATUS, ABNORMAL and
+			// friends describe how a connection ended locally and RFC 6455 7.4.1 says they must
+			// never go on the wire, so echoing one back would be our protocol violation, not theirs.
+			Close(IsCloseCodeSendable(reason) ? WebSocketClose(reason) : WebSocketClose::PROTOCOL_ERROR);
 		} else {
 			Close(WebSocketClose::NO_STATUS);
 		}