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