@@ -640,6 +640,7 @@ typedef struct PicklerObject {
640640 int fast_nesting ;
641641 int fix_imports ; /* Indicate whether Pickler should fix
642642 the name of globals for Python 2.x. */
643+ int running ; /* True when a method of Pickler is executing. */
643644 PyObject * fast_memo ;
644645 PyObject * buffer_callback ; /* Callback for out-of-band buffers, or NULL */
645646} PicklerObject ;
@@ -683,6 +684,8 @@ typedef struct UnpicklerObject {
683684 int proto ; /* Protocol of the pickle loaded. */
684685 int fix_imports ; /* Indicate whether Unpickler should fix
685686 the name of globals pickled by Python 2.x. */
687+ int running ; /* True when a method of Unpickler is executing. */
688+
686689} UnpicklerObject ;
687690
688691typedef struct {
@@ -695,6 +698,32 @@ typedef struct {
695698 UnpicklerObject * unpickler ;
696699} UnpicklerMemoProxyObject ;
697700
701+ #define BEGIN_USING_PICKLER (SELF , RET ) do { \
702+ if ((SELF)->running) { \
703+ PyErr_SetString(PyExc_RuntimeError, \
704+ "Pickler object is already used"); \
705+ return (RET); \
706+ } \
707+ (SELF)->running = 1; \
708+ } while (0)
709+
710+ #define END_USING_PICKLER (SELF ) do { \
711+ (SELF)->running = 0; \
712+ } while (0)
713+
714+ #define BEGIN_USING_UNPICKLER (SELF , RET ) do { \
715+ if ((SELF)->running) { \
716+ PyErr_SetString(PyExc_RuntimeError, \
717+ "Unpickler object is already used"); \
718+ return (RET); \
719+ } \
720+ (SELF)->running = 1; \
721+ } while (0)
722+
723+ #define END_USING_UNPICKLER (SELF ) do { \
724+ (SELF)->running = 0; \
725+ } while (0)
726+
698727/* Forward declarations */
699728static int save (PickleState * state , PicklerObject * , PyObject * , int );
700729static int save_reduce (PickleState * , PicklerObject * , PyObject * , PyObject * );
@@ -1124,6 +1153,7 @@ _Pickler_New(PickleState *st)
11241153 self -> fast = 0 ;
11251154 self -> fast_nesting = 0 ;
11261155 self -> fix_imports = 0 ;
1156+ self -> running = 0 ;
11271157 self -> fast_memo = NULL ;
11281158 self -> buffer_callback = NULL ;
11291159
@@ -1627,6 +1657,7 @@ _Unpickler_New(PyObject *module)
16271657 self -> marks_size = 0 ;
16281658 self -> proto = 0 ;
16291659 self -> fix_imports = 0 ;
1660+ self -> running = 0 ;
16301661
16311662 PyObject_GC_Track (self );
16321663 return self ;
@@ -4611,17 +4642,23 @@ _pickle_Pickler_dump_impl(PicklerObject *self, PyTypeObject *cls,
46114642 Py_TYPE (self )-> tp_name );
46124643 return NULL ;
46134644 }
4645+ BEGIN_USING_PICKLER (self , NULL );
46144646
4615- if (_Pickler_ClearBuffer (self ) < 0 )
4616- return NULL ;
4617-
4618- if (dump (st , self , obj ) < 0 )
4619- return NULL ;
4620-
4621- if (_Pickler_FlushToFile (self ) < 0 )
4622- return NULL ;
4623-
4647+ if (_Pickler_ClearBuffer (self ) < 0 ) {
4648+ goto error ;
4649+ }
4650+ if (dump (st , self , obj ) < 0 ) {
4651+ goto error ;
4652+ }
4653+ if (_Pickler_FlushToFile (self ) < 0 ) {
4654+ goto error ;
4655+ }
4656+ END_USING_PICKLER (self );
46244657 Py_RETURN_NONE ;
4658+
4659+ error :
4660+ END_USING_PICKLER (self );
4661+ return NULL ;
46254662}
46264663
46274664/*[clinic input]
@@ -4760,47 +4797,54 @@ _pickle_Pickler___init___impl(PicklerObject *self, PyObject *file,
47604797 PyObject * buffer_callback )
47614798/*[clinic end generated code: output=0abedc50590d259b input=a7c969699bf5dad3]*/
47624799{
4800+ BEGIN_USING_PICKLER (self , -1 );
47634801 /* In case of multiple __init__() calls, clear previous content. */
47644802 if (self -> write != NULL )
47654803 (void )Pickler_clear (self );
47664804
4767- if (_Pickler_SetProtocol (self , protocol , fix_imports ) < 0 )
4768- return -1 ;
4769-
4770- if (_Pickler_SetOutputStream (self , file ) < 0 )
4771- return -1 ;
4772-
4773- if (_Pickler_SetBufferCallback (self , buffer_callback ) < 0 )
4774- return -1 ;
4775-
4805+ if (_Pickler_SetProtocol (self , protocol , fix_imports ) < 0 ) {
4806+ goto error ;
4807+ }
4808+ if (_Pickler_SetOutputStream (self , file ) < 0 ) {
4809+ goto error ;
4810+ }
4811+ if (_Pickler_SetBufferCallback (self , buffer_callback ) < 0 ) {
4812+ goto error ;
4813+ }
47764814 /* memo and output_buffer may have already been created in _Pickler_New */
47774815 if (self -> memo == NULL ) {
47784816 self -> memo = PyMemoTable_New ();
4779- if (self -> memo == NULL )
4780- return -1 ;
4817+ if (self -> memo == NULL ) {
4818+ goto error ;
4819+ }
47814820 }
47824821 self -> output_len = 0 ;
47834822 if (self -> output_buffer == NULL ) {
47844823 self -> max_output_len = WRITE_BUF_SIZE ;
47854824 self -> output_buffer = PyBytes_FromStringAndSize (NULL ,
47864825 self -> max_output_len );
4787- if (self -> output_buffer == NULL )
4788- return -1 ;
4826+ if (self -> output_buffer == NULL ) {
4827+ goto error ;
4828+ }
47894829 }
47904830
47914831 self -> fast = 0 ;
47924832 self -> fast_nesting = 0 ;
47934833 self -> fast_memo = NULL ;
47944834
4795- if (self -> dispatch_table != NULL ) {
4796- return 0 ;
4797- }
4798- if (PyObject_GetOptionalAttr ((PyObject * )self , & _Py_ID (dispatch_table ),
4799- & self -> dispatch_table ) < 0 ) {
4800- return -1 ;
4835+ if (self -> dispatch_table == NULL ) {
4836+ if (PyObject_GetOptionalAttr ((PyObject * )self , & _Py_ID (dispatch_table ),
4837+ & self -> dispatch_table ) < 0 ) {
4838+ goto error ;
4839+ }
48014840 }
48024841
4842+ END_USING_PICKLER (self );
48034843 return 0 ;
4844+
4845+ error :
4846+ END_USING_PICKLER (self );
4847+ return -1 ;
48044848}
48054849
48064850
@@ -6987,22 +7031,22 @@ static PyObject *
69877031_pickle_Unpickler_load_impl (UnpicklerObject * self , PyTypeObject * cls )
69887032/*[clinic end generated code: output=cc88168f608e3007 input=f5d2f87e61d5f07f]*/
69897033{
6990- UnpicklerObject * unpickler = (UnpicklerObject * )self ;
6991-
69927034 PickleState * st = _Pickle_GetStateByClass (cls );
69937035
69947036 /* Check whether the Unpickler was initialized correctly. This prevents
69957037 segfaulting if a subclass overridden __init__ with a function that does
69967038 not call Unpickler.__init__(). Here, we simply ensure that self->read
69977039 is not NULL. */
6998- if (unpickler -> read == NULL ) {
7040+ if (self -> read == NULL ) {
69997041 PyErr_Format (st -> UnpicklingError ,
70007042 "Unpickler.__init__() was not called by %s.__init__()" ,
7001- Py_TYPE (unpickler )-> tp_name );
7043+ Py_TYPE (self )-> tp_name );
70027044 return NULL ;
70037045 }
7004-
7005- return load (st , unpickler );
7046+ BEGIN_USING_UNPICKLER (self , NULL );
7047+ PyObject * res = load (st , self );
7048+ END_USING_UNPICKLER (self );
7049+ return res ;
70067050}
70077051
70087052/* The name of find_class() is misleading. In newer pickle protocols, this
@@ -7243,35 +7287,41 @@ _pickle_Unpickler___init___impl(UnpicklerObject *self, PyObject *file,
72437287 const char * errors , PyObject * buffers )
72447288/*[clinic end generated code: output=09f0192649ea3f85 input=ca4c1faea9553121]*/
72457289{
7290+ BEGIN_USING_UNPICKLER (self , -1 );
72467291 /* In case of multiple __init__() calls, clear previous content. */
72477292 if (self -> read != NULL )
72487293 (void )Unpickler_clear (self );
72497294
7250- if (_Unpickler_SetInputStream (self , file ) < 0 )
7251- return -1 ;
7252-
7253- if (_Unpickler_SetInputEncoding (self , encoding , errors ) < 0 )
7254- return -1 ;
7255-
7256- if (_Unpickler_SetBuffers (self , buffers ) < 0 )
7257- return -1 ;
7258-
7295+ if (_Unpickler_SetInputStream (self , file ) < 0 ) {
7296+ goto error ;
7297+ }
7298+ if (_Unpickler_SetInputEncoding (self , encoding , errors ) < 0 ) {
7299+ goto error ;
7300+ }
7301+ if (_Unpickler_SetBuffers (self , buffers ) < 0 ) {
7302+ goto error ;
7303+ }
72597304 self -> fix_imports = fix_imports ;
72607305
72617306 PyTypeObject * tp = Py_TYPE (self );
72627307 PickleState * state = _Pickle_FindStateByType (tp );
72637308 self -> stack = (Pdata * )Pdata_New (state );
7264- if (self -> stack == NULL )
7265- return -1 ;
7266-
7309+ if (self -> stack == NULL ) {
7310+ goto error ;
7311+ }
72677312 self -> memo_size = 32 ;
72687313 self -> memo = _Unpickler_NewMemo (self -> memo_size );
7269- if (self -> memo == NULL )
7270- return -1 ;
7271-
7314+ if (self -> memo == NULL ) {
7315+ goto error ;
7316+ }
72727317 self -> proto = 0 ;
72737318
7319+ END_USING_UNPICKLER (self );
72747320 return 0 ;
7321+
7322+ error :
7323+ END_USING_UNPICKLER (self );
7324+ return -1 ;
72757325}
72767326
72777327
0 commit comments