No issues found
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"