Index: Modules/socketmodule.c =================================================================== --- Modules/socketmodule.c (revision 79791) +++ Modules/socketmodule.c (working copy) @@ -433,6 +433,7 @@ /* Global variable holding the exception type for errors detected by this module (but not argument type or memory errors, etc.). */ static PyObject *socket_error; +static PyObject *socket_partialdataerror; static PyObject *socket_herror; static PyObject *socket_gaierror; static PyObject *socket_timeout; @@ -2674,6 +2675,234 @@ Like recv_into(buffer[, nbytes[, flags]]) but also return the sender's address info."); +#ifndef MSG_WAITALL +/* Helper function to set a special error if the recvall function aborted + halfway. It returns an exception with the partially received data in it, + so your application might use that if required. If the receive buffer is + still totally empty, the regular error is raised instead. */ +static PyObject* set_partialreceived_error(PySocketSockObject *sockobj, PyObject *buffer, Py_ssize_t bytes_got) +{ + PyObject *p_type; + PyObject *p_value; + PyObject *p_traceback; + PyObject *partialbuffer; + PyObject *errno_obj; + PyObject *error_message; + PyObject *new_tuple; + + /* use the regular error when the buffer is still empty */ + if(buffer==NULL) { + return sockobj->errorhandler(); + } + if(bytes_got==0) { + Py_DECREF(buffer); + return sockobj->errorhandler(); + } + + partialbuffer=PyString_FromStringAndSize(PyString_AsString(buffer), bytes_got); + Py_DECREF(buffer); + if(partialbuffer==NULL) { + return NULL; + } + sockobj->errorhandler(); /* generate the regular exception */ + PyErr_Fetch(&p_type, &p_value, &p_traceback); + if(p_value!=NULL) { + if(PyTuple_Check(p_value)) { + /* get the first and second arguments from the original exception */ + Py_ssize_t tuplesize=PyTuple_Size(p_value); + if(tuplesize>0) errno_obj=PyTuple_GetItem(p_value,0); + if(tuplesize>1) error_message=PyTuple_GetItem(p_value,1); + } else if(PyString_Check(p_value)) { + /* just get the error message */ + Py_INCREF(p_value); + error_message=p_value; + } + } + + /* replace the exception with socket.partialdataerror */ + new_tuple=PyTuple_New(3); /* new exception: (errno, message, partialdata) */ + if(new_tuple==NULL) { + Py_DECREF(partialbuffer); + return NULL; + } + if(errno_obj!=NULL) { + PyTuple_SET_ITEM(new_tuple, 0, errno_obj); + } else { + Py_INCREF(Py_None); + PyTuple_SET_ITEM(new_tuple,0,Py_None); + } + if(error_message!=NULL) { + PyTuple_SET_ITEM(new_tuple, 1, error_message); + } else { + Py_INCREF(Py_None); + PyTuple_SET_ITEM(new_tuple,1,Py_None); + } + PyTuple_SET_ITEM(new_tuple, 2, partialbuffer); + /* overwrite the original exception with ours: */ + PyErr_Clear(); + PyErr_SetObject(socket_partialdataerror, new_tuple); + return NULL; +} +#endif + +/* s.recvall(nbytes [,flags]) method */ + +static PyObject* +sock_recvall(PySocketSockObject *s, PyObject *args) +{ + int total_size; /* #bytes to receive in total (=arg2) */ + int flags=0; /* optional flags for recv() */ +#ifndef MSG_WAITALL + PyObject *buffer_PyString; /* result buffer (allocated here) */ + char *buffer; /* pointer to start of raw buffer area inside buffer_PyString */ + char *p_Buffer; /* cursor pointing into raw buffer area */ + int bytes_got; /* # bytes already got */ + int n; +#else + PyObject *recv_newargs; /* to pass new args to sock_recv() */ + PyObject *recv_result; /* the result of sock_recv() */ +#endif + + if (!PyArg_ParseTuple(args, "i|i:recvall", &total_size, &flags)) + return NULL; + + if (total_size < 0) { + PyErr_SetString(PyExc_ValueError, "negative buffersize in recvall"); + return NULL; + } + +#ifdef MSG_WAITALL + /* just use sock_recv() but with the MSG_WAITALL flag added. */ + flags |= MSG_WAITALL; + recv_newargs = Py_BuildValue("ii",total_size,flags); + recv_result = sock_recv(s, recv_newargs); + Py_DECREF(recv_newargs); + return recv_result; +#else + /* no support for MSG_WAITALL, so simulate it by doing a loop ourselves. */ + + /* allocate uninitialized string of given length, to use as receive buffer */ + buffer_PyString = PyString_FromStringAndSize(NULL,total_size); + if (buffer_PyString == NULL) + return NULL; + + buffer = PyString_AS_STRING(buffer_PyString); + p_Buffer = buffer; + bytes_got = 0; + n = 1; + + while ((bytes_got 0)) { + int bytes_get = total_size - bytes_got; + n = sock_recv_guts(s, p_Buffer, bytes_get, flags); + if (n<0) + break; + p_Buffer += n; + bytes_got += n; + } + + if(bytes_got==total_size) { + return buffer_PyString; + } else { + return set_partialreceived_error(s,buffer_PyString, bytes_got); + } +#endif /* MSG_WAITALL */ +} + +PyDoc_STRVAR(recvall_doc, +"recvall(buffersize[, flags]) -> data\n\ +\n\ +Receive a piece of data with a well-known size from the socket in one call.\n\ +When the system supports the MSG_WAITALL socket flag, it will use that.\n\ +Otherwise, it will use an internal loop, until all bytes have been\n\ +received. Also see recv()."); + +/* s.recvall_into(buffer, [nbytes [,flags]]) method */ + +static PyObject* +sock_recvall_into(PySocketSockObject *s, PyObject *args, PyObject *kwds) +{ + static char *kwlist[] = {"buffer", "nbytes", "flags", 0}; + + int recvlen = 0, flags = 0; + ssize_t readlen; + Py_buffer buf; + Py_ssize_t buflen; +#ifndef MSG_WAITALL + char *p_Buffer; /* cursor pointing into raw buffer area */ + int bytes_got; /* # bytes already got */ + int n; +#endif + + /* Get the buffer's memory */ + if (!PyArg_ParseTupleAndKeywords(args, kwds, "w*|ii:recv_into", kwlist, + &buf, &recvlen, &flags)) + return NULL; + buflen = buf.len; + assert(buf.buf != 0 && buflen > 0); + + if (recvlen < 0) { + PyErr_SetString(PyExc_ValueError, + "negative buffersize in recv_into"); + goto error; + } + 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"); + goto error; + } + +#ifdef MSG_WAITALL + /* just use the guts but with the MSG_WAITALL flag added. */ + readlen = sock_recv_guts(s, buf.buf, recvlen, flags|MSG_WAITALL); +#else + p_Buffer = buf.buf; + readlen = 0; + bytes_got = 0; + n = 1; + + while ((bytes_got 0)) { + int bytes_get = recvlen - bytes_got; + n = sock_recv_guts(s, p_Buffer, bytes_get, flags); + if (n<0) { + readlen=n; /* cause an error to be returned */ + break; + } + p_Buffer += n; + readlen += n; + } +#endif + + if (readlen < 0) { + /* Return an error. */ + goto error; + } + + PyBuffer_Release(&buf); + /* Return the number of bytes read. Note that we do not do anything + special here in the case that readlen < recvlen. */ + return PyInt_FromSsize_t(readlen); + +error: + PyBuffer_Release(&buf); + return NULL; +} + +PyDoc_STRVAR(recvall_into_doc, +"recvall_into(buffer, [nbytes[, flags]]) -> nbytes_read\n\ +\n\ +Receive a piece of data with a well-known size from the socket in one call.\n\ +When the system supports the MSG_WAITALL socket flag, it will use that.\n\ +Otherwise, it will use an internal loop, until all bytes have been\n\ +received. The data is stored into a buffer rather than creating a\n\ +new string. If buffersize is not specified (or 0), receive exactly the\n\ +size of the given buffer. See recv() for documentation about the flags."); + /* s.send(data [,flags]) method */ static PyObject * @@ -2964,6 +3193,10 @@ recvfrom_doc}, {"recvfrom_into", (PyCFunction)sock_recvfrom_into, METH_VARARGS | METH_KEYWORDS, recvfrom_into_doc}, + {"recvall", (PyCFunction)sock_recvall, METH_VARARGS, + recvall_doc}, + {"recvall_into", (PyCFunction)sock_recvall_into, METH_VARARGS | METH_KEYWORDS, + recvall_into_doc}, {"send", (PyCFunction)sock_send, METH_VARARGS, send_doc}, {"sendall", (PyCFunction)sock_sendall, METH_VARARGS, @@ -4459,6 +4692,12 @@ PySocketModuleAPI.error = socket_error; Py_INCREF(socket_error); PyModule_AddObject(m, "error", socket_error); + socket_partialdataerror = PyErr_NewException("socket.partialdataerror", + PyExc_StandardError, NULL); /* IOError doesn't accept additional arguments in a meaningful way... */ + if (socket_partialdataerror == NULL) + return; + Py_INCREF(socket_partialdataerror); + PyModule_AddObject(m, "partialdataerror", socket_partialdataerror); socket_herror = PyErr_NewException("socket.herror", socket_error, NULL); if (socket_herror == NULL)