net: emulate Windows PEEK and WAITALL receives

Co-authored-by: Firas Ben Mbarek <firasbenmbarek@outlook.com> Co-authored-by: rudy-07 <kanhaiffco2007@gmail.com> Co-authored-by: Matthew Hurley <171001439+theantipopau@users.noreply.github.com> Co-authored-by: usopp-sama <go4av05.pvtx@gmail.com> Co-authored-by: wruwami <wruwami@users.noreply.github.com>
2026-10-25 04:16:55 +01:00 · 2026-10-03 17:13:48 +02:00
parent 19a9464b7d
commit e7f6327068
2 changed files with 429 additions and 39 deletions
@@ -849,20 +849,38 @@ using SocketIoLength                                = size_t;
 #endif

 struct SocketTransport {
-	explicit SocketTransport(NativeSocket value): socket(value) {}
+	SocketTransport(NativeSocket value, int socket_type): socket(value), type(socket_type) {}
 	~SocketTransport() { Close(); }
+	bool HasPendingReceive() const {
+#if defined(_WIN32)
+		return buffered_bytes.load(std::memory_order_relaxed) != 0 ||
+		       receive_shutdown.load(std::memory_order_relaxed);
+#else
+		return false;
+#endif
+	}
 	int Close() {
-		if (socket == INVALID_NATIVE_SOCKET) {
+#if defined(_WIN32)
+		receive_shutdown.store(true, std::memory_order_relaxed);
+#endif
+		const auto value = socket.exchange(INVALID_NATIVE_SOCKET);
+		if (value == INVALID_NATIVE_SOCKET) {
 			return 0;
 		}
-		const auto value = std::exchange(socket, INVALID_NATIVE_SOCKET);
 #if defined(_WIN32)
 		return closesocket(value);
 #else
 		return ::close(value);
 #endif
 	}
-	NativeSocket socket;
+	std::atomic<NativeSocket> socket;
+	const int                type;
+#if defined(_WIN32)
+	std::mutex         receive_mutex;
+	std::vector<char>  receive_buffer;
+	std::atomic_size_t buffered_bytes   = 0;
+	std::atomic_bool   receive_shutdown = false;
+#endif
 };

 struct P2pEndpoint {
@@ -1164,7 +1182,8 @@ static int ConvertHostSockaddr(const sockaddr_storage* addr, SocketLength addrle
 	return 0;
 }

-static bool GetSocketBackend(int guest_fd, NativeSocket* out, SocketSlot* state = nullptr) {
+static bool GetSocketBackend(int guest_fd, NativeSocket* out, SocketSlot* state = nullptr,
+                             std::shared_ptr<SocketTransport>* transport = nullptr) {
 	EXIT_IF(out == nullptr);

 	if (guest_fd < 0 || guest_fd >= SOCKET_FD_MAX) {
@@ -1188,6 +1207,9 @@ static bool GetSocketBackend(int guest_fd, NativeSocket* out, SocketSlot* state
 	if (state != nullptr) {
 		*state = slot;
 	}
+	if (transport != nullptr) {
+		*transport = slot.transport;
+	}
 	return true;
 }

@@ -1234,6 +1256,63 @@ static int AllocSocketFd(const std::shared_ptr<SocketTransport>& transport, bool
 	return -1;
 }

+static int SelectSockets(int nfds, fd_set* readfds, fd_set* writefds, fd_set* exceptfds,
+                         timeval* timeout,
+                         const std::vector<std::shared_ptr<SocketTransport>>& stream_reads) {
+#if defined(_WIN32)
+	if (!stream_reads.empty()) {
+		const auto requested_read   = *readfds;
+		const auto requested_write  = writefds != nullptr ? *writefds : fd_set {};
+		const auto requested_except = exceptfds != nullptr ? *exceptfds : fd_set {};
+		std::optional<std::chrono::steady_clock::time_point> deadline;
+		if (timeout != nullptr) {
+			deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout->tv_sec) +
+			           std::chrono::microseconds(timeout->tv_usec);
+		}
+		for (;;) {
+			*readfds = requested_read;
+			if (writefds != nullptr) {
+				*writefds = requested_write;
+			}
+			if (exceptfds != nullptr) {
+				*exceptfds = requested_except;
+			}
+			const bool ready = std::any_of(stream_reads.begin(), stream_reads.end(),
+			    [](const auto& transport) { return transport->HasPendingReceive(); });
+			// A concurrent peek can move native bytes into the guest buffer after this check.
+			// Bound native waits so both readiness sources are rechecked without holding receive locks.
+			auto wait = std::chrono::microseconds(ready ? 0 : 10'000);
+			if (deadline) {
+				const auto remaining = std::chrono::duration_cast<std::chrono::microseconds>(
+				    *deadline - std::chrono::steady_clock::now());
+				wait = std::min(wait, std::max(remaining, std::chrono::microseconds::zero()));
+			}
+			timeval slice {};
+			slice.tv_usec = static_cast<long>(wait.count());
+			int result = ::select(nfds, readfds, writefds, exceptfds, &slice);
+			if (result < 0) {
+				return result;
+			}
+			for (const auto& transport: stream_reads) {
+				const auto socket = transport->socket.load();
+				if (socket == INVALID_NATIVE_SOCKET) {
+					WSASetLastError(WSAENOTSOCK);
+					return SOCKET_ERROR;
+				}
+				if (transport->HasPendingReceive() && !FD_ISSET(socket, readfds)) {
+					FD_SET(socket, readfds);
+					result++;
+				}
+			}
+			if (result != 0 || (deadline && std::chrono::steady_clock::now() >= *deadline)) {
+				return result;
+			}
+		}
+	}
+#endif
+	return ::select(nfds, readfds, writefds, exceptfds, timeout);
+}
+
 static bool GuestFdIsSet(const void* fds, int fd) {
 	if (fds == nullptr || fd < 0 || fd >= SOCKET_FD_MAX) {
 		return false;
@@ -1621,6 +1700,7 @@ int KYTY_SYSV_ABI EpollWait(int eid, NetEpollEvent* events, int maxevents, int t
 	FD_ZERO(&host_except);

 	std::vector<HostRegistration> host_registrations;
+	std::vector<std::shared_ptr<SocketTransport>> stream_reads;
 	host_registrations.reserve(registrations.size());
 	int host_nfds = 0;
 	for (const auto& registration: registrations) {
@@ -1641,6 +1721,11 @@ int KYTY_SYSV_ABI EpollWait(int eid, NetEpollEvent* events, int maxevents, int t
 #endif
 		if ((registration.event.events & EPOLL_IN) != 0) {
 			FD_SET(socket, &host_read);
+#if defined(_WIN32)
+			if (state.transport->type == SOCK_STREAM) {
+				stream_reads.push_back(state.transport);
+			}
+#endif
 		}
 		if ((registration.event.events & EPOLL_OUT) != 0) {
 			FD_SET(socket, &host_write);
@@ -1660,8 +1745,8 @@ int KYTY_SYSV_ABI EpollWait(int eid, NetEpollEvent* events, int maxevents, int t
 		host_timeout.tv_usec = timeout % 1000000;
 		host_timeout_ptr     = &host_timeout;
 	}
-
-	const int result = ::select(host_nfds, &host_read, &host_write, &host_except, host_timeout_ptr);
+	const int result = SelectSockets(host_nfds, &host_read, &host_write, &host_except,
+	                                 host_timeout_ptr, stream_reads);
 	if (result < 0) {
 		return SetHostSocketError();
 	}
@@ -1727,7 +1812,12 @@ int KYTY_SYSV_ABI SocketClose(int s) {
 	}
 	RemoveSocketFromEpolls(s);

-	if (slot.transport.use_count() == 1 && slot.transport->Close() != 0) {
+	bool close_transport = slot.transport.use_count() == 1;
+#if defined(_WIN32)
+	// Retained receive/readiness state must not defer native TCP close cancellation.
+	close_transport |= !slot.p2p && slot.transport->type == SOCK_STREAM;
+#endif
+	if (close_transport && slot.transport->Close() != 0) {
 		return NET_ERROR_EBADF;
 	}

@@ -1762,7 +1852,7 @@ int KYTY_SYSV_ABI Socket(int family, int type, int protocol) {
 		return SetHostSocketError();
 	}

-	auto transport = std::make_shared<SocketTransport>(socket);
+	auto transport = std::make_shared<SocketTransport>(socket, p2p ? SOCK_DGRAM : type);
 	if (p2p) {
 		// Logical sockets keep their own blocking mode; the shared transport never blocks.
 #if defined(_WIN32)
@@ -1945,7 +2035,7 @@ int KYTY_SYSV_ABI Accept(int s, void* addr, uint32_t* addrlen) {
 		return SetHostSocketError();
 	}

-	auto      transport = std::make_shared<SocketTransport>(accepted);
+	auto      transport = std::make_shared<SocketTransport>(accepted, SOCK_STREAM);
 	const int fd        = AllocSocketFd(transport);
 	if (fd < 0) {
 		*Posix::GetErrorAddr() = Posix::POSIX_EMFILE;
@@ -1974,8 +2064,9 @@ int KYTY_SYSV_ABI Shutdown(int s, int how) {
 	     "\t how = %d\n",
 	     s, how);

-	NativeSocket socket = INVALID_NATIVE_SOCKET;
-	if (!GetSocketBackend(s, &socket)) {
+	NativeSocket                    socket = INVALID_NATIVE_SOCKET;
+	std::shared_ptr<SocketTransport> transport;
+	if (!GetSocketBackend(s, &socket, nullptr, &transport)) {
 		return -1;
 	}

@@ -1988,6 +2079,9 @@ int KYTY_SYSV_ABI Shutdown(int s, int how) {
 	if (::shutdown(socket, how) == SOCKET_ERROR) {
 		return SetHostSocketError();
 	}
+	if (how == SD_RECEIVE || how == SD_BOTH) {
+		transport->receive_shutdown.store(true, std::memory_order_relaxed);
+	}

 	return 0;
 #else
@@ -2205,6 +2299,77 @@ int64_t KYTY_SYSV_ABI Sendto(int s, const void* buf, uint64_t len, int flags, co
 	return result;
 }

+#if defined(_WIN32)
+static int ReceiveStream(SocketTransport& transport, char* output, int length, int flags,
+                         bool dontwait) {
+	std::unique_lock lock(transport.receive_mutex, std::defer_lock);
+	if (dontwait) {
+		if (!lock.try_lock()) {
+			return SetGuestSocketError(Posix::POSIX_EWOULDBLOCK);
+		}
+	} else {
+		lock.lock();
+	}
+	if (transport.receive_shutdown.load(std::memory_order_relaxed)) {
+		return SetGuestSocketError(Posix::POSIX_ESHUTDOWN);
+	}
+	auto&           buffered = transport.receive_buffer;
+	const bool      peek     = (flags & MSG_PEEK) != 0;
+	const bool      waitall  = (flags & MSG_WAITALL) != 0;
+	int             copied   = static_cast<int>(std::min<size_t>(buffered.size(), length));
+	if (copied != 0) {
+		std::memcpy(output, buffered.data(), copied);
+		if (!peek) {
+			buffered.erase(buffered.begin(), buffered.begin() + copied);
+			transport.buffered_bytes.store(buffered.size(), std::memory_order_relaxed);
+		}
+	}
+	if (copied == length || (copied != 0 && !waitall)) {
+		return copied;
+	}
+
+	// Winsock rejects PEEK | WAITALL. Consume into transport-owned storage when peeking,
+	// and let ordinary recv report EOF, timeouts, or nonblocking partial progress.
+	flags &= ~(MSG_PEEK | MSG_WAITALL);
+	do {
+		if (transport.receive_shutdown.load(std::memory_order_relaxed)) {
+			return copied != 0 ? copied : SetGuestSocketError(Posix::POSIX_ESHUTDOWN);
+		}
+		const auto socket = transport.socket.load();
+		if (dontwait) {
+			fd_set readable {};
+			FD_SET(socket, &readable);
+			timeval immediate {};
+			const int ready = ::select(0, &readable, nullptr, nullptr, &immediate);
+			if (ready <= 0) {
+				return copied != 0 ? copied : (ready < 0 ? SetHostSocketError() :
+				    SetGuestSocketError(Posix::POSIX_EWOULDBLOCK));
+			}
+		}
+		const auto previous_size = buffered.size();
+		const int  requested     = peek ? std::min(length - copied, 64 * 1024) : length - copied;
+		char*      destination   = output + copied;
+		if (peek) {
+			buffered.resize(previous_size + requested);
+			destination = buffered.data() + previous_size;
+		}
+		const int received = ::recv(socket, destination, requested, flags);
+		if (peek) {
+			if (received > 0) {
+				std::memcpy(output + copied, destination, received);
+			}
+			buffered.resize(previous_size + std::max(received, 0));
+			transport.buffered_bytes.store(buffered.size(), std::memory_order_relaxed);
+		}
+		if (received <= 0) {
+			return copied != 0 || received == 0 ? copied : SetHostSocketError();
+		}
+		copied += received;
+	} while (waitall && copied < length);
+	return copied;
+}
+#endif
+
 int64_t KYTY_SYSV_ABI Recv(int s, void* buf, uint64_t len, int flags) {
 	return Recvfrom(s, buf, len, flags, nullptr, nullptr);
 }
@@ -2227,8 +2392,9 @@ int64_t KYTY_SYSV_ABI Recvfrom(int s, void* buf, uint64_t len, int flags, void*
 		return -1;
 	}

-	NativeSocket socket = INVALID_NATIVE_SOCKET;
-	if (!GetSocketBackend(s, &socket)) {
+	NativeSocket                    socket = INVALID_NATIVE_SOCKET;
+	std::shared_ptr<SocketTransport> transport;
+	if (!GetSocketBackend(s, &socket, nullptr, &transport)) {
 		return -1;
 	}

@@ -2242,11 +2408,27 @@ int64_t KYTY_SYSV_ABI Recvfrom(int s, void* buf, uint64_t len, int flags, void*
 	sockaddr_storage host_addr {};
 	SocketLength     host_addrlen = sizeof(host_addr);
 	int64_t          result       = 0;
-	if (addr == nullptr) {
-		result = ::recv(socket, static_cast<char*>(buf), host_len, host_flags);
-	} else {
-		result = ::recvfrom(socket, static_cast<char*>(buf), host_len, host_flags,
-		                    reinterpret_cast<sockaddr*>(&host_addr), &host_addrlen);
+#if defined(_WIN32)
+	if (transport->type == SOCK_STREAM) {
+		if (addr != nullptr &&
+		    ::getpeername(socket, reinterpret_cast<sockaddr*>(&host_addr), &host_addrlen) != 0) {
+			return SetHostSocketError();
+		}
+		result = ReceiveStream(*transport, static_cast<char*>(buf), host_len, host_flags,
+		                       (flags & 0x80) != 0);
+		if (result < 0) {
+			return result;
+		}
+	} else
+#endif
+	{
+		transport.reset();
+		if (addr == nullptr) {
+			result = ::recv(socket, static_cast<char*>(buf), host_len, host_flags);
+		} else {
+			result = ::recvfrom(socket, static_cast<char*>(buf), host_len, host_flags,
+			                    reinterpret_cast<sockaddr*>(&host_addr), &host_addrlen);
+		}
 	}
 #if defined(_WIN32)
 	if (result < 0 && WSAGetLastError() == WSAEMSGSIZE) {
@@ -2303,6 +2485,7 @@ int KYTY_SYSV_ABI Select(int nfds, void* readfds, void* writefds, void* exceptfd
 	FD_ZERO(&host_except);

 	std::vector<std::pair<int, NativeSocket>> descriptors;
+	std::vector<std::shared_ptr<SocketTransport>> stream_reads;
 	int host_nfds = 0;
 	for (int fd = 0; fd < nfds; fd++) {
 		const bool read = GuestFdIsSet(readfds, fd);
@@ -2311,13 +2494,14 @@ int KYTY_SYSV_ABI Select(int nfds, void* readfds, void* writefds, void* exceptfd
 		if (!read && !write && !except) {
 			continue;
 		}
-		NativeSocket socket = INVALID_NATIVE_SOCKET;
+		NativeSocket                    socket = INVALID_NATIVE_SOCKET;
+		std::shared_ptr<SocketTransport> transport;
 #if !defined(_WIN32)
 		if (fd < 3) {
 			socket = fd;
 		} else
 #endif
-		if (!GetSocketBackend(fd, &socket)) {
+		if (!GetSocketBackend(fd, &socket, nullptr, &transport)) {
 			return -1;
 		}
 #if defined(_WIN32)
@@ -2335,6 +2519,11 @@ int KYTY_SYSV_ABI Select(int nfds, void* readfds, void* writefds, void* exceptfd
 		descriptors.emplace_back(fd, socket);
 		if (read) {
 			FD_SET(socket, &host_read);
+#if defined(_WIN32)
+			if (transport->type == SOCK_STREAM) {
+				stream_reads.push_back(std::move(transport));
+			}
+#endif
 		}
 		if (write) {
 			FD_SET(socket, &host_write);
@@ -2365,9 +2554,10 @@ int KYTY_SYSV_ABI Select(int nfds, void* readfds, void* writefds, void* exceptfd
 	} else
 #endif
 	{
-		result = ::select(host_nfds, readfds != nullptr ? &host_read : nullptr,
-		                  writefds != nullptr ? &host_write : nullptr,
-		                  exceptfds != nullptr ? &host_except : nullptr, host_timeout_ptr);
+		result = SelectSockets(host_nfds, readfds != nullptr ? &host_read : nullptr,
+		                       writefds != nullptr ? &host_write : nullptr,
+		                       exceptfds != nullptr ? &host_except : nullptr, host_timeout_ptr,
+		                       stream_reads);
 		if (result < 0) {
 			return SetHostSocketError();
 		}
@@ -2376,15 +2566,19 @@ int KYTY_SYSV_ABI Select(int nfds, void* readfds, void* writefds, void* exceptfd
 	GuestFdZero(readfds, nfds);
 	GuestFdZero(writefds, nfds);
 	GuestFdZero(exceptfds, nfds);
+	result = 0;
 	for (const auto& [fd, socket]: descriptors) {
 		if (FD_ISSET(socket, &host_read)) {
 			GuestFdSet(readfds, fd);
+			result++;
 		}
 		if (FD_ISSET(socket, &host_write)) {
 			GuestFdSet(writefds, fd);
+			result++;
 		}
 		if (FD_ISSET(socket, &host_except)) {
 			GuestFdSet(exceptfds, fd);
+			result++;
 		}
 	}
 	return result;