From 12826c9586f19b63198879c4bb93eb6395d0bfe1 Mon Sep 17 00:00:00 2001
From: Viktor Dukhovni <viktor@twosigma.com>
Date: Wed, 6 Nov 2019 20:27:58 -0500
Subject: [PATCH] Handle partial writes on non-blocking sockets

Now that we're using krb5_net_write() with non-blocking sockets in
ipropd_master, we MUST correctly account for partial writes.

Therefore, roken net_write() called from krb5_net_write() now
returns the number of bytes written when the socket error was
EWOULDBLOCK (or EAGAIN).

Also, fix potential issue on Windows, where errno was used instead
of rk_SOCKET_ERRNO whether or not we used _write() or send().
---
 lib/roken/net_write.c | 38 +++++++++++++++++++++++++++++++-------
 1 file changed, 31 insertions(+), 7 deletions(-)

diff --git a/lib/roken/net_write.c b/lib/roken/net_write.c
index e1cfa9907..1e8361999 100644
--- a/lib/roken/net_write.c
+++ b/lib/roken/net_write.c
@@ -36,7 +36,9 @@
 #include "roken.h"
 
 /*
- * Like write but never return partial data.
+ * Like write but blocking sockets never return partial data, i.e. we retry on
+ * EINTR.  With non-blocking sockets (EWOULDBLOCK or EAGAIN) we return the
+ * number of bytes written.
  */
 
 #ifndef _WIN32
@@ -51,10 +53,17 @@ net_write (rk_socket_t fd, const void *buf, size_t nbytes)
     while (rem > 0) {
 	count = write (fd, cbuf, rem);
 	if (count < 0) {
-	    if (errno == EINTR)
+            switch (errno) {
+            case EINTR:
 		continue;
-	    else
+#if defined(EAGAIN) && EAGAIN != EWOULDBLOCK
+            case EAGAIN:
+#endif
+            case EWOULDBLOCK:
+                return nbytes - rem;
+            default:
 		return count;
+            }
 	}
 	cbuf += count;
 	rem -= count;
@@ -93,10 +102,25 @@ net_write(rk_socket_t sock, const void *buf, size_t nbytes)
 	count = send (sock, cbuf, rem, 0);
 #endif
 	if (count < 0) {
-	    if (errno == EINTR)
-		continue;
-	    else
-		return count;
+            if (!use_write) {
+                switch (rk_SOCK_ERRNO) {
+                case WSAEINTR:
+                    continue;
+                case WSAEWOULDBLOCK:
+                    return nbytes - rem;
+                default:
+                    return count;
+                }
+            } else {
+                switch (errno) {
+                case EINTR:
+                    continue;
+                case EWOULDBLOCK:
+                    return nbytes - rem;
+                default:
+                    return count;
+                }
+            }
 	}
 	cbuf += count;
 	rem -= count;