diff options
Diffstat (limited to 'Modules')
| -rw-r--r-- | Modules/overlapped.c | 214 |
1 files changed, 156 insertions, 58 deletions
diff --git a/Modules/overlapped.c b/Modules/overlapped.c index 6099d46..9355ce6 100644 --- a/Modules/overlapped.c +++ b/Modules/overlapped.c @@ -37,8 +37,8 @@ #define T_HANDLE T_POINTER -enum {TYPE_NONE, TYPE_NOT_STARTED, TYPE_READ, TYPE_WRITE, TYPE_ACCEPT, - TYPE_CONNECT, TYPE_DISCONNECT, TYPE_CONNECT_NAMED_PIPE, +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}; typedef struct { @@ -51,10 +51,10 @@ typedef struct { /* Type of operation */ DWORD type; union { - /* Buffer used for reading: TYPE_READ and TYPE_ACCEPT */ - PyObject *read_buffer; - /* Buffer used for writing: TYPE_WRITE */ - Py_buffer write_buffer; + /* Buffer allocated by us: TYPE_READ and TYPE_ACCEPT */ + PyObject *allocated_buffer; + /* Buffer passed by the user: TYPE_WRITE and TYPE_READINTO */ + Py_buffer user_buffer; }; } OverlappedObject; @@ -550,9 +550,9 @@ Overlapped_new(PyTypeObject *type, PyObject *args, PyObject *kwds) self->handle = NULL; self->error = 0; self->type = TYPE_NONE; - self->read_buffer = NULL; + self->allocated_buffer = NULL; memset(&self->overlapped, 0, sizeof(OVERLAPPED)); - memset(&self->write_buffer, 0, sizeof(Py_buffer)); + memset(&self->user_buffer, 0, sizeof(Py_buffer)); if (event) self->overlapped.hEvent = event; return (PyObject *)self; @@ -597,11 +597,12 @@ Overlapped_dealloc(OverlappedObject *self) switch (self->type) { case TYPE_READ: case TYPE_ACCEPT: - Py_CLEAR(self->read_buffer); + Py_CLEAR(self->allocated_buffer); break; case TYPE_WRITE: - if (self->write_buffer.obj) - PyBuffer_Release(&self->write_buffer); + case TYPE_READINTO: + if (self->user_buffer.obj) + PyBuffer_Release(&self->user_buffer); break; } PyObject_Del(self); @@ -676,7 +677,7 @@ Overlapped_getresult(OverlappedObject *self, PyObject *args) case ERROR_MORE_DATA: break; case ERROR_BROKEN_PIPE: - if ((self->type == TYPE_READ || self->type == TYPE_ACCEPT) && self->read_buffer != NULL) + if (self->type == TYPE_READ || self->type == TYPE_READINTO) break; /* fall through */ default: @@ -685,17 +686,45 @@ Overlapped_getresult(OverlappedObject *self, PyObject *args) switch (self->type) { case TYPE_READ: - assert(PyBytes_CheckExact(self->read_buffer)); - if (transferred != PyBytes_GET_SIZE(self->read_buffer) && - _PyBytes_Resize(&self->read_buffer, transferred)) + 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->read_buffer); - return self->read_buffer; + Py_INCREF(self->allocated_buffer); + return self->allocated_buffer; 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: + self->type = TYPE_NOT_STARTED; + return SetFromWindowsErr(err); + } +} + PyDoc_STRVAR( Overlapped_ReadFile_doc, "ReadFile(handle, size) -> Overlapped[message]\n\n" @@ -706,10 +735,7 @@ Overlapped_ReadFile(OverlappedObject *self, PyObject *args) { HANDLE handle; DWORD size; - DWORD nread; PyObject *buf; - BOOL ret; - DWORD err; if (!PyArg_ParseTuple(args, F_HANDLE F_DWORD, &handle, &size)) return NULL; @@ -728,14 +754,66 @@ Overlapped_ReadFile(OverlappedObject *self, PyObject *args) self->type = TYPE_READ; self->handle = handle; - self->read_buffer = buf; + self->allocated_buffer = buf; + + return do_ReadFile(self, handle, PyBytes_AS_STRING(buf), size); +} + +PyDoc_STRVAR( + Overlapped_ReadFileInto_doc, + "ReadFileInto(handle, buf) -> Overlapped[bytes_transferred]\n\n" + "Start overlapped receive"); + +static PyObject * +Overlapped_ReadFileInto(OverlappedObject *self, PyObject *args) +{ + HANDLE handle; + PyObject *bufobj; + + if (!PyArg_ParseTuple(args, F_HANDLE "O", &handle, &bufobj)) + return NULL; + + if (self->type != TYPE_NONE) { + PyErr_SetString(PyExc_ValueError, "operation already attempted"); + return NULL; + } + + if (!PyArg_Parse(bufobj, "y*", &self->user_buffer)) + return NULL; + +#if SIZEOF_SIZE_T > SIZEOF_LONG + if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) { + PyBuffer_Release(&self->user_buffer); + PyErr_SetString(PyExc_ValueError, "buffer too large"); + return NULL; + } +#endif + + self->type = TYPE_READINTO; + self->handle = handle; + + return do_ReadFile(self, handle, self->user_buffer.buf, + (DWORD)self->user_buffer.len); +} + +static PyObject * +do_WSARecv(OverlappedObject *self, HANDLE handle, + char *bufstart, DWORD buflen, DWORD flags) +{ + DWORD nread; + WSABUF wsabuf; + int ret; + DWORD err; + + wsabuf.buf = bufstart; + wsabuf.len = buflen; Py_BEGIN_ALLOW_THREADS - ret = ReadFile(handle, PyBytes_AS_STRING(buf), size, &nread, - &self->overlapped); + ret = WSARecv((SOCKET)handle, &wsabuf, 1, &nread, &flags, + &self->overlapped, NULL); Py_END_ALLOW_THREADS - self->error = err = ret ? ERROR_SUCCESS : GetLastError(); + self->error = err = (ret < 0 ? WSAGetLastError() : ERROR_SUCCESS); switch (err) { case ERROR_BROKEN_PIPE: mark_as_completed(&self->overlapped); @@ -761,11 +839,7 @@ Overlapped_WSARecv(OverlappedObject *self, PyObject *args) HANDLE handle; DWORD size; DWORD flags = 0; - DWORD nread; PyObject *buf; - WSABUF wsabuf; - int ret; - DWORD err; if (!PyArg_ParseTuple(args, F_HANDLE F_DWORD "|" F_DWORD, &handle, &size, &flags)) @@ -785,28 +859,48 @@ Overlapped_WSARecv(OverlappedObject *self, PyObject *args) self->type = TYPE_READ; self->handle = handle; - self->read_buffer = buf; - wsabuf.len = size; - wsabuf.buf = PyBytes_AS_STRING(buf); + self->allocated_buffer = buf; - Py_BEGIN_ALLOW_THREADS - ret = WSARecv((SOCKET)handle, &wsabuf, 1, &nread, &flags, - &self->overlapped, NULL); - Py_END_ALLOW_THREADS + return do_WSARecv(self, handle, PyBytes_AS_STRING(buf), size, flags); +} - self->error = err = (ret < 0 ? WSAGetLastError() : ERROR_SUCCESS); - 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: - self->type = TYPE_NOT_STARTED; - return SetFromWindowsErr(err); +PyDoc_STRVAR( + Overlapped_WSARecvInto_doc, + "WSARecvInto(handle, buf, flags) -> Overlapped[bytes_transferred]\n\n" + "Start overlapped receive"); + +static PyObject * +Overlapped_WSARecvInto(OverlappedObject *self, PyObject *args) +{ + HANDLE handle; + PyObject *bufobj; + DWORD flags; + + if (!PyArg_ParseTuple(args, F_HANDLE "O" F_DWORD, + &handle, &bufobj, &flags)) + return NULL; + + if (self->type != TYPE_NONE) { + PyErr_SetString(PyExc_ValueError, "operation already attempted"); + return NULL; + } + + if (!PyArg_Parse(bufobj, "y*", &self->user_buffer)) + return NULL; + +#if SIZEOF_SIZE_T > SIZEOF_LONG + if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) { + PyBuffer_Release(&self->user_buffer); + PyErr_SetString(PyExc_ValueError, "buffer too large"); + return NULL; } +#endif + + self->type = TYPE_READINTO; + self->handle = handle; + + return do_WSARecv(self, handle, self->user_buffer.buf, + (DWORD)self->user_buffer.len, flags); } PyDoc_STRVAR( @@ -831,12 +925,12 @@ Overlapped_WriteFile(OverlappedObject *self, PyObject *args) return NULL; } - if (!PyArg_Parse(bufobj, "y*", &self->write_buffer)) + if (!PyArg_Parse(bufobj, "y*", &self->user_buffer)) return NULL; #if SIZEOF_SIZE_T > SIZEOF_LONG - if (self->write_buffer.len > (Py_ssize_t)ULONG_MAX) { - PyBuffer_Release(&self->write_buffer); + if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) { + PyBuffer_Release(&self->user_buffer); PyErr_SetString(PyExc_ValueError, "buffer too large"); return NULL; } @@ -846,8 +940,8 @@ Overlapped_WriteFile(OverlappedObject *self, PyObject *args) self->handle = handle; Py_BEGIN_ALLOW_THREADS - ret = WriteFile(handle, self->write_buffer.buf, - (DWORD)self->write_buffer.len, + ret = WriteFile(handle, self->user_buffer.buf, + (DWORD)self->user_buffer.len, &written, &self->overlapped); Py_END_ALLOW_THREADS @@ -887,12 +981,12 @@ Overlapped_WSASend(OverlappedObject *self, PyObject *args) return NULL; } - if (!PyArg_Parse(bufobj, "y*", &self->write_buffer)) + if (!PyArg_Parse(bufobj, "y*", &self->user_buffer)) return NULL; #if SIZEOF_SIZE_T > SIZEOF_LONG - if (self->write_buffer.len > (Py_ssize_t)ULONG_MAX) { - PyBuffer_Release(&self->write_buffer); + if (self->user_buffer.len > (Py_ssize_t)ULONG_MAX) { + PyBuffer_Release(&self->user_buffer); PyErr_SetString(PyExc_ValueError, "buffer too large"); return NULL; } @@ -900,8 +994,8 @@ Overlapped_WSASend(OverlappedObject *self, PyObject *args) self->type = TYPE_WRITE; self->handle = handle; - wsabuf.len = (DWORD)self->write_buffer.len; - wsabuf.buf = self->write_buffer.buf; + wsabuf.len = (DWORD)self->user_buffer.len; + wsabuf.buf = self->user_buffer.buf; Py_BEGIN_ALLOW_THREADS ret = WSASend((SOCKET)handle, &wsabuf, 1, &written, flags, @@ -951,7 +1045,7 @@ Overlapped_AcceptEx(OverlappedObject *self, PyObject *args) self->type = TYPE_ACCEPT; self->handle = (HANDLE)ListenSocket; - self->read_buffer = buf; + self->allocated_buffer = buf; Py_BEGIN_ALLOW_THREADS ret = Py_AcceptEx(ListenSocket, AcceptSocket, PyBytes_AS_STRING(buf), @@ -1193,8 +1287,12 @@ static PyMethodDef Overlapped_methods[] = { METH_NOARGS, Overlapped_cancel_doc}, {"ReadFile", (PyCFunction) Overlapped_ReadFile, METH_VARARGS, Overlapped_ReadFile_doc}, + {"ReadFileInto", (PyCFunction) Overlapped_ReadFileInto, + METH_VARARGS, Overlapped_ReadFileInto_doc}, {"WSARecv", (PyCFunction) Overlapped_WSARecv, METH_VARARGS, Overlapped_WSARecv_doc}, + {"WSARecvInto", (PyCFunction) Overlapped_WSARecvInto, + METH_VARARGS, Overlapped_WSARecvInto_doc}, {"WriteFile", (PyCFunction) Overlapped_WriteFile, METH_VARARGS, Overlapped_WriteFile_doc}, {"WSASend", (PyCFunction) Overlapped_WSASend, |