1 /*
2  * Support for overlapped IO
3  *
4  * Some code borrowed from Modules/_winapi.c of CPython
5  */
6 
7 /* XXX check overflow and DWORD <-> Py_ssize_t conversions
8    Check itemsize */
9 
10 #include "Python.h"
11 #include "structmember.h"
12 
13 #define WINDOWS_LEAN_AND_MEAN
14 #include <winsock2.h>
15 #include <ws2tcpip.h>
16 #include <mswsock.h>
17 
18 #if defined(MS_WIN32) && !defined(MS_WIN64)
19 #  define F_POINTER "k"
20 #  define T_POINTER T_ULONG
21 #else
22 #  define F_POINTER "K"
23 #  define T_POINTER T_ULONGLONG
24 #endif
25 
26 /* Compatibility with Python 3.3 */
27 #if PY_VERSION_HEX < 0x03040000
28 #    define PyMem_RawMalloc PyMem_Malloc
29 #    define PyMem_RawFree PyMem_Free
30 #endif
31 
32 #define F_HANDLE F_POINTER
33 #define F_ULONG_PTR F_POINTER
34 #define F_DWORD "k"
35 #define F_BOOL "i"
36 #define F_UINT "I"
37 
38 #define T_HANDLE T_POINTER
39 
40 enum {TYPE_NONE, TYPE_NOT_STARTED, TYPE_READ, TYPE_READINTO, TYPE_WRITE,
41       TYPE_ACCEPT, TYPE_CONNECT, TYPE_DISCONNECT, TYPE_CONNECT_NAMED_PIPE,
42       TYPE_WAIT_NAMED_PIPE_AND_CONNECT, TYPE_TRANSMIT_FILE, TYPE_READ_FROM,
43       TYPE_WRITE_TO};
44 
45 typedef struct {
46     PyObject_HEAD
47     OVERLAPPED overlapped;
48     /* For convenience, we store the file handle too */
49     HANDLE handle;
50     /* Error returned by last method call */
51     DWORD error;
52     /* Type of operation */
53     DWORD type;
54     union {
55         /* Buffer allocated by us: TYPE_READ and TYPE_ACCEPT */
56         PyObject *allocated_buffer;
57         /* Buffer passed by the user: TYPE_WRITE, TYPE_WRITE_TO, and TYPE_READINTO */
58         Py_buffer user_buffer;
59 
60         /* Data used for reading from a connectionless socket:
61            TYPE_READ_FROM */
62         struct {
63             // A (buffer, (host, port)) tuple
64             PyObject *result;
65             // The actual read buffer
66             PyObject *allocated_buffer;
67             struct sockaddr_in6 address;
68             int address_length;
69         } read_from;
70     };
71 } OverlappedObject;
72 
73 /*
74  * Map Windows error codes to subclasses of OSError
75  */
76 
77 static PyObject *
SetFromWindowsErr(DWORD err)78 SetFromWindowsErr(DWORD err)
79 {
80     PyObject *exception_type;
81 
82     if (err == 0)
83         err = GetLastError();
84     switch (err) {
85         case ERROR_CONNECTION_REFUSED:
86             exception_type = PyExc_ConnectionRefusedError;
87             break;
88         case ERROR_CONNECTION_ABORTED:
89             exception_type = PyExc_ConnectionAbortedError;
90             break;
91         default:
92             exception_type = PyExc_OSError;
93     }
94     return PyErr_SetExcFromWindowsErr(exception_type, err);
95 }
96 
97 /*
98  * Some functions should be loaded at runtime
99  */
100 
101 static LPFN_ACCEPTEX Py_AcceptEx = NULL;
102 static LPFN_CONNECTEX Py_ConnectEx = NULL;
103 static LPFN_DISCONNECTEX Py_DisconnectEx = NULL;
104 static LPFN_TRANSMITFILE Py_TransmitFile = NULL;
105 static BOOL (CALLBACK *Py_CancelIoEx)(HANDLE, LPOVERLAPPED) = NULL;
106 
107 #define GET_WSA_POINTER(s, x)                                           \
108     (SOCKET_ERROR != WSAIoctl(s, SIO_GET_EXTENSION_FUNCTION_POINTER,    \
109                               &Guid##x, sizeof(Guid##x), &Py_##x,       \
110                               sizeof(Py_##x), &dwBytes, NULL, NULL))
111 
112 static int
initialize_function_pointers(void)113 initialize_function_pointers(void)
114 {
115     GUID GuidAcceptEx = WSAID_ACCEPTEX;
116     GUID GuidConnectEx = WSAID_CONNECTEX;
117     GUID GuidDisconnectEx = WSAID_DISCONNECTEX;
118     GUID GuidTransmitFile = WSAID_TRANSMITFILE;
119     HINSTANCE hKernel32;
120     SOCKET s;
121     DWORD dwBytes;
122 
123     s = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
124     if (s == INVALID_SOCKET) {
125         SetFromWindowsErr(WSAGetLastError());
126         return -1;
127     }
128 
129     if (!GET_WSA_POINTER(s, AcceptEx) ||
130         !GET_WSA_POINTER(s, ConnectEx) ||
131         !GET_WSA_POINTER(s, DisconnectEx) ||
132         !GET_WSA_POINTER(s, TransmitFile))
133     {
134         closesocket(s);
135         SetFromWindowsErr(WSAGetLastError());
136         return -1;
137     }
138 
139     closesocket(s);
140 
141     /* On WinXP we will have Py_CancelIoEx == NULL */
142     Py_BEGIN_ALLOW_THREADS
143     hKernel32 = GetModuleHandle("KERNEL32");
144     *(FARPROC *)&Py_CancelIoEx = GetProcAddress(hKernel32, "CancelIoEx");
145     Py_END_ALLOW_THREADS
146     return 0;
147 }
148 
149 /*
150  * Completion port stuff
151  */
152 
153 PyDoc_STRVAR(
154     CreateIoCompletionPort_doc,
155     "CreateIoCompletionPort(handle, port, key, concurrency) -> port\n\n"
156     "Create a completion port or register a handle with a port.");
157 
158 static PyObject *
overlapped_CreateIoCompletionPort(PyObject * self,PyObject * args)159 overlapped_CreateIoCompletionPort(PyObject *self, PyObject *args)
160 {
161     HANDLE FileHandle;
162     HANDLE ExistingCompletionPort;
163     ULONG_PTR CompletionKey;
164     DWORD NumberOfConcurrentThreads;
165     HANDLE ret;
166 
167     if (!PyArg_ParseTuple(args, F_HANDLE F_HANDLE F_ULONG_PTR F_DWORD,
168                           &FileHandle, &ExistingCompletionPort, &CompletionKey,
169                           &NumberOfConcurrentThreads))
170         return NULL;
171 
172     Py_BEGIN_ALLOW_THREADS
173     ret = CreateIoCompletionPort(FileHandle, ExistingCompletionPort,
174                                  CompletionKey, NumberOfConcurrentThreads);
175     Py_END_ALLOW_THREADS
176 
177     if (ret == NULL)
178         return SetFromWindowsErr(0);
179     return Py_BuildValue(F_HANDLE, ret);
180 }
181 
182 PyDoc_STRVAR(
183     GetQueuedCompletionStatus_doc,
184     "GetQueuedCompletionStatus(port, msecs) -> (err, bytes, key, address)\n\n"
185     "Get a message from completion port.  Wait for up to msecs milliseconds.");
186 
187 static PyObject *
overlapped_GetQueuedCompletionStatus(PyObject * self,PyObject * args)188 overlapped_GetQueuedCompletionStatus(PyObject *self, PyObject *args)
189 {
190     HANDLE CompletionPort = NULL;
191     DWORD NumberOfBytes = 0;
192     ULONG_PTR CompletionKey = 0;
193     OVERLAPPED *Overlapped = NULL;
194     DWORD Milliseconds;
195     DWORD err;
196     BOOL ret;
197 
198     if (!PyArg_ParseTuple(args, F_HANDLE F_DWORD,
199                           &CompletionPort, &Milliseconds))
200         return NULL;
201 
202     Py_BEGIN_ALLOW_THREADS
203     ret = GetQueuedCompletionStatus(CompletionPort, &NumberOfBytes,
204                                     &CompletionKey, &Overlapped, Milliseconds);
205     Py_END_ALLOW_THREADS
206 
207     err = ret ? ERROR_SUCCESS : GetLastError();
208     if (Overlapped == NULL) {
209         if (err == WAIT_TIMEOUT)
210             Py_RETURN_NONE;
211         else
212             return SetFromWindowsErr(err);
213     }
214     return Py_BuildValue(F_DWORD F_DWORD F_ULONG_PTR F_POINTER,
215                          err, NumberOfBytes, CompletionKey, Overlapped);
216 }
217 
218 PyDoc_STRVAR(
219     PostQueuedCompletionStatus_doc,
220     "PostQueuedCompletionStatus(port, bytes, key, address) -> None\n\n"
221     "Post a message to completion port.");
222 
223 static PyObject *
overlapped_PostQueuedCompletionStatus(PyObject * self,PyObject * args)224 overlapped_PostQueuedCompletionStatus(PyObject *self, PyObject *args)
225 {
226     HANDLE CompletionPort;
227     DWORD NumberOfBytes;
228     ULONG_PTR CompletionKey;
229     OVERLAPPED *Overlapped;
230     BOOL ret;
231 
232     if (!PyArg_ParseTuple(args, F_HANDLE F_DWORD F_ULONG_PTR F_POINTER,
233                           &CompletionPort, &NumberOfBytes, &CompletionKey,
234                           &Overlapped))
235         return NULL;
236 
237     Py_BEGIN_ALLOW_THREADS
238     ret = PostQueuedCompletionStatus(CompletionPort, NumberOfBytes,
239                                      CompletionKey, Overlapped);
240     Py_END_ALLOW_THREADS
241 
242     if (!ret)
243         return SetFromWindowsErr(0);
244     Py_RETURN_NONE;
245 }
246 
247 /*
248  * Wait for a handle
249  */
250 
251 struct PostCallbackData {
252     HANDLE CompletionPort;
253     LPOVERLAPPED Overlapped;
254 };
255 
256 static VOID CALLBACK
PostToQueueCallback(PVOID lpParameter,BOOLEAN TimerOrWaitFired)257 PostToQueueCallback(PVOID lpParameter, BOOLEAN TimerOrWaitFired)
258 {
259     struct PostCallbackData *p = (struct PostCallbackData*) lpParameter;
260 
261     PostQueuedCompletionStatus(p->CompletionPort, TimerOrWaitFired,
262                                0, p->Overlapped);
263     /* ignore possible error! */
264     PyMem_RawFree(p);
265 }
266 
267 PyDoc_STRVAR(
268     RegisterWaitWithQueue_doc,
269     "RegisterWaitWithQueue(Object, CompletionPort, Overlapped, Timeout)\n"
270     "    -> WaitHandle\n\n"
271     "Register wait for Object; when complete CompletionPort is notified.\n");
272 
273 static PyObject *
overlapped_RegisterWaitWithQueue(PyObject * self,PyObject * args)274 overlapped_RegisterWaitWithQueue(PyObject *self, PyObject *args)
275 {
276     HANDLE NewWaitObject;
277     HANDLE Object;
278     ULONG Milliseconds;
279     struct PostCallbackData data, *pdata;
280 
281     if (!PyArg_ParseTuple(args, F_HANDLE F_HANDLE F_POINTER F_DWORD,
282                           &Object,
283                           &data.CompletionPort,
284                           &data.Overlapped,
285                           &Milliseconds))
286         return NULL;
287 
288     /* Use PyMem_RawMalloc() rather than PyMem_Malloc(), since
289        PostToQueueCallback() will call PyMem_Free() from a new C thread
290        which doesn't hold the GIL. */
291     pdata = PyMem_RawMalloc(sizeof(struct PostCallbackData));
292     if (pdata == NULL)
293         return SetFromWindowsErr(0);
294 
295     *pdata = data;
296 
297     if (!RegisterWaitForSingleObject(
298             &NewWaitObject, Object, PostToQueueCallback, pdata, Milliseconds,
299             WT_EXECUTEINWAITTHREAD | WT_EXECUTEONLYONCE))
300     {
301         PyMem_RawFree(pdata);
302         return SetFromWindowsErr(0);
303     }
304 
305     return Py_BuildValue(F_HANDLE, NewWaitObject);
306 }
307 
308 PyDoc_STRVAR(
309     UnregisterWait_doc,
310     "UnregisterWait(WaitHandle) -> None\n\n"
311     "Unregister wait handle.\n");
312 
313 static PyObject *
overlapped_UnregisterWait(PyObject * self,PyObject * args)314 overlapped_UnregisterWait(PyObject *self, PyObject *args)
315 {
316     HANDLE WaitHandle;
317     BOOL ret;
318 
319     if (!PyArg_ParseTuple(args, F_HANDLE, &WaitHandle))
320         return NULL;
321 
322     Py_BEGIN_ALLOW_THREADS
323     ret = UnregisterWait(WaitHandle);
324     Py_END_ALLOW_THREADS
325 
326     if (!ret)
327         return SetFromWindowsErr(0);
328     Py_RETURN_NONE;
329 }
330 
331 PyDoc_STRVAR(
332     UnregisterWaitEx_doc,
333     "UnregisterWaitEx(WaitHandle, Event) -> None\n\n"
334     "Unregister wait handle.\n");
335 
336 static PyObject *
overlapped_UnregisterWaitEx(PyObject * self,PyObject * args)337 overlapped_UnregisterWaitEx(PyObject *self, PyObject *args)
338 {
339     HANDLE WaitHandle, Event;
340     BOOL ret;
341 
342     if (!PyArg_ParseTuple(args, F_HANDLE F_HANDLE, &WaitHandle, &Event))
343         return NULL;
344 
345     Py_BEGIN_ALLOW_THREADS
346     ret = UnregisterWaitEx(WaitHandle, Event);
347     Py_END_ALLOW_THREADS
348 
349     if (!ret)
350         return SetFromWindowsErr(0);
351     Py_RETURN_NONE;
352 }
353 
354 /*
355  * Event functions -- currently only used by tests
356  */
357 
358 PyDoc_STRVAR(
359     CreateEvent_doc,
360     "CreateEvent(EventAttributes, ManualReset, InitialState, Name)"
361     " -> Handle\n\n"
362     "Create an event.  EventAttributes must be None.\n");
363 
364 static PyObject *
overlapped_CreateEvent(PyObject * self,PyObject * args)365 overlapped_CreateEvent(PyObject *self, PyObject *args)
366 {
367     PyObject *EventAttributes;
368     BOOL ManualReset;
369     BOOL InitialState;
370     Py_UNICODE *Name;
371     HANDLE Event;
372 
373     if (!PyArg_ParseTuple(args, "O" F_BOOL F_BOOL "Z",
374                           &EventAttributes, &ManualReset,
375                           &InitialState, &Name))
376         return NULL;
377 
378     if (EventAttributes != Py_None) {
379         PyErr_SetString(PyExc_ValueError, "EventAttributes must be None");
380         return NULL;
381     }
382 
383     Py_BEGIN_ALLOW_THREADS
384     Event = CreateEventW(NULL, ManualReset, InitialState, Name);
385     Py_END_ALLOW_THREADS
386 
387     if (Event == NULL)
388         return SetFromWindowsErr(0);
389     return Py_BuildValue(F_HANDLE, Event);
390 }
391 
392 PyDoc_STRVAR(
393     SetEvent_doc,
394     "SetEvent(Handle) -> None\n\n"
395     "Set event.\n");
396 
397 static PyObject *
overlapped_SetEvent(PyObject * self,PyObject * args)398 overlapped_SetEvent(PyObject *self, PyObject *args)
399 {
400     HANDLE Handle;
401     BOOL ret;
402 
403     if (!PyArg_ParseTuple(args, F_HANDLE, &Handle))
404         return NULL;
405 
406     Py_BEGIN_ALLOW_THREADS
407     ret = SetEvent(Handle);
408     Py_END_ALLOW_THREADS
409 
410     if (!ret)
411         return SetFromWindowsErr(0);
412     Py_RETURN_NONE;
413 }
414 
415 PyDoc_STRVAR(
416     ResetEvent_doc,
417     "ResetEvent(Handle) -> None\n\n"
418     "Reset event.\n");
419 
420 static PyObject *
overlapped_ResetEvent(PyObject * self,PyObject * args)421 overlapped_ResetEvent(PyObject *self, PyObject *args)
422 {
423     HANDLE Handle;
424     BOOL ret;
425 
426     if (!PyArg_ParseTuple(args, F_HANDLE, &Handle))
427         return NULL;
428 
429     Py_BEGIN_ALLOW_THREADS
430     ret = ResetEvent(Handle);
431     Py_END_ALLOW_THREADS
432 
433     if (!ret)
434         return SetFromWindowsErr(0);
435     Py_RETURN_NONE;
436 }
437 
438 /*
439  * Bind socket handle to local port without doing slow getaddrinfo()
440  */
441 
442 PyDoc_STRVAR(
443     BindLocal_doc,
444     "BindLocal(handle, family) -> None\n\n"
445     "Bind a socket handle to an arbitrary local port.\n"
446     "family should AF_INET or AF_INET6.\n");
447 
448 static PyObject *
overlapped_BindLocal(PyObject * self,PyObject * args)449 overlapped_BindLocal(PyObject *self, PyObject *args)
450 {
451     SOCKET Socket;
452     int Family;
453     BOOL ret;
454 
455     if (!PyArg_ParseTuple(args, F_HANDLE "i", &Socket, &Family))
456         return NULL;
457 
458     if (Family == AF_INET) {
459         struct sockaddr_in addr;
460         memset(&addr, 0, sizeof(addr));
461         addr.sin_family = AF_INET;
462         addr.sin_port = 0;
463         addr.sin_addr.S_un.S_addr = INADDR_ANY;
464         ret = bind(Socket, (SOCKADDR*)&addr, sizeof(addr)) != SOCKET_ERROR;
465     } else if (Family == AF_INET6) {
466         struct sockaddr_in6 addr;
467         memset(&addr, 0, sizeof(addr));
468         addr.sin6_family = AF_INET6;
469         addr.sin6_port = 0;
470         addr.sin6_addr = in6addr_any;
471         ret = bind(Socket, (SOCKADDR*)&addr, sizeof(addr)) != SOCKET_ERROR;
472     } else {
473         PyErr_SetString(PyExc_ValueError, "expected tuple of length 2 or 4");
474         return NULL;
475     }
476 
477     if (!ret)
478         return SetFromWindowsErr(WSAGetLastError());
479     Py_RETURN_NONE;
480 }
481 
482 /*
483  * Windows equivalent of os.strerror() -- compare _ctypes/callproc.c
484  */
485 
486 PyDoc_STRVAR(
487     FormatMessage_doc,
488     "FormatMessage(error_code) -> error_message\n\n"
489     "Return error message for an error code.");
490 
491 static PyObject *
overlapped_FormatMessage(PyObject * ignore,PyObject * args)492 overlapped_FormatMessage(PyObject *ignore, PyObject *args)
493 {
494     DWORD code, n;
495     WCHAR *lpMsgBuf;
496     PyObject *res;
497 
498     if (!PyArg_ParseTuple(args, F_DWORD, &code))
499         return NULL;
500 
501     n = FormatMessageW(FORMAT_MESSAGE_ALLOCATE_BUFFER |
502                        FORMAT_MESSAGE_FROM_SYSTEM |
503                        FORMAT_MESSAGE_IGNORE_INSERTS,
504                        NULL,
505                        code,
506                        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
507                        (LPWSTR) &lpMsgBuf,
508                        0,
509                        NULL);
510     if (n) {
511         while (iswspace(lpMsgBuf[n-1]))
512             --n;
513         lpMsgBuf[n] = L'\0';
514         res = Py_BuildValue("u", lpMsgBuf);
515     } else {
516         res = PyUnicode_FromFormat("unknown error code %u", code);
517     }
518     LocalFree(lpMsgBuf);
519     return res;
520 }
521 
522 
523 /*
524  * Mark operation as completed - used when reading produces ERROR_BROKEN_PIPE
525  */
526 
527 static void
mark_as_completed(OVERLAPPED * ov)528 mark_as_completed(OVERLAPPED *ov)
529 {
530     ov->Internal = 0;
531     if (ov->hEvent != NULL)
532         SetEvent(ov->hEvent);
533 }
534 
535 /*
536  * A Python object wrapping an OVERLAPPED structure and other useful data
537  * for overlapped I/O
538  */
539 
540 PyDoc_STRVAR(
541     Overlapped_doc,
542     "Overlapped object");
543 
544 static PyObject *
Overlapped_new(PyTypeObject * type,PyObject * args,PyObject * kwds)545 Overlapped_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
546 {
547     OverlappedObject *self;
548     HANDLE event = INVALID_HANDLE_VALUE;
549     static char *kwlist[] = {"event", NULL};
550 
551     if (!PyArg_ParseTupleAndKeywords(args, kwds, "|" F_HANDLE, kwlist, &event))
552         return NULL;
553 
554     if (event == INVALID_HANDLE_VALUE) {
555         event = CreateEvent(NULL, TRUE, FALSE, NULL);
556         if (event == NULL)
557             return SetFromWindowsErr(0);
558     }
559 
560     self = PyObject_New(OverlappedObject, type);
561     if (self == NULL) {
562         if (event != NULL)
563             CloseHandle(event);
564         return NULL;
565     }
566 
567     self->handle = NULL;
568     self->error = 0;
569     self->type = TYPE_NONE;
570     self->allocated_buffer = NULL;
571     memset(&self->overlapped, 0, sizeof(OVERLAPPED));
572     memset(&self->user_buffer, 0, sizeof(Py_buffer));
573     if (event)
574         self->overlapped.hEvent = event;
575     return (PyObject *)self;
576 }
577 
578 
579 /* Note (bpo-32710): OverlappedType.tp_clear is not defined to not release
580    buffers while overlapped are still running, to prevent a crash. */
581 static int
Overlapped_clear(OverlappedObject * self)582 Overlapped_clear(OverlappedObject *self)
583 {
584     switch (self->type) {
585         case TYPE_READ:
586         case TYPE_ACCEPT: {
587             Py_CLEAR(self->allocated_buffer);
588             break;
589         }
590         case TYPE_READ_FROM: {
591             // An initial call to WSARecvFrom will only allocate the buffer.
592             // The result tuple of (message, address) is only
593             // allocated _after_ a message has been received.
594             if(self->read_from.result) {
595                 // We've received a message, free the result tuple.
596                 Py_CLEAR(self->read_from.result);
597             }
598             if(self->read_from.allocated_buffer) {
599                 Py_CLEAR(self->read_from.allocated_buffer);
600             }
601             break;
602         }
603         case TYPE_WRITE:
604         case TYPE_WRITE_TO:
605         case TYPE_READINTO: {
606             if (self->user_buffer.obj) {
607                 PyBuffer_Release(&self->user_buffer);
608             }
609             break;
610         }
611     }
612     self->type = TYPE_NOT_STARTED;
613     return 0;
614 }
615 
616 static void
Overlapped_dealloc(OverlappedObject * self)617 Overlapped_dealloc(OverlappedObject *self)
618 {
619     DWORD bytes;
620     DWORD olderr = GetLastError();
621     BOOL wait = FALSE;
622     BOOL ret;
623 
624     if (!HasOverlappedIoCompleted(&self->overlapped) &&
625         self->type != TYPE_NOT_STARTED)
626     {
627         if (Py_CancelIoEx && Py_CancelIoEx(self->handle, &self->overlapped))
628             wait = TRUE;
629 
630         Py_BEGIN_ALLOW_THREADS
631         ret = GetOverlappedResult(self->handle, &self->overlapped,
632                                   &bytes, wait);
633         Py_END_ALLOW_THREADS
634 
635         switch (ret ? ERROR_SUCCESS : GetLastError()) {
636             case ERROR_SUCCESS:
637             case ERROR_NOT_FOUND:
638             case ERROR_OPERATION_ABORTED:
639                 break;
640             default:
641                 PyErr_Format(
642                     PyExc_RuntimeError,
643                     "%R still has pending operation at "
644                     "deallocation, the process may crash", self);
645                 PyErr_WriteUnraisable(NULL);
646         }
647     }
648 
649     if (self->overlapped.hEvent != NULL) {
650         CloseHandle(self->overlapped.hEvent);
651     }
652 
653     Overlapped_clear(self);
654     PyObject_Del(self);
655     SetLastError(olderr);
656 }
657 
658 
659 /* Convert IPv4 sockaddr to a Python str. */
660 
661 static PyObject *
make_ipv4_addr(const struct sockaddr_in * addr)662 make_ipv4_addr(const struct sockaddr_in *addr)
663 {
664         char buf[INET_ADDRSTRLEN];
665         if (inet_ntop(AF_INET, &addr->sin_addr, buf, sizeof(buf)) == NULL) {
666                 PyErr_SetFromErrno(PyExc_OSError);
667                 return NULL;
668         }
669         return PyUnicode_FromString(buf);
670 }
671 
672 /* Convert IPv6 sockaddr to a Python str. */
673 
674 static PyObject *
make_ipv6_addr(const struct sockaddr_in6 * addr)675 make_ipv6_addr(const struct sockaddr_in6 *addr)
676 {
677         char buf[INET6_ADDRSTRLEN];
678         if (inet_ntop(AF_INET6, &addr->sin6_addr, buf, sizeof(buf)) == NULL) {
679                 PyErr_SetFromErrno(PyExc_OSError);
680                 return NULL;
681         }
682         return PyUnicode_FromString(buf);
683 }
684 
685 static PyObject*
unparse_address(LPSOCKADDR Address,DWORD Length)686 unparse_address(LPSOCKADDR Address, DWORD Length)
687 {
688         /* The function is adopted from mocketmodule.c makesockaddr()*/
689 
690     switch(Address->sa_family) {
691         case AF_INET: {
692             const struct sockaddr_in *a = (const struct sockaddr_in *)Address;
693             PyObject *addrobj = make_ipv4_addr(a);
694             PyObject *ret = NULL;
695             if (addrobj) {
696                 ret = Py_BuildValue("Oi", addrobj, ntohs(a->sin_port));
697                 Py_DECREF(addrobj);
698             }
699             return ret;
700         }
701         case AF_INET6: {
702             const struct sockaddr_in6 *a = (const struct sockaddr_in6 *)Address;
703             PyObject *addrobj = make_ipv6_addr(a);
704             PyObject *ret = NULL;
705             if (addrobj) {
706                 ret = Py_BuildValue("OiII",
707                                     addrobj,
708                                     ntohs(a->sin6_port),
709                                     ntohl(a->sin6_flowinfo),
710                                     a->sin6_scope_id);
711                 Py_DECREF(addrobj);
712             }
713             return ret;
714         }
715         default: {
716             PyErr_SetString(PyExc_ValueError, "recvfrom returned unsupported address family");
717             return NULL;
718         }
719     }
720 }
721 
722 PyDoc_STRVAR(
723     Overlapped_cancel_doc,
724     "cancel() -> None\n\n"
725     "Cancel overlapped operation");
726 
727 static PyObject *
Overlapped_cancel(OverlappedObject * self,PyObject * Py_UNUSED (ignored))728 Overlapped_cancel(OverlappedObject *self, PyObject *Py_UNUSED(ignored))
729 {
730     BOOL ret = TRUE;
731 
732     if (self->type == TYPE_NOT_STARTED
733         || self->type == TYPE_WAIT_NAMED_PIPE_AND_CONNECT)
734         Py_RETURN_NONE;
735 
736     if (!HasOverlappedIoCompleted(&self->overlapped)) {
737         Py_BEGIN_ALLOW_THREADS
738         if (Py_CancelIoEx)
739             ret = Py_CancelIoEx(self->handle, &self->overlapped);
740         else
741             ret = CancelIo(self->handle);
742         Py_END_ALLOW_THREADS
743     }
744 
745     /* CancelIoEx returns ERROR_NOT_FOUND if the I/O completed in-between */
746     if (!ret && GetLastError() != ERROR_NOT_FOUND)
747         return SetFromWindowsErr(0);
748     Py_RETURN_NONE;
749 }
750 
751 PyDoc_STRVAR(
752     Overlapped_getresult_doc,
753     "getresult(wait=False) -> result\n\n"
754     "Retrieve result of operation.  If wait is true then it blocks\n"
755     "until the operation is finished.  If wait is false and the\n"
756     "operation is still pending then an error is raised.");
757 
758 static PyObject *
Overlapped_getresult(OverlappedObject * self,PyObject * args)759 Overlapped_getresult(OverlappedObject *self, PyObject *args)
760 {
761     BOOL wait = FALSE;
762     DWORD transferred = 0;
763     BOOL ret;
764     DWORD err;
765     PyObject *addr;
766 
767     if (!PyArg_ParseTuple(args, "|" F_BOOL, &wait))
768         return NULL;
769 
770     if (self->type == TYPE_NONE) {
771         PyErr_SetString(PyExc_ValueError, "operation not yet attempted");
772         return NULL;
773     }
774 
775     if (self->type == TYPE_NOT_STARTED) {
776         PyErr_SetString(PyExc_ValueError, "operation failed to start");
777         return NULL;
778     }
779 
780     Py_BEGIN_ALLOW_THREADS
781     ret = GetOverlappedResult(self->handle, &self->overlapped, &transferred,
782                               wait);
783     Py_END_ALLOW_THREADS
784 
785     self->error = err = ret ? ERROR_SUCCESS : GetLastError();
786     switch (err) {
787         case ERROR_SUCCESS:
788         case ERROR_MORE_DATA:
789             break;
790         case ERROR_BROKEN_PIPE:
791             if (self->type == TYPE_READ || self->type == TYPE_READINTO) {
792                 break;
793             }
794             else if (self->type == TYPE_READ_FROM &&
795                      (self->read_from.result != NULL ||
796                       self->read_from.allocated_buffer != NULL))
797             {
798                 break;
799             }
800             /* fall through */
801         default:
802             return SetFromWindowsErr(err);
803     }
804 
805     switch (self->type) {
806         case TYPE_READ:
807             assert(PyBytes_CheckExact(self->allocated_buffer));
808             if (transferred != PyBytes_GET_SIZE(self->allocated_buffer) &&
809                 _PyBytes_Resize(&self->allocated_buffer, transferred))
810                 return NULL;
811 
812             Py_INCREF(self->allocated_buffer);
813             return self->allocated_buffer;
814         case TYPE_READ_FROM:
815             assert(PyBytes_CheckExact(self->read_from.allocated_buffer));
816 
817             if (transferred != PyBytes_GET_SIZE(
818                     self->read_from.allocated_buffer) &&
819                 _PyBytes_Resize(&self->read_from.allocated_buffer, transferred))
820             {
821                 return NULL;
822             }
823 
824             // unparse the address
825             addr = unparse_address((SOCKADDR*)&self->read_from.address,
826                                    self->read_from.address_length);
827 
828             if (addr == NULL) {
829                 return NULL;
830             }
831 
832             // The result is a two item tuple: (message, address)
833             self->read_from.result = PyTuple_New(2);
834             if (self->read_from.result == NULL) {
835                 Py_CLEAR(addr);
836                 return NULL;
837             }
838 
839             // first item: message
840             Py_INCREF(self->read_from.allocated_buffer);
841             PyTuple_SET_ITEM(self->read_from.result, 0,
842                              self->read_from.allocated_buffer);
843             // second item: address
844             PyTuple_SET_ITEM(self->read_from.result, 1, addr);
845 
846             Py_INCREF(self->read_from.result);
847             return self->read_from.result;
848         default:
849             return PyLong_FromUnsignedLong((unsigned long) transferred);
850     }
851 }
852 
853 static PyObject *
do_ReadFile(OverlappedObject * self,HANDLE handle,char * bufstart,DWORD buflen)854 do_ReadFile(OverlappedObject *self, HANDLE handle,
855             char *bufstart, DWORD buflen)
856 {
857     DWORD nread;
858     int ret;
859     DWORD err;
860 
861     Py_BEGIN_ALLOW_THREADS
862     ret = ReadFile(handle, bufstart, buflen, &nread,
863                    &self->overlapped);
864     Py_END_ALLOW_THREADS
865 
866     self->error = err = ret ? ERROR_SUCCESS : GetLastError();
867     switch (err) {
868         case ERROR_BROKEN_PIPE:
869             mark_as_completed(&self->overlapped);
870             return SetFromWindowsErr(err);
871         case ERROR_SUCCESS:
872         case ERROR_MORE_DATA:
873         case ERROR_IO_PENDING:
874             Py_RETURN_NONE;
875         default:
876             Overlapped_clear(self);
877             return SetFromWindowsErr(err);
878     }
879 }
880 
881 PyDoc_STRVAR(
882     Overlapped_ReadFile_doc,
883     "ReadFile(handle, size) -> Overlapped[message]\n\n"
884     "Start overlapped read");
885 
886 static PyObject *
Overlapped_ReadFile(OverlappedObject * self,PyObject * args)887 Overlapped_ReadFile(OverlappedObject *self, PyObject *args)
888 {
889     HANDLE handle;
890     DWORD size;
891     PyObject *buf;
892 
893     if (!PyArg_ParseTuple(args, F_HANDLE F_DWORD, &handle, &size))
894         return NULL;
895 
896     if (self->type != TYPE_NONE) {
897         PyErr_SetString(PyExc_ValueError, "operation already attempted");
898         return NULL;
899     }
900 
901 #if SIZEOF_SIZE_T <= SIZEOF_LONG
902     size = Py_MIN(size, (DWORD)PY_SSIZE_T_MAX);
903 #endif
904     buf = PyBytes_FromStringAndSize(NULL, Py_MAX(size, 1));
905     if (buf == NULL)
906         return NULL;
907 
908     self->type = TYPE_READ;
909     self->handle = handle;
910     self->allocated_buffer = buf;
911 
912     return do_ReadFile(self, handle, PyBytes_AS_STRING(buf), size);
913 }
914 
915 PyDoc_STRVAR(
916     Overlapped_ReadFileInto_doc,
917     "ReadFileInto(handle, buf) -> Overlapped[bytes_transferred]\n\n"
918     "Start overlapped receive");
919 
920 static PyObject *
Overlapped_ReadFileInto(OverlappedObject * self,PyObject * args)921 Overlapped_ReadFileInto(OverlappedObject *self, PyObject *args)
922 {
923     HANDLE handle;
924     PyObject *bufobj;
925 
926     if (!PyArg_ParseTuple(args, F_HANDLE "O", &handle, &bufobj))
927         return NULL;
928 
929     if (self->type != TYPE_NONE) {
930         PyErr_SetString(PyExc_ValueError, "operation already attempted");
931         return NULL;
932     }
933 
934     if (!PyArg_Parse(bufobj, "y*", &self->user_buffer))
935         return NULL;
936 
937 #if SIZEOF_SIZE_T > SIZEOF_LONG
938     if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) {
939         PyBuffer_Release(&self->user_buffer);
940         PyErr_SetString(PyExc_ValueError, "buffer too large");
941         return NULL;
942     }
943 #endif
944 
945     self->type = TYPE_READINTO;
946     self->handle = handle;
947 
948     return do_ReadFile(self, handle, self->user_buffer.buf,
949                        (DWORD)self->user_buffer.len);
950 }
951 
952 static PyObject *
do_WSARecv(OverlappedObject * self,HANDLE handle,char * bufstart,DWORD buflen,DWORD flags)953 do_WSARecv(OverlappedObject *self, HANDLE handle,
954            char *bufstart, DWORD buflen, DWORD flags)
955 {
956     DWORD nread;
957     WSABUF wsabuf;
958     int ret;
959     DWORD err;
960 
961     wsabuf.buf = bufstart;
962     wsabuf.len = buflen;
963 
964     Py_BEGIN_ALLOW_THREADS
965     ret = WSARecv((SOCKET)handle, &wsabuf, 1, &nread, &flags,
966                   &self->overlapped, NULL);
967     Py_END_ALLOW_THREADS
968 
969     self->error = err = (ret < 0 ? WSAGetLastError() : ERROR_SUCCESS);
970     switch (err) {
971         case ERROR_BROKEN_PIPE:
972             mark_as_completed(&self->overlapped);
973             return SetFromWindowsErr(err);
974         case ERROR_SUCCESS:
975         case ERROR_MORE_DATA:
976         case ERROR_IO_PENDING:
977             Py_RETURN_NONE;
978         default:
979             Overlapped_clear(self);
980             return SetFromWindowsErr(err);
981     }
982 }
983 
984 PyDoc_STRVAR(
985     Overlapped_WSARecv_doc,
986     "RecvFile(handle, size, flags) -> Overlapped[message]\n\n"
987     "Start overlapped receive");
988 
989 static PyObject *
Overlapped_WSARecv(OverlappedObject * self,PyObject * args)990 Overlapped_WSARecv(OverlappedObject *self, PyObject *args)
991 {
992     HANDLE handle;
993     DWORD size;
994     DWORD flags = 0;
995     PyObject *buf;
996 
997     if (!PyArg_ParseTuple(args, F_HANDLE F_DWORD "|" F_DWORD,
998                           &handle, &size, &flags))
999         return NULL;
1000 
1001     if (self->type != TYPE_NONE) {
1002         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1003         return NULL;
1004     }
1005 
1006 #if SIZEOF_SIZE_T <= SIZEOF_LONG
1007     size = Py_MIN(size, (DWORD)PY_SSIZE_T_MAX);
1008 #endif
1009     buf = PyBytes_FromStringAndSize(NULL, Py_MAX(size, 1));
1010     if (buf == NULL)
1011         return NULL;
1012 
1013     self->type = TYPE_READ;
1014     self->handle = handle;
1015     self->allocated_buffer = buf;
1016 
1017     return do_WSARecv(self, handle, PyBytes_AS_STRING(buf), size, flags);
1018 }
1019 
1020 PyDoc_STRVAR(
1021     Overlapped_WSARecvInto_doc,
1022     "WSARecvInto(handle, buf, flags) -> Overlapped[bytes_transferred]\n\n"
1023     "Start overlapped receive");
1024 
1025 static PyObject *
Overlapped_WSARecvInto(OverlappedObject * self,PyObject * args)1026 Overlapped_WSARecvInto(OverlappedObject *self, PyObject *args)
1027 {
1028     HANDLE handle;
1029     PyObject *bufobj;
1030     DWORD flags;
1031 
1032     if (!PyArg_ParseTuple(args, F_HANDLE "O" F_DWORD,
1033                           &handle, &bufobj, &flags))
1034         return NULL;
1035 
1036     if (self->type != TYPE_NONE) {
1037         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1038         return NULL;
1039     }
1040 
1041     if (!PyArg_Parse(bufobj, "y*", &self->user_buffer))
1042         return NULL;
1043 
1044 #if SIZEOF_SIZE_T > SIZEOF_LONG
1045     if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) {
1046         PyBuffer_Release(&self->user_buffer);
1047         PyErr_SetString(PyExc_ValueError, "buffer too large");
1048         return NULL;
1049     }
1050 #endif
1051 
1052     self->type = TYPE_READINTO;
1053     self->handle = handle;
1054 
1055     return do_WSARecv(self, handle, self->user_buffer.buf,
1056                       (DWORD)self->user_buffer.len, flags);
1057 }
1058 
1059 PyDoc_STRVAR(
1060     Overlapped_WriteFile_doc,
1061     "WriteFile(handle, buf) -> Overlapped[bytes_transferred]\n\n"
1062     "Start overlapped write");
1063 
1064 static PyObject *
Overlapped_WriteFile(OverlappedObject * self,PyObject * args)1065 Overlapped_WriteFile(OverlappedObject *self, PyObject *args)
1066 {
1067     HANDLE handle;
1068     PyObject *bufobj;
1069     DWORD written;
1070     BOOL ret;
1071     DWORD err;
1072 
1073     if (!PyArg_ParseTuple(args, F_HANDLE "O", &handle, &bufobj))
1074         return NULL;
1075 
1076     if (self->type != TYPE_NONE) {
1077         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1078         return NULL;
1079     }
1080 
1081     if (!PyArg_Parse(bufobj, "y*", &self->user_buffer))
1082         return NULL;
1083 
1084 #if SIZEOF_SIZE_T > SIZEOF_LONG
1085     if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) {
1086         PyBuffer_Release(&self->user_buffer);
1087         PyErr_SetString(PyExc_ValueError, "buffer too large");
1088         return NULL;
1089     }
1090 #endif
1091 
1092     self->type = TYPE_WRITE;
1093     self->handle = handle;
1094 
1095     Py_BEGIN_ALLOW_THREADS
1096     ret = WriteFile(handle, self->user_buffer.buf,
1097                     (DWORD)self->user_buffer.len,
1098                     &written, &self->overlapped);
1099     Py_END_ALLOW_THREADS
1100 
1101     self->error = err = ret ? ERROR_SUCCESS : GetLastError();
1102     switch (err) {
1103         case ERROR_SUCCESS:
1104         case ERROR_IO_PENDING:
1105             Py_RETURN_NONE;
1106         default:
1107             Overlapped_clear(self);
1108             return SetFromWindowsErr(err);
1109     }
1110 }
1111 
1112 PyDoc_STRVAR(
1113     Overlapped_WSASend_doc,
1114     "WSASend(handle, buf, flags) -> Overlapped[bytes_transferred]\n\n"
1115     "Start overlapped send");
1116 
1117 static PyObject *
Overlapped_WSASend(OverlappedObject * self,PyObject * args)1118 Overlapped_WSASend(OverlappedObject *self, PyObject *args)
1119 {
1120     HANDLE handle;
1121     PyObject *bufobj;
1122     DWORD flags;
1123     DWORD written;
1124     WSABUF wsabuf;
1125     int ret;
1126     DWORD err;
1127 
1128     if (!PyArg_ParseTuple(args, F_HANDLE "O" F_DWORD,
1129                           &handle, &bufobj, &flags))
1130         return NULL;
1131 
1132     if (self->type != TYPE_NONE) {
1133         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1134         return NULL;
1135     }
1136 
1137     if (!PyArg_Parse(bufobj, "y*", &self->user_buffer))
1138         return NULL;
1139 
1140 #if SIZEOF_SIZE_T > SIZEOF_LONG
1141     if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) {
1142         PyBuffer_Release(&self->user_buffer);
1143         PyErr_SetString(PyExc_ValueError, "buffer too large");
1144         return NULL;
1145     }
1146 #endif
1147 
1148     self->type = TYPE_WRITE;
1149     self->handle = handle;
1150     wsabuf.len = (DWORD)self->user_buffer.len;
1151     wsabuf.buf = self->user_buffer.buf;
1152 
1153     Py_BEGIN_ALLOW_THREADS
1154     ret = WSASend((SOCKET)handle, &wsabuf, 1, &written, flags,
1155                   &self->overlapped, NULL);
1156     Py_END_ALLOW_THREADS
1157 
1158     self->error = err = (ret < 0 ? WSAGetLastError() : ERROR_SUCCESS);
1159     switch (err) {
1160         case ERROR_SUCCESS:
1161         case ERROR_IO_PENDING:
1162             Py_RETURN_NONE;
1163         default:
1164             Overlapped_clear(self);
1165             return SetFromWindowsErr(err);
1166     }
1167 }
1168 
1169 PyDoc_STRVAR(
1170     Overlapped_AcceptEx_doc,
1171     "AcceptEx(listen_handle, accept_handle) -> Overlapped[address_as_bytes]\n\n"
1172     "Start overlapped wait for client to connect");
1173 
1174 static PyObject *
Overlapped_AcceptEx(OverlappedObject * self,PyObject * args)1175 Overlapped_AcceptEx(OverlappedObject *self, PyObject *args)
1176 {
1177     SOCKET ListenSocket;
1178     SOCKET AcceptSocket;
1179     DWORD BytesReceived;
1180     DWORD size;
1181     PyObject *buf;
1182     BOOL ret;
1183     DWORD err;
1184 
1185     if (!PyArg_ParseTuple(args, F_HANDLE F_HANDLE,
1186                           &ListenSocket, &AcceptSocket))
1187         return NULL;
1188 
1189     if (self->type != TYPE_NONE) {
1190         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1191         return NULL;
1192     }
1193 
1194     size = sizeof(struct sockaddr_in6) + 16;
1195     buf = PyBytes_FromStringAndSize(NULL, size*2);
1196     if (!buf)
1197         return NULL;
1198 
1199     self->type = TYPE_ACCEPT;
1200     self->handle = (HANDLE)ListenSocket;
1201     self->allocated_buffer = buf;
1202 
1203     Py_BEGIN_ALLOW_THREADS
1204     ret = Py_AcceptEx(ListenSocket, AcceptSocket, PyBytes_AS_STRING(buf),
1205                       0, size, size, &BytesReceived, &self->overlapped);
1206     Py_END_ALLOW_THREADS
1207 
1208     self->error = err = ret ? ERROR_SUCCESS : WSAGetLastError();
1209     switch (err) {
1210         case ERROR_SUCCESS:
1211         case ERROR_IO_PENDING:
1212             Py_RETURN_NONE;
1213         default:
1214             Overlapped_clear(self);
1215             return SetFromWindowsErr(err);
1216     }
1217 }
1218 
1219 
1220 static int
parse_address(PyObject * obj,SOCKADDR * Address,int Length)1221 parse_address(PyObject *obj, SOCKADDR *Address, int Length)
1222 {
1223     Py_UNICODE *Host;
1224     unsigned short Port;
1225     unsigned long FlowInfo;
1226     unsigned long ScopeId;
1227 
1228     memset(Address, 0, Length);
1229 
1230     if (PyArg_ParseTuple(obj, "uH", &Host, &Port))
1231     {
1232         Address->sa_family = AF_INET;
1233         if (WSAStringToAddressW(Host, AF_INET, NULL, Address, &Length) < 0) {
1234             SetFromWindowsErr(WSAGetLastError());
1235             return -1;
1236         }
1237         ((SOCKADDR_IN*)Address)->sin_port = htons(Port);
1238         return Length;
1239     }
1240     else if (PyArg_ParseTuple(obj,
1241                               "uHkk;ConnectEx(): illegal address_as_bytes "
1242                               "argument", &Host, &Port, &FlowInfo, &ScopeId))
1243     {
1244         PyErr_Clear();
1245         Address->sa_family = AF_INET6;
1246         if (WSAStringToAddressW(Host, AF_INET6, NULL, Address, &Length) < 0) {
1247             SetFromWindowsErr(WSAGetLastError());
1248             return -1;
1249         }
1250         ((SOCKADDR_IN6*)Address)->sin6_port = htons(Port);
1251         ((SOCKADDR_IN6*)Address)->sin6_flowinfo = FlowInfo;
1252         ((SOCKADDR_IN6*)Address)->sin6_scope_id = ScopeId;
1253         return Length;
1254     }
1255 
1256     return -1;
1257 }
1258 
1259 PyDoc_STRVAR(
1260     Overlapped_ConnectEx_doc,
1261     "ConnectEx(client_handle, address_as_bytes) -> Overlapped[None]\n\n"
1262     "Start overlapped connect.  client_handle should be unbound.");
1263 
1264 static PyObject *
Overlapped_ConnectEx(OverlappedObject * self,PyObject * args)1265 Overlapped_ConnectEx(OverlappedObject *self, PyObject *args)
1266 {
1267     SOCKET ConnectSocket;
1268     PyObject *AddressObj;
1269     char AddressBuf[sizeof(struct sockaddr_in6)];
1270     SOCKADDR *Address = (SOCKADDR*)AddressBuf;
1271     int Length;
1272     BOOL ret;
1273     DWORD err;
1274 
1275     if (!PyArg_ParseTuple(args, F_HANDLE "O!:ConnectEx",
1276                           &ConnectSocket, &PyTuple_Type, &AddressObj))
1277     {
1278         return NULL;
1279     }
1280 
1281     if (self->type != TYPE_NONE) {
1282         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1283         return NULL;
1284     }
1285 
1286     Length = sizeof(AddressBuf);
1287     Length = parse_address(AddressObj, Address, Length);
1288     if (Length < 0)
1289         return NULL;
1290 
1291     self->type = TYPE_CONNECT;
1292     self->handle = (HANDLE)ConnectSocket;
1293 
1294     Py_BEGIN_ALLOW_THREADS
1295     ret = Py_ConnectEx(ConnectSocket, Address, Length,
1296                        NULL, 0, NULL, &self->overlapped);
1297     Py_END_ALLOW_THREADS
1298 
1299     self->error = err = ret ? ERROR_SUCCESS : WSAGetLastError();
1300     switch (err) {
1301         case ERROR_SUCCESS:
1302         case ERROR_IO_PENDING:
1303             Py_RETURN_NONE;
1304         default:
1305             Overlapped_clear(self);
1306             return SetFromWindowsErr(err);
1307     }
1308 }
1309 
1310 PyDoc_STRVAR(
1311     Overlapped_DisconnectEx_doc,
1312     "DisconnectEx(handle, flags) -> Overlapped[None]\n\n"
1313     "Start overlapped connect.  client_handle should be unbound.");
1314 
1315 static PyObject *
Overlapped_DisconnectEx(OverlappedObject * self,PyObject * args)1316 Overlapped_DisconnectEx(OverlappedObject *self, PyObject *args)
1317 {
1318     SOCKET Socket;
1319     DWORD flags;
1320     BOOL ret;
1321     DWORD err;
1322 
1323     if (!PyArg_ParseTuple(args, F_HANDLE F_DWORD, &Socket, &flags))
1324         return NULL;
1325 
1326     if (self->type != TYPE_NONE) {
1327         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1328         return NULL;
1329     }
1330 
1331     self->type = TYPE_DISCONNECT;
1332     self->handle = (HANDLE)Socket;
1333 
1334     Py_BEGIN_ALLOW_THREADS
1335     ret = Py_DisconnectEx(Socket, &self->overlapped, flags, 0);
1336     Py_END_ALLOW_THREADS
1337 
1338     self->error = err = ret ? ERROR_SUCCESS : WSAGetLastError();
1339     switch (err) {
1340         case ERROR_SUCCESS:
1341         case ERROR_IO_PENDING:
1342             Py_RETURN_NONE;
1343         default:
1344             Overlapped_clear(self);
1345             return SetFromWindowsErr(err);
1346     }
1347 }
1348 
1349 PyDoc_STRVAR(
1350     Overlapped_TransmitFile_doc,
1351     "TransmitFile(socket, file, offset, offset_high, "
1352     "count_to_write, count_per_send, flags) "
1353     "-> Overlapped[None]\n\n"
1354     "Transmit file data over a connected socket.");
1355 
1356 static PyObject *
Overlapped_TransmitFile(OverlappedObject * self,PyObject * args)1357 Overlapped_TransmitFile(OverlappedObject *self, PyObject *args)
1358 {
1359     SOCKET Socket;
1360     HANDLE File;
1361     DWORD offset;
1362     DWORD offset_high;
1363     DWORD count_to_write;
1364     DWORD count_per_send;
1365     DWORD flags;
1366     BOOL ret;
1367     DWORD err;
1368 
1369     if (!PyArg_ParseTuple(args,
1370                           F_HANDLE F_HANDLE F_DWORD F_DWORD
1371                           F_DWORD F_DWORD F_DWORD,
1372                           &Socket, &File, &offset, &offset_high,
1373                           &count_to_write, &count_per_send,
1374                           &flags))
1375         return NULL;
1376 
1377     if (self->type != TYPE_NONE) {
1378         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1379         return NULL;
1380     }
1381 
1382     self->type = TYPE_TRANSMIT_FILE;
1383     self->handle = (HANDLE)Socket;
1384     self->overlapped.Offset = offset;
1385     self->overlapped.OffsetHigh = offset_high;
1386 
1387     Py_BEGIN_ALLOW_THREADS
1388     ret = Py_TransmitFile(Socket, File, count_to_write, count_per_send,
1389                           &self->overlapped,
1390                           NULL, flags);
1391     Py_END_ALLOW_THREADS
1392 
1393     self->error = err = ret ? ERROR_SUCCESS : WSAGetLastError();
1394     switch (err) {
1395         case ERROR_SUCCESS:
1396         case ERROR_IO_PENDING:
1397             Py_RETURN_NONE;
1398         default:
1399             Overlapped_clear(self);
1400             return SetFromWindowsErr(err);
1401     }
1402 }
1403 
1404 PyDoc_STRVAR(
1405     Overlapped_ConnectNamedPipe_doc,
1406     "ConnectNamedPipe(handle) -> Overlapped[None]\n\n"
1407     "Start overlapped wait for a client to connect.");
1408 
1409 static PyObject *
Overlapped_ConnectNamedPipe(OverlappedObject * self,PyObject * args)1410 Overlapped_ConnectNamedPipe(OverlappedObject *self, PyObject *args)
1411 {
1412     HANDLE Pipe;
1413     BOOL ret;
1414     DWORD err;
1415 
1416     if (!PyArg_ParseTuple(args, F_HANDLE, &Pipe))
1417         return NULL;
1418 
1419     if (self->type != TYPE_NONE) {
1420         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1421         return NULL;
1422     }
1423 
1424     self->type = TYPE_CONNECT_NAMED_PIPE;
1425     self->handle = Pipe;
1426 
1427     Py_BEGIN_ALLOW_THREADS
1428     ret = ConnectNamedPipe(Pipe, &self->overlapped);
1429     Py_END_ALLOW_THREADS
1430 
1431     self->error = err = ret ? ERROR_SUCCESS : GetLastError();
1432     switch (err) {
1433         case ERROR_PIPE_CONNECTED:
1434             mark_as_completed(&self->overlapped);
1435             Py_RETURN_TRUE;
1436         case ERROR_SUCCESS:
1437         case ERROR_IO_PENDING:
1438             Py_RETURN_FALSE;
1439         default:
1440             Overlapped_clear(self);
1441             return SetFromWindowsErr(err);
1442     }
1443 }
1444 
1445 PyDoc_STRVAR(
1446     ConnectPipe_doc,
1447     "ConnectPipe(addr) -> pipe_handle\n\n"
1448     "Connect to the pipe for asynchronous I/O (overlapped).");
1449 
1450 static PyObject *
overlapped_ConnectPipe(PyObject * self,PyObject * args)1451 overlapped_ConnectPipe(PyObject *self, PyObject *args)
1452 {
1453     PyObject *AddressObj;
1454     wchar_t *Address;
1455     HANDLE PipeHandle;
1456 
1457     if (!PyArg_ParseTuple(args, "U",  &AddressObj))
1458         return NULL;
1459 
1460     Address = PyUnicode_AsWideCharString(AddressObj, NULL);
1461     if (Address == NULL)
1462         return NULL;
1463 
1464     Py_BEGIN_ALLOW_THREADS
1465     PipeHandle = CreateFileW(Address,
1466                              GENERIC_READ | GENERIC_WRITE,
1467                              0, NULL, OPEN_EXISTING,
1468                              FILE_FLAG_OVERLAPPED, NULL);
1469     Py_END_ALLOW_THREADS
1470 
1471     PyMem_Free(Address);
1472     if (PipeHandle == INVALID_HANDLE_VALUE)
1473         return SetFromWindowsErr(0);
1474     return Py_BuildValue(F_HANDLE, PipeHandle);
1475 }
1476 
1477 static PyObject*
Overlapped_getaddress(OverlappedObject * self)1478 Overlapped_getaddress(OverlappedObject *self)
1479 {
1480     return PyLong_FromVoidPtr(&self->overlapped);
1481 }
1482 
1483 static PyObject*
Overlapped_getpending(OverlappedObject * self)1484 Overlapped_getpending(OverlappedObject *self)
1485 {
1486     return PyBool_FromLong(!HasOverlappedIoCompleted(&self->overlapped) &&
1487                            self->type != TYPE_NOT_STARTED);
1488 }
1489 
1490 static int
Overlapped_traverse(OverlappedObject * self,visitproc visit,void * arg)1491 Overlapped_traverse(OverlappedObject *self, visitproc visit, void *arg)
1492 {
1493     switch (self->type) {
1494     case TYPE_READ:
1495     case TYPE_ACCEPT:
1496         Py_VISIT(self->allocated_buffer);
1497         break;
1498     case TYPE_WRITE:
1499     case TYPE_WRITE_TO:
1500     case TYPE_READINTO:
1501         if (self->user_buffer.obj) {
1502             Py_VISIT(&self->user_buffer.obj);
1503         }
1504         break;
1505     case TYPE_READ_FROM:
1506         if(self->read_from.result) {
1507             Py_VISIT(self->read_from.result);
1508         }
1509         if(self->read_from.allocated_buffer) {
1510             Py_VISIT(self->read_from.allocated_buffer);
1511         }
1512     }
1513     return 0;
1514 }
1515 
1516 // UDP functions
1517 
1518 PyDoc_STRVAR(
1519     WSAConnect_doc,
1520     "WSAConnect(client_handle, address_as_bytes) -> Overlapped[None]\n\n"
1521     "Bind a remote address to a connectionless (UDP) socket");
1522 
1523 /*
1524  * Note: WSAConnect does not support Overlapped I/O so this function should
1525  * _only_ be used for connectionless sockets (UDP).
1526  */
1527 static PyObject *
overlapped_WSAConnect(PyObject * self,PyObject * args)1528 overlapped_WSAConnect(PyObject *self, PyObject *args)
1529 {
1530     SOCKET ConnectSocket;
1531     PyObject *AddressObj;
1532     char AddressBuf[sizeof(struct sockaddr_in6)];
1533     SOCKADDR *Address = (SOCKADDR*)AddressBuf;
1534     int Length;
1535     int err;
1536 
1537     if (!PyArg_ParseTuple(args, F_HANDLE "O", &ConnectSocket, &AddressObj)) {
1538         return NULL;
1539     }
1540 
1541     Length = sizeof(AddressBuf);
1542     Length = parse_address(AddressObj, Address, Length);
1543     if (Length < 0) {
1544         return NULL;
1545     }
1546 
1547     Py_BEGIN_ALLOW_THREADS
1548     // WSAConnect does not support overlapped I/O so this call will
1549     // successfully complete immediately.
1550     err = WSAConnect(ConnectSocket, Address, Length,
1551                         NULL, NULL, NULL, NULL);
1552     Py_END_ALLOW_THREADS
1553 
1554     if (err == 0) {
1555         Py_RETURN_NONE;
1556     }
1557     else {
1558         return SetFromWindowsErr(WSAGetLastError());
1559     }
1560 }
1561 
1562 PyDoc_STRVAR(
1563     Overlapped_WSASendTo_doc,
1564     "WSASendTo(handle, buf, flags, address_as_bytes) -> "
1565     "Overlapped[bytes_transferred]\n\n"
1566     "Start overlapped sendto over a connectionless (UDP) socket");
1567 
1568 static PyObject *
Overlapped_WSASendTo(OverlappedObject * self,PyObject * args)1569 Overlapped_WSASendTo(OverlappedObject *self, PyObject *args)
1570 {
1571     HANDLE handle;
1572     PyObject *bufobj;
1573     DWORD flags;
1574     PyObject *AddressObj;
1575     char AddressBuf[sizeof(struct sockaddr_in6)];
1576     SOCKADDR *Address = (SOCKADDR*)AddressBuf;
1577     int AddressLength;
1578     DWORD written;
1579     WSABUF wsabuf;
1580     int ret;
1581     DWORD err;
1582 
1583     if (!PyArg_ParseTuple(args, F_HANDLE "O" F_DWORD "O",
1584                           &handle, &bufobj, &flags, &AddressObj))
1585     {
1586         return NULL;
1587     }
1588 
1589     // Parse the "to" address
1590     AddressLength = sizeof(AddressBuf);
1591     AddressLength = parse_address(AddressObj, Address, AddressLength);
1592     if (AddressLength < 0) {
1593         return NULL;
1594     }
1595 
1596     if (self->type != TYPE_NONE) {
1597         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1598         return NULL;
1599     }
1600 
1601     if (!PyArg_Parse(bufobj, "y*", &self->user_buffer)) {
1602         return NULL;
1603     }
1604 
1605 #if SIZEOF_SIZE_T > SIZEOF_LONG
1606     if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) {
1607         PyBuffer_Release(&self->user_buffer);
1608         PyErr_SetString(PyExc_ValueError, "buffer too large");
1609         return NULL;
1610     }
1611 #endif
1612 
1613     self->type = TYPE_WRITE_TO;
1614     self->handle = handle;
1615     wsabuf.len = (DWORD)self->user_buffer.len;
1616     wsabuf.buf = self->user_buffer.buf;
1617 
1618     Py_BEGIN_ALLOW_THREADS
1619     ret = WSASendTo((SOCKET)handle, &wsabuf, 1, &written, flags,
1620                     Address, AddressLength, &self->overlapped, NULL);
1621     Py_END_ALLOW_THREADS
1622 
1623     self->error = err = (ret == SOCKET_ERROR ? WSAGetLastError() :
1624                                                ERROR_SUCCESS);
1625 
1626     switch(err) {
1627         case ERROR_SUCCESS:
1628         case ERROR_IO_PENDING:
1629             Py_RETURN_NONE;
1630         default:
1631             self->type = TYPE_NOT_STARTED;
1632             return SetFromWindowsErr(err);
1633     }
1634 }
1635 
1636 
1637 
1638 PyDoc_STRVAR(
1639     Overlapped_WSARecvFrom_doc,
1640     "RecvFile(handle, size, flags) -> Overlapped[(message, (host, port))]\n\n"
1641     "Start overlapped receive");
1642 
1643 static PyObject *
Overlapped_WSARecvFrom(OverlappedObject * self,PyObject * args)1644 Overlapped_WSARecvFrom(OverlappedObject *self, PyObject *args)
1645 {
1646     HANDLE handle;
1647     DWORD size;
1648     DWORD flags = 0;
1649     DWORD nread;
1650     PyObject *buf;
1651     WSABUF wsabuf;
1652     int ret;
1653     DWORD err;
1654 
1655     if (!PyArg_ParseTuple(args, F_HANDLE F_DWORD "|" F_DWORD,
1656                           &handle, &size, &flags))
1657     {
1658         return NULL;
1659     }
1660 
1661     if (self->type != TYPE_NONE) {
1662         PyErr_SetString(PyExc_ValueError, "operation already attempted");
1663         return NULL;
1664     }
1665 
1666 #if SIZEOF_SIZE_T <= SIZEOF_LONG
1667     size = Py_MIN(size, (DWORD)PY_SSIZE_T_MAX);
1668 #endif
1669     buf = PyBytes_FromStringAndSize(NULL, Py_MAX(size, 1));
1670     if (buf == NULL) {
1671         return NULL;
1672     }
1673 
1674     wsabuf.len = size;
1675     wsabuf.buf = PyBytes_AS_STRING(buf);
1676 
1677     self->type = TYPE_READ_FROM;
1678     self->handle = handle;
1679     self->read_from.allocated_buffer = buf;
1680     memset(&self->read_from.address, 0, sizeof(self->read_from.address));
1681     self->read_from.address_length = sizeof(self->read_from.address);
1682 
1683     Py_BEGIN_ALLOW_THREADS
1684     ret = WSARecvFrom((SOCKET)handle, &wsabuf, 1, &nread, &flags,
1685                       (SOCKADDR*)&self->read_from.address,
1686                       &self->read_from.address_length,
1687                       &self->overlapped, NULL);
1688     Py_END_ALLOW_THREADS
1689 
1690     self->error = err = (ret < 0 ? WSAGetLastError() : ERROR_SUCCESS);
1691 
1692     switch(err) {
1693     case ERROR_BROKEN_PIPE:
1694         mark_as_completed(&self->overlapped);
1695         return SetFromWindowsErr(err);
1696     case ERROR_SUCCESS:
1697     case ERROR_MORE_DATA:
1698     case ERROR_IO_PENDING:
1699         Py_RETURN_NONE;
1700     default:
1701         self->type = TYPE_NOT_STARTED;
1702         return SetFromWindowsErr(err);
1703     }
1704 }
1705 
1706 
1707 static PyMethodDef Overlapped_methods[] = {
1708     {"getresult", (PyCFunction) Overlapped_getresult,
1709      METH_VARARGS, Overlapped_getresult_doc},
1710     {"cancel", (PyCFunction) Overlapped_cancel,
1711      METH_NOARGS, Overlapped_cancel_doc},
1712     {"ReadFile", (PyCFunction) Overlapped_ReadFile,
1713      METH_VARARGS, Overlapped_ReadFile_doc},
1714     {"ReadFileInto", (PyCFunction) Overlapped_ReadFileInto,
1715      METH_VARARGS, Overlapped_ReadFileInto_doc},
1716     {"WSARecv", (PyCFunction) Overlapped_WSARecv,
1717      METH_VARARGS, Overlapped_WSARecv_doc},
1718     {"WSARecvInto", (PyCFunction) Overlapped_WSARecvInto,
1719      METH_VARARGS, Overlapped_WSARecvInto_doc},
1720     {"WriteFile", (PyCFunction) Overlapped_WriteFile,
1721      METH_VARARGS, Overlapped_WriteFile_doc},
1722     {"WSASend", (PyCFunction) Overlapped_WSASend,
1723      METH_VARARGS, Overlapped_WSASend_doc},
1724     {"AcceptEx", (PyCFunction) Overlapped_AcceptEx,
1725      METH_VARARGS, Overlapped_AcceptEx_doc},
1726     {"ConnectEx", (PyCFunction) Overlapped_ConnectEx,
1727      METH_VARARGS, Overlapped_ConnectEx_doc},
1728     {"DisconnectEx", (PyCFunction) Overlapped_DisconnectEx,
1729      METH_VARARGS, Overlapped_DisconnectEx_doc},
1730     {"TransmitFile", (PyCFunction) Overlapped_TransmitFile,
1731      METH_VARARGS, Overlapped_TransmitFile_doc},
1732     {"ConnectNamedPipe", (PyCFunction) Overlapped_ConnectNamedPipe,
1733      METH_VARARGS, Overlapped_ConnectNamedPipe_doc},
1734     {"WSARecvFrom", (PyCFunction) Overlapped_WSARecvFrom,
1735      METH_VARARGS, Overlapped_WSARecvFrom_doc },
1736     {"WSASendTo", (PyCFunction) Overlapped_WSASendTo,
1737      METH_VARARGS, Overlapped_WSASendTo_doc },
1738     {NULL}
1739 };
1740 
1741 static PyMemberDef Overlapped_members[] = {
1742     {"error", T_ULONG,
1743      offsetof(OverlappedObject, error),
1744      READONLY, "Error from last operation"},
1745     {"event", T_HANDLE,
1746      offsetof(OverlappedObject, overlapped) + offsetof(OVERLAPPED, hEvent),
1747      READONLY, "Overlapped event handle"},
1748     {NULL}
1749 };
1750 
1751 static PyGetSetDef Overlapped_getsets[] = {
1752     {"address", (getter)Overlapped_getaddress, NULL,
1753      "Address of overlapped structure"},
1754     {"pending", (getter)Overlapped_getpending, NULL,
1755      "Whether the operation is pending"},
1756     {NULL},
1757 };
1758 
1759 PyTypeObject OverlappedType = {
1760     PyVarObject_HEAD_INIT(NULL, 0)
1761     /* tp_name           */ "_overlapped.Overlapped",
1762     /* tp_basicsize      */ sizeof(OverlappedObject),
1763     /* tp_itemsize       */ 0,
1764     /* tp_dealloc        */ (destructor) Overlapped_dealloc,
1765     /* tp_vectorcall_offset */ 0,
1766     /* tp_getattr        */ 0,
1767     /* tp_setattr        */ 0,
1768     /* tp_as_async       */ 0,
1769     /* tp_repr           */ 0,
1770     /* tp_as_number      */ 0,
1771     /* tp_as_sequence    */ 0,
1772     /* tp_as_mapping     */ 0,
1773     /* tp_hash           */ 0,
1774     /* tp_call           */ 0,
1775     /* tp_str            */ 0,
1776     /* tp_getattro       */ 0,
1777     /* tp_setattro       */ 0,
1778     /* tp_as_buffer      */ 0,
1779     /* tp_flags          */ Py_TPFLAGS_DEFAULT,
1780     /* tp_doc            */ "OVERLAPPED structure wrapper",
1781     /* tp_traverse       */ (traverseproc)Overlapped_traverse,
1782     /* tp_clear          */ 0,
1783     /* tp_richcompare    */ 0,
1784     /* tp_weaklistoffset */ 0,
1785     /* tp_iter           */ 0,
1786     /* tp_iternext       */ 0,
1787     /* tp_methods        */ Overlapped_methods,
1788     /* tp_members        */ Overlapped_members,
1789     /* tp_getset         */ Overlapped_getsets,
1790     /* tp_base           */ 0,
1791     /* tp_dict           */ 0,
1792     /* tp_descr_get      */ 0,
1793     /* tp_descr_set      */ 0,
1794     /* tp_dictoffset     */ 0,
1795     /* tp_init           */ 0,
1796     /* tp_alloc          */ 0,
1797     /* tp_new            */ Overlapped_new,
1798 };
1799 
1800 static PyMethodDef overlapped_functions[] = {
1801     {"CreateIoCompletionPort", overlapped_CreateIoCompletionPort,
1802      METH_VARARGS, CreateIoCompletionPort_doc},
1803     {"GetQueuedCompletionStatus", overlapped_GetQueuedCompletionStatus,
1804      METH_VARARGS, GetQueuedCompletionStatus_doc},
1805     {"PostQueuedCompletionStatus", overlapped_PostQueuedCompletionStatus,
1806      METH_VARARGS, PostQueuedCompletionStatus_doc},
1807     {"FormatMessage", overlapped_FormatMessage,
1808      METH_VARARGS, FormatMessage_doc},
1809     {"BindLocal", overlapped_BindLocal,
1810      METH_VARARGS, BindLocal_doc},
1811     {"RegisterWaitWithQueue", overlapped_RegisterWaitWithQueue,
1812      METH_VARARGS, RegisterWaitWithQueue_doc},
1813     {"UnregisterWait", overlapped_UnregisterWait,
1814      METH_VARARGS, UnregisterWait_doc},
1815     {"UnregisterWaitEx", overlapped_UnregisterWaitEx,
1816      METH_VARARGS, UnregisterWaitEx_doc},
1817     {"CreateEvent", overlapped_CreateEvent,
1818      METH_VARARGS, CreateEvent_doc},
1819     {"SetEvent", overlapped_SetEvent,
1820      METH_VARARGS, SetEvent_doc},
1821     {"ResetEvent", overlapped_ResetEvent,
1822      METH_VARARGS, ResetEvent_doc},
1823     {"ConnectPipe", overlapped_ConnectPipe,
1824      METH_VARARGS, ConnectPipe_doc},
1825     {"WSAConnect", overlapped_WSAConnect,
1826      METH_VARARGS, WSAConnect_doc},
1827     {NULL}
1828 };
1829 
1830 static struct PyModuleDef overlapped_module = {
1831     PyModuleDef_HEAD_INIT,
1832     "_overlapped",
1833     NULL,
1834     -1,
1835     overlapped_functions,
1836     NULL,
1837     NULL,
1838     NULL,
1839     NULL
1840 };
1841 
1842 #define WINAPI_CONSTANT(fmt, con) \
1843     PyDict_SetItemString(d, #con, Py_BuildValue(fmt, con))
1844 
1845 PyMODINIT_FUNC
PyInit__overlapped(void)1846 PyInit__overlapped(void)
1847 {
1848     PyObject *m, *d;
1849 
1850     /* Ensure WSAStartup() called before initializing function pointers */
1851     m = PyImport_ImportModule("_socket");
1852     if (!m)
1853         return NULL;
1854     Py_DECREF(m);
1855 
1856     if (initialize_function_pointers() < 0)
1857         return NULL;
1858 
1859     if (PyType_Ready(&OverlappedType) < 0)
1860         return NULL;
1861 
1862     m = PyModule_Create(&overlapped_module);
1863     if (PyModule_AddObject(m, "Overlapped", (PyObject *)&OverlappedType) < 0)
1864         return NULL;
1865 
1866     d = PyModule_GetDict(m);
1867 
1868     WINAPI_CONSTANT(F_DWORD,  ERROR_IO_PENDING);
1869     WINAPI_CONSTANT(F_DWORD,  ERROR_NETNAME_DELETED);
1870     WINAPI_CONSTANT(F_DWORD,  ERROR_OPERATION_ABORTED);
1871     WINAPI_CONSTANT(F_DWORD,  ERROR_SEM_TIMEOUT);
1872     WINAPI_CONSTANT(F_DWORD,  ERROR_PIPE_BUSY);
1873     WINAPI_CONSTANT(F_DWORD,  INFINITE);
1874     WINAPI_CONSTANT(F_HANDLE, INVALID_HANDLE_VALUE);
1875     WINAPI_CONSTANT(F_HANDLE, NULL);
1876     WINAPI_CONSTANT(F_DWORD,  SO_UPDATE_ACCEPT_CONTEXT);
1877     WINAPI_CONSTANT(F_DWORD,  SO_UPDATE_CONNECT_CONTEXT);
1878     WINAPI_CONSTANT(F_DWORD,  TF_REUSE_SOCKET);
1879 
1880     return m;
1881 }
1882