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)

   157 {

   158 	int res;

   159 	fd_set rfds;

   160 

   161 	FD_ZERO(&rfds);

   162 	FD_SET((SOCKET)conn->handle, &rfds);

   163 

   164 	if (timeout < 0.0) {

   165 		res = select((int)conn->handle+1, &rfds, NULL, NULL, NULL);

   166 	} else {

   167 		struct timeval tv;

   168 		tv.tv_sec = (long)timeout;

   169 		tv.tv_usec = (long)((timeout - tv.tv_sec) * 1e6 + 0.5);

   170 		res = select((int)conn->handle+1, &rfds, NULL, NULL, &tv);

   171 	}

   172 

   173 	if (res < 0) {

   174 		return MP_SOCKET_ERROR;

   175 	} else if (FD_ISSET(conn->handle, &rfds)) {

   176 		return TRUE;

   177 	} else {

   178 		assert(res == 0);

   179 		return FALSE;

   180 	}

   181 }

   182 

   183 /*

   184  * "connection.h" defines the Connection type using defs above

   185  */

   186 

   187 #define CONNECTION_NAME "Connection"

   188 #define CONNECTION_TYPE ConnectionType

   189 

   190 #include "connection.h"