fix: retry partial writes in packet and push paths

attach.c

@@ -27,58 +27,63 @@ char const *clear_csi_data(void)
 	return "\033[999H\r\n";
 }
 
-/* Write buf to fd handling partial writes. Exit on failure. */
-void write_buf_or_fail(int fd, const void *buf, size_t count)
+/* Write all of buf to fd, retrying on short writes and EINTR.
+** Returns 0 on success, -1 on failure (errno is set). */
+static int write_all(int fd, const void *buf, size_t count)
 {
 	while (count != 0) {
 		ssize_t ret = write(fd, buf, count);
 
-		if (ret >= 0) {
+		if (ret > 0) {
 			buf = (const char *)buf + ret;
 			count -= ret;
 		} else if (ret < 0 && errno == EINTR)
 			continue;
 		else {
-			if (session_start) {
-				char age[32];
-				session_age(age, sizeof(age));
-				printf
-				    ("%s[%s: session '%s' write failed after %s]\r\n",
-				     clear_csi_data(), progname,
-				     session_shortname(), age);
-			} else {
-				printf("%s[%s: write failed]\r\n",
-				       clear_csi_data(), progname);
-			}
-			exit(1);
+			/* ret == 0 (no progress) or ret < 0 (real error) */
+			if (ret == 0)
+				errno = EIO;
+			return -1;
 		}
 	}
+	return 0;
+}
+
+/* Write buf to fd handling partial writes. Exit on failure. */
+void write_buf_or_fail(int fd, const void *buf, size_t count)
+{
+	if (write_all(fd, buf, count) < 0) {
+		if (session_start) {
+			char age[32];
+			session_age(age, sizeof(age));
+			printf
+			    ("%s[%s: session '%s' write failed after %s]\r\n",
+			     clear_csi_data(), progname,
+			     session_shortname(), age);
+		} else {
+			printf("%s[%s: write failed]\r\n",
+			       clear_csi_data(), progname);
+		}
+		exit(1);
+	}
 }
 
 /* Write pkt to fd. Exit on failure. */
 void write_packet_or_fail(int fd, const struct packet *pkt)
 {
-	while (1) {
-		ssize_t ret = write(fd, pkt, sizeof(struct packet));
-
-		if (ret == sizeof(struct packet))
-			return;
-		else if (ret < 0 && errno == EINTR)
-			continue;
-		else {
-			if (session_start) {
-				char age[32];
-				session_age(age, sizeof(age));
-				printf
-				    ("%s[%s: session '%s' write failed after %s]\r\n",
-				     clear_csi_data(), progname,
-				     session_shortname(), age);
-			} else {
-				printf("%s[%s: write failed]\r\n",
-				       clear_csi_data(), progname);
-			}
-			exit(1);
+	if (write_all(fd, pkt, sizeof(struct packet)) < 0) {
+		if (session_start) {
+			char age[32];
+			session_age(age, sizeof(age));
+			printf
+			    ("%s[%s: session '%s' write failed after %s]\r\n",
+			     clear_csi_data(), progname,
+			     session_shortname(), age);
+		} else {
+			printf("%s[%s: write failed]\r\n",
+			       clear_csi_data(), progname);
 		}
+		exit(1);
 	}
 }
 
@@ -497,11 +502,7 @@ int push_main()
 		}
 
 		pkt.len = len;
-		len = write(s, &pkt, sizeof(struct packet));
-		if (len != sizeof(struct packet)) {
-			if (len >= 0)
-				errno = EPIPE;
-
+		if (write_all(s, &pkt, sizeof(struct packet)) < 0) {
 			printf("%s: %s: %s\n", progname, sockname,
 			       strerror(errno));
 			return 1;
@@ -521,9 +522,9 @@ static int send_kill(int sig)
 	memset(&pkt, 0, sizeof(pkt));
 	pkt.type = MSG_KILL;
 	pkt.len = (unsigned char)sig;
-	ret = write(s, &pkt, sizeof(pkt));
+	ret = write_all(s, &pkt, sizeof(pkt));
 	close(s);
-	return (ret == sizeof(pkt)) ? 0 : -1;
+	return ret;
 }
 
 static int session_gone(void)

tests/preload_short_write.c

@@ -0,0 +1,58 @@
+#include <stdlib.h>
+#include <sys/stat.h>
+#include <sys/syscall.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+static int did_inject;
+
+static ssize_t real_write(int fd, const void *buf, size_t count)
+{
+	return syscall(SYS_write, fd, buf, count);
+}
+
+static int should_inject(int fd, size_t count)
+{
+	struct stat st;
+
+	if (did_inject || count <= 1)
+		return 0;
+	if (!getenv("ATCH_FAULT_SHORT_WRITE_ONCE"))
+		return 0;
+	if (fstat(fd, &st) < 0)
+		return 0;
+	return S_ISSOCK(st.st_mode);
+}
+
+static ssize_t short_write_impl(int fd, const void *buf, size_t count)
+{
+	if (should_inject(fd, count)) {
+		did_inject = 1;
+		return real_write(fd, buf, 1);
+	}
+	return real_write(fd, buf, count);
+}
+
+#ifdef __APPLE__
+#define DYLD_INTERPOSE(_replacement, _replacee) \
+	__attribute__((used)) static struct { \
+		const void *replacement; \
+		const void *replacee; \
+	} _interpose_##_replacee \
+	__attribute__((section("__DATA,__interpose"))) = { \
+		(const void *)(unsigned long)&_replacement, \
+		(const void *)(unsigned long)&_replacee \
+	}
+
+ssize_t interposed_write(int fd, const void *buf, size_t count)
+{
+	return short_write_impl(fd, buf, count);
+}
+
+DYLD_INTERPOSE(interposed_write, write);
+#else
+ssize_t write(int fd, const void *buf, size_t count)
+{
+	return short_write_impl(fd, buf, count);
+}
+#endif

tests/test.sh

@@ -720,6 +720,63 @@ assert_contains "no args: shows Usage:"              "Usage:" "$out"
 run "$ATCH" --help
 assert_contains "help: shows tail command"           "tail" "$out"
 
+# ── 23. fault injection: short socket writes are retried ───────────────────
+# Force the first packet write to a socket to complete with 1 byte.
+# Verifies write_all() retries correctly instead of treating short writes
+# as fatal.
+
+TESTS_DIR=$(CDPATH= cd -- "$(dirname "$0")" && pwd)
+OS_NAME=$(uname -s)
+
+FAULT_LIB=
+build_short_write_injector() {
+    [ -n "$FAULT_LIB" ] && return 0
+    case "$OS_NAME" in
+        Darwin)
+            FAULT_LIB="$TESTDIR/libshortwrite.dylib"
+            cc -dynamiclib -O2 -Wall -o "$FAULT_LIB" \
+                "$TESTS_DIR/preload_short_write.c" >/dev/null 2>&1 ;;
+        *)
+            FAULT_LIB="$TESTDIR/libshortwrite.so"
+            cc -shared -fPIC -O2 -Wall -o "$FAULT_LIB" \
+                "$TESTS_DIR/preload_short_write.c" -ldl >/dev/null 2>&1 ;;
+    esac
+}
+
+with_short_socket_write() {
+    build_short_write_injector || return 1
+    case "$OS_NAME" in
+        Darwin)
+            env DYLD_INSERT_LIBRARIES="$FAULT_LIB" \
+                DYLD_FORCE_FLAT_NAMESPACE=1 \
+                ATCH_FAULT_SHORT_WRITE_ONCE=1 "$@" ;;
+        *)
+            env LD_PRELOAD="$FAULT_LIB" \
+                ATCH_FAULT_SHORT_WRITE_ONCE=1 "$@" ;;
+    esac
+}
+
+"$ATCH" start short-push sh -c 'cat'
+wait_socket short-push
+out=$(printf 'short-write-marker\n' | with_short_socket_write \
+    "$ATCH" push short-push 2>&1)
+prc=$?
+assert_exit "fault: push retries short socket write" 0 "$prc"
+sleep 0.2
+assert_contains "fault: push data reaches session after short write" \
+    "short-write-marker" "$(cat "$HOME/.cache/atch/short-push.log" 2>/dev/null)"
+tidy short-push
+
+"$ATCH" start short-kill sleep 999
+wait_socket short-kill
+out=$(with_short_socket_write "$ATCH" kill short-kill 2>&1)
+krc=$?
+assert_exit "fault: kill retries short socket write" 0 "$krc"
+run "$ATCH" list
+assert_not_contains "fault: session is gone after short-write kill" \
+    "short-kill" "$out"
+"$ATCH" kill -f short-kill >/dev/null 2>&1 || true
+
 # ── summary ──────────────────────────────────────────────────────────────────
 
 printf "\n1..%d\n" "$T"