init/aws-nitro: support partial read/write ethernet packet forwarding

Ensure the requested number of bytes is forwarded to/from the host
vsock providing the network access and the guest TAP device routing the
application network traffic.

To write this code I used Cursor and the claude-4.6-opus-high model.

Signed-off-by: Jake Correnti <jakecorrenti+github@proton.me>

Author: jakecorrenti

init/aws-nitro/device/net_tap_afvsock.c

@@ -31,6 +31,58 @@
 #define TUN_DEV_MAJOR 10
 #define TUN_DEV_MINOR 200
 
+#define ETH_HEADER_LEN 14
+
+/*
+ * Read exactly n bytes into the buffer, retrying on partial reads.
+ * Returns n on success, 0 on clean EOF, or -1 on error.
+ */
+static ssize_t read_exact(int fd, void *buf, size_t n)
+{
+    size_t total = 0;
+
+    while (total < n) {
+        ssize_t r = read(fd, (char *)buf + total, n - total);
+        if (r < 0) {
+            if (errno == EINTR)
+                continue;
+            return -1;
+        } else if (r == 0) {
+            if (total > 0) {
+                errno = EIO;
+                return -1;
+            }
+            return 0;
+        }
+        total += r;
+    }
+    return (ssize_t)total;
+}
+
+/*
+ * Write exactly n bytes from the buffer to the fd, retrying on partial writes.
+ * Returns n on success, or -1 on error.
+ */
+static ssize_t write_all(int fd, const void *buf, size_t n)
+{
+    size_t total = 0;
+
+    while (total < n) {
+        ssize_t w = write(fd, (const char *)buf + total, n - total);
+        if (w <= 0) {
+            if (w < 0 && errno == EINTR)
+                continue;
+
+            if (w == 0)
+                errno = EIO;
+
+            return -1;
+        }
+        total += w;
+    }
+    return (ssize_t)total;
+}
+
 /*
  * Forward ethernet packets to/from the host vsock providing network access and
  * the guest TAP device routing application network traffic.
@@ -43,12 +95,12 @@ static int tap_vsock_forward(int tun_fd, int vsock_fd, int shutdown_fd,
     bool event_found;
     struct ifreq ifr;
     int ret, sock_fd;
-    unsigned int sz;
+    uint32_t sz;
     ssize_t nread;
 
     /*
      * Fetch the TAP device's Maximum Transfer Unit (MTU) and allocate a buffer
-     * in that size to transfer ethernet frames to/from the host.
+     * large enough for a full eth frame (MTU + eth header).
      */
     sock_fd = socket(AF_INET, SOCK_DGRAM, 0);
     if (sock_fd < 0) {
@@ -68,17 +120,19 @@ static int tap_vsock_forward(int tun_fd, int vsock_fd, int shutdown_fd,
 
     close(sock_fd);
 
-    buf = (unsigned char *)malloc(ifr.ifr_mtu);
+    unsigned int buffer_size = ifr.ifr_mtu + ETH_HEADER_LEN;
+    buf = (unsigned char *)malloc(buffer_size);
     if (buf == NULL) {
         perror("allocate buffer for TAP/vsock communication");
         exit(-1);
     }
 
     // Forward the MTU to the host for it to allocate a corresponding buffer.
-    ret = write(vsock_fd, (void *)&ifr.ifr_mtu, sizeof(int));
-    if (ret < sizeof(int)) {
+    // The host is expecting a 32-bit header.
+    uint32_t mtu = ifr.ifr_mtu;
+    if (write_all(vsock_fd, &mtu, sizeof(mtu)) < 0) {
         perror("write TAP device MTU to host");
-        exit(-errno);
+        exit(EXIT_FAILURE);
     }
 
     pfds[0].fd = vsock_fd;
@@ -97,25 +151,62 @@ static int tap_vsock_forward(int tun_fd, int vsock_fd, int shutdown_fd,
         event_found = false;
         // Event on vsock. Read the frame and write it to the TAP device.
         if (pfds[0].revents & POLLIN) {
-            nread = read(vsock_fd, &sz, 4);
-            if (nread != 4)
+            nread = read_exact(vsock_fd, &sz, sizeof(sz));
+            if (nread < 0) {
+                perror("failure to read header");
+                exit(EXIT_FAILURE);
+            } else if (nread == 0) {
+                // vsock connection is closed
                 exit(0);
+            }
+
+            unsigned int len = ntohl(sz);
+            if (len > buffer_size) {
+                fprintf(stderr, "frame size %u exceeds buffer %u\n", len,
+                        buffer_size);
+                exit(EXIT_FAILURE);
+            }
 
-            unsigned int len = htonl(sz);
+            nread = read_exact(vsock_fd, buf, len);
+            if (nread != (ssize_t)len) {
+                if (nread == 0)
+                    errno = EIO;
+                perror("failure to read from buffer");
+                exit(EXIT_FAILURE);
+            }
 
-            nread = read(vsock_fd, buf, len);
-            write(tun_fd, buf, nread);
+            // TAP device is expected to write exactly one eth frame
+            ssize_t w;
+            do {
+                w = write(tun_fd, buf, nread);
+            } while (w < 0 && errno == EINTR);
+
+            if (w != nread) {
+                if (w >= 0)
+                    errno = EIO;
+                perror("failure to write buffer to tap device");
+                exit(EXIT_FAILURE);
+            }
 
             event_found = true;
         }
 
         // Event on the TAP device. Read the frame and write it to the vsock.
         if (pfds[1].revents & POLLIN) {
-            nread = read(tun_fd, buf, ifr.ifr_mtu);
+            do {
+                nread = read(tun_fd, buf, buffer_size);
+            } while (nread < 0 && errno == EINTR);
+
             if (nread > 0) {
                 sz = htonl(nread);
-                write(vsock_fd, (void *)&sz, 4);
-                write(vsock_fd, buf, nread);
+                if (write_all(vsock_fd, &sz, sizeof(sz)) < 0) {
+                    perror("unable to write header to vsock");
+                    exit(EXIT_FAILURE);
+                }
+                if (write_all(vsock_fd, buf, nread) < 0) {
+                    perror("unable to write buffer to vsock");
+                    exit(EXIT_FAILURE);
+                }
             }
 
             event_found = true;