#include "Python.h"

#include "structmember.h" // PyMemberDef

#define WINDOWS_LEAN_AND_MEAN

#include <winsock2.h>

#include <ws2tcpip.h>

#include <mswsock.h>

#if defined(MS_WIN32) && !defined(MS_WIN64)

# define F_POINTER "k"

# define T_POINTER T_ULONG

# define F_POINTER "K"

# define T_POINTER T_ULONGLONG

#define F_HANDLE F_POINTER

#define F_ULONG_PTR F_POINTER

#define F_DWORD "k"

#define F_BOOL "i"

#define F_UINT "I"

#define T_HANDLE T_POINTER

enum {TYPE_NONE, TYPE_NOT_STARTED, TYPE_READ, TYPE_READINTO, TYPE_WRITE,

TYPE_ACCEPT, TYPE_CONNECT, TYPE_DISCONNECT, TYPE_CONNECT_NAMED_PIPE,

TYPE_WAIT_NAMED_PIPE_AND_CONNECT, TYPE_TRANSMIT_FILE, TYPE_READ_FROM,

TYPE_WRITE_TO};

typedef struct {

PyObject_HEAD

OVERLAPPED overlapped;

/* For convenience, we store the file handle too */

HANDLE handle;

/* Error returned by last method call */

DWORD error;

/* Type of operation */

DWORD type;

union {

/* Buffer allocated by us: TYPE_READ and TYPE_ACCEPT */

PyObject *allocated_buffer;

/* Buffer passed by the user: TYPE_WRITE, TYPE_WRITE_TO, and TYPE_READINTO */

Py_buffer user_buffer;

/* Data used for reading from a connectionless socket:

struct {

// A (buffer, (host, port)) tuple

PyObject *result;

// The actual read buffer

PyObject *allocated_buffer;

struct sockaddr_in6 address;

int address_length;

} read_from;

};

} OverlappedObject;

typedef struct {

PyTypeObject *overlapped_type;

} OverlappedState;

static inline OverlappedState*

overlapped_get_state(PyObject *module)

{

void *state = PyModule_GetState(module);

assert(state != NULL);

return (OverlappedState *)state;

}

static PyObject *

SetFromWindowsErr(DWORD err)

{

PyObject *exception_type;

if (err == 0)

err = GetLastError();

switch (err) {

case ERROR_CONNECTION_REFUSED:

exception_type = PyExc_ConnectionRefusedError;

break;

case ERROR_CONNECTION_ABORTED:

exception_type = PyExc_ConnectionAbortedError;

break;

default:

exception_type = PyExc_OSError;

}

return PyErr_SetExcFromWindowsErr(exception_type, err);

}

static LPFN_ACCEPTEX Py_AcceptEx = NULL;

static LPFN_CONNECTEX Py_ConnectEx = NULL;

static LPFN_DISCONNECTEX Py_DisconnectEx = NULL;

static LPFN_TRANSMITFILE Py_TransmitFile = NULL;

static BOOL (CALLBACK *Py_CancelIoEx)(HANDLE, LPOVERLAPPED) = NULL;

#define GET_WSA_POINTER(s, x) \

(SOCKET_ERROR != WSAIoctl(s, SIO_GET_EXTENSION_FUNCTION_POINTER, \

&Guid##x, sizeof(Guid##x), &Py_##x, \

sizeof(Py_##x), &dwBytes, NULL, NULL))

static int

initialize_function_pointers(void)

{

GUID GuidAcceptEx = WSAID_ACCEPTEX;

GUID GuidConnectEx = WSAID_CONNECTEX;

GUID GuidDisconnectEx = WSAID_DISCONNECTEX;

GUID GuidTransmitFile = WSAID_TRANSMITFILE;

HINSTANCE hKernel32;

SOCKET s;

DWORD dwBytes;

s = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);

if (s == INVALID_SOCKET) {

SetFromWindowsErr(WSAGetLastError());

return -1;

}

if (!GET_WSA_POINTER(s, AcceptEx) ||

!GET_WSA_POINTER(s, ConnectEx) ||

!GET_WSA_POINTER(s, DisconnectEx) ||

!GET_WSA_POINTER(s, TransmitFile))

{

closesocket(s);

SetFromWindowsErr(WSAGetLastError());

return -1;

}

closesocket(s);

/* On WinXP we will have Py_CancelIoEx == NULL */

Py_BEGIN_ALLOW_THREADS

hKernel32 = GetModuleHandle("KERNEL32");

*(FARPROC *)&Py_CancelIoEx = GetProcAddress(hKernel32, "CancelIoEx");

Py_END_ALLOW_THREADS

return 0;

}

static PyObject *

_overlapped_CreateIoCompletionPort_impl(PyObject *module, HANDLE FileHandle,

HANDLE ExistingCompletionPort,

ULONG_PTR CompletionKey,

DWORD NumberOfConcurrentThreads)

{

HANDLE ret;

Py_BEGIN_ALLOW_THREADS

ret = CreateIoCompletionPort(FileHandle, ExistingCompletionPort,

CompletionKey, NumberOfConcurrentThreads);

Py_END_ALLOW_THREADS

if (ret == NULL)

return SetFromWindowsErr(0);

return Py_BuildValue(F_HANDLE, ret);

}

static PyObject *

_overlapped_GetQueuedCompletionStatus_impl(PyObject *module,

HANDLE CompletionPort,

DWORD Milliseconds)

{

DWORD NumberOfBytes = 0;

ULONG_PTR CompletionKey = 0;

OVERLAPPED *Overlapped = NULL;

DWORD err;

BOOL ret;

Py_BEGIN_ALLOW_THREADS

ret = GetQueuedCompletionStatus(CompletionPort, &NumberOfBytes,

&CompletionKey, &Overlapped, Milliseconds);

Py_END_ALLOW_THREADS

err = ret ? ERROR_SUCCESS : GetLastError();

if (Overlapped == NULL) {

if (err == WAIT_TIMEOUT)

Py_RETURN_NONE;

else

return SetFromWindowsErr(err);

}

return Py_BuildValue(F_DWORD F_DWORD F_ULONG_PTR F_POINTER,

err, NumberOfBytes, CompletionKey, Overlapped);

}

static PyObject *

_overlapped_PostQueuedCompletionStatus_impl(PyObject *module,

OVERLAPPED *Overlapped)

{

BOOL ret;

Py_BEGIN_ALLOW_THREADS

ret = PostQueuedCompletionStatus(CompletionPort, NumberOfBytes,

CompletionKey, Overlapped);

Py_END_ALLOW_THREADS

if (!ret)

return SetFromWindowsErr(0);

Py_RETURN_NONE;

}

struct PostCallbackData {

HANDLE CompletionPort;

LPOVERLAPPED Overlapped;

};

static VOID CALLBACK

PostToQueueCallback(PVOID lpParameter, BOOLEAN TimerOrWaitFired)

{

struct PostCallbackData *p = (struct PostCallbackData*) lpParameter;

PostQueuedCompletionStatus(p->CompletionPort, TimerOrWaitFired,

0, p->Overlapped);

/* ignore possible error! */

PyMem_RawFree(p);

}

static PyObject *

_overlapped_RegisterWaitWithQueue_impl(PyObject *module, HANDLE Object,

HANDLE CompletionPort,

OVERLAPPED *Overlapped,

DWORD Milliseconds)

{

HANDLE NewWaitObject;

struct PostCallbackData data = {CompletionPort, Overlapped}, *pdata;

/* Use PyMem_RawMalloc() rather than PyMem_Malloc(), since

pdata = PyMem_RawMalloc(sizeof(struct PostCallbackData));

if (pdata == NULL)

return SetFromWindowsErr(0);

*pdata = data;

if (!RegisterWaitForSingleObject(

&NewWaitObject, Object, PostToQueueCallback, pdata, Milliseconds,

WT_EXECUTEINWAITTHREAD | WT_EXECUTEONLYONCE))

{

PyMem_RawFree(pdata);

return SetFromWindowsErr(0);

}

return Py_BuildValue(F_HANDLE, NewWaitObject);

}

static PyObject *

_overlapped_UnregisterWait_impl(PyObject *module, HANDLE WaitHandle)

{

BOOL ret;

Py_BEGIN_ALLOW_THREADS

ret = UnregisterWait(WaitHandle);

Py_END_ALLOW_THREADS

if (!ret)

return SetFromWindowsErr(0);

Py_RETURN_NONE;

}

static PyObject *

_overlapped_UnregisterWaitEx_impl(PyObject *module, HANDLE WaitHandle,

HANDLE Event)

{

BOOL ret;

Py_BEGIN_ALLOW_THREADS

ret = UnregisterWaitEx(WaitHandle, Event);

Py_END_ALLOW_THREADS

if (!ret)

return SetFromWindowsErr(0);

Py_RETURN_NONE;

}

static PyObject *

_overlapped_CreateEvent_impl(PyObject *module, PyObject *EventAttributes,

BOOL ManualReset, BOOL InitialState,

const Py_UNICODE *Name)

{

HANDLE Event;

if (EventAttributes != Py_None) {

PyErr_SetString(PyExc_ValueError, "EventAttributes must be None");

return NULL;

}

Py_BEGIN_ALLOW_THREADS

Event = CreateEventW(NULL, ManualReset, InitialState, Name);

Py_END_ALLOW_THREADS

if (Event == NULL)

return SetFromWindowsErr(0);

return Py_BuildValue(F_HANDLE, Event);

}

static PyObject *

_overlapped_SetEvent_impl(PyObject *module, HANDLE Handle)

{

BOOL ret;

Py_BEGIN_ALLOW_THREADS

ret = SetEvent(Handle);

Py_END_ALLOW_THREADS

if (!ret)

return SetFromWindowsErr(0);

Py_RETURN_NONE;

}

static PyObject *

_overlapped_ResetEvent_impl(PyObject *module, HANDLE Handle)

{

BOOL ret;

Py_BEGIN_ALLOW_THREADS

ret = ResetEvent(Handle);

Py_END_ALLOW_THREADS

if (!ret)

return SetFromWindowsErr(0);

Py_RETURN_NONE;

}

static PyObject *

_overlapped_BindLocal_impl(PyObject *module, HANDLE Socket, int Family)

{

BOOL ret;

if (Family == AF_INET) {

struct sockaddr_in addr;

memset(&addr, 0, sizeof(addr));

addr.sin_family = AF_INET;

addr.sin_port = 0;

addr.sin_addr.S_un.S_addr = INADDR_ANY;

ret = bind((SOCKET)Socket, (SOCKADDR*)&addr, sizeof(addr))

!= SOCKET_ERROR;

} else if (Family == AF_INET6) {

struct sockaddr_in6 addr;

memset(&addr, 0, sizeof(addr));

addr.sin6_family = AF_INET6;

addr.sin6_port = 0;

addr.sin6_addr = in6addr_any;

ret = bind((SOCKET)Socket, (SOCKADDR*)&addr, sizeof(addr))

!= SOCKET_ERROR;

} else {

PyErr_SetString(PyExc_ValueError, "expected tuple of length 2 or 4");

return NULL;

}

if (!ret)

return SetFromWindowsErr(WSAGetLastError());

Py_RETURN_NONE;

}

static PyObject *

_overlapped_FormatMessage_impl(PyObject *module, DWORD code)

{

DWORD n;

WCHAR *lpMsgBuf;

PyObject *res;

n = FormatMessageW(FORMAT_MESSAGE_ALLOCATE_BUFFER |

FORMAT_MESSAGE_FROM_SYSTEM |

FORMAT_MESSAGE_IGNORE_INSERTS,

NULL,

code,

MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),

(LPWSTR) &lpMsgBuf,

0,

NULL);

if (n) {

while (iswspace(lpMsgBuf[n-1]))

--n;

lpMsgBuf[n] = L'\0';

res = Py_BuildValue("u", lpMsgBuf);

} else {

res = PyUnicode_FromFormat("unknown error code %u", code);

}

LocalFree(lpMsgBuf);

return res;

}

static void

mark_as_completed(OVERLAPPED *ov)

{

ov->Internal = 0;

if (ov->hEvent != NULL)

SetEvent(ov->hEvent);

}

static PyObject *

_overlapped_Overlapped_impl(PyTypeObject *type, HANDLE event)

{

OverlappedObject *self;

if (event == INVALID_HANDLE_VALUE) {

event = CreateEvent(NULL, TRUE, FALSE, NULL);

if (event == NULL)

return SetFromWindowsErr(0);

}

self = PyObject_New(OverlappedObject, type);

if (self == NULL) {

if (event != NULL)

CloseHandle(event);

return NULL;

}

self->handle = NULL;

self->error = 0;

self->type = TYPE_NONE;

self->allocated_buffer = NULL;

memset(&self->overlapped, 0, sizeof(OVERLAPPED));

memset(&self->user_buffer, 0, sizeof(Py_buffer));

if (event)

self->overlapped.hEvent = event;

return (PyObject *)self;

}

static int

Overlapped_clear(OverlappedObject *self)

{

switch (self->type) {

case TYPE_READ:

case TYPE_ACCEPT: {

Py_CLEAR(self->allocated_buffer);

break;

}

case TYPE_READ_FROM: {

// An initial call to WSARecvFrom will only allocate the buffer.

// The result tuple of (message, address) is only

// allocated _after_ a message has been received.

if(self->read_from.result) {

// We've received a message, free the result tuple.

Py_CLEAR(self->read_from.result);

}

if(self->read_from.allocated_buffer) {

Py_CLEAR(self->read_from.allocated_buffer);

}

break;

}

case TYPE_WRITE:

case TYPE_WRITE_TO:

case TYPE_READINTO: {

if (self->user_buffer.obj) {

PyBuffer_Release(&self->user_buffer);

}

break;

}

}

self->type = TYPE_NOT_STARTED;

return 0;

}

static void

Overlapped_dealloc(OverlappedObject *self)

{

DWORD bytes;

DWORD olderr = GetLastError();

BOOL wait = FALSE;

BOOL ret;

if (!HasOverlappedIoCompleted(&self->overlapped) &&

self->type != TYPE_NOT_STARTED)

{

if (Py_CancelIoEx && Py_CancelIoEx(self->handle, &self->overlapped))

wait = TRUE;

Py_BEGIN_ALLOW_THREADS

ret = GetOverlappedResult(self->handle, &self->overlapped,

&bytes, wait);

Py_END_ALLOW_THREADS

switch (ret ? ERROR_SUCCESS : GetLastError()) {

case ERROR_SUCCESS:

case ERROR_NOT_FOUND:

case ERROR_OPERATION_ABORTED:

break;

default:

PyErr_Format(

PyExc_RuntimeError,

"%R still has pending operation at "

"deallocation, the process may crash", self);

PyErr_WriteUnraisable(NULL);

}

}

if (self->overlapped.hEvent != NULL) {

CloseHandle(self->overlapped.hEvent);

}

Overlapped_clear(self);

SetLastError(olderr);

PyTypeObject *tp = Py_TYPE(self);

PyObject_Free(self);

Py_DECREF(tp);

}

static PyObject *

make_ipv4_addr(const struct sockaddr_in *addr)

{

char buf[INET_ADDRSTRLEN];

if (inet_ntop(AF_INET, &addr->sin_addr, buf, sizeof(buf)) == NULL) {

PyErr_SetFromErrno(PyExc_OSError);

return NULL;

}

return PyUnicode_FromString(buf);

}

static PyObject *

make_ipv6_addr(const struct sockaddr_in6 *addr)

{

char buf[INET6_ADDRSTRLEN];

if (inet_ntop(AF_INET6, &addr->sin6_addr, buf, sizeof(buf)) == NULL) {

PyErr_SetFromErrno(PyExc_OSError);

return NULL;

}

return PyUnicode_FromString(buf);

}

static PyObject*

unparse_address(LPSOCKADDR Address, DWORD Length)

{

/* The function is adopted from mocketmodule.c makesockaddr()*/

switch(Address->sa_family) {

case AF_INET: {

const struct sockaddr_in *a = (const struct sockaddr_in *)Address;

PyObject *addrobj = make_ipv4_addr(a);

PyObject *ret = NULL;

if (addrobj) {

ret = Py_BuildValue("Oi", addrobj, ntohs(a->sin_port));

Py_DECREF(addrobj);

}

return ret;

}

case AF_INET6: {

const struct sockaddr_in6 *a = (const struct sockaddr_in6 *)Address;

PyObject *addrobj = make_ipv6_addr(a);

PyObject *ret = NULL;

if (addrobj) {

ret = Py_BuildValue("OiII",

addrobj,

ntohs(a->sin6_port),

ntohl(a->sin6_flowinfo),

a->sin6_scope_id);

Py_DECREF(addrobj);

}

return ret;

}

default: {

PyErr_SetString(PyExc_ValueError, "recvfrom returned unsupported address family");

return NULL;

}

}

}

static PyObject *

_overlapped_Overlapped_cancel_impl(OverlappedObject *self)

{

BOOL ret = TRUE;

if (self->type == TYPE_NOT_STARTED

|| self->type == TYPE_WAIT_NAMED_PIPE_AND_CONNECT)

Py_RETURN_NONE;

if (!HasOverlappedIoCompleted(&self->overlapped)) {

Py_BEGIN_ALLOW_THREADS

if (Py_CancelIoEx)

ret = Py_CancelIoEx(self->handle, &self->overlapped);

else

ret = CancelIo(self->handle);

Py_END_ALLOW_THREADS

}

/* CancelIoEx returns ERROR_NOT_FOUND if the I/O completed in-between */

if (!ret && GetLastError() != ERROR_NOT_FOUND)

return SetFromWindowsErr(0);

Py_RETURN_NONE;

}

static PyObject *

_overlapped_Overlapped_getresult_impl(OverlappedObject *self, BOOL wait)

{

DWORD transferred = 0;

BOOL ret;

DWORD err;

PyObject *addr;

if (self->type == TYPE_NONE) {

PyErr_SetString(PyExc_ValueError, "operation not yet attempted");

return NULL;

}

if (self->type == TYPE_NOT_STARTED) {

PyErr_SetString(PyExc_ValueError, "operation failed to start");

return NULL;

}

Py_BEGIN_ALLOW_THREADS

ret = GetOverlappedResult(self->handle, &self->overlapped, &transferred,

wait);

Py_END_ALLOW_THREADS

self->error = err = ret ? ERROR_SUCCESS : GetLastError();

switch (err) {

case ERROR_SUCCESS:

case ERROR_MORE_DATA:

break;

case ERROR_BROKEN_PIPE:

if (self->type == TYPE_READ || self->type == TYPE_READINTO) {

break;

}

else if (self->type == TYPE_READ_FROM &&

(self->read_from.result != NULL ||

self->read_from.allocated_buffer != NULL))

{

break;

}

/* fall through */

default:

return SetFromWindowsErr(err);

}

switch (self->type) {

case TYPE_READ:

assert(PyBytes_CheckExact(self->allocated_buffer));

if (transferred != PyBytes_GET_SIZE(self->allocated_buffer) &&

_PyBytes_Resize(&self->allocated_buffer, transferred))

return NULL;

Py_INCREF(self->allocated_buffer);

return self->allocated_buffer;

case TYPE_READ_FROM:

assert(PyBytes_CheckExact(self->read_from.allocated_buffer));

if (transferred != PyBytes_GET_SIZE(

self->read_from.allocated_buffer) &&

_PyBytes_Resize(&self->read_from.allocated_buffer, transferred))

{

return NULL;

}

// unparse the address

addr = unparse_address((SOCKADDR*)&self->read_from.address,

self->read_from.address_length);

if (addr == NULL) {

return NULL;

}

// The result is a two item tuple: (message, address)

self->read_from.result = PyTuple_New(2);

if (self->read_from.result == NULL) {

Py_CLEAR(addr);

return NULL;

}

// first item: message

Py_INCREF(self->read_from.allocated_buffer);

PyTuple_SET_ITEM(self->read_from.result, 0,

self->read_from.allocated_buffer);

// second item: address

PyTuple_SET_ITEM(self->read_from.result, 1, addr);

Py_INCREF(self->read_from.result);

return self->read_from.result;

default:

return PyLong_FromUnsignedLong((unsigned long) transferred);

}

}

static PyObject *

do_ReadFile(OverlappedObject *self, HANDLE handle,

char *bufstart, DWORD buflen)

{

DWORD nread;

int ret;

DWORD err;

Py_BEGIN_ALLOW_THREADS

ret = ReadFile(handle, bufstart, buflen, &nread,

&self->overlapped);

Py_END_ALLOW_THREADS

self->error = err = ret ? ERROR_SUCCESS : GetLastError();

switch (err) {

case ERROR_BROKEN_PIPE:

mark_as_completed(&self->overlapped);

return SetFromWindowsErr(err);

case ERROR_SUCCESS:

case ERROR_MORE_DATA:

case ERROR_IO_PENDING:

Py_RETURN_NONE;

default:

Overlapped_clear(self);

return SetFromWindowsErr(err);

}

}

static PyObject *

_overlapped_Overlapped_ReadFile_impl(OverlappedObject *self, HANDLE handle,

DWORD size)

{

PyObject *buf;

if (self->type != TYPE_NONE) {

PyErr_SetString(PyExc_ValueError, "operation already attempted");

return NULL;

}

#if SIZEOF_SIZE_T <= SIZEOF_LONG

size = Py_MIN(size, (DWORD)PY_SSIZE_T_MAX);

buf = PyBytes_FromStringAndSize(NULL, Py_MAX(size, 1));

if (buf == NULL)

return NULL;

self->type = TYPE_READ;

self->handle = handle;

self->allocated_buffer = buf;

return do_ReadFile(self, handle, PyBytes_AS_STRING(buf), size);

}

static PyObject *

_overlapped_Overlapped_ReadFileInto_impl(OverlappedObject *self,

HANDLE handle, PyObject *bufobj)

{