From c6069c562b8e3a01a2bef13d0cd113c85ccc4b16 Mon Sep 17 00:00:00 2001
From: Brian Coleman <brianfcoleman@gmail.com>
Date: Wed, 1 Mar 2017 08:53:34 +0000
Subject: [PATCH] bpo-29683 - Fixes to _PyCode_SetExtra when
 co_extra->ce->extras is allocated.

On PyMem_Realloc failure, _PyCode_SetExtra should free co_extra if
co_extra->ce_extras could not be allocated.
On PyMem_Realloc success, _PyCode_SetExtra should set all unused slots in
co_extra->ce_extras to NULL.
---
 Misc/NEWS            |  3 +++
 Objects/codeobject.c | 24 ++++++++++++++----------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/Misc/NEWS b/Misc/NEWS
index 0fe1f9b9d651bf..3b081418c37f3d 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -10,6 +10,9 @@ What's New in Python 3.6.1 release candidate 1?
 Core and Builtins
 -----------------
 
+- bpo-29683: Fixes to memory allocation in _PyCode_SetExtra.  Patch by
+  Brian Coleman.
+
 - bpo-29438: Fixed use-after-free problem in key sharing dict.
 
 - Issue #29319: Prevent RunMainFromImporter overwriting sys.path[0].
diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index 0d8a675f9fbfc6..df8b9538fe9ed5 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -856,16 +856,15 @@ _PyCode_SetExtra(PyObject *code, Py_ssize_t index, void *extra)
     _PyCodeObjectExtra *co_extra = (_PyCodeObjectExtra *) o->co_extra;
 
     if (co_extra == NULL) {
-        o->co_extra = (_PyCodeObjectExtra*) PyMem_Malloc(
-            sizeof(_PyCodeObjectExtra));
-        if (o->co_extra == NULL) {
+        co_extra = PyMem_Malloc(sizeof(_PyCodeObjectExtra));
+        if (co_extra == NULL) {
             return -1;
         }
-        co_extra = (_PyCodeObjectExtra *) o->co_extra;
 
         co_extra->ce_extras = PyMem_Malloc(
             tstate->co_extra_user_count * sizeof(void*));
         if (co_extra->ce_extras == NULL) {
+            PyMem_Free(co_extra);
             return -1;
         }
 
@@ -874,20 +873,25 @@ _PyCode_SetExtra(PyObject *code, Py_ssize_t index, void *extra)
         for (Py_ssize_t i = 0; i < co_extra->ce_size; i++) {
             co_extra->ce_extras[i] = NULL;
         }
+
+        o->co_extra = co_extra;
     }
     else if (co_extra->ce_size <= index) {
-        co_extra->ce_extras = PyMem_Realloc(
+        void** ce_extras = PyMem_Realloc(
             co_extra->ce_extras, tstate->co_extra_user_count * sizeof(void*));
 
-        if (co_extra->ce_extras == NULL) {
+        if (ce_extras == NULL) {
             return -1;
         }
 
-        co_extra->ce_size = tstate->co_extra_user_count;
-
-        for (Py_ssize_t i = co_extra->ce_size; i < co_extra->ce_size; i++) {
-            co_extra->ce_extras[i] = NULL;
+        for (Py_ssize_t i = co_extra->ce_size;
+             i < tstate->co_extra_user_count;
+             i++) {
+            ce_extras[i] = NULL;
         }
+
+        co_extra->ce_extras = ce_extras;
+        co_extra->ce_size = tstate->co_extra_user_count;
     }
 
     co_extra->ce_extras[index] = extra;