@@ -668,6 +668,10 @@ typedef struct UnpicklerObject {
668668 Py_ssize_t input_len ;
669669 Py_ssize_t next_read_idx ;
670670 Py_ssize_t prefetched_idx ; /* index of first prefetched byte */
671+ Py_ssize_t frame_end ; /* End of the current frame, or -1 if not in a
672+ frame. While in a frame, input_len is capped
673+ here so reads can't cross it (PEP 3154). */
674+ Py_ssize_t saved_input_len ; /* input_len to restore when the frame ends. */
671675
672676 PyObject * read ; /* read() method of the input stream. */
673677 PyObject * readinto ; /* readinto() method of the input stream. */
@@ -1249,6 +1253,7 @@ _Unpickler_SetStringInput(UnpicklerObject *self, PyObject *input)
12491253 self -> input_len = self -> buffer .len ;
12501254 self -> next_read_idx = 0 ;
12511255 self -> prefetched_idx = self -> input_len ;
1256+ self -> frame_end = -1 ;
12521257 return self -> input_len ;
12531258}
12541259
@@ -1259,6 +1264,30 @@ bad_readline(PickleState *st)
12591264 return -1 ;
12601265}
12611266
1267+ /* End the current frame, uncapping input_len. The frame must be fully read. */
1268+ static void
1269+ _Unpickler_EndFrame (UnpicklerObject * self )
1270+ {
1271+ assert (self -> frame_end >= 0 );
1272+ assert (self -> next_read_idx == self -> frame_end );
1273+ self -> input_len = self -> saved_input_len ;
1274+ self -> frame_end = -1 ;
1275+ }
1276+
1277+ /* Reached the end of the current frame. If straddle, the read crosses the
1278+ frame boundary (PEP 3154) and fails; otherwise end the frame. */
1279+ static int
1280+ _Unpickler_LeaveFrame (PickleState * st , UnpicklerObject * self , int straddle )
1281+ {
1282+ if (straddle ) {
1283+ PyErr_SetString (st -> UnpicklingError ,
1284+ "pickle exhausted before end of frame" );
1285+ return -1 ;
1286+ }
1287+ _Unpickler_EndFrame (self );
1288+ return 0 ;
1289+ }
1290+
12621291/* Skip any consumed data that was only prefetched using peek() */
12631292static int
12641293_Unpickler_SkipConsumed (UnpicklerObject * self )
@@ -1443,6 +1472,19 @@ _Unpickler_ReadImpl(UnpicklerObject *self, PickleState *st, char **s, Py_ssize_t
14431472 return -1 ;
14441473 }
14451474
1475+ if (self -> frame_end >= 0 ) {
1476+ if (_Unpickler_LeaveFrame (st , self ,
1477+ self -> next_read_idx < self -> frame_end ) < 0 ) {
1478+ return -1 ;
1479+ }
1480+ /* Frame ended; the read may now be satisfied from the buffer. */
1481+ if (n <= self -> input_len - self -> next_read_idx ) {
1482+ * s = self -> input_buffer + self -> next_read_idx ;
1483+ self -> next_read_idx += n ;
1484+ return n ;
1485+ }
1486+ }
1487+
14461488 /* This case is handled by the _Unpickler_Read() macro for efficiency */
14471489 assert (self -> next_read_idx + n > self -> input_len );
14481490
@@ -1488,6 +1530,25 @@ _Unpickler_ReadInto(PickleState *state, UnpicklerObject *self, char *buf,
14881530 }
14891531 }
14901532
1533+ if (self -> frame_end >= 0 ) {
1534+ /* in_buffer > 0 is next_read_idx < frame_end on entry: frame data was
1535+ consumed above, so the read crosses the frame boundary. */
1536+ if (_Unpickler_LeaveFrame (state , self , in_buffer > 0 ) < 0 ) {
1537+ return -1 ;
1538+ }
1539+ in_buffer = self -> input_len - self -> next_read_idx ;
1540+ if (in_buffer > 0 ) {
1541+ Py_ssize_t to_read = Py_MIN (in_buffer , n );
1542+ memcpy (buf , self -> input_buffer + self -> next_read_idx , to_read );
1543+ self -> next_read_idx += to_read ;
1544+ buf += to_read ;
1545+ n -= to_read ;
1546+ if (n == 0 ) {
1547+ return n ;
1548+ }
1549+ }
1550+ }
1551+
14911552 /* Read from file */
14921553 if (!self -> read ) {
14931554 /* We're unpickling memory, this means the input is truncated */
@@ -1546,6 +1607,7 @@ _Unpickler_Readline(PickleState *state, UnpicklerObject *self, char **result)
15461607{
15471608 Py_ssize_t i , num_read ;
15481609
1610+ rescan :
15491611 for (i = self -> next_read_idx ; i < self -> input_len ; i ++ ) {
15501612 if (self -> input_buffer [i ] == '\n' ) {
15511613 char * line_start = self -> input_buffer + self -> next_read_idx ;
@@ -1554,6 +1616,14 @@ _Unpickler_Readline(PickleState *state, UnpicklerObject *self, char **result)
15541616 return _Unpickler_CopyLine (self , line_start , num_read , result );
15551617 }
15561618 }
1619+ if (self -> frame_end >= 0 ) {
1620+ if (_Unpickler_LeaveFrame (state , self ,
1621+ self -> next_read_idx < self -> frame_end ) < 0 ) {
1622+ return -1 ;
1623+ }
1624+ /* Frame ended; continue the line past its end. */
1625+ goto rescan ;
1626+ }
15571627 if (!self -> read )
15581628 return bad_readline (state );
15591629
@@ -1728,6 +1798,8 @@ _Unpickler_New(PyObject *module)
17281798 self -> input_len = 0 ;
17291799 self -> next_read_idx = 0 ;
17301800 self -> prefetched_idx = 0 ;
1801+ self -> frame_end = -1 ;
1802+ self -> saved_input_len = 0 ;
17311803 self -> read = NULL ;
17321804 self -> readinto = NULL ;
17331805 self -> readline = NULL ;
@@ -7079,6 +7151,17 @@ load_frame(PickleState *state, UnpicklerObject *self)
70797151 if (_Unpickler_Read (self , state , & s , 8 ) < 0 )
70807152 return -1 ;
70817153
7154+ /* A new frame must not begin before the current one ends (PEP 3154). Its
7155+ header may lie at the tail of the current frame if it drains it. */
7156+ if (self -> frame_end >= 0 ) {
7157+ if (self -> next_read_idx < self -> frame_end ) {
7158+ PyErr_SetString (state -> UnpicklingError ,
7159+ "beginning of a new frame before end of current frame" );
7160+ return -1 ;
7161+ }
7162+ _Unpickler_EndFrame (self );
7163+ }
7164+
70827165 frame_len = calc_binsize (s , 8 );
70837166 if (frame_len < 0 ) {
70847167 PyErr_Format (PyExc_OverflowError ,
@@ -7092,6 +7175,12 @@ load_frame(PickleState *state, UnpicklerObject *self)
70927175
70937176 /* Rewind to start of frame */
70947177 self -> next_read_idx -= frame_len ;
7178+
7179+ /* Cap input_len at the frame end so reads can't cross it (PEP 3154);
7180+ _Unpickler_EndFrame() restores it when the frame is fully read. */
7181+ self -> frame_end = self -> next_read_idx + frame_len ;
7182+ self -> saved_input_len = self -> input_len ;
7183+ self -> input_len = self -> frame_end ;
70957184 return 0 ;
70967185}
70977186
0 commit comments