Author: martin.blais Date: Fri May 26 14:03:27 2006 New Revision: 46300 Modified: python/trunk/Lib/socket.py python/trunk/Lib/test/test_socket.py python/trunk/Lib/test/test_struct.py python/trunk/Modules/_struct.c python/trunk/Modules/arraymodule.c python/trunk/Modules/socketmodule.c Log: Support for buffer protocol for socket and struct. * Added socket.recv_buf() and socket.recvfrom_buf() methods, that use the buffer protocol (send and sendto already did). * Added struct.pack_to(), that is the corresponding buffer compatible method to unpack_from(). * Fixed minor typos in arraymodule. Modified: python/trunk/Lib/socket.py ============================================================================== --- python/trunk/Lib/socket.py (original) +++ python/trunk/Lib/socket.py Fri May 26 14:03:27 2006 @@ -140,7 +140,9 @@ __doc__ = _realsocket.__doc__ - __slots__ = ["_sock", "send", "recv", "sendto", "recvfrom", + __slots__ = ["_sock", + "recv", "recv_buf", "recvfrom_buf", + "send", "sendto", "recvfrom", "__weakref__"] def __init__(self, family=AF_INET, type=SOCK_STREAM, proto=0, _sock=None): @@ -149,8 +151,10 @@ self._sock = _sock self.send = self._sock.send self.recv = self._sock.recv + self.recv_buf = self._sock.recv_buf self.sendto = self._sock.sendto self.recvfrom = self._sock.recvfrom + self.recvfrom_buf = self._sock.recvfrom_buf def close(self): self._sock = _closedsocket() Modified: python/trunk/Lib/test/test_socket.py ============================================================================== --- python/trunk/Lib/test/test_socket.py (original) +++ python/trunk/Lib/test/test_socket.py Fri May 26 14:03:27 2006 @@ -9,6 +9,7 @@ import thread, threading import Queue import sys +import array from weakref import proxy PORT = 50007 @@ -852,8 +853,38 @@ self.assertRaises(socket.error, s.bind, address) +class BufferIOTest(SocketConnectedTest): + """ + Test the buffer versions of socket.recv() and socket.send(). + """ + def __init__(self, methodName='runTest'): + SocketConnectedTest.__init__(self, methodName=methodName) + + def testRecvBuf(self): + buf = array.array('c', ' '*1024) + nbytes = self.cli_conn.recv_buf(buf) + self.assertEqual(nbytes, len(MSG)) + msg = buf.tostring()[:len(MSG)] + self.assertEqual(msg, MSG) + + def _testRecvBuf(self): + buf = buffer(MSG) + self.serv_conn.send(buf) + + def testRecvFromBuf(self): + buf = array.array('c', ' '*1024) + nbytes, addr = self.cli_conn.recvfrom_buf(buf) + self.assertEqual(nbytes, len(MSG)) + msg = buf.tostring()[:len(MSG)] + self.assertEqual(msg, MSG) + + def _testRecvFromBuf(self): + buf = buffer(MSG) + self.serv_conn.send(buf) + def test_main(): - tests = [GeneralModuleTests, BasicTCPTest, TCPTimeoutTest, TestExceptions] + tests = [GeneralModuleTests, BasicTCPTest, TCPTimeoutTest, TestExceptions, + BufferIOTest] if sys.platform != 'mac': tests.extend([ BasicUDPTest, UDPTimeoutTest ]) Modified: python/trunk/Lib/test/test_struct.py ============================================================================== --- python/trunk/Lib/test/test_struct.py (original) +++ python/trunk/Lib/test/test_struct.py Fri May 26 14:03:27 2006 @@ -1,5 +1,8 @@ from test.test_support import TestFailed, verbose, verify +import test.test_support import struct +import array +import unittest import sys ISBIGENDIAN = sys.byteorder == "big" @@ -438,31 +441,6 @@ test_705836() -def test_unpack_from(): - test_string = 'abcd01234' - fmt = '4s' - s = struct.Struct(fmt) - for cls in (str, buffer): - data = cls(test_string) - assert s.unpack_from(data) == ('abcd',) - assert s.unpack_from(data, 2) == ('cd01',) - assert s.unpack_from(data, 4) == ('0123',) - for i in xrange(6): - assert s.unpack_from(data, i) == (data[i:i+4],) - for i in xrange(6, len(test_string) + 1): - simple_err(s.unpack_from, data, i) - for cls in (str, buffer): - data = cls(test_string) - assert struct.unpack_from(fmt, data) == ('abcd',) - assert struct.unpack_from(fmt, data, 2) == ('cd01',) - assert struct.unpack_from(fmt, data, 4) == ('0123',) - for i in xrange(6): - assert struct.unpack_from(fmt, data, i) == (data[i:i+4],) - for i in xrange(6, len(test_string) + 1): - simple_err(struct.unpack_from, fmt, data, i) - -test_unpack_from() - def test_1229380(): for endian in ('', '>', '<'): for cls in (int, long): @@ -478,3 +456,60 @@ if 0: # TODO: bug #1229380 test_1229380() + +class PackBufferTestCase(unittest.TestCase): + """ + Test the packing methods that work on buffers. + """ + + def test_unpack_from( self ): + test_string = 'abcd01234' + fmt = '4s' + s = struct.Struct(fmt) + for cls in (str, buffer): + data = cls(test_string) + self.assertEquals(s.unpack_from(data), ('abcd',)) + self.assertEquals(s.unpack_from(data, 2), ('cd01',)) + self.assertEquals(s.unpack_from(data, 4), ('0123',)) + for i in xrange(6): + self.assertEquals(s.unpack_from(data, i), (data[i:i+4],)) + for i in xrange(6, len(test_string) + 1): + simple_err(s.unpack_from, data, i) + for cls in (str, buffer): + data = cls(test_string) + self.assertEquals(struct.unpack_from(fmt, data), ('abcd',)) + self.assertEquals(struct.unpack_from(fmt, data, 2), ('cd01',)) + self.assertEquals(struct.unpack_from(fmt, data, 4), ('0123',)) + for i in xrange(6): + self.assertEquals(struct.unpack_from(fmt, data, i), + (data[i:i+4],)) + for i in xrange(6, len(test_string) + 1): + simple_err(struct.unpack_from, fmt, data, i) + + def test_pack_to( self ): + test_string = 'Reykjavik rocks, eow!' + writable_buf = array.array('c', ' '*100) + fmt = '21s' + s = struct.Struct(fmt) + + # Test without offset + s.pack_to(writable_buf, 0, test_string) + from_buf = writable_buf.tostring()[:len(test_string)] + self.assertEquals(from_buf, test_string) + + # Test with offset. + s.pack_to(writable_buf, 10, test_string) + from_buf = writable_buf.tostring()[:len(test_string)+10] + self.assertEquals(from_buf, (test_string[:10] + test_string)) + + # Go beyond boundaries. + small_buf = array.array('c', ' '*10) + self.assertRaises(struct.error, s.pack_to, small_buf, 0, test_string) + self.assertRaises(struct.error, s.pack_to, small_buf, 2, test_string) + +def test_main(): + test.test_support.run_unittest(PackBufferTestCase) + +if __name__ == "__main__": + test_main() + Modified: python/trunk/Modules/_struct.c ============================================================================== --- python/trunk/Modules/_struct.c (original) +++ python/trunk/Modules/_struct.c Fri May 26 14:03:27 2006 @@ -15,6 +15,7 @@ typedef int Py_ssize_t; #endif + /* PY_USE_INT_WHEN_POSSIBLE is an experimental flag that changes the struct API to return int instead of long when possible. This is often a significant performance improvement. */ @@ -24,7 +25,6 @@ /* The translation function for each format character is table driven */ - typedef struct _formatdef { char format; int size; @@ -1315,50 +1315,36 @@ return s_unpack_internal(soself, buffer + offset); } -PyDoc_STRVAR(s_pack__doc__, -"pack(v1, v2, ...) -> string\n\ -\n\ -Return a string containing values v1, v2, ... packed according to this\n\ -Struct's format. See struct.__doc__ for more on format strings."); -static PyObject * -s_pack(PyObject *self, PyObject *args) +/* + * Guts of the pack function. + * + * Takes a struct object, a tuple of arguments, and offset in that tuple of + * argument for where to start processing the arguments for packing, and a + * character buffer for writing the packed string. The caller must insure + * that the buffer may contain the required length for packing the arguments. + * 0 is returned on success, 1 is returned if there is an error. + * + */ +static int +s_pack_internal(PyStructObject *soself, PyObject *args, int offset, char* buf) { - PyStructObject *soself; - PyObject *result; - char *restart; formatcode *code; Py_ssize_t i; - - soself = (PyStructObject *)self; - assert(PyStruct_Check(self)); - assert(soself->s_codes != NULL); - if (args == NULL || !PyTuple_Check(args) || - PyTuple_GET_SIZE(args) != soself->s_len) - { - PyErr_Format(StructError, - "pack requires exactly %d arguments", soself->s_len); - return NULL; - } - - result = PyString_FromStringAndSize((char *)NULL, soself->s_size); - if (result == NULL) - return NULL; - - restart = PyString_AS_STRING(result); - memset(restart, '\0', soself->s_size); - i = 0; + + memset(buf, '\0', soself->s_size); + i = offset; for (code = soself->s_codes; code->fmtdef != NULL; code++) { Py_ssize_t n; PyObject *v; const formatdef *e = code->fmtdef; - char *res = restart + code->offset; + char *res = buf + code->offset; if (e->format == 's') { v = PyTuple_GET_ITEM(args, i++); if (!PyString_Check(v)) { PyErr_SetString(StructError, "argument for 's' must be a string"); - goto fail; + return -1; } n = PyString_GET_SIZE(v); if (n > code->size) @@ -1370,7 +1356,7 @@ if (!PyString_Check(v)) { PyErr_SetString(StructError, "argument for 'p' must be a string"); - goto fail; + return -1; } n = PyString_GET_SIZE(v); if (n > (code->size - 1)) @@ -1383,16 +1369,109 @@ } else { v = PyTuple_GET_ITEM(args, i++); if (e->pack(res, v, e) < 0) - goto fail; + return -1; } } + /* Success */ + return 0; +} + + +PyDoc_STRVAR(s_pack__doc__, +"pack(v1, v2, ...) -> string\n\ +\n\ +Return a string containing values v1, v2, ... packed according to this\n\ +Struct's format. See struct.__doc__ for more on format strings."); + +static PyObject * +s_pack(PyObject *self, PyObject *args) +{ + PyStructObject *soself; + PyObject *result; + + /* Validate arguments. */ + soself = (PyStructObject *)self; + assert(PyStruct_Check(self)); + assert(soself->s_codes != NULL); + if (args == NULL || !PyTuple_Check(args) || + PyTuple_GET_SIZE(args) != soself->s_len) + { + PyErr_Format(StructError, + "pack requires exactly %d arguments", soself->s_len); + return NULL; + } + + /* Allocate a new string */ + result = PyString_FromStringAndSize((char *)NULL, soself->s_size); + if (result == NULL) + return NULL; + + /* Call the guts */ + if ( s_pack_internal(soself, args, 0, PyString_AS_STRING(result)) != 0 ) { + Py_DECREF(result); + return NULL; + } + return result; +} -fail: - Py_DECREF(result); - return NULL; +PyDoc_STRVAR(s_pack_to__doc__, +"pack_to(buffer, offset, v1, v2, ...)\n\ +\n\ +Pack the values v2, v2, ... according to this Struct's format, write \n\ +the packed bytes into the given buffer at the given offset. Note that \n\ +the offset is not an optional argument. See struct.__doc__ for \n\ +more on format strings."); + +static PyObject * +s_pack_to(PyObject *self, PyObject *args) +{ + PyStructObject *soself; + char *buffer; + Py_ssize_t buffer_len, offset; + + /* Validate arguments. +1 is for the first arg as buffer. */ + soself = (PyStructObject *)self; + assert(PyStruct_Check(self)); + assert(soself->s_codes != NULL); + if (args == NULL || !PyTuple_Check(args) || + PyTuple_GET_SIZE(args) != (soself->s_len + 2)) + { + PyErr_Format(StructError, + "pack_to requires exactly %d arguments", + (soself->s_len + 2)); + return NULL; + } + + /* Extract a writable memory buffer from the first argument */ + if ( PyObject_AsWriteBuffer(PyTuple_GET_ITEM(args, 0), + (void**)&buffer, &buffer_len) == -1 ) { + return NULL; + } + assert( buffer_len >= 0 ); + + /* Extract the offset from the first argument */ + offset = PyInt_AsLong(PyTuple_GET_ITEM(args, 1)); + + /* Support negative offsets. */ + if (offset < 0) + offset += buffer_len; + + /* Check boundaries */ + if (offset < 0 || (buffer_len - offset) < soself->s_size) { + PyErr_Format(StructError, + "pack_to requires a buffer of at least %d bytes", + soself->s_size); + return NULL; + } + /* Call the guts */ + if ( s_pack_internal(soself, args, 2, buffer + offset) != 0 ) { + return NULL; + } + + return Py_None; } @@ -1400,6 +1479,7 @@ static struct PyMethodDef s_methods[] = { {"pack", (PyCFunction)s_pack, METH_VARARGS, s_pack__doc__}, + {"pack_to", (PyCFunction)s_pack_to, METH_VARARGS, s_pack_to__doc__}, {"unpack", (PyCFunction)s_unpack, METH_O, s_unpack__doc__}, {"unpack_from", (PyCFunction)s_unpack_from, METH_KEYWORDS, s_unpack_from__doc__}, {NULL, NULL} /* sentinel */ Modified: python/trunk/Modules/arraymodule.c ============================================================================== --- python/trunk/Modules/arraymodule.c (original) +++ python/trunk/Modules/arraymodule.c Fri May 26 14:03:27 2006 @@ -1975,9 +1975,9 @@ 0, /* tp_setattr */ 0, /* tp_compare */ (reprfunc)array_repr, /* tp_repr */ - 0, /* tp_as _number*/ - &array_as_sequence, /* tp_as _sequence*/ - &array_as_mapping, /* tp_as _mapping*/ + 0, /* tp_as_number*/ + &array_as_sequence, /* tp_as_sequence*/ + &array_as_mapping, /* tp_as_mapping*/ 0, /* tp_hash */ 0, /* tp_call */ 0, /* tp_str */ Modified: python/trunk/Modules/socketmodule.c ============================================================================== --- python/trunk/Modules/socketmodule.c (original) +++ python/trunk/Modules/socketmodule.c Fri May 26 14:03:27 2006 @@ -104,7 +104,10 @@ listen(n) -- start listening for incoming connections\n\ makefile([mode, [bufsize]]) -- return a file object for the socket [*]\n\ recv(buflen[, flags]) -- receive data\n\ -recvfrom(buflen[, flags]) -- receive data and sender's address\n\ +recv_buf(buffer[, nbytes[, flags]]) -- receive data (into a buffer)\n\ +recvfrom(buflen[, flags]) -- receive data and sender\'s address\n\ +recvfrom_buf(buffer[, nbytes, [, flags])\n\ + -- receive data and sender\'s address (into a buffer)\n\ sendall(data[, flags]) -- send all data\n\ send(data[, flags]) -- send data, may not send all of it\n\ sendto(data[, flags], addr) -- send data to a given address\n\ @@ -205,7 +208,7 @@ functions are declared correctly if compiling with MIPSPro 7.x in ANSI C mode (default) */ -/* XXX Using _SGIAPI is the wrong thing, +/* XXX Using _SGIAPI is the wrong thing, but I don't know what the right thing is. */ #undef _SGIAPI /* to avoid warning */ #define _SGIAPI 1 @@ -223,8 +226,8 @@ #include <netdb.h> #endif -/* Irix 6.5 fails to define this variable at all. This is needed - for both GCC and SGI's compiler. I'd say that the SGI headers +/* Irix 6.5 fails to define this variable at all. This is needed + for both GCC and SGI's compiler. I'd say that the SGI headers are just busted. Same thing for Solaris. */ #if (defined(__sgi) || defined(sun)) && !defined(INET_ADDRSTRLEN) #define INET_ADDRSTRLEN 16 @@ -1194,10 +1197,10 @@ args->ob_type->tp_name); return 0; } - if (!PyArg_ParseTuple(args, "eti:getsockaddrarg", + if (!PyArg_ParseTuple(args, "eti:getsockaddrarg", "idna", &host, &port)) return 0; - result = setipaddr(host, (struct sockaddr *)addr, + result = setipaddr(host, (struct sockaddr *)addr, sizeof(*addr), AF_INET); PyMem_Free(host); if (result < 0) @@ -1225,12 +1228,12 @@ args->ob_type->tp_name); return 0; } - if (!PyArg_ParseTuple(args, "eti|ii", + if (!PyArg_ParseTuple(args, "eti|ii", "idna", &host, &port, &flowinfo, &scope_id)) { return 0; } - result = setipaddr(host, (struct sockaddr *)addr, + result = setipaddr(host, (struct sockaddr *)addr, sizeof(*addr), AF_INET6); PyMem_Free(host); if (result < 0) @@ -1839,7 +1842,7 @@ int res_size = sizeof res; /* It must be in the exception set */ assert(FD_ISSET(s->sock_fd, &fds_exc)); - if (0 == getsockopt(s->sock_fd, SOL_SOCKET, SO_ERROR, + if (0 == getsockopt(s->sock_fd, SOL_SOCKET, SO_ERROR, (char *)&res, &res_size)) /* getsockopt also clears WSAGetLastError, so reset it back. */ @@ -2135,95 +2138,126 @@ #endif /* NO_DUP */ - -/* s.recv(nbytes [,flags]) method */ - -static PyObject * -sock_recv(PySocketSockObject *s, PyObject *args) +/* + * This is the guts of the recv() and recv_buf() methods, which reads into a + * char buffer. If you have any inc/def ref to do to the objects that contain + * the buffer, do it in the caller. This function returns the number of bytes + * succesfully read. If there was an error, it returns -1. Note that it is + * also possible that we return a number of bytes smaller than the request + * bytes. + */ +static int +sock_recv_guts(PySocketSockObject *s, char* cbuf, int len, int flags) { - int len, n = 0, flags = 0, timeout; - PyObject *buf; + int timeout, outlen = 0; #ifdef __VMS - int read_length; + int remaining, nread; char *read_buf; #endif - if (!PyArg_ParseTuple(args, "i|i:recv", &len, &flags)) - return NULL; - - if (len < 0) { - PyErr_SetString(PyExc_ValueError, - "negative buffersize in recv"); - return NULL; + if (!IS_SELECTABLE(s)) { + select_error(); + return -1; } - buf = PyString_FromStringAndSize((char *) 0, len); - if (buf == NULL) - return NULL; - - if (!IS_SELECTABLE(s)) - return select_error(); - #ifndef __VMS Py_BEGIN_ALLOW_THREADS timeout = internal_select(s, 0); if (!timeout) - n = recv(s->sock_fd, PyString_AS_STRING(buf), len, flags); + outlen = recv(s->sock_fd, cbuf, len, flags); Py_END_ALLOW_THREADS if (timeout) { - Py_DECREF(buf); PyErr_SetString(socket_timeout, "timed out"); - return NULL; + return -1; } - if (n < 0) { - Py_DECREF(buf); - return s->errorhandler(); + if (outlen < 0) { + /* Note: the call to errorhandler() ALWAYS indirectly returned + NULL, so ignore its return value */ + s->errorhandler(); + return -1; } - if (n != len) - _PyString_Resize(&buf, n); #else - read_buf = PyString_AsString(buf); - read_length = len; - while (read_length != 0) { + read_buf = cbuf; + remaining = len; + while (remaining != 0) { unsigned int segment; - segment = read_length /SEGMENT_SIZE; + segment = remaining /SEGMENT_SIZE; if (segment != 0) { segment = SEGMENT_SIZE; } else { - segment = read_length; + segment = remaining; } Py_BEGIN_ALLOW_THREADS timeout = internal_select(s, 0); if (!timeout) - n = recv(s->sock_fd, read_buf, segment, flags); + nread = recv(s->sock_fd, read_buf, segment, flags); Py_END_ALLOW_THREADS if (timeout) { - Py_DECREF(buf); PyErr_SetString(socket_timeout, "timed out"); - return NULL; + return -1; } - if (n < 0) { - Py_DECREF(buf); - return s->errorhandler(); + if (nread < 0) { + s->errorhandler(); + return -1; } - if (n != read_length) { - read_buf += n; + if (nread != remaining) { + read_buf += nread; break; } - read_length -= segment; + remaining -= segment; read_buf += segment; } - if (_PyString_Resize(&buf, (read_buf - PyString_AsString(buf))) < 0) - { - return NULL; - } + outlen = read_buf - cbuf; #endif /* !__VMS */ + + return outlen; +} + + +/* s.recv(nbytes [,flags]) method */ + +static PyObject * +sock_recv(PySocketSockObject *s, PyObject *args) +{ + int recvlen, flags = 0, outlen; + PyObject *buf; + + if (!PyArg_ParseTuple(args, "i|i:recv", &recvlen, &flags)) + return NULL; + + if (recvlen < 0) { + PyErr_SetString(PyExc_ValueError, + "negative buffersize in recv"); + return NULL; + } + + /* Allocate a new string. */ + buf = PyString_FromStringAndSize((char *) 0, recvlen); + if (buf == NULL) + return NULL; + + /* Call the guts */ + outlen = sock_recv_guts(s, PyString_AsString(buf), recvlen, flags); + if (outlen < 0) { + /* An error occured, release the string and return an + error. */ + Py_DECREF(buf); + return NULL; + } + if (outlen != recvlen) { + /* We did not read as many bytes as we anticipated, resize the + string if possible and be succesful. */ + if (_PyString_Resize(&buf, outlen) < 0) + /* Oopsy, not so succesful after all. */ + return NULL; + } + return buf; } @@ -2236,29 +2270,90 @@ the remote end is closed and all data is read, return the empty string."); -/* s.recvfrom(nbytes [,flags]) method */ +/* s.recv_buf(buffer, [nbytes [,flags]]) method */ -static PyObject * -sock_recvfrom(PySocketSockObject *s, PyObject *args) +static PyObject* +sock_recv_buf(PySocketSockObject *s, PyObject *args, PyObject *kwds) { - sock_addr_t addrbuf; - PyObject *buf = NULL; - PyObject *addr = NULL; - PyObject *ret = NULL; - int len, n = 0, flags = 0, timeout; - socklen_t addrlen; + static char *kwlist[] = {"buffer", "nbytes", "flags", 0}; - if (!PyArg_ParseTuple(args, "i|i:recvfrom", &len, &flags)) + int recvlen = 0, flags = 0, readlen; + char *buf; + int buflen; + + /* Get the buffer's memory */ + if (!PyArg_ParseTupleAndKeywords(args, kwds, "s#|ii:recv", kwlist, + &buf, &buflen, &recvlen, &flags)) return NULL; + assert(buf != 0 && buflen > 0); - if (!getsockaddrlen(s, &addrlen)) + if (recvlen < 0) { + PyErr_SetString(PyExc_ValueError, + "negative buffersize in recv"); return NULL; - buf = PyString_FromStringAndSize((char *) 0, len); - if (buf == NULL) + } + if (recvlen == 0) { + /* If nbytes was not specified, use the buffer's length */ + recvlen = buflen; + } + + /* Check if the buffer is large enough */ + if (buflen < recvlen) { + PyErr_SetString(PyExc_ValueError, + "buffer too small for requested bytes"); return NULL; + } - if (!IS_SELECTABLE(s)) - return select_error(); + /* Call the guts */ + readlen = sock_recv_guts(s, buf, recvlen, flags); + if (readlen < 0) { + /* Return an error. */ + return NULL; + } + + /* Return the number of bytes read. Note that we do not do anything + special here in the case that readlen < recvlen. */ + return PyInt_FromLong(readlen); +} + +PyDoc_STRVAR(recv_buf_doc, +"recv_buf(buffer, [nbytes[, flags]]) -> nbytes_read\n\ +\n\ +A version of recv() that stores its data into a buffer rather than creating \n\ +a new string. Receive up to buffersize bytes from the socket. If buffersize \n\ +is not specified (or 0), receive up to the size available in the given buffer.\n\ +\n\ +See recv() for documentation about the flags."); + + +/* + * This is the guts of the recv() and recv_buf() methods, which reads into a + * char buffer. If you have any inc/def ref to do to the objects that contain + * the buffer, do it in the caller. This function returns the number of bytes + * succesfully read. If there was an error, it returns -1. Note that it is + * also possible that we return a number of bytes smaller than the request + * bytes. + * + * 'addr' is a return value for the address object. Note that you must decref + * it yourself. + */ +static int +sock_recvfrom_guts(PySocketSockObject *s, char* cbuf, int len, int flags, + PyObject** addr) +{ + sock_addr_t addrbuf; + int n = 0, timeout; + socklen_t addrlen; + + *addr = NULL; + + if (!getsockaddrlen(s, &addrlen)) + return -1; + + if (!IS_SELECTABLE(s)) { + select_error(); + return -1; + } Py_BEGIN_ALLOW_THREADS memset(&addrbuf, 0, addrlen); @@ -2266,41 +2361,71 @@ if (!timeout) { #ifndef MS_WINDOWS #if defined(PYOS_OS2) && !defined(PYCC_GCC) - n = recvfrom(s->sock_fd, PyString_AS_STRING(buf), len, flags, + n = recvfrom(s->sock_fd, cbuf, len, flags, (struct sockaddr *) &addrbuf, &addrlen); #else - n = recvfrom(s->sock_fd, PyString_AS_STRING(buf), len, flags, + n = recvfrom(s->sock_fd, cbuf, len, flags, (void *) &addrbuf, &addrlen); #endif #else - n = recvfrom(s->sock_fd, PyString_AS_STRING(buf), len, flags, + n = recvfrom(s->sock_fd, cbuf, len, flags, (struct sockaddr *) &addrbuf, &addrlen); #endif } Py_END_ALLOW_THREADS if (timeout) { - Py_DECREF(buf); PyErr_SetString(socket_timeout, "timed out"); - return NULL; + return -1; } if (n < 0) { - Py_DECREF(buf); - return s->errorhandler(); + s->errorhandler(); + return -1; } - if (n != len && _PyString_Resize(&buf, n) < 0) + if (!(*addr = makesockaddr(s->sock_fd, (struct sockaddr *) &addrbuf, + addrlen, s->sock_proto))) + return -1; + + return n; +} + +/* s.recvfrom(nbytes [,flags]) method */ + +static PyObject * +sock_recvfrom(PySocketSockObject *s, PyObject *args) +{ + PyObject *buf = NULL; + PyObject *addr = NULL; + PyObject *ret = NULL; + int recvlen, outlen, flags = 0; + + if (!PyArg_ParseTuple(args, "i|i:recvfrom", &recvlen, &flags)) + return NULL; + + buf = PyString_FromStringAndSize((char *) 0, recvlen); + if (buf == NULL) return NULL; - if (!(addr = makesockaddr(s->sock_fd, (struct sockaddr *) &addrbuf, - addrlen, s->sock_proto))) + outlen = sock_recvfrom_guts(s, PyString_AS_STRING(buf), + recvlen, flags, &addr); + if (outlen < 0) { goto finally; + } + + if (outlen != recvlen) { + /* We did not read as many bytes as we anticipated, resize the + string if possible and be succesful. */ + if (_PyString_Resize(&buf, outlen) < 0) + /* Oopsy, not so succesful after all. */ + goto finally; + } ret = PyTuple_Pack(2, buf, addr); finally: - Py_XDECREF(addr); Py_XDECREF(buf); + Py_XDECREF(addr); return ret; } @@ -2309,6 +2434,57 @@ \n\ Like recv(buffersize, flags) but also return the sender's address info."); + +/* s.recvfrom_buf(buffer[, nbytes [,flags]]) method */ + +static PyObject * +sock_recvfrom_buf(PySocketSockObject *s, PyObject *args, PyObject* kwds) +{ + static char *kwlist[] = {"buffer", "nbytes", "flags", 0}; + + int recvlen = 0, flags = 0, readlen; + char *buf; + int buflen; + + PyObject *addr = NULL; + PyObject *ret = NULL; + + if (!PyArg_ParseTupleAndKeywords(args, kwds, "s#|ii:recvfrom", kwlist, + &buf, &buflen, &recvlen, &flags)) + return NULL; + assert(buf != 0 && buflen > 0); + + if (recvlen < 0) { + PyErr_SetString(PyExc_ValueError, + "negative buffersize in recv"); + return NULL; + } + if (recvlen == 0) { + /* If nbytes was not specified, use the buffer's length */ + recvlen = buflen; + } + + readlen = sock_recvfrom_guts(s, buf, recvlen, flags, &addr); + if (readlen < 0) { + /* Return an error */ + goto finally; + } + + /* Return the number of bytes read and the address. Note that we do + not do anything special here in the case that readlen < recvlen. */ + ret = PyTuple_Pack(2, PyInt_FromLong(readlen), addr); + +finally: + Py_XDECREF(addr); + return ret; +} + +PyDoc_STRVAR(recvfrom_buf_doc, +"recvfrom_buf(buffer[, nbytes[, flags]]) -> (nbytes, address info)\n\ +\n\ +Like recv_buf(buffer[, nbytes[, flags]]) but also return the sender's address info."); + + /* s.send(data [,flags]) method */ static PyObject * @@ -2503,59 +2679,63 @@ /* List of methods for socket objects */ static PyMethodDef sock_methods[] = { - {"accept", (PyCFunction)sock_accept, METH_NOARGS, - accept_doc}, - {"bind", (PyCFunction)sock_bind, METH_O, - bind_doc}, - {"close", (PyCFunction)sock_close, METH_NOARGS, - close_doc}, - {"connect", (PyCFunction)sock_connect, METH_O, - connect_doc}, - {"connect_ex", (PyCFunction)sock_connect_ex, METH_O, - connect_ex_doc}, + {"accept", (PyCFunction)sock_accept, METH_NOARGS, + accept_doc}, + {"bind", (PyCFunction)sock_bind, METH_O, + bind_doc}, + {"close", (PyCFunction)sock_close, METH_NOARGS, + close_doc}, + {"connect", (PyCFunction)sock_connect, METH_O, + connect_doc}, + {"connect_ex", (PyCFunction)sock_connect_ex, METH_O, + connect_ex_doc}, #ifndef NO_DUP - {"dup", (PyCFunction)sock_dup, METH_NOARGS, - dup_doc}, + {"dup", (PyCFunction)sock_dup, METH_NOARGS, + dup_doc}, #endif - {"fileno", (PyCFunction)sock_fileno, METH_NOARGS, - fileno_doc}, + {"fileno", (PyCFunction)sock_fileno, METH_NOARGS, + fileno_doc}, #ifdef HAVE_GETPEERNAME - {"getpeername", (PyCFunction)sock_getpeername, - METH_NOARGS, getpeername_doc}, + {"getpeername", (PyCFunction)sock_getpeername, + METH_NOARGS, getpeername_doc}, #endif - {"getsockname", (PyCFunction)sock_getsockname, - METH_NOARGS, getsockname_doc}, - {"getsockopt", (PyCFunction)sock_getsockopt, METH_VARARGS, - getsockopt_doc}, - {"listen", (PyCFunction)sock_listen, METH_O, - listen_doc}, + {"getsockname", (PyCFunction)sock_getsockname, + METH_NOARGS, getsockname_doc}, + {"getsockopt", (PyCFunction)sock_getsockopt, METH_VARARGS, + getsockopt_doc}, + {"listen", (PyCFunction)sock_listen, METH_O, + listen_doc}, #ifndef NO_DUP - {"makefile", (PyCFunction)sock_makefile, METH_VARARGS, - makefile_doc}, + {"makefile", (PyCFunction)sock_makefile, METH_VARARGS, + makefile_doc}, #endif - {"recv", (PyCFunction)sock_recv, METH_VARARGS, - recv_doc}, - {"recvfrom", (PyCFunction)sock_recvfrom, METH_VARARGS, - recvfrom_doc}, - {"send", (PyCFunction)sock_send, METH_VARARGS, - send_doc}, - {"sendall", (PyCFunction)sock_sendall, METH_VARARGS, - sendall_doc}, - {"sendto", (PyCFunction)sock_sendto, METH_VARARGS, - sendto_doc}, - {"setblocking", (PyCFunction)sock_setblocking, METH_O, - setblocking_doc}, - {"settimeout", (PyCFunction)sock_settimeout, METH_O, - settimeout_doc}, - {"gettimeout", (PyCFunction)sock_gettimeout, METH_NOARGS, - gettimeout_doc}, - {"setsockopt", (PyCFunction)sock_setsockopt, METH_VARARGS, - setsockopt_doc}, - {"shutdown", (PyCFunction)sock_shutdown, METH_O, - shutdown_doc}, + {"recv", (PyCFunction)sock_recv, METH_VARARGS, + recv_doc}, + {"recv_buf", (PyCFunction)sock_recv_buf, METH_VARARGS | METH_KEYWORDS, + recv_buf_doc}, + {"recvfrom", (PyCFunction)sock_recvfrom, METH_VARARGS, + recvfrom_doc}, + {"recvfrom_buf", (PyCFunction)sock_recvfrom_buf, METH_VARARGS | METH_KEYWORDS, + recvfrom_buf_doc}, + {"send", (PyCFunction)sock_send, METH_VARARGS, + send_doc}, + {"sendall", (PyCFunction)sock_sendall, METH_VARARGS, + sendall_doc}, + {"sendto", (PyCFunction)sock_sendto, METH_VARARGS, + sendto_doc}, + {"setblocking", (PyCFunction)sock_setblocking, METH_O, + setblocking_doc}, + {"settimeout", (PyCFunction)sock_settimeout, METH_O, + settimeout_doc}, + {"gettimeout", (PyCFunction)sock_gettimeout, METH_NOARGS, + gettimeout_doc}, + {"setsockopt", (PyCFunction)sock_setsockopt, METH_VARARGS, + setsockopt_doc}, + {"shutdown", (PyCFunction)sock_shutdown, METH_O, + shutdown_doc}, #ifdef RISCOS - {"sleeptaskw", (PyCFunction)sock_sleeptaskw, METH_O, - sleeptaskw_doc}, + {"sleeptaskw", (PyCFunction)sock_sleeptaskw, METH_O, + sleeptaskw_doc}, #endif {NULL, NULL} /* sentinel */ }; @@ -3401,7 +3581,7 @@ if (strcmp(ip_addr, "255.255.255.255") == 0) { packed_addr = 0xFFFFFFFF; } else { - + packed_addr = inet_addr(ip_addr); if (packed_addr == INADDR_NONE) { /* invalid address */ @@ -3476,7 +3656,7 @@ "can't use AF_INET6, IPv6 is disabled"); return NULL; } -#endif +#endif retval = inet_pton(af, ip, packed); if (retval < 0) { @@ -3499,7 +3679,7 @@ return NULL; } } - + PyDoc_STRVAR(inet_ntop_doc, "inet_ntop(af, packed_ip) -> string formatted IP address\n\ \n\ @@ -3517,7 +3697,7 @@ #else char ip[INET_ADDRSTRLEN + 1]; #endif - + /* Guarantee NUL-termination for PyString_FromString() below */ memset((void *) &ip[0], '\0', sizeof(ip)); @@ -3595,7 +3775,7 @@ } else if (PyString_Check(hobj)) { hptr = PyString_AsString(hobj); } else { - PyErr_SetString(PyExc_TypeError, + PyErr_SetString(PyExc_TypeError, "getaddrinfo() argument 1 must be string or None"); return NULL; }