Modules/selectmodule.c

   1 /* select - Module containing unix select(2) call.
   2    Under Unix, the file descriptors are small integers.
   3    Under Win32, select only exists for sockets, and sockets may
   4    have any value except INVALID_SOCKET.
   5    Under BeOS, we suffer the same dichotomy as Win32; sockets can be anything
   6    >= 0.
   7 */
   8
   9 #include "Python.h"
  10
  11 /* Windows #defines FD_SETSIZE to 64 if FD_SETSIZE isn't already defined.
  12    64 is too small (too many people have bumped into that limit).
  13    Here we boost it.
  14    Users who want even more than the boosted limit should #define
  15    FD_SETSIZE higher before this; e.g., via compiler /D switch.
  16 */
  17 #if defined(MS_WINDOWS) && !defined(FD_SETSIZE)
  18 #define FD_SETSIZE 512
  19 #endif
  20
  21 #ifdef HAVE_UNISTD_H
  22 #include <unistd.h>
  23 #endif
  24 #if defined(HAVE_POLL_H)
  25 #include <poll.h>
  26 #elif defined(HAVE_SYS_POLL_H)
  27 #include <sys/poll.h>
  28 #endif
  29
  30 #ifdef __sgi
  31 /* This is missing from unistd.h */
  32 extern void bzero(void *, int);
  33 #endif
  34
  35 #ifndef DONT_HAVE_SYS_TYPES_H
  36 #include <sys/types.h>
  37 #endif
  38
  39 #if defined(PYOS_OS2)
  40 #include <sys/time.h>
  41 #include <utils.h>
  42 #endif
  43
  44 #ifdef MS_WINDOWS
  45 #include <winsock.h>
  46 #else
  47 #ifdef __BEOS__
  48 #include <net/socket.h>
  49 #define SOCKET int
  50 #else
  51 #define SOCKET int
  52 #endif
  53 #endif
  54
  55 #ifdef RISCOS
  56 #define NO_DUP
  57 #undef off_t
  58 #undef uid_t
  59 #undef gid_t
  60 #undef errno
  61 #include "socklib.h"
  62 #endif /* RISCOS */
  63
  64 static PyObject *SelectError;
  65
  66 /* list of Python objects and their file descriptor */
  67 typedef struct {
  68         PyObject *obj;                       /* owned reference */
  69         SOCKET fd;
  70         int sentinel;                        /* -1 == sentinel */
  71 } pylist;
  72
  73 static void
  74 reap_obj(pylist fd2obj[FD_SETSIZE + 3])
  75 {
  76         int i;
  77         for (i = 0; i < FD_SETSIZE + 3 && fd2obj[i].sentinel >= 0; i++) {
  78                 Py_XDECREF(fd2obj[i].obj);
  79                 fd2obj[i].obj = NULL;
  80         }
  81         fd2obj[0].sentinel = -1;
  82 }
  83
  84
  85 /* returns -1 and sets the Python exception if an error occurred, otherwise
  86    returns a number >= 0
  87 */
  88 static int
  89 list2set(PyObject *list, fd_set *set, pylist fd2obj[FD_SETSIZE + 3])
  90 {
  91         int i;
  92         int max = -1;
  93         int index = 0;
  94         int len = PyList_Size(list);
  95         PyObject* o = NULL;
  96
  97         fd2obj[0].obj = (PyObject*)0;        /* set list to zero size */
  98         FD_ZERO(set);
  99
 100         for (i = 0; i < len; i++)  {
 101                 SOCKET v;
 102
 103                 /* any intervening fileno() calls could decr this refcnt */
 104                 if (!(o = PyList_GetItem(list, i)))
 105                     return -1;
 106
 107                 Py_INCREF(o);
 108                 v = PyObject_AsFileDescriptor( o );
 109                 if (v == -1) goto finally;
 110
 111 #if defined(_MSC_VER)
 112                 max = 0;                     /* not used for Win32 */
 113 #else  /* !_MSC_VER */
 114                 if (v < 0 || v >= FD_SETSIZE) {
 115                         PyErr_SetString(PyExc_ValueError,
 116                                     "filedescriptor out of range in select()");
 117                         goto finally;
 118                 }
 119                 if (v > max)
 120                         max = v;
 121 #endif /* _MSC_VER */
 122                 FD_SET(v, set);
 123
 124                 /* add object and its file descriptor to the list */
 125                 if (index >= FD_SETSIZE) {
 126                         PyErr_SetString(PyExc_ValueError,
 127                                       "too many file descriptors in select()");
 128                         goto finally;
 129                 }
 130                 fd2obj[index].obj = o;
 131                 fd2obj[index].fd = v;
 132                 fd2obj[index].sentinel = 0;
 133                 fd2obj[++index].sentinel = -1;
 134         }
 135         return max+1;
 136
 137   finally:
 138         Py_XDECREF(o);
 139         return -1;
 140 }
 141
 142 /* returns NULL and sets the Python exception if an error occurred */
 143 static PyObject *
 144 set2list(fd_set *set, pylist fd2obj[FD_SETSIZE + 3])
 145 {
 146         int i, j, count=0;
 147         PyObject *list, *o;
 148         SOCKET fd;
 149
 150         for (j = 0; fd2obj[j].sentinel >= 0; j++) {
 151                 if (FD_ISSET(fd2obj[j].fd, set))
 152                         count++;
 153         }
 154         list = PyList_New(count);
 155         if (!list)
 156                 return NULL;
 157
 158         i = 0;
 159         for (j = 0; fd2obj[j].sentinel >= 0; j++) {
 160                 fd = fd2obj[j].fd;
 161                 if (FD_ISSET(fd, set)) {
 162 #ifndef _MSC_VER
 163                         if (fd > FD_SETSIZE) {
 164                                 PyErr_SetString(PyExc_SystemError,
 165                            "filedescriptor out of range returned in select()");
 166                                 goto finally;
 167                         }
 168 #endif
 169                         o = fd2obj[j].obj;
 170                         fd2obj[j].obj = NULL;
 171                         /* transfer ownership */
 172                         if (PyList_SetItem(list, i, o) < 0)
 173                                 goto finally;
 174
 175                         i++;
 176                 }
 177         }
 178         return list;
 179   finally:
 180         Py_DECREF(list);
 181         return NULL;
 182 }
 183
 184
 185 static PyObject *
 186 select_select(PyObject *self, PyObject *args)
 187 {
 188 #ifdef MS_WINDOWS
 189         /* This would be an awful lot of stack space on Windows! */
 190         pylist *rfd2obj, *wfd2obj, *efd2obj;
 191 #else
 192         pylist rfd2obj[FD_SETSIZE + 3];
 193         pylist wfd2obj[FD_SETSIZE + 3];
 194         pylist efd2obj[FD_SETSIZE + 3];
 195 #endif
 196         PyObject *ifdlist, *ofdlist, *efdlist;
 197         PyObject *ret = NULL;
 198         PyObject *tout = Py_None;
 199         fd_set ifdset, ofdset, efdset;
 200         double timeout;
 201         struct timeval tv, *tvp;
 202         long seconds;
 203         int imax, omax, emax, max;
 204         int n;
 205
 206         /* convert arguments */
 207         if (!PyArg_ParseTuple(args, "OOO|O:select",
 208                               &ifdlist, &ofdlist, &efdlist, &tout))
 209                 return NULL;
 210
 211         if (tout == Py_None)
 212                 tvp = (struct timeval *)0;
 213         else if (!PyArg_Parse(tout, "d", &timeout)) {
 214                 PyErr_SetString(PyExc_TypeError,
 215                                 "timeout must be a float or None");
 216                 return NULL;
 217         }
 218         else {
 219                 if (timeout > (double)LONG_MAX) {
 220                         PyErr_SetString(PyExc_OverflowError, "timeout period too long");
 221                         return NULL;
 222                 }
 223                 seconds = (long)timeout;
 224                 timeout = timeout - (double)seconds;
 225                 tv.tv_sec = seconds;
 226                 tv.tv_usec = (long)(timeout*1000000.0);
 227                 tvp = &tv;
 228         }
 229
 230         /* sanity check first three arguments */
 231         if (!PyList_Check(ifdlist) ||
 232             !PyList_Check(ofdlist) ||
 233             !PyList_Check(efdlist))
 234         {
 235                 PyErr_SetString(PyExc_TypeError,
 236                                 "arguments 1-3 must be lists");
 237                 return NULL;
 238         }
 239
 240 #ifdef MS_WINDOWS
 241         /* Allocate memory for the lists */
 242         rfd2obj = PyMem_NEW(pylist, FD_SETSIZE + 3);
 243         wfd2obj = PyMem_NEW(pylist, FD_SETSIZE + 3);
 244         efd2obj = PyMem_NEW(pylist, FD_SETSIZE + 3);
 245         if (rfd2obj == NULL || wfd2obj == NULL || efd2obj == NULL) {
 246                 if (rfd2obj) PyMem_DEL(rfd2obj);
 247                 if (wfd2obj) PyMem_DEL(wfd2obj);
 248                 if (efd2obj) PyMem_DEL(efd2obj);
 249                 return NULL;
 250         }
 251 #endif
 252         /* Convert lists to fd_sets, and get maximum fd number
 253          * propagates the Python exception set in list2set()
 254          */
 255         rfd2obj[0].sentinel = -1;
 256         wfd2obj[0].sentinel = -1;
 257         efd2obj[0].sentinel = -1;
 258         if ((imax=list2set(ifdlist, &ifdset, rfd2obj)) < 0)
 259                 goto finally;
 260         if ((omax=list2set(ofdlist, &ofdset, wfd2obj)) < 0)
 261                 goto finally;
 262         if ((emax=list2set(efdlist, &efdset, efd2obj)) < 0)
 263                 goto finally;
 264         max = imax;
 265         if (omax > max) max = omax;
 266         if (emax > max) max = emax;
 267
 268         Py_BEGIN_ALLOW_THREADS
 269         n = select(max, &ifdset, &ofdset, &efdset, tvp);
 270         Py_END_ALLOW_THREADS
 271
 272         if (n < 0) {
 273                 PyErr_SetFromErrno(SelectError);
 274         }
 275         else if (n == 0) {
 276                 /* optimization */
 277                 ifdlist = PyList_New(0);
 278                 if (ifdlist) {
 279                         ret = Py_BuildValue("OOO", ifdlist, ifdlist, ifdlist);
 280                         Py_DECREF(ifdlist);
 281                 }
 282         }
 283         else {
 284                 /* any of these three calls can raise an exception.  it's more
 285                    convenient to test for this after all three calls... but
 286                    is that acceptable?
 287                 */
 288                 ifdlist = set2list(&ifdset, rfd2obj);
 289                 ofdlist = set2list(&ofdset, wfd2obj);
 290                 efdlist = set2list(&efdset, efd2obj);
 291                 if (PyErr_Occurred())
 292                         ret = NULL;
 293                 else
 294                         ret = Py_BuildValue("OOO", ifdlist, ofdlist, efdlist);
 295
 296                 Py_DECREF(ifdlist);
 297                 Py_DECREF(ofdlist);
 298                 Py_DECREF(efdlist);
 299         }
 300
 301   finally:
 302         reap_obj(rfd2obj);
 303         reap_obj(wfd2obj);
 304         reap_obj(efd2obj);
 305 #ifdef MS_WINDOWS
 306         PyMem_DEL(rfd2obj);
 307         PyMem_DEL(wfd2obj);
 308         PyMem_DEL(efd2obj);
 309 #endif
 310         return ret;
 311 }
 312
 313 #ifdef HAVE_POLL
 314 /*
 315  * poll() support
 316  */
 317
 318 typedef struct {
 319         PyObject_HEAD
 320         PyObject *dict;
 321         int ufd_uptodate;
 322         int ufd_len;
 323         struct pollfd *ufds;
 324 } pollObject;
 325
 326 staticforward PyTypeObject poll_Type;
 327
 328 /* Update the malloc'ed array of pollfds to match the dictionary
 329    contained within a pollObject.  Return 1 on success, 0 on an error.
 330 */
 331
 332 static int
 333 update_ufd_array(pollObject *self)
 334 {
 335         int i, pos;
 336         PyObject *key, *value;
 337
 338         self->ufd_len = PyDict_Size(self->dict);
 339         PyMem_Resize(self->ufds, struct pollfd, self->ufd_len);
 340         if (self->ufds == NULL) {
 341                 PyErr_NoMemory();
 342                 return 0;
 343         }
 344
 345         i = pos = 0;
 346         while (PyDict_Next(self->dict, &pos, &key, &value)) {
 347                 self->ufds[i].fd = PyInt_AsLong(key);
 348                 self->ufds[i].events = (short)PyInt_AsLong(value);
 349                 i++;
 350         }
 351         self->ufd_uptodate = 1;
 352         return 1;
 353 }
 354
 355 static char poll_register_doc[] =
 356 "register(fd [, eventmask] ) -> None\n\n\
 357 Register a file descriptor with the polling object.\n\
 358 fd -- either an integer, or an object with a fileno() method returning an int.\n\
 359 events -- an optional bitmask describing the type of events to check for";
 360
 361 static PyObject *
 362 poll_register(pollObject *self, PyObject *args)
 363 {
 364         PyObject *o, *key, *value;
 365         int fd, events = POLLIN | POLLPRI | POLLOUT;
 366
 367         if (!PyArg_ParseTuple(args, "O|i:register", &o, &events)) {
 368                 return NULL;
 369         }
 370
 371         fd = PyObject_AsFileDescriptor(o);
 372         if (fd == -1) return NULL;
 373
 374         /* Add entry to the internal dictionary: the key is the
 375            file descriptor, and the value is the event mask. */
 376         if ( (NULL == (key = PyInt_FromLong(fd))) ||
 377              (NULL == (value = PyInt_FromLong(events))) ||
 378              (PyDict_SetItem(self->dict, key, value)) == -1) {
 379                 return NULL;
 380         }
 381         self->ufd_uptodate = 0;
 382
 383         Py_INCREF(Py_None);
 384         return Py_None;
 385 }
 386
 387 static char poll_unregister_doc[] =
 388 "unregister(fd) -> None\n\n\
 389 Remove a file descriptor being tracked by the polling object.";
 390
 391 static PyObject *
 392 poll_unregister(pollObject *self, PyObject *args)
 393 {
 394         PyObject *o, *key;
 395         int fd;
 396
 397         if (!PyArg_ParseTuple(args, "O:unregister", &o)) {
 398                 return NULL;
 399         }
 400
 401         fd = PyObject_AsFileDescriptor( o );
 402         if (fd == -1)
 403                 return NULL;
 404
 405         /* Check whether the fd is already in the array */
 406         key = PyInt_FromLong(fd);
 407         if (key == NULL)
 408                 return NULL;
 409
 410         if (PyDict_DelItem(self->dict, key) == -1) {
 411                 Py_DECREF(key);
 412                 /* This will simply raise the KeyError set by PyDict_DelItem
 413                    if the file descriptor isn't registered. */
 414                 return NULL;
 415         }
 416
 417         Py_DECREF(key);
 418         self->ufd_uptodate = 0;
 419
 420         Py_INCREF(Py_None);
 421         return Py_None;
 422 }
 423
 424 static char poll_poll_doc[] =
 425 "poll( [timeout] ) -> list of (fd, event) 2-tuples\n\n\
 426 Polls the set of registered file descriptors, returning a list containing \n\
 427 any descriptors that have events or errors to report.";
 428
 429 static PyObject *
 430 poll_poll(pollObject *self, PyObject *args)
 431 {
 432         PyObject *result_list = NULL, *tout = NULL;
 433         int timeout = 0, poll_result, i, j;
 434         PyObject *value = NULL, *num = NULL;
 435
 436         if (!PyArg_ParseTuple(args, "|O:poll", &tout)) {
 437                 return NULL;
 438         }
 439
 440         /* Check values for timeout */
 441         if (tout == NULL || tout == Py_None)
 442                 timeout = -1;
 443         else if (!PyArg_Parse(tout, "i", &timeout)) {
 444                 PyErr_SetString(PyExc_TypeError,
 445                                 "timeout must be an integer or None");
 446                 return NULL;
 447         }
 448
 449         /* Ensure the ufd array is up to date */
 450         if (!self->ufd_uptodate)
 451                 if (update_ufd_array(self) == 0)
 452                         return NULL;
 453
 454         /* call poll() */
 455         Py_BEGIN_ALLOW_THREADS;
 456         poll_result = poll(self->ufds, self->ufd_len, timeout);
 457         Py_END_ALLOW_THREADS;
 458
 459         if (poll_result < 0) {
 460                 PyErr_SetFromErrno(SelectError);
 461                 return NULL;
 462         }
 463
 464         /* build the result list */
 465
 466         result_list = PyList_New(poll_result);
 467         if (!result_list)
 468                 return NULL;
 469         else {
 470                 for (i = 0, j = 0; j < poll_result; j++) {
 471                         /* skip to the next fired descriptor */
 472                         while (!self->ufds[i].revents) {
 473                                 i++;
 474                         }
 475                         /* if we hit a NULL return, set value to NULL
 476                            and break out of loop; code at end will
 477                            clean up result_list */
 478                         value = PyTuple_New(2);
 479                         if (value == NULL)
 480                                 goto error;
 481                         num = PyInt_FromLong(self->ufds[i].fd);
 482                         if (num == NULL) {
 483                                 Py_DECREF(value);
 484                                 goto error;
 485                         }
 486                         PyTuple_SET_ITEM(value, 0, num);
 487
 488                         num = PyInt_FromLong(self->ufds[i].revents);
 489                         if (num == NULL) {
 490                                 Py_DECREF(value);
 491                                 goto error;
 492                         }
 493                         PyTuple_SET_ITEM(value, 1, num);
 494                         if ((PyList_SetItem(result_list, j, value)) == -1) {
 495                                 Py_DECREF(value);
 496                                 goto error;
 497                         }
 498                         i++;
 499                 }
 500         }
 501         return result_list;
 502
 503   error:
 504         Py_DECREF(result_list);
 505         return NULL;
 506 }
 507
 508 static PyMethodDef poll_methods[] = {
 509         {"register",    (PyCFunction)poll_register,
 510          METH_VARARGS,  poll_register_doc},
 511         {"unregister",  (PyCFunction)poll_unregister,
 512          METH_VARARGS,  poll_unregister_doc},
 513         {"poll",        (PyCFunction)poll_poll,
 514          METH_VARARGS,  poll_poll_doc},
 515         {NULL,          NULL}           /* sentinel */
 516 };
 517
 518 static pollObject *
 519 newPollObject(void)
 520 {
 521         pollObject *self;
 522         self = PyObject_New(pollObject, &poll_Type);
 523         if (self == NULL)
 524                 return NULL;
 525         /* ufd_uptodate is a Boolean, denoting whether the
 526            array pointed to by ufds matches the contents of the dictionary. */
 527         self->ufd_uptodate = 0;
 528         self->ufds = NULL;
 529         self->dict = PyDict_New();
 530         if (self->dict == NULL) {
 531                 Py_DECREF(self);
 532                 return NULL;
 533         }
 534         return self;
 535 }
 536
 537 static void
 538 poll_dealloc(pollObject *self)
 539 {
 540         if (self->ufds != NULL)
 541                 PyMem_DEL(self->ufds);
 542         Py_XDECREF(self->dict);
 543         PyObject_Del(self);
 544 }
 545
 546 static PyObject *
 547 poll_getattr(pollObject *self, char *name)
 548 {
 549         return Py_FindMethod(poll_methods, (PyObject *)self, name);
 550 }
 551
 552 statichere PyTypeObject poll_Type = {
 553         /* The ob_type field must be initialized in the module init function
 554          * to be portable to Windows without using C++. */
 555         PyObject_HEAD_INIT(NULL)
 556         0,                      /*ob_size*/
 557         "poll",                 /*tp_name*/
 558         sizeof(pollObject),     /*tp_basicsize*/
 559         0,                      /*tp_itemsize*/
 560         /* methods */
 561         (destructor)poll_dealloc, /*tp_dealloc*/
 562         0,                      /*tp_print*/
 563         (getattrfunc)poll_getattr, /*tp_getattr*/
 564         0,                      /*tp_setattr*/
 565         0,                      /*tp_compare*/
 566         0,                      /*tp_repr*/
 567         0,                      /*tp_as_number*/
 568         0,                      /*tp_as_sequence*/
 569         0,                      /*tp_as_mapping*/
 570         0,                      /*tp_hash*/
 571 };
 572
 573 static char poll_doc[] =
 574 "Returns a polling object, which supports registering and\n\
 575 unregistering file descriptors, and then polling them for I/O events.";
 576
 577 static PyObject *
 578 select_poll(PyObject *self, PyObject *args)
 579 {
 580         pollObject *rv;
 581
 582         if (!PyArg_ParseTuple(args, ":poll"))
 583                 return NULL;
 584         rv = newPollObject();
 585         if ( rv == NULL )
 586                 return NULL;
 587         return (PyObject *)rv;
 588 }
 589 #endif /* HAVE_POLL */
 590
 591 static char select_doc[] =
 592 "select(rlist, wlist, xlist[, timeout]) -> (rlist, wlist, xlist)\n\
 593 \n\
 594 Wait until one or more file descriptors are ready for some kind of I/O.\n\
 595 The first three arguments are lists of file descriptors to be waited for:\n\
 596 rlist -- wait until ready for reading\n\
 597 wlist -- wait until ready for writing\n\
 598 xlist -- wait for an ``exceptional condition''\n\
 599 If only one kind of condition is required, pass [] for the other lists.\n\
 600 A file descriptor is either a socket or file object, or a small integer\n\
 601 gotten from a fileno() method call on one of those.\n\
 602 \n\
 603 The optional 4th argument specifies a timeout in seconds; it may be\n\
 604 a floating point number to specify fractions of seconds.  If it is absent\n\
 605 or None, the call will never time out.\n\
 606 \n\
 607 The return value is a tuple of three lists corresponding to the first three\n\
 608 arguments; each contains the subset of the corresponding file descriptors\n\
 609 that are ready.\n\
 610 \n\
 611 *** IMPORTANT NOTICE ***\n\
 612 On Windows, only sockets are supported; on Unix, all file descriptors.";
 613
 614 static PyMethodDef select_methods[] = {
 615     {"select",  select_select, METH_VARARGS, select_doc},
 616 #ifdef HAVE_POLL
 617     {"poll",    select_poll,   METH_VARARGS, poll_doc},
 618 #endif /* HAVE_POLL */
 619     {0,         0},                          /* sentinel */
 620 };
 621
 622 static char module_doc[] =
 623 "This module supports asynchronous I/O on multiple file descriptors.\n\
 624 \n\
 625 *** IMPORTANT NOTICE ***\n\
 626 On Windows, only sockets are supported; on Unix, all file descriptors.";
 627
 628 /*
 629  * Convenience routine to export an integer value.
 630  * For simplicity, errors (which are unlikely anyway) are ignored.
 631  */
 632
 633 static void
 634 insint(PyObject *d, char *name, int value)
 635 {
 636         PyObject *v = PyInt_FromLong((long) value);
 637         if (v == NULL) {
 638                 /* Don't bother reporting this error */
 639                 PyErr_Clear();
 640         }
 641         else {
 642                 PyDict_SetItemString(d, name, v);
 643                 Py_DECREF(v);
 644         }
 645 }
 646
 647 DL_EXPORT(void)
 648 initselect(void)
 649 {
 650         PyObject *m, *d;
 651         m = Py_InitModule3("select", select_methods, module_doc);
 652         d = PyModule_GetDict(m);
 653         SelectError = PyErr_NewException("select.error", NULL, NULL);
 654         PyDict_SetItemString(d, "error", SelectError);
 655 #ifdef HAVE_POLL
 656         poll_Type.ob_type = &PyType_Type;
 657         insint(d, "POLLIN", POLLIN);
 658         insint(d, "POLLPRI", POLLPRI);
 659         insint(d, "POLLOUT", POLLOUT);
 660         insint(d, "POLLERR", POLLERR);
 661         insint(d, "POLLHUP", POLLHUP);
 662         insint(d, "POLLNVAL", POLLNVAL);
 663
 664 #ifdef POLLRDNORM
 665         insint(d, "POLLRDNORM", POLLRDNORM);
 666 #endif
 667 #ifdef POLLRDBAND
 668         insint(d, "POLLRDBAND", POLLRDBAND);
 669 #endif
 670 #ifdef POLLWRNORM
 671         insint(d, "POLLWRNORM", POLLWRNORM);
 672 #endif
 673 #ifdef POLLWRBAND
 674         insint(d, "POLLWRBAND", POLLWRBAND);
 675 #endif
 676 #ifdef POLLMSG
 677         insint(d, "POLLMSG", POLLMSG);
 678 #endif
 679 #endif /* HAVE_POLL */
 680 }