From a8cee0e3e20f050bf9562d243d0d636e66bfd59c Mon Sep 17 00:00:00 2001
From: Gavin Hayes <gavin@computoid.com>
Date: Wed, 8 May 2024 01:07:55 -0400
Subject: [PATCH] fix nt connect not initializing with
 SO_UPDATE_CONNECT_CONTEXT

---
 libc/sock/connect-nt.c            | 11 +++++-
 test/libc/sock/getpeername_test.c | 62 +++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/libc/sock/connect-nt.c b/libc/sock/connect-nt.c
index 3065c51c4..ce1951452 100644
--- a/libc/sock/connect-nt.c
+++ b/libc/sock/connect-nt.c
@@ -35,11 +35,14 @@
 #include "libc/sock/syscall_fd.internal.h"
 #include "libc/sock/wsaid.internal.h"
 #include "libc/sysv/consts/o.h"
+#include "libc/sysv/consts/sol.h"
 #include "libc/sysv/errfuns.h"
 
 #ifdef __x86_64__
 #include "libc/sock/yoink.inc"
 
+__msabi extern typeof(__sys_setsockopt_nt) *const __imp_setsockopt;
+
 struct ConnectArgs {
   const void *addr;
   uint32_t addrsize;
@@ -113,6 +116,8 @@ static textwindows int sys_connect_nt_impl(struct Fd *f, const void *addr,
       // return ETIMEDOUT if SO_SNDTIMEO elapsed
       // note that Linux will return EINPROGRESS
       errno = etimedout();
+    } else if (!rc) {
+      __imp_setsockopt(f->handle, SOL_SOCKET, kNtSoUpdateConnectContext, 0, 0);
     }
     return rc;
   }
@@ -131,7 +136,11 @@ static textwindows int sys_connect_nt_impl(struct Fd *f, const void *addr,
     ok = WSAGetOverlappedResult(f->handle, overlap, &dwBytes, false, &dwFlags);
     WSACloseEvent(overlap->hEvent);
     free(overlap);
-    return ok ? 0 : __winsockerr();
+    if (!ok) {
+      return __winsockerr();
+    }
+    __imp_setsockopt(f->handle, SOL_SOCKET, kNtSoUpdateConnectContext, 0, 0);
+    return 0;
   } else if (WSAGetLastError() == kNtErrorIoPending) {
     f->connect_op = overlap;
     return einprogress();
diff --git a/test/libc/sock/getpeername_test.c b/test/libc/sock/getpeername_test.c
index ae17983ce..e46f00612 100644
--- a/test/libc/sock/getpeername_test.c
+++ b/test/libc/sock/getpeername_test.c
@@ -90,3 +90,65 @@ TEST(getpeername, worksAfterAcceptingOnParent) {
   ASSERT_SYS(0, 0, close(3));
   WAIT(exit, 0);
 }
+
+TEST(getpeername, worksAfterConnectingOnFork) {
+  char buf[16] = {0};
+  uint32_t addrsize = sizeof(struct sockaddr_in);
+  struct sockaddr_in addr = {
+      .sin_family = AF_INET,
+      .sin_addr.s_addr = htonl(0x7f000001),
+  };
+  ASSERT_SYS(0, 3, socket(AF_INET, SOCK_STREAM, IPPROTO_TCP));
+  ASSERT_SYS(0, 0, bind(3, (struct sockaddr *)&addr, sizeof(addr)));
+  ASSERT_SYS(0, 0, getsockname(3, (struct sockaddr *)&addr, &addrsize));
+  ASSERT_SYS(0, 0, listen(3, SOMAXCONN));
+  ASSERT_SYS(0, 4, socket(AF_INET, SOCK_STREAM, IPPROTO_TCP));
+  SPAWN(fork);
+  ASSERT_SYS(0, 0, close(3));
+  ASSERT_SYS(0, 0, connect(4, (struct sockaddr *)&addr, sizeof(addr)));
+  struct sockaddr_storage out = {0};
+  uint32_t out_size = sizeof(out);
+  ASSERT_SYS(0, 0, getpeername(4, (struct sockaddr *)&out, &out_size));
+  EXPECT_GE(sizeof(struct sockaddr_in), out_size);
+  EXPECT_EQ(AF_INET, ((struct sockaddr_in *)&out)->sin_family);
+  EXPECT_EQ(htonl(0x7f000001), ((struct sockaddr_in *)&out)->sin_addr.s_addr);
+  ASSERT_SYS(0, 5, send(4, "hello", 5, 0));
+  PARENT();
+  ASSERT_SYS(0, 0, close(4));
+  ASSERT_SYS(0, 4, accept(3, (struct sockaddr *)&addr, &addrsize));
+  ASSERT_SYS(0, 5, read(4, buf, 16));
+  ASSERT_STREQ("hello", buf);
+  ASSERT_SYS(0, 0, close(4));
+  ASSERT_SYS(0, 0, close(3));
+  WAIT(exit, 0);
+}
+
+TEST(getpeername, worksAfterConnectingOnParent) {
+  char buf[16] = {0};
+  uint32_t addrsize = sizeof(struct sockaddr_in);
+  struct sockaddr_in addr = {
+      .sin_family = AF_INET,
+      .sin_addr.s_addr = htonl(0x7f000001),
+  };
+  ASSERT_SYS(0, 3, socket(AF_INET, SOCK_STREAM, IPPROTO_TCP));
+  ASSERT_SYS(0, 0, bind(3, (struct sockaddr *)&addr, sizeof(addr)));
+  ASSERT_SYS(0, 0, getsockname(3, (struct sockaddr *)&addr, &addrsize));
+  ASSERT_SYS(0, 0, listen(3, SOMAXCONN));
+  SPAWN(fork);
+  ASSERT_SYS(0, 4, accept(3, (struct sockaddr *)&addr, &addrsize));
+  ASSERT_SYS(0, 5, read(4, buf, 16));
+  ASSERT_STREQ("hello", buf);
+  PARENT();
+  ASSERT_SYS(0, 0, close(3));
+  ASSERT_SYS(0, 3, socket(AF_INET, SOCK_STREAM, IPPROTO_TCP));
+  ASSERT_SYS(0, 0, connect(3, (struct sockaddr *)&addr, sizeof(addr)));
+  struct sockaddr_storage out = {0};
+  uint32_t out_size = sizeof(out);
+  ASSERT_SYS(0, 0, getpeername(3, (struct sockaddr *)&out, &out_size));
+  EXPECT_GE(sizeof(struct sockaddr_in), out_size);
+  EXPECT_EQ(AF_INET, ((struct sockaddr_in *)&out)->sin_family);
+  EXPECT_EQ(htonl(0x7f000001), ((struct sockaddr_in *)&out)->sin_addr.s_addr);
+  ASSERT_SYS(0, 5, send(3, "hello", 5, 0));
+  ASSERT_SYS(0, 0, close(3));
+  WAIT(exit, 0);
+}