@@ -669,6 +669,10 @@ typedef struct UnpicklerObject {
669669 Py_ssize_t input_len ;
670670 Py_ssize_t next_read_idx ;
671671 Py_ssize_t prefetched_idx ; /* index of first prefetched byte */
672+ Py_ssize_t frame_end ; /* End of the current frame, or -1 if not in a
673+ frame. While in a frame, input_len is capped
674+ here so reads can't cross it (PEP 3154). */
675+ Py_ssize_t saved_input_len ; /* input_len to restore when the frame ends. */
672676
673677 PyObject * read ; /* read() method of the input stream. */
674678 PyObject * readinto ; /* readinto() method of the input stream. */
@@ -1250,6 +1254,7 @@ _Unpickler_SetStringInput(UnpicklerObject *self, PyObject *input)
12501254 self -> input_len = self -> buffer .len ;
12511255 self -> next_read_idx = 0 ;
12521256 self -> prefetched_idx = self -> input_len ;
1257+ self -> frame_end = -1 ;
12531258 return self -> input_len ;
12541259}
12551260
@@ -1260,6 +1265,30 @@ bad_readline(PickleState *st)
12601265 return -1 ;
12611266}
12621267
1268+ /* End the current frame, uncapping input_len. The frame must be fully read. */
1269+ static void
1270+ _Unpickler_EndFrame (UnpicklerObject * self )
1271+ {
1272+ assert (self -> frame_end >= 0 );
1273+ assert (self -> next_read_idx == self -> frame_end );
1274+ self -> input_len = self -> saved_input_len ;
1275+ self -> frame_end = -1 ;
1276+ }
1277+
1278+ /* Reached the end of the current frame. If straddle, the read crosses the
1279+ frame boundary (PEP 3154) and fails; otherwise end the frame. */
1280+ static int
1281+ _Unpickler_LeaveFrame (PickleState * st , UnpicklerObject * self , int straddle )
1282+ {
1283+ if (straddle ) {
1284+ PyErr_SetString (st -> UnpicklingError ,
1285+ "pickle exhausted before end of frame" );
1286+ return -1 ;
1287+ }
1288+ _Unpickler_EndFrame (self );
1289+ return 0 ;
1290+ }
1291+
12631292/* Skip any consumed data that was only prefetched using peek() */
12641293static int
12651294_Unpickler_SkipConsumed (UnpicklerObject * self )
@@ -1444,6 +1473,19 @@ _Unpickler_ReadImpl(UnpicklerObject *self, PickleState *st, char **s, Py_ssize_t
14441473 return -1 ;
14451474 }
14461475
1476+ if (self -> frame_end >= 0 ) {
1477+ if (_Unpickler_LeaveFrame (st , self ,
1478+ self -> next_read_idx < self -> frame_end ) < 0 ) {
1479+ return -1 ;
1480+ }
1481+ /* Frame ended; the read may now be satisfied from the buffer. */
1482+ if (n <= self -> input_len - self -> next_read_idx ) {
1483+ * s = self -> input_buffer + self -> next_read_idx ;
1484+ self -> next_read_idx += n ;
1485+ return n ;
1486+ }
1487+ }
1488+
14471489 /* This case is handled by the _Unpickler_Read() macro for efficiency */
14481490 assert (self -> next_read_idx + n > self -> input_len );
14491491
@@ -1489,6 +1531,25 @@ _Unpickler_ReadInto(PickleState *state, UnpicklerObject *self, char *buf,
14891531 }
14901532 }
14911533
1534+ if (self -> frame_end >= 0 ) {
1535+ /* in_buffer > 0 is next_read_idx < frame_end on entry: frame data was
1536+ consumed above, so the read crosses the frame boundary. */
1537+ if (_Unpickler_LeaveFrame (state , self , in_buffer > 0 ) < 0 ) {
1538+ return -1 ;
1539+ }
1540+ in_buffer = self -> input_len - self -> next_read_idx ;
1541+ if (in_buffer > 0 ) {
1542+ Py_ssize_t to_read = Py_MIN (in_buffer , n );
1543+ memcpy (buf , self -> input_buffer + self -> next_read_idx , to_read );
1544+ self -> next_read_idx += to_read ;
1545+ buf += to_read ;
1546+ n -= to_read ;
1547+ if (n == 0 ) {
1548+ return n ;
1549+ }
1550+ }
1551+ }
1552+
14921553 /* Read from file */
14931554 if (!self -> read ) {
14941555 /* We're unpickling memory, this means the input is truncated */
@@ -1547,6 +1608,7 @@ _Unpickler_Readline(PickleState *state, UnpicklerObject *self, char **result)
15471608{
15481609 Py_ssize_t i , num_read ;
15491610
1611+ rescan :
15501612 for (i = self -> next_read_idx ; i < self -> input_len ; i ++ ) {
15511613 if (self -> input_buffer [i ] == '\n' ) {
15521614 char * line_start = self -> input_buffer + self -> next_read_idx ;
@@ -1555,6 +1617,14 @@ _Unpickler_Readline(PickleState *state, UnpicklerObject *self, char **result)
15551617 return _Unpickler_CopyLine (self , line_start , num_read , result );
15561618 }
15571619 }
1620+ if (self -> frame_end >= 0 ) {
1621+ if (_Unpickler_LeaveFrame (state , self ,
1622+ self -> next_read_idx < self -> frame_end ) < 0 ) {
1623+ return -1 ;
1624+ }
1625+ /* Frame ended; continue the line past its end. */
1626+ goto rescan ;
1627+ }
15581628 if (!self -> read )
15591629 return bad_readline (state );
15601630
@@ -1729,6 +1799,8 @@ _Unpickler_New(PyObject *module)
17291799 self -> input_len = 0 ;
17301800 self -> next_read_idx = 0 ;
17311801 self -> prefetched_idx = 0 ;
1802+ self -> frame_end = -1 ;
1803+ self -> saved_input_len = 0 ;
17321804 self -> read = NULL ;
17331805 self -> readinto = NULL ;
17341806 self -> readline = NULL ;
@@ -7078,6 +7150,17 @@ load_frame(PickleState *state, UnpicklerObject *self)
70787150 if (_Unpickler_Read (self , state , & s , 8 ) < 0 )
70797151 return -1 ;
70807152
7153+ /* A new frame must not begin before the current one ends (PEP 3154). Its
7154+ header may lie at the tail of the current frame if it drains it. */
7155+ if (self -> frame_end >= 0 ) {
7156+ if (self -> next_read_idx < self -> frame_end ) {
7157+ PyErr_SetString (state -> UnpicklingError ,
7158+ "beginning of a new frame before end of current frame" );
7159+ return -1 ;
7160+ }
7161+ _Unpickler_EndFrame (self );
7162+ }
7163+
70817164 frame_len = calc_binsize (s , 8 );
70827165 if (frame_len < 0 ) {
70837166 PyErr_Format (PyExc_OverflowError ,
@@ -7091,6 +7174,12 @@ load_frame(PickleState *state, UnpicklerObject *self)
70917174
70927175 /* Rewind to start of frame */
70937176 self -> next_read_idx -= frame_len ;
7177+
7178+ /* Cap input_len at the frame end so reads can't cross it (PEP 3154);
7179+ _Unpickler_EndFrame() restores it when the frame is fully read. */
7180+ self -> frame_end = self -> next_read_idx + frame_len ;
7181+ self -> saved_input_len = self -> input_len ;
7182+ self -> input_len = self -> frame_end ;
70947183 return 0 ;
70957184}
70967185
0 commit comments