@@ -664,6 +664,10 @@ typedef struct UnpicklerObject {
664664 Py_ssize_t input_len ;
665665 Py_ssize_t next_read_idx ;
666666 Py_ssize_t prefetched_idx ; /* index of first prefetched byte */
667+ Py_ssize_t frame_end ; /* End of the current frame, or -1 if not in a
668+ frame. While in a frame, input_len is capped
669+ here so reads can't cross it (PEP 3154). */
670+ Py_ssize_t saved_input_len ; /* input_len to restore when the frame ends. */
667671
668672 PyObject * read ; /* read() method of the input stream. */
669673 PyObject * readinto ; /* readinto() method of the input stream. */
@@ -1240,6 +1244,7 @@ _Unpickler_SetStringInput(UnpicklerObject *self, PyObject *input)
12401244 self -> input_len = self -> buffer .len ;
12411245 self -> next_read_idx = 0 ;
12421246 self -> prefetched_idx = self -> input_len ;
1247+ self -> frame_end = -1 ;
12431248 return self -> input_len ;
12441249}
12451250
@@ -1250,6 +1255,30 @@ bad_readline(PickleState *st)
12501255 return -1 ;
12511256}
12521257
1258+ /* End the current frame, uncapping input_len. The frame must be fully read. */
1259+ static void
1260+ _Unpickler_EndFrame (UnpicklerObject * self )
1261+ {
1262+ assert (self -> frame_end >= 0 );
1263+ assert (self -> next_read_idx == self -> frame_end );
1264+ self -> input_len = self -> saved_input_len ;
1265+ self -> frame_end = -1 ;
1266+ }
1267+
1268+ /* Reached the end of the current frame. If straddle, the read crosses the
1269+ frame boundary (PEP 3154) and fails; otherwise end the frame. */
1270+ static int
1271+ _Unpickler_LeaveFrame (PickleState * st , UnpicklerObject * self , int straddle )
1272+ {
1273+ if (straddle ) {
1274+ PyErr_SetString (st -> UnpicklingError ,
1275+ "pickle exhausted before end of frame" );
1276+ return -1 ;
1277+ }
1278+ _Unpickler_EndFrame (self );
1279+ return 0 ;
1280+ }
1281+
12531282/* Skip any consumed data that was only prefetched using peek() */
12541283static int
12551284_Unpickler_SkipConsumed (UnpicklerObject * self )
@@ -1355,6 +1384,19 @@ _Unpickler_ReadImpl(UnpicklerObject *self, PickleState *st, char **s, Py_ssize_t
13551384 return -1 ;
13561385 }
13571386
1387+ if (self -> frame_end >= 0 ) {
1388+ if (_Unpickler_LeaveFrame (st , self ,
1389+ self -> next_read_idx < self -> frame_end ) < 0 ) {
1390+ return -1 ;
1391+ }
1392+ /* Frame ended; the read may now be satisfied from the buffer. */
1393+ if (n <= self -> input_len - self -> next_read_idx ) {
1394+ * s = self -> input_buffer + self -> next_read_idx ;
1395+ self -> next_read_idx += n ;
1396+ return n ;
1397+ }
1398+ }
1399+
13581400 /* This case is handled by the _Unpickler_Read() macro for efficiency */
13591401 assert (self -> next_read_idx + n > self -> input_len );
13601402
@@ -1400,6 +1442,25 @@ _Unpickler_ReadInto(PickleState *state, UnpicklerObject *self, char *buf,
14001442 }
14011443 }
14021444
1445+ if (self -> frame_end >= 0 ) {
1446+ /* in_buffer > 0 is next_read_idx < frame_end on entry: frame data was
1447+ consumed above, so the read crosses the frame boundary. */
1448+ if (_Unpickler_LeaveFrame (state , self , in_buffer > 0 ) < 0 ) {
1449+ return -1 ;
1450+ }
1451+ in_buffer = self -> input_len - self -> next_read_idx ;
1452+ if (in_buffer > 0 ) {
1453+ Py_ssize_t to_read = Py_MIN (in_buffer , n );
1454+ memcpy (buf , self -> input_buffer + self -> next_read_idx , to_read );
1455+ self -> next_read_idx += to_read ;
1456+ buf += to_read ;
1457+ n -= to_read ;
1458+ if (n == 0 ) {
1459+ return n ;
1460+ }
1461+ }
1462+ }
1463+
14031464 /* Read from file */
14041465 if (!self -> read ) {
14051466 /* We're unpickling memory, this means the input is truncated */
@@ -1508,6 +1569,7 @@ _Unpickler_Readline(PickleState *state, UnpicklerObject *self, char **result)
15081569{
15091570 Py_ssize_t i , num_read ;
15101571
1572+ rescan :
15111573 for (i = self -> next_read_idx ; i < self -> input_len ; i ++ ) {
15121574 if (self -> input_buffer [i ] == '\n' ) {
15131575 char * line_start = self -> input_buffer + self -> next_read_idx ;
@@ -1516,6 +1578,14 @@ _Unpickler_Readline(PickleState *state, UnpicklerObject *self, char **result)
15161578 return _Unpickler_CopyLine (self , line_start , num_read , result );
15171579 }
15181580 }
1581+ if (self -> frame_end >= 0 ) {
1582+ if (_Unpickler_LeaveFrame (state , self ,
1583+ self -> next_read_idx < self -> frame_end ) < 0 ) {
1584+ return -1 ;
1585+ }
1586+ /* Frame ended; continue the line past its end. */
1587+ goto rescan ;
1588+ }
15191589 if (!self -> read )
15201590 return bad_readline (state );
15211591
@@ -1645,6 +1715,8 @@ _Unpickler_New(PyObject *module)
16451715 self -> input_len = 0 ;
16461716 self -> next_read_idx = 0 ;
16471717 self -> prefetched_idx = 0 ;
1718+ self -> frame_end = -1 ;
1719+ self -> saved_input_len = 0 ;
16481720 self -> read = NULL ;
16491721 self -> readinto = NULL ;
16501722 self -> readline = NULL ;
@@ -6844,6 +6916,17 @@ load_frame(PickleState *state, UnpicklerObject *self)
68446916 if (_Unpickler_Read (self , state , & s , 8 ) < 0 )
68456917 return -1 ;
68466918
6919+ /* A new frame must not begin before the current one ends (PEP 3154). Its
6920+ header may lie at the tail of the current frame if it drains it. */
6921+ if (self -> frame_end >= 0 ) {
6922+ if (self -> next_read_idx < self -> frame_end ) {
6923+ PyErr_SetString (state -> UnpicklingError ,
6924+ "beginning of a new frame before end of current frame" );
6925+ return -1 ;
6926+ }
6927+ _Unpickler_EndFrame (self );
6928+ }
6929+
68476930 frame_len = calc_binsize (s , 8 );
68486931 if (frame_len < 0 ) {
68496932 PyErr_Format (PyExc_OverflowError ,
@@ -6857,6 +6940,12 @@ load_frame(PickleState *state, UnpicklerObject *self)
68576940
68586941 /* Rewind to start of frame */
68596942 self -> next_read_idx -= frame_len ;
6943+
6944+ /* Cap input_len at the frame end so reads can't cross it (PEP 3154);
6945+ _Unpickler_EndFrame() restores it when the frame is fully read. */
6946+ self -> frame_end = self -> next_read_idx + frame_len ;
6947+ self -> saved_input_len = self -> input_len ;
6948+ self -> input_len = self -> frame_end ;
68606949 return 0 ;
68616950}
68626951
0 commit comments