Modules/_multiprocessing/socket_connection.c

   1 /*
   2  * A type which wraps a socket
   3  *
   4  * socket_connection.c
   5  *
   6  * Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt
   7  */
   8
   9 #include "multiprocessing.h"
  10
  11 #ifdef MS_WINDOWS
  12 #  define WRITE(h, buffer, length) send((SOCKET)h, buffer, length, 0)
  13 #  define READ(h, buffer, length) recv((SOCKET)h, buffer, length, 0)
  14 #  define CLOSE(h) closesocket((SOCKET)h)
  15 #else
  16 #  define WRITE(h, buffer, length) write(h, buffer, length)
  17 #  define READ(h, buffer, length) read(h, buffer, length)
  18 #  define CLOSE(h) close(h)
  19 #endif
  20
  21 /*
  22  * Send string to file descriptor
  23  */
  24
  25 static Py_ssize_t
  26 _conn_sendall(HANDLE h, char *string, size_t length)
  27 {
  28         char *p = string;
  29         Py_ssize_t res;
  30
  31         while (length > 0) {
  32                 res = WRITE(h, p, length);
  33                 if (res < 0)
  34                         return MP_SOCKET_ERROR;
  35                 length -= res;
  36                 p += res;
  37         }
  38
  39         return MP_SUCCESS;
  40 }
  41
  42 /*
  43  * Receive string of exact length from file descriptor
  44  */
  45
  46 static Py_ssize_t
  47 _conn_recvall(HANDLE h, char *buffer, size_t length)
  48 {
  49         size_t remaining = length;
  50         Py_ssize_t temp;
  51         char *p = buffer;
  52
  53         while (remaining > 0) {
  54                 temp = READ(h, p, remaining);
  55                 if (temp <= 0) {
  56                         if (temp == 0)
  57                                 return remaining == length ?
  58                                         MP_END_OF_FILE : MP_EARLY_END_OF_FILE;
  59                         else
  60                                 return temp;
  61                 }
  62                 remaining -= temp;
  63                 p += temp;
  64         }
  65
  66         return MP_SUCCESS;
  67 }
  68
  69 /*
  70  * Send a string prepended by the string length in network byte order
  71  */
  72
  73 static Py_ssize_t
  74 conn_send_string(ConnectionObject *conn, char *string, size_t length)
  75 {
  76         Py_ssize_t res;
  77         /* The "header" of the message is a 32 bit unsigned number (in
  78            network order) which specifies the length of the "body".  If
  79            the message is shorter than about 16kb then it is quicker to
  80            combine the "header" and the "body" of the message and send
  81            them at once. */
  82         if (length < (16*1024)) {
  83                 char *message;
  84
  85                 message = PyMem_Malloc(length+4);
  86                 if (message == NULL)
  87                         return MP_MEMORY_ERROR;
  88
  89                 *(UINT32*)message = htonl((UINT32)length);
  90                 memcpy(message+4, string, length);
  91                 Py_BEGIN_ALLOW_THREADS
  92                 res = _conn_sendall(conn->handle, message, length+4);
  93                 Py_END_ALLOW_THREADS
  94                 PyMem_Free(message);
  95         } else {
  96                 UINT32 lenbuff;
  97
  98                 if (length > MAX_MESSAGE_LENGTH)
  99                         return MP_BAD_MESSAGE_LENGTH;
 100
 101                 lenbuff = htonl((UINT32)length);
 102                 Py_BEGIN_ALLOW_THREADS
 103                 res = _conn_sendall(conn->handle, (char*)&lenbuff, 4) ||
 104                         _conn_sendall(conn->handle, string, length);
 105                 Py_END_ALLOW_THREADS
 106         }
 107         return res;
 108 }
 109
 110 /*
 111  * Attempts to read into buffer, or failing that into *newbuffer
 112  *
 113  * Returns number of bytes read.
 114  */
 115
 116 static Py_ssize_t
 117 conn_recv_string(ConnectionObject *conn, char *buffer,
 118                  size_t buflength, char **newbuffer, size_t maxlength)
 119 {
 120         int res;
 121         UINT32 ulength;
 122
 123         *newbuffer = NULL;
 124
 125         Py_BEGIN_ALLOW_THREADS
 126         res = _conn_recvall(conn->handle, (char*)&ulength, 4);
 127         Py_END_ALLOW_THREADS
 128         if (res < 0)
 129                 return res;
 130
 131         ulength = ntohl(ulength);
 132         if (ulength > maxlength)
 133                 return MP_BAD_MESSAGE_LENGTH;
 134
 135         if (ulength <= buflength) {
 136                 Py_BEGIN_ALLOW_THREADS
 137                 res = _conn_recvall(conn->handle, buffer, (size_t)ulength);
 138                 Py_END_ALLOW_THREADS
 139                 return res < 0 ? res : ulength;
 140         } else {
 141                 *newbuffer = PyMem_Malloc((size_t)ulength);
 142                 if (*newbuffer == NULL)
 143                         return MP_MEMORY_ERROR;
 144                 Py_BEGIN_ALLOW_THREADS
 145                 res = _conn_recvall(conn->handle, *newbuffer, (size_t)ulength);
 146                 Py_END_ALLOW_THREADS
 147                 return res < 0 ? (Py_ssize_t)res : (Py_ssize_t)ulength;
 148         }
 149 }
 150
 151 /*
 152  * Check whether any data is available for reading -- neg timeout blocks
 153  */
 154
 155 static int
 156 conn_poll(ConnectionObject *conn, double timeout, PyThreadState *_save)
 157 {
 158         int res;
 159         fd_set rfds;
 160
 161         /*
 162          * Verify the handle, issue 3321. Not required for windows.
 163          */
 164         #ifndef MS_WINDOWS
 165                 if (((int)conn->handle) < 0 || ((int)conn->handle) >= FD_SETSIZE) {
 166                         Py_BLOCK_THREADS
 167                         PyErr_SetString(PyExc_IOError, "handle out of range in select()");
 168                         Py_UNBLOCK_THREADS
 169                         return MP_EXCEPTION_HAS_BEEN_SET;
 170                 }
 171         #endif
 172
 173         FD_ZERO(&rfds);
 174         FD_SET((SOCKET)conn->handle, &rfds);
 175
 176         if (timeout < 0.0) {
 177                 res = select((int)conn->handle+1, &rfds, NULL, NULL, NULL);
 178         } else {
 179                 struct timeval tv;
 180                 tv.tv_sec = (long)timeout;
 181                 tv.tv_usec = (long)((timeout - tv.tv_sec) * 1e6 + 0.5);
 182                 res = select((int)conn->handle+1, &rfds, NULL, NULL, &tv);
 183         }
 184
 185         if (res < 0) {
 186                 return MP_SOCKET_ERROR;
 187         } else if (FD_ISSET(conn->handle, &rfds)) {
 188                 return TRUE;
 189         } else {
 190                 assert(res == 0);
 191                 return FALSE;
 192         }
 193 }
 194
 195 /*
 196  * "connection.h" defines the Connection type using defs above
 197  */
 198
 199 #define CONNECTION_NAME "Connection"
 200 #define CONNECTION_TYPE ConnectionType
 201
 202 #include "connection.h"