1 | /*
|
---|
2 | * Definition of a `Connection` type.
|
---|
3 | * Used by `socket_connection.c` and `pipe_connection.c`.
|
---|
4 | *
|
---|
5 | * connection.h
|
---|
6 | *
|
---|
7 | * Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt
|
---|
8 | */
|
---|
9 |
|
---|
10 | #ifndef CONNECTION_H
|
---|
11 | #define CONNECTION_H
|
---|
12 |
|
---|
13 | /*
|
---|
14 | * Read/write flags
|
---|
15 | */
|
---|
16 |
|
---|
17 | #define READABLE 1
|
---|
18 | #define WRITABLE 2
|
---|
19 |
|
---|
20 | #define CHECK_READABLE(self) \
|
---|
21 | if (!(self->flags & READABLE)) { \
|
---|
22 | PyErr_SetString(PyExc_IOError, "connection is write-only"); \
|
---|
23 | return NULL; \
|
---|
24 | }
|
---|
25 |
|
---|
26 | #define CHECK_WRITABLE(self) \
|
---|
27 | if (!(self->flags & WRITABLE)) { \
|
---|
28 | PyErr_SetString(PyExc_IOError, "connection is read-only"); \
|
---|
29 | return NULL; \
|
---|
30 | }
|
---|
31 |
|
---|
32 | /*
|
---|
33 | * Allocation and deallocation
|
---|
34 | */
|
---|
35 |
|
---|
36 | static PyObject *
|
---|
37 | connection_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
|
---|
38 | {
|
---|
39 | ConnectionObject *self;
|
---|
40 | HANDLE handle;
|
---|
41 | BOOL readable = TRUE, writable = TRUE;
|
---|
42 |
|
---|
43 | static char *kwlist[] = {"handle", "readable", "writable", NULL};
|
---|
44 |
|
---|
45 | if (!PyArg_ParseTupleAndKeywords(args, kwds, F_HANDLE "|ii", kwlist,
|
---|
46 | &handle, &readable, &writable))
|
---|
47 | return NULL;
|
---|
48 |
|
---|
49 | if (handle == INVALID_HANDLE_VALUE || (Py_ssize_t)handle < 0) {
|
---|
50 | PyErr_Format(PyExc_IOError, "invalid handle %zd",
|
---|
51 | (Py_ssize_t)handle);
|
---|
52 | return NULL;
|
---|
53 | }
|
---|
54 |
|
---|
55 | if (!readable && !writable) {
|
---|
56 | PyErr_SetString(PyExc_ValueError,
|
---|
57 | "either readable or writable must be true");
|
---|
58 | return NULL;
|
---|
59 | }
|
---|
60 |
|
---|
61 | self = PyObject_New(ConnectionObject, type);
|
---|
62 | if (self == NULL)
|
---|
63 | return NULL;
|
---|
64 |
|
---|
65 | self->weakreflist = NULL;
|
---|
66 | self->handle = handle;
|
---|
67 | self->flags = 0;
|
---|
68 |
|
---|
69 | if (readable)
|
---|
70 | self->flags |= READABLE;
|
---|
71 | if (writable)
|
---|
72 | self->flags |= WRITABLE;
|
---|
73 | assert(self->flags >= 1 && self->flags <= 3);
|
---|
74 |
|
---|
75 | return (PyObject*)self;
|
---|
76 | }
|
---|
77 |
|
---|
78 | static void
|
---|
79 | connection_dealloc(ConnectionObject* self)
|
---|
80 | {
|
---|
81 | if (self->weakreflist != NULL)
|
---|
82 | PyObject_ClearWeakRefs((PyObject*)self);
|
---|
83 |
|
---|
84 | if (self->handle != INVALID_HANDLE_VALUE) {
|
---|
85 | Py_BEGIN_ALLOW_THREADS
|
---|
86 | CLOSE(self->handle);
|
---|
87 | Py_END_ALLOW_THREADS
|
---|
88 | }
|
---|
89 | PyObject_Del(self);
|
---|
90 | }
|
---|
91 |
|
---|
92 | /*
|
---|
93 | * Functions for transferring buffers
|
---|
94 | */
|
---|
95 |
|
---|
96 | static PyObject *
|
---|
97 | connection_sendbytes(ConnectionObject *self, PyObject *args)
|
---|
98 | {
|
---|
99 | char *buffer;
|
---|
100 | Py_ssize_t length, offset=0, size=PY_SSIZE_T_MIN;
|
---|
101 | int res;
|
---|
102 |
|
---|
103 | if (!PyArg_ParseTuple(args, F_RBUFFER "#|" F_PY_SSIZE_T F_PY_SSIZE_T,
|
---|
104 | &buffer, &length, &offset, &size))
|
---|
105 | return NULL;
|
---|
106 |
|
---|
107 | CHECK_WRITABLE(self);
|
---|
108 |
|
---|
109 | if (offset < 0) {
|
---|
110 | PyErr_SetString(PyExc_ValueError, "offset is negative");
|
---|
111 | return NULL;
|
---|
112 | }
|
---|
113 | if (length < offset) {
|
---|
114 | PyErr_SetString(PyExc_ValueError, "buffer length < offset");
|
---|
115 | return NULL;
|
---|
116 | }
|
---|
117 |
|
---|
118 | if (size == PY_SSIZE_T_MIN) {
|
---|
119 | size = length - offset;
|
---|
120 | } else {
|
---|
121 | if (size < 0) {
|
---|
122 | PyErr_SetString(PyExc_ValueError, "size is negative");
|
---|
123 | return NULL;
|
---|
124 | }
|
---|
125 | if (offset + size > length) {
|
---|
126 | PyErr_SetString(PyExc_ValueError,
|
---|
127 | "buffer length < offset + size");
|
---|
128 | return NULL;
|
---|
129 | }
|
---|
130 | }
|
---|
131 |
|
---|
132 | res = conn_send_string(self, buffer + offset, size);
|
---|
133 |
|
---|
134 | if (res < 0) {
|
---|
135 | if (PyErr_Occurred())
|
---|
136 | return NULL;
|
---|
137 | else
|
---|
138 | return mp_SetError(PyExc_IOError, res);
|
---|
139 | }
|
---|
140 |
|
---|
141 | Py_RETURN_NONE;
|
---|
142 | }
|
---|
143 |
|
---|
144 | static PyObject *
|
---|
145 | connection_recvbytes(ConnectionObject *self, PyObject *args)
|
---|
146 | {
|
---|
147 | char *freeme = NULL;
|
---|
148 | Py_ssize_t res, maxlength = PY_SSIZE_T_MAX;
|
---|
149 | PyObject *result = NULL;
|
---|
150 |
|
---|
151 | if (!PyArg_ParseTuple(args, "|" F_PY_SSIZE_T, &maxlength))
|
---|
152 | return NULL;
|
---|
153 |
|
---|
154 | CHECK_READABLE(self);
|
---|
155 |
|
---|
156 | if (maxlength < 0) {
|
---|
157 | PyErr_SetString(PyExc_ValueError, "maxlength < 0");
|
---|
158 | return NULL;
|
---|
159 | }
|
---|
160 |
|
---|
161 | res = conn_recv_string(self, self->buffer, CONNECTION_BUFFER_SIZE,
|
---|
162 | &freeme, maxlength);
|
---|
163 |
|
---|
164 | if (res < 0) {
|
---|
165 | if (res == MP_BAD_MESSAGE_LENGTH) {
|
---|
166 | if ((self->flags & WRITABLE) == 0) {
|
---|
167 | Py_BEGIN_ALLOW_THREADS
|
---|
168 | CLOSE(self->handle);
|
---|
169 | Py_END_ALLOW_THREADS
|
---|
170 | self->handle = INVALID_HANDLE_VALUE;
|
---|
171 | } else {
|
---|
172 | self->flags = WRITABLE;
|
---|
173 | }
|
---|
174 | }
|
---|
175 | mp_SetError(PyExc_IOError, res);
|
---|
176 | } else {
|
---|
177 | if (freeme == NULL) {
|
---|
178 | result = PyString_FromStringAndSize(self->buffer, res);
|
---|
179 | } else {
|
---|
180 | result = PyString_FromStringAndSize(freeme, res);
|
---|
181 | PyMem_Free(freeme);
|
---|
182 | }
|
---|
183 | }
|
---|
184 |
|
---|
185 | return result;
|
---|
186 | }
|
---|
187 |
|
---|
188 | static PyObject *
|
---|
189 | connection_recvbytes_into(ConnectionObject *self, PyObject *args)
|
---|
190 | {
|
---|
191 | char *freeme = NULL, *buffer = NULL;
|
---|
192 | Py_ssize_t res, length, offset = 0;
|
---|
193 | PyObject *result = NULL;
|
---|
194 | Py_buffer pbuf;
|
---|
195 |
|
---|
196 | CHECK_READABLE(self);
|
---|
197 |
|
---|
198 | if (!PyArg_ParseTuple(args, "w*|" F_PY_SSIZE_T,
|
---|
199 | &pbuf, &offset))
|
---|
200 | return NULL;
|
---|
201 |
|
---|
202 | buffer = pbuf.buf;
|
---|
203 | length = pbuf.len;
|
---|
204 |
|
---|
205 | if (offset < 0) {
|
---|
206 | PyErr_SetString(PyExc_ValueError, "negative offset");
|
---|
207 | goto _error;
|
---|
208 | }
|
---|
209 |
|
---|
210 | if (offset > length) {
|
---|
211 | PyErr_SetString(PyExc_ValueError, "offset too large");
|
---|
212 | goto _error;
|
---|
213 | }
|
---|
214 |
|
---|
215 | res = conn_recv_string(self, buffer+offset, length-offset,
|
---|
216 | &freeme, PY_SSIZE_T_MAX);
|
---|
217 |
|
---|
218 | if (res < 0) {
|
---|
219 | if (res == MP_BAD_MESSAGE_LENGTH) {
|
---|
220 | if ((self->flags & WRITABLE) == 0) {
|
---|
221 | Py_BEGIN_ALLOW_THREADS
|
---|
222 | CLOSE(self->handle);
|
---|
223 | Py_END_ALLOW_THREADS
|
---|
224 | self->handle = INVALID_HANDLE_VALUE;
|
---|
225 | } else {
|
---|
226 | self->flags = WRITABLE;
|
---|
227 | }
|
---|
228 | }
|
---|
229 | mp_SetError(PyExc_IOError, res);
|
---|
230 | } else {
|
---|
231 | if (freeme == NULL) {
|
---|
232 | result = PyInt_FromSsize_t(res);
|
---|
233 | } else {
|
---|
234 | result = PyObject_CallFunction(BufferTooShort,
|
---|
235 | F_RBUFFER "#",
|
---|
236 | freeme, res);
|
---|
237 | PyMem_Free(freeme);
|
---|
238 | if (result) {
|
---|
239 | PyErr_SetObject(BufferTooShort, result);
|
---|
240 | Py_DECREF(result);
|
---|
241 | }
|
---|
242 | goto _error;
|
---|
243 | }
|
---|
244 | }
|
---|
245 |
|
---|
246 | _cleanup:
|
---|
247 | PyBuffer_Release(&pbuf);
|
---|
248 | return result;
|
---|
249 |
|
---|
250 | _error:
|
---|
251 | result = NULL;
|
---|
252 | goto _cleanup;
|
---|
253 | }
|
---|
254 |
|
---|
255 | /*
|
---|
256 | * Functions for transferring objects
|
---|
257 | */
|
---|
258 |
|
---|
259 | static PyObject *
|
---|
260 | connection_send_obj(ConnectionObject *self, PyObject *obj)
|
---|
261 | {
|
---|
262 | char *buffer;
|
---|
263 | int res;
|
---|
264 | Py_ssize_t length;
|
---|
265 | PyObject *pickled_string = NULL;
|
---|
266 |
|
---|
267 | CHECK_WRITABLE(self);
|
---|
268 |
|
---|
269 | pickled_string = PyObject_CallFunctionObjArgs(pickle_dumps, obj,
|
---|
270 | pickle_protocol, NULL);
|
---|
271 | if (!pickled_string)
|
---|
272 | goto failure;
|
---|
273 |
|
---|
274 | if (PyString_AsStringAndSize(pickled_string, &buffer, &length) < 0)
|
---|
275 | goto failure;
|
---|
276 |
|
---|
277 | res = conn_send_string(self, buffer, (int)length);
|
---|
278 |
|
---|
279 | if (res < 0) {
|
---|
280 | mp_SetError(PyExc_IOError, res);
|
---|
281 | goto failure;
|
---|
282 | }
|
---|
283 |
|
---|
284 | Py_XDECREF(pickled_string);
|
---|
285 | Py_RETURN_NONE;
|
---|
286 |
|
---|
287 | failure:
|
---|
288 | Py_XDECREF(pickled_string);
|
---|
289 | return NULL;
|
---|
290 | }
|
---|
291 |
|
---|
292 | static PyObject *
|
---|
293 | connection_recv_obj(ConnectionObject *self)
|
---|
294 | {
|
---|
295 | char *freeme = NULL;
|
---|
296 | Py_ssize_t res;
|
---|
297 | PyObject *temp = NULL, *result = NULL;
|
---|
298 |
|
---|
299 | CHECK_READABLE(self);
|
---|
300 |
|
---|
301 | res = conn_recv_string(self, self->buffer, CONNECTION_BUFFER_SIZE,
|
---|
302 | &freeme, PY_SSIZE_T_MAX);
|
---|
303 |
|
---|
304 | if (res < 0) {
|
---|
305 | if (res == MP_BAD_MESSAGE_LENGTH) {
|
---|
306 | if ((self->flags & WRITABLE) == 0) {
|
---|
307 | Py_BEGIN_ALLOW_THREADS
|
---|
308 | CLOSE(self->handle);
|
---|
309 | Py_END_ALLOW_THREADS
|
---|
310 | self->handle = INVALID_HANDLE_VALUE;
|
---|
311 | } else {
|
---|
312 | self->flags = WRITABLE;
|
---|
313 | }
|
---|
314 | }
|
---|
315 | mp_SetError(PyExc_IOError, res);
|
---|
316 | } else {
|
---|
317 | if (freeme == NULL) {
|
---|
318 | temp = PyString_FromStringAndSize(self->buffer, res);
|
---|
319 | } else {
|
---|
320 | temp = PyString_FromStringAndSize(freeme, res);
|
---|
321 | PyMem_Free(freeme);
|
---|
322 | }
|
---|
323 | }
|
---|
324 |
|
---|
325 | if (temp)
|
---|
326 | result = PyObject_CallFunctionObjArgs(pickle_loads,
|
---|
327 | temp, NULL);
|
---|
328 | Py_XDECREF(temp);
|
---|
329 | return result;
|
---|
330 | }
|
---|
331 |
|
---|
332 | /*
|
---|
333 | * Other functions
|
---|
334 | */
|
---|
335 |
|
---|
336 | static PyObject *
|
---|
337 | connection_poll(ConnectionObject *self, PyObject *args)
|
---|
338 | {
|
---|
339 | PyObject *timeout_obj = NULL;
|
---|
340 | double timeout = 0.0;
|
---|
341 | int res;
|
---|
342 |
|
---|
343 | CHECK_READABLE(self);
|
---|
344 |
|
---|
345 | if (!PyArg_ParseTuple(args, "|O", &timeout_obj))
|
---|
346 | return NULL;
|
---|
347 |
|
---|
348 | if (timeout_obj == NULL) {
|
---|
349 | timeout = 0.0;
|
---|
350 | } else if (timeout_obj == Py_None) {
|
---|
351 | timeout = -1.0; /* block forever */
|
---|
352 | } else {
|
---|
353 | timeout = PyFloat_AsDouble(timeout_obj);
|
---|
354 | if (PyErr_Occurred())
|
---|
355 | return NULL;
|
---|
356 | if (timeout < 0.0)
|
---|
357 | timeout = 0.0;
|
---|
358 | }
|
---|
359 |
|
---|
360 | Py_BEGIN_ALLOW_THREADS
|
---|
361 | res = conn_poll(self, timeout, _save);
|
---|
362 | Py_END_ALLOW_THREADS
|
---|
363 |
|
---|
364 | switch (res) {
|
---|
365 | case TRUE:
|
---|
366 | Py_RETURN_TRUE;
|
---|
367 | case FALSE:
|
---|
368 | Py_RETURN_FALSE;
|
---|
369 | default:
|
---|
370 | return mp_SetError(PyExc_IOError, res);
|
---|
371 | }
|
---|
372 | }
|
---|
373 |
|
---|
374 | static PyObject *
|
---|
375 | connection_fileno(ConnectionObject* self)
|
---|
376 | {
|
---|
377 | if (self->handle == INVALID_HANDLE_VALUE) {
|
---|
378 | PyErr_SetString(PyExc_IOError, "handle is invalid");
|
---|
379 | return NULL;
|
---|
380 | }
|
---|
381 | return PyInt_FromLong((long)self->handle);
|
---|
382 | }
|
---|
383 |
|
---|
384 | static PyObject *
|
---|
385 | connection_close(ConnectionObject *self)
|
---|
386 | {
|
---|
387 | if (self->handle != INVALID_HANDLE_VALUE) {
|
---|
388 | Py_BEGIN_ALLOW_THREADS
|
---|
389 | CLOSE(self->handle);
|
---|
390 | Py_END_ALLOW_THREADS
|
---|
391 | self->handle = INVALID_HANDLE_VALUE;
|
---|
392 | }
|
---|
393 |
|
---|
394 | Py_RETURN_NONE;
|
---|
395 | }
|
---|
396 |
|
---|
397 | static PyObject *
|
---|
398 | connection_repr(ConnectionObject *self)
|
---|
399 | {
|
---|
400 | static char *conn_type[] = {"read-only", "write-only", "read-write"};
|
---|
401 |
|
---|
402 | assert(self->flags >= 1 && self->flags <= 3);
|
---|
403 | return FROM_FORMAT("<%s %s, handle %zd>",
|
---|
404 | conn_type[self->flags - 1],
|
---|
405 | CONNECTION_NAME, (Py_ssize_t)self->handle);
|
---|
406 | }
|
---|
407 |
|
---|
408 | /*
|
---|
409 | * Getters and setters
|
---|
410 | */
|
---|
411 |
|
---|
412 | static PyObject *
|
---|
413 | connection_closed(ConnectionObject *self, void *closure)
|
---|
414 | {
|
---|
415 | return PyBool_FromLong((long)(self->handle == INVALID_HANDLE_VALUE));
|
---|
416 | }
|
---|
417 |
|
---|
418 | static PyObject *
|
---|
419 | connection_readable(ConnectionObject *self, void *closure)
|
---|
420 | {
|
---|
421 | return PyBool_FromLong((long)(self->flags & READABLE));
|
---|
422 | }
|
---|
423 |
|
---|
424 | static PyObject *
|
---|
425 | connection_writable(ConnectionObject *self, void *closure)
|
---|
426 | {
|
---|
427 | return PyBool_FromLong((long)(self->flags & WRITABLE));
|
---|
428 | }
|
---|
429 |
|
---|
430 | /*
|
---|
431 | * Tables
|
---|
432 | */
|
---|
433 |
|
---|
434 | static PyMethodDef connection_methods[] = {
|
---|
435 | {"send_bytes", (PyCFunction)connection_sendbytes, METH_VARARGS,
|
---|
436 | "send the byte data from a readable buffer-like object"},
|
---|
437 | {"recv_bytes", (PyCFunction)connection_recvbytes, METH_VARARGS,
|
---|
438 | "receive byte data as a string"},
|
---|
439 | {"recv_bytes_into",(PyCFunction)connection_recvbytes_into,METH_VARARGS,
|
---|
440 | "receive byte data into a writeable buffer-like object\n"
|
---|
441 | "returns the number of bytes read"},
|
---|
442 |
|
---|
443 | {"send", (PyCFunction)connection_send_obj, METH_O,
|
---|
444 | "send a (picklable) object"},
|
---|
445 | {"recv", (PyCFunction)connection_recv_obj, METH_NOARGS,
|
---|
446 | "receive a (picklable) object"},
|
---|
447 |
|
---|
448 | {"poll", (PyCFunction)connection_poll, METH_VARARGS,
|
---|
449 | "whether there is any input available to be read"},
|
---|
450 | {"fileno", (PyCFunction)connection_fileno, METH_NOARGS,
|
---|
451 | "file descriptor or handle of the connection"},
|
---|
452 | {"close", (PyCFunction)connection_close, METH_NOARGS,
|
---|
453 | "close the connection"},
|
---|
454 |
|
---|
455 | {NULL} /* Sentinel */
|
---|
456 | };
|
---|
457 |
|
---|
458 | static PyGetSetDef connection_getset[] = {
|
---|
459 | {"closed", (getter)connection_closed, NULL,
|
---|
460 | "True if the connection is closed", NULL},
|
---|
461 | {"readable", (getter)connection_readable, NULL,
|
---|
462 | "True if the connection is readable", NULL},
|
---|
463 | {"writable", (getter)connection_writable, NULL,
|
---|
464 | "True if the connection is writable", NULL},
|
---|
465 | {NULL}
|
---|
466 | };
|
---|
467 |
|
---|
468 | /*
|
---|
469 | * Connection type
|
---|
470 | */
|
---|
471 |
|
---|
472 | PyDoc_STRVAR(connection_doc,
|
---|
473 | "Connection type whose constructor signature is\n\n"
|
---|
474 | " Connection(handle, readable=True, writable=True).\n\n"
|
---|
475 | "The constructor does *not* duplicate the handle.");
|
---|
476 |
|
---|
477 | PyTypeObject CONNECTION_TYPE = {
|
---|
478 | PyVarObject_HEAD_INIT(NULL, 0)
|
---|
479 | /* tp_name */ "_multiprocessing." CONNECTION_NAME,
|
---|
480 | /* tp_basicsize */ sizeof(ConnectionObject),
|
---|
481 | /* tp_itemsize */ 0,
|
---|
482 | /* tp_dealloc */ (destructor)connection_dealloc,
|
---|
483 | /* tp_print */ 0,
|
---|
484 | /* tp_getattr */ 0,
|
---|
485 | /* tp_setattr */ 0,
|
---|
486 | /* tp_compare */ 0,
|
---|
487 | /* tp_repr */ (reprfunc)connection_repr,
|
---|
488 | /* tp_as_number */ 0,
|
---|
489 | /* tp_as_sequence */ 0,
|
---|
490 | /* tp_as_mapping */ 0,
|
---|
491 | /* tp_hash */ 0,
|
---|
492 | /* tp_call */ 0,
|
---|
493 | /* tp_str */ 0,
|
---|
494 | /* tp_getattro */ 0,
|
---|
495 | /* tp_setattro */ 0,
|
---|
496 | /* tp_as_buffer */ 0,
|
---|
497 | /* tp_flags */ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE |
|
---|
498 | Py_TPFLAGS_HAVE_WEAKREFS,
|
---|
499 | /* tp_doc */ connection_doc,
|
---|
500 | /* tp_traverse */ 0,
|
---|
501 | /* tp_clear */ 0,
|
---|
502 | /* tp_richcompare */ 0,
|
---|
503 | /* tp_weaklistoffset */ offsetof(ConnectionObject, weakreflist),
|
---|
504 | /* tp_iter */ 0,
|
---|
505 | /* tp_iternext */ 0,
|
---|
506 | /* tp_methods */ connection_methods,
|
---|
507 | /* tp_members */ 0,
|
---|
508 | /* tp_getset */ connection_getset,
|
---|
509 | /* tp_base */ 0,
|
---|
510 | /* tp_dict */ 0,
|
---|
511 | /* tp_descr_get */ 0,
|
---|
512 | /* tp_descr_set */ 0,
|
---|
513 | /* tp_dictoffset */ 0,
|
---|
514 | /* tp_init */ 0,
|
---|
515 | /* tp_alloc */ 0,
|
---|
516 | /* tp_new */ connection_new,
|
---|
517 | };
|
---|
518 |
|
---|
519 | #endif /* CONNECTION_H */
|
---|