@@ -645,6 +645,7 @@ typedef struct PicklerObject {
645645 int fast_nesting ;
646646 int fix_imports ; /* Indicate whether Pickler should fix
647647 the name of globals for Python 2.x. */
648+ int running ; /* True when a method of Pickler is executing. */
648649 PyObject * fast_memo ;
649650 PyObject * buffer_callback ; /* Callback for out-of-band buffers, or NULL */
650651} PicklerObject ;
@@ -688,6 +689,8 @@ typedef struct UnpicklerObject {
688689 int proto ; /* Protocol of the pickle loaded. */
689690 int fix_imports ; /* Indicate whether Unpickler should fix
690691 the name of globals pickled by Python 2.x. */
692+ int running ; /* True when a method of Unpickler is executing. */
693+
691694} UnpicklerObject ;
692695
693696typedef struct {
@@ -705,6 +708,32 @@ typedef struct {
705708#define PicklerMemoProxyObject_CAST (op ) ((PicklerMemoProxyObject *)(op))
706709#define UnpicklerMemoProxyObject_CAST (op ) ((UnpicklerMemoProxyObject *)(op))
707710
711+ #define BEGIN_USING_PICKLER (SELF , RET ) do { \
712+ if ((SELF)->running) { \
713+ PyErr_SetString(PyExc_RuntimeError, \
714+ "Pickler object is already used"); \
715+ return (RET); \
716+ } \
717+ (SELF)->running = 1; \
718+ } while (0)
719+
720+ #define END_USING_PICKLER (SELF ) do { \
721+ (SELF)->running = 0; \
722+ } while (0)
723+
724+ #define BEGIN_USING_UNPICKLER (SELF , RET ) do { \
725+ if ((SELF)->running) { \
726+ PyErr_SetString(PyExc_RuntimeError, \
727+ "Unpickler object is already used"); \
728+ return (RET); \
729+ } \
730+ (SELF)->running = 1; \
731+ } while (0)
732+
733+ #define END_USING_UNPICKLER (SELF ) do { \
734+ (SELF)->running = 0; \
735+ } while (0)
736+
708737/* Forward declarations */
709738static int save (PickleState * state , PicklerObject * , PyObject * , int );
710739static int save_reduce (PickleState * , PicklerObject * , PyObject * , PyObject * );
@@ -1134,6 +1163,7 @@ _Pickler_New(PickleState *st)
11341163 self -> fast = 0 ;
11351164 self -> fast_nesting = 0 ;
11361165 self -> fix_imports = 0 ;
1166+ self -> running = 0 ;
11371167 self -> fast_memo = NULL ;
11381168 self -> buffer_callback = NULL ;
11391169
@@ -1637,6 +1667,7 @@ _Unpickler_New(PyObject *module)
16371667 self -> marks_size = 0 ;
16381668 self -> proto = 0 ;
16391669 self -> fix_imports = 0 ;
1670+ self -> running = 0 ;
16401671
16411672 PyObject_GC_Track (self );
16421673 return self ;
@@ -4693,17 +4724,23 @@ _pickle_Pickler_dump_impl(PicklerObject *self, PyTypeObject *cls,
46934724 Py_TYPE (self )-> tp_name );
46944725 return NULL ;
46954726 }
4727+ BEGIN_USING_PICKLER (self , NULL );
46964728
4697- if (_Pickler_ClearBuffer (self ) < 0 )
4698- return NULL ;
4699-
4700- if (dump (st , self , obj ) < 0 )
4701- return NULL ;
4702-
4703- if (_Pickler_FlushToFile (self ) < 0 )
4704- return NULL ;
4705-
4729+ if (_Pickler_ClearBuffer (self ) < 0 ) {
4730+ goto error ;
4731+ }
4732+ if (dump (st , self , obj ) < 0 ) {
4733+ goto error ;
4734+ }
4735+ if (_Pickler_FlushToFile (self ) < 0 ) {
4736+ goto error ;
4737+ }
4738+ END_USING_PICKLER (self );
47064739 Py_RETURN_NONE ;
4740+
4741+ error :
4742+ END_USING_PICKLER (self );
4743+ return NULL ;
47074744}
47084745
47094746/*[clinic input]
@@ -4844,47 +4881,54 @@ _pickle_Pickler___init___impl(PicklerObject *self, PyObject *file,
48444881 PyObject * buffer_callback )
48454882/*[clinic end generated code: output=0abedc50590d259b input=cddc50f66b770002]*/
48464883{
4884+ BEGIN_USING_PICKLER (self , -1 );
48474885 /* In case of multiple __init__() calls, clear previous content. */
48484886 if (self -> write != NULL )
48494887 (void )Pickler_clear ((PyObject * )self );
48504888
4851- if (_Pickler_SetProtocol (self , protocol , fix_imports ) < 0 )
4852- return -1 ;
4853-
4854- if (_Pickler_SetOutputStream (self , file ) < 0 )
4855- return -1 ;
4856-
4857- if (_Pickler_SetBufferCallback (self , buffer_callback ) < 0 )
4858- return -1 ;
4859-
4889+ if (_Pickler_SetProtocol (self , protocol , fix_imports ) < 0 ) {
4890+ goto error ;
4891+ }
4892+ if (_Pickler_SetOutputStream (self , file ) < 0 ) {
4893+ goto error ;
4894+ }
4895+ if (_Pickler_SetBufferCallback (self , buffer_callback ) < 0 ) {
4896+ goto error ;
4897+ }
48604898 /* memo and output_buffer may have already been created in _Pickler_New */
48614899 if (self -> memo == NULL ) {
48624900 self -> memo = PyMemoTable_New ();
4863- if (self -> memo == NULL )
4864- return -1 ;
4901+ if (self -> memo == NULL ) {
4902+ goto error ;
4903+ }
48654904 }
48664905 self -> output_len = 0 ;
48674906 if (self -> output_buffer == NULL ) {
48684907 self -> max_output_len = WRITE_BUF_SIZE ;
48694908 self -> output_buffer = PyBytes_FromStringAndSize (NULL ,
48704909 self -> max_output_len );
4871- if (self -> output_buffer == NULL )
4872- return -1 ;
4910+ if (self -> output_buffer == NULL ) {
4911+ goto error ;
4912+ }
48734913 }
48744914
48754915 self -> fast = 0 ;
48764916 self -> fast_nesting = 0 ;
48774917 self -> fast_memo = NULL ;
48784918
4879- if (self -> dispatch_table != NULL ) {
4880- return 0 ;
4881- }
4882- if (PyObject_GetOptionalAttr ((PyObject * )self , & _Py_ID (dispatch_table ),
4883- & self -> dispatch_table ) < 0 ) {
4884- return -1 ;
4919+ if (self -> dispatch_table == NULL ) {
4920+ if (PyObject_GetOptionalAttr ((PyObject * )self , & _Py_ID (dispatch_table ),
4921+ & self -> dispatch_table ) < 0 ) {
4922+ goto error ;
4923+ }
48854924 }
48864925
4926+ END_USING_PICKLER (self );
48874927 return 0 ;
4928+
4929+ error :
4930+ END_USING_PICKLER (self );
4931+ return -1 ;
48884932}
48894933
48904934
@@ -7073,22 +7117,22 @@ static PyObject *
70737117_pickle_Unpickler_load_impl (UnpicklerObject * self , PyTypeObject * cls )
70747118/*[clinic end generated code: output=cc88168f608e3007 input=f5d2f87e61d5f07f]*/
70757119{
7076- UnpicklerObject * unpickler = (UnpicklerObject * )self ;
7077-
70787120 PickleState * st = _Pickle_GetStateByClass (cls );
70797121
70807122 /* Check whether the Unpickler was initialized correctly. This prevents
70817123 segfaulting if a subclass overridden __init__ with a function that does
70827124 not call Unpickler.__init__(). Here, we simply ensure that self->read
70837125 is not NULL. */
7084- if (unpickler -> read == NULL ) {
7126+ if (self -> read == NULL ) {
70857127 PyErr_Format (st -> UnpicklingError ,
70867128 "Unpickler.__init__() was not called by %s.__init__()" ,
7087- Py_TYPE (unpickler )-> tp_name );
7129+ Py_TYPE (self )-> tp_name );
70887130 return NULL ;
70897131 }
7090-
7091- return load (st , unpickler );
7132+ BEGIN_USING_UNPICKLER (self , NULL );
7133+ PyObject * res = load (st , self );
7134+ END_USING_UNPICKLER (self );
7135+ return res ;
70927136}
70937137
70947138/* The name of find_class() is misleading. In newer pickle protocols, this
@@ -7350,35 +7394,41 @@ _pickle_Unpickler___init___impl(UnpicklerObject *self, PyObject *file,
73507394 const char * errors , PyObject * buffers )
73517395/*[clinic end generated code: output=09f0192649ea3f85 input=ca4c1faea9553121]*/
73527396{
7397+ BEGIN_USING_UNPICKLER (self , -1 );
73537398 /* In case of multiple __init__() calls, clear previous content. */
73547399 if (self -> read != NULL )
73557400 (void )Unpickler_clear ((PyObject * )self );
73567401
7357- if (_Unpickler_SetInputStream (self , file ) < 0 )
7358- return -1 ;
7359-
7360- if (_Unpickler_SetInputEncoding (self , encoding , errors ) < 0 )
7361- return -1 ;
7362-
7363- if (_Unpickler_SetBuffers (self , buffers ) < 0 )
7364- return -1 ;
7365-
7402+ if (_Unpickler_SetInputStream (self , file ) < 0 ) {
7403+ goto error ;
7404+ }
7405+ if (_Unpickler_SetInputEncoding (self , encoding , errors ) < 0 ) {
7406+ goto error ;
7407+ }
7408+ if (_Unpickler_SetBuffers (self , buffers ) < 0 ) {
7409+ goto error ;
7410+ }
73667411 self -> fix_imports = fix_imports ;
73677412
73687413 PyTypeObject * tp = Py_TYPE (self );
73697414 PickleState * state = _Pickle_FindStateByType (tp );
73707415 self -> stack = (Pdata * )Pdata_New (state );
7371- if (self -> stack == NULL )
7372- return -1 ;
7373-
7416+ if (self -> stack == NULL ) {
7417+ goto error ;
7418+ }
73747419 self -> memo_size = 32 ;
73757420 self -> memo = _Unpickler_NewMemo (self -> memo_size );
7376- if (self -> memo == NULL )
7377- return -1 ;
7378-
7421+ if (self -> memo == NULL ) {
7422+ goto error ;
7423+ }
73797424 self -> proto = 0 ;
73807425
7426+ END_USING_UNPICKLER (self );
73817427 return 0 ;
7428+
7429+ error :
7430+ END_USING_UNPICKLER (self );
7431+ return -1 ;
73827432}
73837433
73847434
0 commit comments