From d2c9ed3d1e7c1a318e6fd018eb0e0f3ba5ee3365 Mon Sep 17 00:00:00 2001
From: Kevin Svetlitski <svetlitski@meta.com>
Date: Thu, 10 Aug 2023 10:43:42 -0700
Subject: [PATCH] Ensure short `read(2)`s/`write(2)`s are properly handled by
 IO utilities

`read(2)` and `write(2)` may read or write fewer bytes than were
requested. In order to robustly ensure that all of the requested bytes
are read/written, these edge-cases must be handled.
---
 include/jemalloc/internal/malloc_io.h | 35 +++++++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/include/jemalloc/internal/malloc_io.h b/include/jemalloc/internal/malloc_io.h
index 0afb0429..91e7b2ba 100644
--- a/include/jemalloc/internal/malloc_io.h
+++ b/include/jemalloc/internal/malloc_io.h
@@ -68,7 +68,7 @@ void malloc_cprintf(write_cb_t *write_cb, void *cbopaque, const char *format,
 void malloc_printf(const char *format, ...) JEMALLOC_FORMAT_PRINTF(1, 2);
 
 static inline ssize_t
-malloc_write_fd(int fd, const void *buf, size_t count) {
+malloc_write_fd_syscall(int fd, const void *buf, size_t count) {
 #if defined(JEMALLOC_USE_SYSCALL) && defined(SYS_write)
 	/*
 	 * Use syscall(2) rather than write(2) when possible in order to avoid
@@ -90,7 +90,22 @@ malloc_write_fd(int fd, const void *buf, size_t count) {
 }
 
 static inline ssize_t
-malloc_read_fd(int fd, void *buf, size_t count) {
+malloc_write_fd(int fd, const void *buf, size_t count) {
+	size_t bytes_written = 0;
+	do {
+		ssize_t result = malloc_write_fd_syscall(fd,
+		    &((const byte_t *)buf)[bytes_written],
+		    count - bytes_written);
+		if (result < 0) {
+			return result;
+		}
+		bytes_written += result;
+	} while (bytes_written < count);
+	return bytes_written;
+}
+
+static inline ssize_t
+malloc_read_fd_syscall(int fd, void *buf, size_t count) {
 #if defined(JEMALLOC_USE_SYSCALL) && defined(SYS_read)
 	long result = syscall(SYS_read, fd, buf, count);
 #else
@@ -103,4 +118,20 @@ malloc_read_fd(int fd, void *buf, size_t count) {
 	return (ssize_t)result;
 }
 
+static inline ssize_t
+malloc_read_fd(int fd, void *buf, size_t count) {
+	size_t bytes_read = 0;
+	do {
+		ssize_t result = malloc_read_fd_syscall(fd,
+		    &((byte_t *)buf)[bytes_read], count - bytes_read);
+		if (result < 0) {
+			return result;
+		} else if (result == 0) {
+			break;
+		}
+		bytes_read += result;
+	} while (bytes_read < count);
+	return bytes_read;
+}
+
 #endif /* JEMALLOC_INTERNAL_MALLOC_IO_H */