From 197651b4e38449c7fe4c13e7c6feebeb5c52d220 Mon Sep 17 00:00:00 2001
From: Richard Oudkerk <shibturn@gmail.com>
Date: Mon, 11 Jun 2012 15:12:12 +0100
Subject: [PATCH] Issue #10133: Make multiprocessing deallocate buffer if
 socket read fails.

Patch by Hallvard B Furuseth.
---
 Misc/NEWS                                    |  3 ++
 Modules/_multiprocessing/socket_connection.c | 29 +++++++++++---------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/Misc/NEWS b/Misc/NEWS
index bd79b257abf..22cf31cb693 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -70,6 +70,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #10133: Make multiprocessing deallocate buffer if socket read
+  fails.  Patch by Hallvard B Furuseth.
+
 - Issue #13854: Make multiprocessing properly handle non-integer
   non-string argument to SystemExit.
 
diff --git a/Modules/_multiprocessing/socket_connection.c b/Modules/_multiprocessing/socket_connection.c
index 7ebf3389b3f..1169c0a9237 100644
--- a/Modules/_multiprocessing/socket_connection.c
+++ b/Modules/_multiprocessing/socket_connection.c
@@ -117,7 +117,7 @@ static Py_ssize_t
 conn_recv_string(ConnectionObject *conn, char *buffer,
                  size_t buflength, char **newbuffer, size_t maxlength)
 {
-    int res;
+    Py_ssize_t res;
     UINT32 ulength;
 
     *newbuffer = NULL;
@@ -132,20 +132,23 @@ conn_recv_string(ConnectionObject *conn, char *buffer,
     if (ulength > maxlength)
         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;
 }
 
 /*