Index: Modules/_multiprocessing/socket_connection.c --- Modules/_multiprocessing/socket_connection.c +++ Modules/_multiprocessing/socket_connection.c @@ -118,5 +118,5 @@ conn_recv_string(ConnectionObject *conn, size_t buflength, char **newbuffer, size_t maxlength) { - int res; + Py_ssize_t res; UINT32 ulength; @@ -133,18 +133,21 @@ conn_recv_string(ConnectionObject *conn, return MP_BAD_MESSAGE_LENGTH; - if (ulength <= buflength) { - Py_BEGIN_ALLOW_THREADS - res = _conn_recvall(conn->handle, buffer, (size_t)ulength); - Py_END_ALLOW_THREADS - return res < 0 ? res : ulength; - } else { - *newbuffer = PyMem_Malloc((size_t)ulength); - if (*newbuffer == NULL) + if (ulength > buflength) { + *newbuffer = buffer = PyMem_Malloc((size_t)ulength); + if (buffer == NULL) return MP_MEMORY_ERROR; - Py_BEGIN_ALLOW_THREADS - res = _conn_recvall(conn->handle, *newbuffer, (size_t)ulength); - Py_END_ALLOW_THREADS - return res < 0 ? (Py_ssize_t)res : (Py_ssize_t)ulength; } + + Py_BEGIN_ALLOW_THREADS + res = _conn_recvall(conn->handle, buffer, (size_t)ulength); + Py_END_ALLOW_THREADS + + if (res >= 0) { + res = (Py_ssize_t)ulength; + } else if (*newbuffer != NULL) { + PyMem_Free(*newbuffer); + *newbuffer = NULL; + } + return res; }