Windows: add uv_msafd_poll, to support overlapped socket polling

joyent · May 2, 2012 · 19aca7a · 19aca7a
1 parent 9f0dc26
commit 19aca7a
Show file tree

Hide file tree

Showing 4 changed files with 138 additions and 1 deletion.
diff --git a/include/uv-private/uv-win.h b/include/uv-private/uv-win.h
@@ -128,6 +128,24 @@ typedef int (WSAAPI* LPFN_WSARECVFROM)
              LPWSAOVERLAPPED overlapped,
              LPWSAOVERLAPPED_COMPLETION_ROUTINE completion_routine);
 
+#ifndef _NTDEF_
+  typedef LONG NTSTATUS;
+  typedef NTSTATUS *PNTSTATUS;
+#endif
+
+typedef struct _AFD_POLL_HANDLE_INFO {
+  HANDLE Handle;
+  ULONG Events;
+  NTSTATUS Status;
+} AFD_POLL_HANDLE_INFO, *PAFD_POLL_HANDLE_INFO;
+
+typedef struct _AFD_POLL_INFO {
+  LARGE_INTEGER Timeout;
+  ULONG NumberOfHandles;
+  ULONG Exclusive;
+  AFD_POLL_HANDLE_INFO Handles[1];
+} AFD_POLL_INFO, *PAFD_POLL_INFO;
+
 
 /**
  * It should be possible to cast uv_buf_t[] to WSABUF[]

diff --git a/src/win/internal.h b/src/win/internal.h
@@ -349,6 +349,9 @@ int WSAAPI uv_wsarecvfrom_workaround(SOCKET socket, WSABUF* buffers,
     int* addr_len, WSAOVERLAPPED *overlapped,
     LPWSAOVERLAPPED_COMPLETION_ROUTINE completion_routine);
 
+int WSAAPI uv_msafd_poll(SOCKET socket, AFD_POLL_INFO* info,
+    OVERLAPPED* overlapped);
+
 /* Whether ipv6 is supported */
 extern int uv_allow_ipv6;
 

diff --git a/src/win/winsock.c b/src/win/winsock.c
@@ -468,3 +468,90 @@ int WSAAPI uv_wsarecvfrom_workaround(SOCKET socket, WSABUF* buffers,
     return SOCKET_ERROR;
   }
 }
+
+
+int WSAAPI uv_msafd_poll(SOCKET socket, AFD_POLL_INFO* info,
+    OVERLAPPED* overlapped) {
+  IO_STATUS_BLOCK iosb;
+  IO_STATUS_BLOCK* iosb_ptr;
+  HANDLE event = NULL;
+  void* apc_context;
+  NTSTATUS status;
+  DWORD error;
+
+  if (overlapped != NULL) {
+    /* Overlapped operation. */
+    iosb_ptr = (IO_STATUS_BLOCK*) &overlapped->Internal;
+    event = overlapped->hEvent;
+
+    /* Do not report iocp completion if hEvent is tagged. */
+    if ((uintptr_t) event & 1) {
+      event = (HANDLE)((uintptr_t) event & ~(uintptr_t) 1);
+      apc_context = NULL;
+    } else {
+      apc_context = overlapped;
+    }
+
+  } else {
+    /* Blocking operation. */
+    iosb_ptr = &iosb;
+    event = CreateEvent(NULL, FALSE, FALSE, NULL);
+    if (event == NULL) {
+      return SOCKET_ERROR;
+    }
+    apc_context = NULL;
+  }
+
+  iosb_ptr->Status = STATUS_PENDING;
+  status = pNtDeviceIoControlFile((HANDLE) socket,
+                                  event,
+                                  NULL,
+                                  apc_context,
+                                  iosb_ptr,
+                                  IOCTL_AFD_POLL,
+                                  info,
+                                  sizeof *info,
+                                  info,
+                                  sizeof *info);
+
+  if (overlapped == NULL) {
+    /* If this is a blocking operation, wait for the event to become */
+    /* signaled, and then grab the real status from the io status block. */
+    if (status == STATUS_PENDING) {
+      DWORD r = WaitForSingleObject(event, INFINITE);
+
+      if (r == WAIT_FAILED) {
+        DWORD saved_error = GetLastError();
+        CloseHandle(event);
+        WSASetLastError(saved_error);
+        return SOCKET_ERROR;
+      }
+
+      status = iosb.Status;
+    }
+
+    CloseHandle(event);
+  }
+
+  switch (status) {
+    case STATUS_SUCCESS:
+      error = ERROR_SUCCESS;
+      break;
+
+    case STATUS_PENDING:
+      error = WSA_IO_PENDING;
+      break;
+
+    default:
+      error = uv_ntstatus_to_winsock_error(status);
+      break;
+  }
+
+  WSASetLastError(error);
+
+  if (error == ERROR_SUCCESS) {
+    return 0;
+  } else {
+    return SOCKET_ERROR;
+  }
+}
diff --git a/src/win/winsock.h b/src/win/winsock.h
@@ -81,6 +81,32 @@
 #define AFD_OVERLAPPED   0x00000002
 #define AFD_IMMEDIATE    0x00000004
 
+#define AFD_POLL_RECEIVE_BIT            0
+#define AFD_POLL_RECEIVE                (1 << AFD_POLL_RECEIVE_BIT)
+#define AFD_POLL_RECEIVE_EXPEDITED_BIT  1
+#define AFD_POLL_RECEIVE_EXPEDITED      (1 << AFD_POLL_RECEIVE_EXPEDITED_BIT)
+#define AFD_POLL_SEND_BIT               2
+#define AFD_POLL_SEND                   (1 << AFD_POLL_SEND_BIT)
+#define AFD_POLL_DISCONNECT_BIT         3
+#define AFD_POLL_DISCONNECT             (1 << AFD_POLL_DISCONNECT_BIT)
+#define AFD_POLL_ABORT_BIT              4
+#define AFD_POLL_ABORT                  (1 << AFD_POLL_ABORT_BIT)
+#define AFD_POLL_LOCAL_CLOSE_BIT        5
+#define AFD_POLL_LOCAL_CLOSE            (1 << AFD_POLL_LOCAL_CLOSE_BIT)
+#define AFD_POLL_CONNECT_BIT            6
+#define AFD_POLL_CONNECT                (1 << AFD_POLL_CONNECT_BIT)
+#define AFD_POLL_ACCEPT_BIT             7
+#define AFD_POLL_ACCEPT                 (1 << AFD_POLL_ACCEPT_BIT)
+#define AFD_POLL_CONNECT_FAIL_BIT       8
+#define AFD_POLL_CONNECT_FAIL           (1 << AFD_POLL_CONNECT_FAIL_BIT)
+#define AFD_POLL_QOS_BIT                9
+#define AFD_POLL_QOS                    (1 << AFD_POLL_QOS_BIT)
+#define AFD_POLL_GROUP_QOS_BIT          10
+#define AFD_POLL_GROUP_QOS              (1 << AFD_POLL_GROUP_QOS_BIT)
+
+#define AFD_NUM_POLL_EVENTS             11
+#define AFD_POLL_ALL                    ((1 << AFD_NUM_POLL_EVENTS) - 1)
+
 typedef struct _AFD_RECV_DATAGRAM_INFO {
     LPWSABUF BufferArray;
     ULONG BufferCount;
@@ -105,15 +131,18 @@ typedef struct _AFD_RECV_INFO {
 
 #define AFD_RECEIVE            5
 #define AFD_RECEIVE_DATAGRAM   6
+#define AFD_POLL               9
 
 #define IOCTL_AFD_RECEIVE \
     _AFD_CONTROL_CODE(AFD_RECEIVE, METHOD_NEITHER)
 
 #define IOCTL_AFD_RECEIVE_DATAGRAM \
     _AFD_CONTROL_CODE(AFD_RECEIVE_DATAGRAM, METHOD_NEITHER)
 
-#if defined(__MINGW32__) && !defined(__MINGW64_VERSION_MAJOR)
+#define IOCTL_AFD_POLL \
+    _AFD_CONTROL_CODE(AFD_POLL, METHOD_BUFFERED)
 
+#if defined(__MINGW32__) && !defined(__MINGW64_VERSION_MAJOR)
 typedef struct _IP_ADAPTER_UNICAST_ADDRESS_XP {
   /* FIXME: __C89_NAMELESS was removed */
   /* __C89_NAMELESS */ union {