Submit fix for issue3393: Memory corruption in multiprocessing module

3 files changed, 28 insertions, 19 deletions

diff --git a/Modules/_multiprocessing/connection.h b/Modules/_multiprocessing/connection.h
index 8e2f7c2..4b475c6 100644
--- a/Modules/_multiprocessing/connection.h
+++ b/Modules/_multiprocessing/connection.h
@@ -129,9 +129,7 @@ connection_sendbytes(ConnectionObject *self, PyObject *args)
 		}
 	}
 
-	Py_BEGIN_ALLOW_THREADS
 	res = conn_send_string(self, buffer + offset, size);
-	Py_END_ALLOW_THREADS
 
 	if (res < 0)
 		return mp_SetError(PyExc_IOError, res);
@@ -156,10 +154,8 @@ connection_recvbytes(ConnectionObject *self, PyObject *args)
 		return NULL;
 	}
 	
-	Py_BEGIN_ALLOW_THREADS
 	res = conn_recv_string(self, self->buffer, CONNECTION_BUFFER_SIZE, 
 			       &freeme, maxlength);
-	Py_END_ALLOW_THREADS
 	
 	if (res < 0) {
 		if (res == MP_BAD_MESSAGE_LENGTH) {
@@ -208,10 +204,8 @@ connection_recvbytes_into(ConnectionObject *self, PyObject *args)
 		return NULL;
 	}
 
-	Py_BEGIN_ALLOW_THREADS
 	res = conn_recv_string(self, buffer+offset, length-offset, 
 			       &freeme, PY_SSIZE_T_MAX);
-	Py_END_ALLOW_THREADS
 
 	if (res < 0) {
 		if (res == MP_BAD_MESSAGE_LENGTH) {
@@ -266,9 +260,7 @@ connection_send_obj(ConnectionObject *self, PyObject *obj)
 	if (PyString_AsStringAndSize(pickled_string, &buffer, &length) < 0)
 		goto failure;
 
-	Py_BEGIN_ALLOW_THREADS
 	res = conn_send_string(self, buffer, (int)length);
-	Py_END_ALLOW_THREADS
 
 	if (res < 0) {
 		mp_SetError(PyExc_IOError, res);
@@ -292,10 +284,8 @@ connection_recv_obj(ConnectionObject *self)
 
 	CHECK_READABLE(self);
 
-	Py_BEGIN_ALLOW_THREADS
 	res = conn_recv_string(self, self->buffer, CONNECTION_BUFFER_SIZE, 
 			       &freeme, PY_SSIZE_T_MAX);
-	Py_END_ALLOW_THREADS
 
 	if (res < 0) {
 		if (res == MP_BAD_MESSAGE_LENGTH) {
diff --git a/Modules/_multiprocessing/pipe_connection.c b/Modules/_multiprocessing/pipe_connection.c
index 1592cca..d655c50 100644
--- a/Modules/_multiprocessing/pipe_connection.c
+++ b/Modules/_multiprocessing/pipe_connection.c
@@ -18,9 +18,12 @@ static Py_ssize_t
 conn_send_string(ConnectionObject *conn, char *string, size_t length)
 {
 	DWORD amount_written;
+	BOOL ret;
 
-	return WriteFile(conn->handle, string, length, &amount_written, NULL)
-		? MP_SUCCESS : MP_STANDARD_ERROR;
+	Py_BEGIN_ALLOW_THREADS
+	ret = WriteFile(conn->handle, string, length, &amount_written, NULL);
+	Py_END_ALLOW_THREADS
+	return ret ? MP_SUCCESS : MP_STANDARD_ERROR;
 }
 
 /*
@@ -34,11 +37,14 @@ conn_recv_string(ConnectionObject *conn, char *buffer,
 		 size_t buflength, char **newbuffer, size_t maxlength)
 {
 	DWORD left, length, full_length, err;
-
+	BOOL ret;
 	*newbuffer = NULL;
 
-	if (ReadFile(conn->handle, buffer, MIN(buflength, maxlength), 
-		     &length, NULL))
+	Py_BEGIN_ALLOW_THREADS
+	ret = ReadFile(conn->handle, buffer, MIN(buflength, maxlength), 
+		      &length, NULL);
+	Py_END_ALLOW_THREADS
+	if (ret)
 		return length;
 
 	err = GetLastError();
@@ -61,7 +67,10 @@ conn_recv_string(ConnectionObject *conn, char *buffer,
 
 	memcpy(*newbuffer, buffer, length);
 
-	if (ReadFile(conn->handle, *newbuffer+length, left, &length, NULL)) {
+	Py_BEGIN_ALLOW_THREADS
+	ret = ReadFile(conn->handle, *newbuffer+length, left, &length, NULL)
+	Py_END_ALLOW_THREADS
+	if (ret) {
 		assert(length == left);
 		return full_length;
 	} else {
diff --git a/Modules/_multiprocessing/socket_connection.c b/Modules/_multiprocessing/socket_connection.c
index a6ff9dd..e5d2d15 100644
--- a/Modules/_multiprocessing/socket_connection.c
+++ b/Modules/_multiprocessing/socket_connection.c
@@ -73,6 +73,7 @@ _conn_recvall(HANDLE h, char *buffer, size_t length)
 static Py_ssize_t
 conn_send_string(ConnectionObject *conn, char *string, size_t length)
 {
+	Py_ssize_t res;
 	/* The "header" of the message is a 32 bit unsigned number (in
 	   network order) which specifies the length of the "body".  If
 	   the message is shorter than about 16kb then it is quicker to
@@ -80,7 +81,6 @@ conn_send_string(ConnectionObject *conn, char *string, size_t length)
 	   them at once. */
 	if (length < (16*1024)) {
 		char *message;
-		int res;
 
 		message = PyMem_Malloc(length+4);
 		if (message == NULL)
@@ -88,9 +88,10 @@ conn_send_string(ConnectionObject *conn, char *string, size_t length)
 
 		*(UINT32*)message = htonl((UINT32)length);     
 		memcpy(message+4, string, length);
+		Py_BEGIN_ALLOW_THREADS
 		res = _conn_sendall(conn->handle, message, length+4);
+		Py_END_ALLOW_THREADS
 		PyMem_Free(message);
-		return res;
 	} else {
 		UINT32 lenbuff;
 
@@ -98,9 +99,12 @@ conn_send_string(ConnectionObject *conn, char *string, size_t length)
 			return MP_BAD_MESSAGE_LENGTH;
 
 		lenbuff = htonl((UINT32)length);
-		return _conn_sendall(conn->handle, (char*)&lenbuff, 4) || 
+		Py_BEGIN_ALLOW_THREADS
+		res = _conn_sendall(conn->handle, (char*)&lenbuff, 4) || 
 			_conn_sendall(conn->handle, string, length);
+		Py_END_ALLOW_THREADS
 	}
+	return res;
 }
 
 /*
@@ -118,7 +122,9 @@ conn_recv_string(ConnectionObject *conn, char *buffer,
 
 	*newbuffer = NULL;
 
+	Py_BEGIN_ALLOW_THREADS
 	res = _conn_recvall(conn->handle, (char*)&ulength, 4);
+	Py_END_ALLOW_THREADS
 	if (res < 0)
 		return res;
 
@@ -127,13 +133,17 @@ conn_recv_string(ConnectionObject *conn, char *buffer,
 		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)
 			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;
 	}
 }