#include "Python.h"

#include <structmember.h>

#ifdef __APPLE__

/* Perform runtime testing for a broken poll on OSX to make it easier

#undef HAVE_BROKEN_POLL

#if defined(MS_WINDOWS) && !defined(FD_SETSIZE)

#define FD_SETSIZE 512

#if defined(HAVE_POLL_H)

#include <poll.h>

#elif defined(HAVE_SYS_POLL_H)

#include <sys/poll.h>

#ifdef __sgi

extern void bzero(void *, int);

#ifdef HAVE_SYS_TYPES_H

#include <sys/types.h>

#if defined(PYOS_OS2) && !defined(PYCC_GCC)

#include <sys/time.h>

#include <utils.h>

#ifdef MS_WINDOWS

# include <winsock2.h>

# define SOCKET int

# ifdef __BEOS__

# include <net/socket.h>

# elif defined(__VMS)

# include <socket.h>

static PyObject *SelectError;

typedef struct {

PyObject *obj; /* owned reference */

SOCKET fd;

int sentinel; /* -1 == sentinel */

} pylist;

static void

reap_obj(pylist fd2obj[FD_SETSIZE + 1])

{

int i;

for (i = 0; i < FD_SETSIZE + 1 && fd2obj[i].sentinel >= 0; i++) {

Py_XDECREF(fd2obj[i].obj);

fd2obj[i].obj = NULL;

}

fd2obj[0].sentinel = -1;

}

static int

seq2set(PyObject *seq, fd_set *set, pylist fd2obj[FD_SETSIZE + 1])

{

int i;

int max = -1;

int index = 0;

PyObject* fast_seq = NULL;

PyObject* o = NULL;

fd2obj[0].obj = (PyObject*)0; /* set list to zero size */

FD_ZERO(set);

fast_seq = PySequence_Fast(seq, "arguments 1-3 must be sequences");

if (!fast_seq)

return -1;

for (i = 0; i < PySequence_Fast_GET_SIZE(fast_seq); i++) {

SOCKET v;

/* any intervening fileno() calls could decr this refcnt */

if (!(o = PySequence_Fast_GET_ITEM(fast_seq, i)))

return -1;

Py_INCREF(o);

v = PyObject_AsFileDescriptor( o );

if (v == -1) goto finally;

#if defined(_MSC_VER)

max = 0; /* not used for Win32 */

#else /* !_MSC_VER */

if (!_PyIsSelectable_fd(v)) {

PyErr_SetString(PyExc_ValueError,

"filedescriptor out of range in select()");

goto finally;

}

if (v > max)

max = v;

#endif /* _MSC_VER */

FD_SET(v, set);

/* add object and its file descriptor to the list */

if (index >= FD_SETSIZE) {

PyErr_SetString(PyExc_ValueError,

"too many file descriptors in select()");

goto finally;

}

fd2obj[index].obj = o;

fd2obj[index].fd = v;

fd2obj[index].sentinel = 0;

fd2obj[++index].sentinel = -1;

}

Py_DECREF(fast_seq);

return max+1;

finally:

Py_XDECREF(o);

Py_DECREF(fast_seq);

return -1;

}

static PyObject *

set2list(fd_set *set, pylist fd2obj[FD_SETSIZE + 1])

{

int i, j, count=0;

PyObject *list, *o;

SOCKET fd;

for (j = 0; fd2obj[j].sentinel >= 0; j++) {

if (FD_ISSET(fd2obj[j].fd, set))

count++;

}

list = PyList_New(count);

if (!list)

return NULL;

i = 0;

for (j = 0; fd2obj[j].sentinel >= 0; j++) {

fd = fd2obj[j].fd;

if (FD_ISSET(fd, set)) {

o = fd2obj[j].obj;

fd2obj[j].obj = NULL;

/* transfer ownership */

if (PyList_SetItem(list, i, o) < 0)

goto finally;

i++;

}

}

return list;

finally:

Py_DECREF(list);

return NULL;

}

#undef SELECT_USES_HEAP

#if FD_SETSIZE > 1024

#define SELECT_USES_HEAP

#endif /* FD_SETSIZE > 1024 */

static PyObject *

select_select(PyObject *self, PyObject *args)

{

#ifdef SELECT_USES_HEAP

pylist *rfd2obj, *wfd2obj, *efd2obj;

#else /* !SELECT_USES_HEAP */

/* XXX: All this should probably be implemented as follows:

pylist rfd2obj[FD_SETSIZE + 1];

pylist wfd2obj[FD_SETSIZE + 1];

pylist efd2obj[FD_SETSIZE + 1];

#endif /* SELECT_USES_HEAP */

PyObject *ifdlist, *ofdlist, *efdlist;

PyObject *ret = NULL;

PyObject *tout = Py_None;

fd_set ifdset, ofdset, efdset;

double timeout;

struct timeval tv, *tvp;

long seconds;

int imax, omax, emax, max;

int n;

/* convert arguments */

if (!PyArg_UnpackTuple(args, "select", 3, 4,

&ifdlist, &ofdlist, &efdlist, &tout))

return NULL;

if (tout == Py_None)

tvp = (struct timeval *)0;

else if (!PyNumber_Check(tout)) {

PyErr_SetString(PyExc_TypeError,

"timeout must be a float or None");

return NULL;

}

else {

timeout = PyFloat_AsDouble(tout);

if (timeout == -1 && PyErr_Occurred())

return NULL;

if (timeout > (double)LONG_MAX) {

PyErr_SetString(PyExc_OverflowError,

"timeout period too long");

return NULL;

}

seconds = (long)timeout;

timeout = timeout - (double)seconds;

tv.tv_sec = seconds;

tv.tv_usec = (long)(timeout * 1E6);

tvp = &tv;

}

#ifdef SELECT_USES_HEAP

/* Allocate memory for the lists */

rfd2obj = PyMem_NEW(pylist, FD_SETSIZE + 1);

wfd2obj = PyMem_NEW(pylist, FD_SETSIZE + 1);

efd2obj = PyMem_NEW(pylist, FD_SETSIZE + 1);

if (rfd2obj == NULL || wfd2obj == NULL || efd2obj == NULL) {

if (rfd2obj) PyMem_DEL(rfd2obj);

if (wfd2obj) PyMem_DEL(wfd2obj);

if (efd2obj) PyMem_DEL(efd2obj);

return PyErr_NoMemory();

}

#endif /* SELECT_USES_HEAP */

/* Convert sequences to fd_sets, and get maximum fd number

rfd2obj[0].sentinel = -1;

wfd2obj[0].sentinel = -1;

efd2obj[0].sentinel = -1;

if ((imax=seq2set(ifdlist, &ifdset, rfd2obj)) < 0)

goto finally;

if ((omax=seq2set(ofdlist, &ofdset, wfd2obj)) < 0)

goto finally;

if ((emax=seq2set(efdlist, &efdset, efd2obj)) < 0)

goto finally;

max = imax;

if (omax > max) max = omax;

if (emax > max) max = emax;

Py_BEGIN_ALLOW_THREADS

n = select(max, &ifdset, &ofdset, &efdset, tvp);

Py_END_ALLOW_THREADS

#ifdef MS_WINDOWS

if (n == SOCKET_ERROR) {

PyErr_SetExcFromWindowsErr(SelectError, WSAGetLastError());

}

if (n < 0) {

PyErr_SetFromErrno(SelectError);

}

else {

/* any of these three calls can raise an exception. it's more

ifdlist = set2list(&ifdset, rfd2obj);

ofdlist = set2list(&ofdset, wfd2obj);

efdlist = set2list(&efdset, efd2obj);

if (PyErr_Occurred())

ret = NULL;

else

ret = PyTuple_Pack(3, ifdlist, ofdlist, efdlist);

Py_DECREF(ifdlist);

Py_DECREF(ofdlist);

Py_DECREF(efdlist);

}

finally:

reap_obj(rfd2obj);

reap_obj(wfd2obj);

reap_obj(efd2obj);

#ifdef SELECT_USES_HEAP

PyMem_DEL(rfd2obj);

PyMem_DEL(wfd2obj);

PyMem_DEL(efd2obj);

#endif /* SELECT_USES_HEAP */

return ret;

}

#if defined(HAVE_POLL) && !defined(HAVE_BROKEN_POLL)

typedef struct {

PyObject_HEAD

PyObject *dict;

int ufd_uptodate;

int ufd_len;

struct pollfd *ufds;

} pollObject;

static PyTypeObject poll_Type;

static int

update_ufd_array(pollObject *self)

{

Py_ssize_t i, pos;

PyObject *key, *value;

struct pollfd *old_ufds = self->ufds;

self->ufd_len = PyDict_Size(self->dict);

PyMem_RESIZE(self->ufds, struct pollfd, self->ufd_len);

if (self->ufds == NULL) {

self->ufds = old_ufds;

PyErr_NoMemory();

return 0;

}

i = pos = 0;

while (PyDict_Next(self->dict, &pos, &key, &value)) {

assert(i < self->ufd_len);

/* Never overflow */

self->ufds[i].fd = (int)PyInt_AsLong(key);

self->ufds[i].events = (short)PyInt_AsLong(value);

i++;

}

assert(i == self->ufd_len);

self->ufd_uptodate = 1;

return 1;

}

PyDoc_STRVAR(poll_register_doc,

events -- an optional bitmask describing the type of events to check for");

static PyObject *

poll_register(pollObject *self, PyObject *args)

{

PyObject *o, *key, *value;

int fd;

short events = POLLIN | POLLPRI | POLLOUT;

int err;

if (!PyArg_ParseTuple(args, "O|h:register", &o, &events)) {

return NULL;

}

fd = PyObject_AsFileDescriptor(o);

if (fd == -1) return NULL;

/* Add entry to the internal dictionary: the key is the

key = PyInt_FromLong(fd);

if (key == NULL)

return NULL;

value = PyInt_FromLong(events);

if (value == NULL) {

Py_DECREF(key);

return NULL;

}

err = PyDict_SetItem(self->dict, key, value);

Py_DECREF(key);

Py_DECREF(value);

if (err < 0)

return NULL;

self->ufd_uptodate = 0;

Py_INCREF(Py_None);

return Py_None;

}

PyDoc_STRVAR(poll_modify_doc,

events -- an optional bitmask describing the type of events to check for");

static PyObject *

poll_modify(pollObject *self, PyObject *args)

{

PyObject *o, *key, *value;

int fd, events;

int err;

if (!PyArg_ParseTuple(args, "Oi:modify", &o, &events)) {

return NULL;

}

fd = PyObject_AsFileDescriptor(o);

if (fd == -1) return NULL;

/* Modify registered fd */

key = PyInt_FromLong(fd);

if (key == NULL)

return NULL;

if (PyDict_GetItem(self->dict, key) == NULL) {

errno = ENOENT;

PyErr_SetFromErrno(PyExc_IOError);

return NULL;

}

value = PyInt_FromLong(events);

if (value == NULL) {

Py_DECREF(key);

return NULL;

}

err = PyDict_SetItem(self->dict, key, value);

Py_DECREF(key);

Py_DECREF(value);

if (err < 0)

return NULL;

self->ufd_uptodate = 0;

Py_INCREF(Py_None);

return Py_None;

}

PyDoc_STRVAR(poll_unregister_doc,

Remove a file descriptor being tracked by the polling object.");

static PyObject *

poll_unregister(pollObject *self, PyObject *o)

{

PyObject *key;

int fd;

fd = PyObject_AsFileDescriptor( o );

if (fd == -1)

return NULL;

/* Check whether the fd is already in the array */

key = PyInt_FromLong(fd);

if (key == NULL)

return NULL;

if (PyDict_DelItem(self->dict, key) == -1) {

Py_DECREF(key);

/* This will simply raise the KeyError set by PyDict_DelItem

return NULL;

}

Py_DECREF(key);

self->ufd_uptodate = 0;

Py_INCREF(Py_None);

return Py_None;

}

PyDoc_STRVAR(poll_poll_doc,

any descriptors that have events or errors to report.");

static PyObject *

poll_poll(pollObject *self, PyObject *args)

{

PyObject *result_list = NULL, *tout = NULL;

int timeout = 0, poll_result, i, j;

PyObject *value = NULL, *num = NULL;

if (!PyArg_UnpackTuple(args, "poll", 0, 1, &tout)) {

return NULL;

}

/* Check values for timeout */

if (tout == NULL || tout == Py_None)

timeout = -1;

else if (!PyNumber_Check(tout)) {

PyErr_SetString(PyExc_TypeError,

"timeout must be an integer or None");

return NULL;

}

else {

tout = PyNumber_Int(tout);

if (!tout)

return NULL;

timeout = _PyInt_AsInt(tout);

Py_DECREF(tout);

if (timeout == -1 && PyErr_Occurred())

return NULL;

}

/* Ensure the ufd array is up to date */

if (!self->ufd_uptodate)

if (update_ufd_array(self) == 0)

return NULL;

/* call poll() */

Py_BEGIN_ALLOW_THREADS

poll_result = poll(self->ufds, self->ufd_len, timeout);

Py_END_ALLOW_THREADS

if (poll_result < 0) {

PyErr_SetFromErrno(SelectError);

return NULL;

}

/* build the result list */

result_list = PyList_New(poll_result);

if (!result_list)

return NULL;

else {

for (i = 0, j = 0; j < poll_result; j++) {

/* skip to the next fired descriptor */

while (!self->ufds[i].revents) {

i++;

}

/* if we hit a NULL return, set value to NULL

value = PyTuple_New(2);

if (value == NULL)

goto error;

num = PyInt_FromLong(self->ufds[i].fd);

if (num == NULL) {

Py_DECREF(value);

goto error;

}

PyTuple_SET_ITEM(value, 0, num);

/* The &0xffff is a workaround for AIX. 'revents'

num = PyInt_FromLong(self->ufds[i].revents & 0xffff);

if (num == NULL) {

Py_DECREF(value);

goto error;

}

PyTuple_SET_ITEM(value, 1, num);

if ((PyList_SetItem(result_list, j, value)) == -1) {

Py_DECREF(value);

goto error;

}

i++;

}

}

return result_list;

error:

Py_DECREF(result_list);

return NULL;

}

static PyMethodDef poll_methods[] = {

{"register", (PyCFunction)poll_register,

METH_VARARGS, poll_register_doc},

{"modify", (PyCFunction)poll_modify,

METH_VARARGS, poll_modify_doc},

{"unregister", (PyCFunction)poll_unregister,

METH_O, poll_unregister_doc},

{"poll", (PyCFunction)poll_poll,

METH_VARARGS, poll_poll_doc},

{NULL, NULL} /* sentinel */

};

static pollObject *

newPollObject(void)

{

pollObject *self;

self = PyObject_New(pollObject, &poll_Type);

if (self == NULL)

return NULL;

/* ufd_uptodate is a Boolean, denoting whether the

self->ufd_uptodate = 0;

self->ufds = NULL;

self->dict = PyDict_New();

if (self->dict == NULL) {

Py_DECREF(self);

return NULL;

}

return self;

}

static void

poll_dealloc(pollObject *self)

{

if (self->ufds != NULL)

PyMem_DEL(self->ufds);

Py_XDECREF(self->dict);

PyObject_Del(self);

}

static PyObject *

poll_getattr(pollObject *self, char *name)

{

return Py_FindMethod(poll_methods, (PyObject *)self, name);

}

static PyTypeObject poll_Type = {

/* The ob_type field must be initialized in the module init function

PyVarObject_HEAD_INIT(NULL, 0)

"select.poll", /*tp_name*/

sizeof(pollObject), /*tp_basicsize*/

0, /*tp_itemsize*/

/* methods */

(destructor)poll_dealloc, /*tp_dealloc*/

0, /*tp_print*/

(getattrfunc)poll_getattr, /*tp_getattr*/

0, /*tp_setattr*/

0, /*tp_compare*/

0, /*tp_repr*/

0, /*tp_as_number*/

0, /*tp_as_sequence*/

0, /*tp_as_mapping*/

0, /*tp_hash*/

};

PyDoc_STRVAR(poll_doc,

unregistering file descriptors, and then polling them for I/O events.");

static PyObject *

select_poll(PyObject *self, PyObject *unused)

{

return (PyObject *)newPollObject();

}

#ifdef __APPLE__

static int select_have_broken_poll(void)

{

int poll_test;

int filedes[2];

struct pollfd poll_struct = { 0, POLLIN|POLLPRI|POLLOUT, 0 };

/* Create a file descriptor to make invalid */

if (pipe(filedes) < 0) {

return 1;

}

poll_struct.fd = filedes[0];

close(filedes[0]);

close(filedes[1]);

poll_test = poll(&poll_struct, 1, 0);

if (poll_test < 0) {

return 1;

} else if (poll_test == 0 && poll_struct.revents != POLLNVAL) {

return 1;

}

return 0;

}

#endif /* __APPLE__ */

#endif /* HAVE_POLL */

#ifdef HAVE_EPOLL

#ifdef HAVE_SYS_EPOLL_H

#include <sys/epoll.h>

typedef struct {

PyObject_HEAD

SOCKET epfd; /* epoll control file descriptor */

} pyEpoll_Object;

static PyTypeObject pyEpoll_Type;

#define pyepoll_CHECK(op) (PyObject_TypeCheck((op), &pyEpoll_Type))

static PyObject *

pyepoll_err_closed(void)

{

PyErr_SetString(PyExc_ValueError, "I/O operation on closed epoll fd");

return NULL;

}

static int

pyepoll_internal_close(pyEpoll_Object *self)

{

int save_errno = 0;

if (self->epfd >= 0) {

int epfd = self->epfd;

self->epfd = -1;

Py_BEGIN_ALLOW_THREADS

if (close(epfd) < 0)

save_errno = errno;

Py_END_ALLOW_THREADS

}

return save_errno;

}

static PyObject *

newPyEpoll_Object(PyTypeObject *type, int sizehint, SOCKET fd)

{

pyEpoll_Object *self;

if (sizehint == -1) {

sizehint = FD_SETSIZE-1;

}

else if (sizehint < 1) {

PyErr_Format(PyExc_ValueError,

"sizehint must be greater zero, got %d",

sizehint);

return NULL;

}

assert(type != NULL && type->tp_alloc != NULL);

self = (pyEpoll_Object *) type->tp_alloc(type, 0);

if (self == NULL)

return NULL;

if (fd == -1) {

Py_BEGIN_ALLOW_THREADS

self->epfd = epoll_create(sizehint);

Py_END_ALLOW_THREADS

}

else {

self->epfd = fd;

}

if (self->epfd < 0) {

Py_DECREF(self);

PyErr_SetFromErrno(PyExc_IOError);

return NULL;

}

return (PyObject *)self;

}

static PyObject *

pyepoll_new(PyTypeObject *type, PyObject *args, PyObject *kwds)

{

int sizehint = -1;

static char *kwlist[] = {"sizehint", NULL};

if (!PyArg_ParseTupleAndKeywords(args, kwds, "|i:epoll", kwlist,

&sizehint))

return NULL;

return newPyEpoll_Object(type, sizehint, -1);

}

static void

pyepoll_dealloc(pyEpoll_Object *self)

{

(void)pyepoll_internal_close(self);

Py_TYPE(self)->tp_free(self);

}

static PyObject*

pyepoll_close(pyEpoll_Object *self)

{

errno = pyepoll_internal_close(self);

if (errno < 0) {

PyErr_SetFromErrno(PyExc_IOError);

return NULL;

}

Py_RETURN_NONE;

}

PyDoc_STRVAR(pyepoll_close_doc,

object will raise an exception.");

static PyObject*

pyepoll_get_closed(pyEpoll_Object *self)

{

if (self->epfd < 0)

Py_RETURN_TRUE;

else

Py_RETURN_FALSE;

}

static PyObject*

pyepoll_fileno(pyEpoll_Object *self)

{

if (self->epfd < 0)

return pyepoll_err_closed();

return PyInt_FromLong(self->epfd);

}

PyDoc_STRVAR(pyepoll_fileno_doc,

Return the epoll control file descriptor.");

static PyObject*

pyepoll_fromfd(PyObject *cls, PyObject *args)

{

SOCKET fd;

if (!PyArg_ParseTuple(args, "i:fromfd", &fd))

return NULL;

return newPyEpoll_Object((PyTypeObject*)cls, -1, fd);

}

PyDoc_STRVAR(pyepoll_fromfd_doc,

Create an epoll object from a given control fd.");

static PyObject *

pyepoll_internal_ctl(int epfd, int op, PyObject *pfd, unsigned int events)

{

struct epoll_event ev;

int result;

int fd;

if (epfd < 0)

return pyepoll_err_closed();

fd = PyObject_AsFileDescriptor(pfd);

if (fd == -1) {

return NULL;

}

switch(op) {

case EPOLL_CTL_ADD:

case EPOLL_CTL_MOD:

ev.events = events;

ev.data.fd = fd;

Py_BEGIN_ALLOW_THREADS

result = epoll_ctl(epfd, op, fd, &ev);

Py_END_ALLOW_THREADS

break;

case EPOLL_CTL_DEL:

/* In kernel versions before 2.6.9, the EPOLL_CTL_DEL

Py_BEGIN_ALLOW_THREADS

result = epoll_ctl(epfd, op, fd, &ev);

if (errno == EBADF) {

/* fd already closed */

result = 0;

errno = 0;

}

Py_END_ALLOW_THREADS

break;

default:

result = -1;

errno = EINVAL;

}

if (result < 0) {

PyErr_SetFromErrno(PyExc_IOError);

return NULL;

}

Py_RETURN_NONE;

}

static PyObject *

pyepoll_register(pyEpoll_Object *self, PyObject *args, PyObject *kwds)

{

PyObject *pfd;

unsigned int events = EPOLLIN | EPOLLOUT | EPOLLPRI;

static char *kwlist[] = {"fd", "eventmask", NULL};

if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|I:register", kwlist,

&pfd, &events)) {

return NULL;

}

return pyepoll_internal_ctl(self->epfd, EPOLL_CTL_ADD, pfd, events);

}

PyDoc_STRVAR(pyepoll_register_doc,

The epoll interface supports all file descriptors that support poll.");

static PyObject *

pyepoll_modify(pyEpoll_Object *self, PyObject *args, PyObject *kwds)

{

PyObject *pfd;

unsigned int events;

static char *kwlist[] = {"fd", "eventmask", NULL};

if (!PyArg_ParseTupleAndKeywords(args, kwds, "OI:modify", kwlist,

&pfd, &events)) {

return NULL;

}

return pyepoll_internal_ctl(self->epfd, EPOLL_CTL_MOD, pfd, events);

}

PyDoc_STRVAR(pyepoll_modify_doc,

events is a bit set composed of the various EPOLL constants");

static PyObject *

pyepoll_unregister(pyEpoll_Object *self, PyObject *args, PyObject *kwds)

{

PyObject *pfd;

static char *kwlist[] = {"fd", NULL};

if (!PyArg_ParseTupleAndKeywords(args, kwds, "O:unregister", kwlist,

&pfd)) {

return NULL;

}

return pyepoll_internal_ctl(self->epfd, EPOLL_CTL_DEL, pfd, 0);

}

PyDoc_STRVAR(pyepoll_unregister_doc,

fd is the target file descriptor of the operation.");

static PyObject *

pyepoll_poll(pyEpoll_Object *self, PyObject *args, PyObject *kwds)

{

double dtimeout = -1.;

int timeout;

int maxevents = -1;

int nfds, i;

PyObject *elist = NULL, *etuple = NULL;

struct epoll_event *evs = NULL;

static char *kwlist[] = {"timeout", "maxevents", NULL};

if (self->epfd < 0)

return pyepoll_err_closed();

if (!PyArg_ParseTupleAndKeywords(args, kwds, "|di:poll", kwlist,

&dtimeout, &maxevents)) {

return NULL;

}

if (dtimeout < 0) {

timeout = -1;

}

else if (dtimeout * 1000.0 > INT_MAX) {

PyErr_SetString(PyExc_OverflowError,

"timeout is too large");

return NULL;

}

else {

timeout = (int)(dtimeout * 1000.0);

}

if (maxevents == -1) {

maxevents = FD_SETSIZE-1;

}

else if (maxevents < 1) {

PyErr_Format(PyExc_ValueError,

"maxevents must be greater than 0, got %d",

maxevents);

return NULL;

}

evs = PyMem_New(struct epoll_event, maxevents);

if (evs == NULL) {

Py_DECREF(self);

PyErr_NoMemory();