gh-153569: bound tokenizer input storage

Author: pablogsal

Lib/test/test_tokenize.py

@@ -2427,6 +2427,31 @@ def test_stop_iteration_skips_encoded_readline_codec_lookup(self):
             (token.ENDMARKER, "", (1, 0), (1, 0), ""),
         )
 
+    def test_fstring_offsets_survive_buffer_reallocation(self):
+        padding = " " * 9000
+        expression_line = ")=:>{2}}\n"
+        physical_lines = [
+            'f"""\n',
+            "{(\n",
+            padding + "1\n",
+            expression_line,
+            '"""\n',
+        ]
+        source = "".join(physical_lines)
+        chunks = iter([
+            "".join(physical_lines[:2]),
+            "".join(physical_lines[2:4]),
+            physical_lines[4],
+            "",
+        ])
+
+        expected = self._get_tokens(source, extra_tokens=True)
+        tokens = list(tokenize._generate_tokens_from_c_tokenizer(
+            chunks.__next__,
+            extra_tokens=True,
+        ))
+        self.assertEqual(tokens, expected)
+
     def test_extra_tokens_relaxes_lexer_errors(self):
         cases = [
             (

Parser/lexer/buffer.c

@@ -1,62 +1,45 @@
 #include "Python.h"
-#include "errcode.h"
-
+#include "buffer.h"
 #include "state.h"
 
-/* Traverse and remember all f-string buffers, in order to be able to restore
-   them after reallocating tok->buf */
 void
-_PyLexer_remember_fstring_buffers(struct tok_state *tok)
+_PyLexer_SnapshotBuffer(struct tok_state *tok, const char *base,
+                        _PyLexer_BufferSnapshot *snapshot)
 {
-    int index;
-    tokenizer_mode *mode;
-
-    for (index = tok->tok_mode_stack_index; index >= 0; --index) {
-        mode = &(tok->tok_mode_stack[index]);
+    snapshot->buf = tok->buf - base;
+    snapshot->cur = tok->cur - tok->buf;
+    snapshot->inp = tok->inp - tok->buf;
+    snapshot->start = tok->start == NULL ? -1 : tok->start - tok->buf;
+    snapshot->line_start = tok->line_start == NULL
+        ? -1 : tok->line_start - tok->buf;
+    snapshot->multi_line_start = tok->multi_line_start == NULL
+        ? -1 : tok->multi_line_start - tok->buf;
+    for (int index = tok->tok_mode_stack_index; index >= 0; --index) {
+        tokenizer_mode *mode = &tok->tok_mode_stack[index];
         mode->start_offset = mode->start == NULL ? -1 : mode->start - tok->buf;
-        mode->multi_line_start_offset = mode->multi_line_start == NULL ? -1 : mode->multi_line_start - tok->buf;
+        mode->multi_line_start_offset = mode->multi_line_start == NULL
+            ? -1 : mode->multi_line_start - tok->buf;
     }
 }
 
-/* Traverse and restore all f-string buffers after reallocating tok->buf */
 void
-_PyLexer_restore_fstring_buffers(struct tok_state *tok)
-{
-    int index;
-    tokenizer_mode *mode;
-
-    for (index = tok->tok_mode_stack_index; index >= 0; --index) {
-        mode = &(tok->tok_mode_stack[index]);
-        mode->start = mode->start_offset < 0 ? NULL : tok->buf + mode->start_offset;
-        mode->multi_line_start = mode->multi_line_start_offset < 0 ? NULL : tok->buf + mode->multi_line_start_offset;
-    }
-}
-
-int
-_PyLexer_tok_reserve_buf(struct tok_state *tok, Py_ssize_t size)
+_PyLexer_RestoreBuffer(struct tok_state *tok, char *base,
+                       const _PyLexer_BufferSnapshot *snapshot)
 {
-    Py_ssize_t cur = tok->cur - tok->buf;
-    Py_ssize_t oldsize = tok->inp - tok->buf;
-    Py_ssize_t newsize = oldsize + Py_MAX(size, oldsize >> 1);
-    if (newsize > tok->end - tok->buf) {
-        char *newbuf = tok->buf;
-        Py_ssize_t start = tok->start == NULL ? -1 : tok->start - tok->buf;
-        Py_ssize_t line_start = tok->start == NULL ? -1 : tok->line_start - tok->buf;
-        Py_ssize_t multi_line_start = tok->multi_line_start - tok->buf;
-        _PyLexer_remember_fstring_buffers(tok);
-        newbuf = (char *)PyMem_Realloc(newbuf, newsize);
-        if (newbuf == NULL) {
-            tok->done = E_NOMEM;
-            return 0;
-        }
-        tok->buf = newbuf;
-        tok->cur = tok->buf + cur;
-        tok->inp = tok->buf + oldsize;
-        tok->end = tok->buf + newsize;
-        tok->start = start < 0 ? NULL : tok->buf + start;
-        tok->line_start = line_start < 0 ? NULL : tok->buf + line_start;
-        tok->multi_line_start = multi_line_start < 0 ? NULL : tok->buf + multi_line_start;
-        _PyLexer_restore_fstring_buffers(tok);
+    tok->buf = base + snapshot->buf;
+    tok->cur = tok->buf + snapshot->cur;
+    tok->inp = tok->buf + snapshot->inp;
+    tok->start = snapshot->start < 0
+        ? NULL : tok->buf + snapshot->start;
+    tok->line_start = snapshot->line_start < 0
+        ? NULL : tok->buf + snapshot->line_start;
+    tok->multi_line_start = snapshot->multi_line_start < 0
+        ? NULL : tok->buf + snapshot->multi_line_start;
+    for (int index = tok->tok_mode_stack_index; index >= 0; --index) {
+        tokenizer_mode *mode = &tok->tok_mode_stack[index];
+        mode->start = mode->start_offset < 0
+            ? NULL : tok->buf + mode->start_offset;
+        mode->multi_line_start = mode->multi_line_start_offset < 0
+            ? NULL : tok->buf + mode->multi_line_start_offset;
     }
-    return 1;
 }

Parser/lexer/buffer.h

@@ -3,8 +3,20 @@
 
 #include "pyport.h"
 
-void _PyLexer_remember_fstring_buffers(struct tok_state *tok);
-void _PyLexer_restore_fstring_buffers(struct tok_state *tok);
-int _PyLexer_tok_reserve_buf(struct tok_state *tok, Py_ssize_t size);
+struct tok_state;
+
+typedef struct {
+    Py_ssize_t buf;
+    Py_ssize_t cur;
+    Py_ssize_t inp;
+    Py_ssize_t start;
+    Py_ssize_t line_start;
+    Py_ssize_t multi_line_start;
+} _PyLexer_BufferSnapshot;
+
+void _PyLexer_SnapshotBuffer(
+    struct tok_state *, const char *, _PyLexer_BufferSnapshot *);
+void _PyLexer_RestoreBuffer(
+    struct tok_state *, char *, const _PyLexer_BufferSnapshot *);
 
 #endif

Parser/lexer/state.c

@@ -26,7 +26,6 @@ _PyTokenizer_tok_new(void)
     tok->interactive_src_start = NULL;
     tok->interactive_src_end = NULL;
     tok->start = NULL;
-    tok->end = NULL;
     tok->done = E_OK;
     tok->fp = NULL;
     tok->tabsize = TABSIZE;
@@ -52,6 +51,7 @@ _PyTokenizer_tok_new(void)
     tok->comment_newline = 0;
     tok->implicit_newline = 0;
     _PyTok_SourceInit(&tok->source);
+    _PyTok_CursorInit(&tok->reader_cursor, &tok->source);
     tok->reader = NULL;
     tok->tok_mode_stack[0] = (tokenizer_mode){.kind =TOK_REGULAR_MODE, .quote='\0', .quote_size = 0, .in_debug=0};
     tok->tok_mode_stack_index = 0;

Parser/lexer/state.h

@@ -2,7 +2,7 @@
 #define _PY_LEXER_H_
 
 #include "object.h"
-#include "../tokenizer/source.h"
+#include "../tokenizer/cursor.h"
 
 #define MAXINDENT 100       /* Max indentation level */
 #define MAXLEVEL 200        /* Max parentheses level */
@@ -67,15 +67,15 @@ typedef struct _tokenizer_mode {
 
 /* Tokenizer state */
 struct tok_state {
-    /* Input state; buf <= cur <= inp <= end */
+    /* Input state; buf <= cur <= inp */
     /* NB an entire line is held in the buffer */
     char *buf;          /* Input buffer, or NULL; malloc'ed if fp != NULL or readline != NULL */
     char *cur;          /* Next character in buffer */
     char *inp;          /* End of data in buffer */
+    _PyTok_Off buf_offset; /* Logical offset of buf[0]. */
     int fp_interactive; /* If the file descriptor is interactive */
     char *interactive_src_start; /* The start of the source parsed so far in interactive mode */
     char *interactive_src_end; /* The end of the source parsed so far in interactive mode */
-    const char *end;    /* End of input buffer if buf != NULL */
     const char *start;  /* Start of current token if not NULL */
     int done;           /* E_OK normally, E_EOF at EOF, otherwise error code */
     /* NB If done != E_OK, cur must be == inp!!! */
@@ -110,6 +110,7 @@ struct tok_state {
     char* str;          /* Source string being tokenized (if tokenizing from a string)*/
 
     _PyTok_SourceText source;
+    _PyTok_Cursor reader_cursor;
     struct _PyTok_Reader *reader;
 
     int type_comments;      /* Whether to look for type comments */

Parser/lexer/string.c

@@ -125,7 +125,8 @@ _PyLexer_update_ftstring_expr(struct tok_state *tok, char cur)
 {
     assert(tok->cur != NULL);
 
-    Py_ssize_t size = strlen(tok->cur);
+    Py_ssize_t size = cur == 0
+        ? tok->inp - tok->cur : (Py_ssize_t)strlen(tok->cur);
     tokenizer_mode *tok_mode = TOK_GET_MODE(tok);
 
     switch (cur) {
@@ -142,7 +143,8 @@ _PyLexer_update_ftstring_expr(struct tok_state *tok, char cur)
                 goto error;
             }
             tok_mode->last_expr_buffer = new_buffer;
-            strncpy(tok_mode->last_expr_buffer + tok_mode->last_expr_size, tok->cur, size);
+            memcpy(tok_mode->last_expr_buffer + tok_mode->last_expr_size,
+                   tok->cur, size);
             tok_mode->last_expr_size += size;
             break;
         case '{':
@@ -155,7 +157,7 @@ _PyLexer_update_ftstring_expr(struct tok_state *tok, char cur)
             }
             tok_mode->last_expr_size = size;
             tok_mode->last_expr_end = -1;
-            strncpy(tok_mode->last_expr_buffer, tok->cur, size);
+            memcpy(tok_mode->last_expr_buffer, tok->cur, size);
             break;
         case '}':
         case '!':