gh-103656: Transfer f-string buffers to parser to avoid use-after-free (GH-103896)

Co-authored-by: Pablo Galindo <pablogsal@gmail.com>

1 files changed, 51 insertions, 41 deletions

diff --git a/Parser/action_helpers.c b/Parser/action_helpers.c
index 55c0f6f..0aaaed6 100644
--- a/Parser/action_helpers.c
+++ b/Parser/action_helpers.c
@@ -965,17 +965,43 @@ _PyPegen_check_legacy_stmt(Parser *p, expr_ty name) {
     return 0;
 }
 
-expr_ty
-_PyPegen_check_fstring_conversion(Parser *p, Token* symbol, expr_ty conv) {
-    if (symbol->lineno != conv->lineno || symbol->end_col_offset != conv->col_offset) {
+static ResultTokenWithMetadata *
+result_token_with_metadata(Parser *p, void *result, PyObject *metadata)
+{
+    ResultTokenWithMetadata *res = _PyArena_Malloc(p->arena, sizeof(ResultTokenWithMetadata));
+    if (res == NULL) {
+        return NULL;
+    }
+    res->metadata = metadata;
+    res->result = result;
+    return res;
+}
+
+ResultTokenWithMetadata *
+_PyPegen_check_fstring_conversion(Parser *p, Token* conv_token, expr_ty conv)
+{
+    if (conv_token->lineno != conv->lineno || conv_token->end_col_offset != conv->col_offset) {
         return RAISE_SYNTAX_ERROR_KNOWN_RANGE(
-            symbol, conv,
+            conv_token, conv,
             "f-string: conversion type must come right after the exclamanation mark"
         );
     }
-    return conv;
+    return result_token_with_metadata(p, conv, conv_token->metadata);
 }
 
+ResultTokenWithMetadata *
+_PyPegen_setup_full_format_spec(Parser *p, Token *colon, asdl_expr_seq *spec, int lineno, int col_offset,
+                                int end_lineno, int end_col_offset, PyArena *arena)
+{
+    if (!spec) {
+        return NULL;
+    }
+    expr_ty res = _PyAST_JoinedStr(spec, lineno, col_offset, end_lineno, end_col_offset, p->arena);
+    if (!res) {
+        return NULL;
+    }
+    return result_token_with_metadata(p, res, colon->metadata);
+}
 
 const char *
 _PyPegen_get_expr_name(expr_ty e)
@@ -1198,27 +1224,6 @@ _PyPegen_nonparen_genexp_in_call(Parser *p, expr_ty args, asdl_comprehension_seq
 // Fstring stuff
 
 static expr_ty
-decode_fstring_buffer(Parser *p, int lineno, int col_offset, int end_lineno,
-                      int end_col_offset)
-{
-    tokenizer_mode *tok_mode = &(p->tok->tok_mode_stack[p->tok->tok_mode_stack_index]);
-    assert(tok_mode->last_expr_buffer != NULL);
-    assert(tok_mode->last_expr_size >= 0 && tok_mode->last_expr_end >= 0);
-
-    PyObject *res = PyUnicode_DecodeUTF8(
-        tok_mode->last_expr_buffer,
-        tok_mode->last_expr_size - tok_mode->last_expr_end,
-        NULL
-    );
-    if (!res || _PyArena_AddPyObject(p->arena, res) < 0) {
-        Py_XDECREF(res);
-        return NULL;
-    }
-
-    return _PyAST_Constant(res, NULL, lineno, col_offset, end_lineno, end_col_offset, p->arena);
-}
-
-static expr_ty
 _PyPegen_decode_fstring_part(Parser* p, int is_raw, expr_ty constant) {
     assert(PyUnicode_CheckExact(constant->v.Constant.value));
 
@@ -1386,19 +1391,20 @@ expr_ty _PyPegen_constant_from_string(Parser* p, Token* tok) {
     return _PyAST_Constant(s, kind, tok->lineno, tok->col_offset, tok->end_lineno, tok->end_col_offset, p->arena);
 }
 
-expr_ty _PyPegen_formatted_value(Parser *p, expr_ty expression, Token *debug, expr_ty conversion,
-                                 expr_ty format, int lineno, int col_offset, int end_lineno, int end_col_offset,
-                                 PyArena *arena) {
+expr_ty _PyPegen_formatted_value(Parser *p, expr_ty expression, Token *debug, ResultTokenWithMetadata *conversion,
+                                 ResultTokenWithMetadata *format, Token *closing_brace, int lineno, int col_offset,
+                                 int end_lineno, int end_col_offset, PyArena *arena) {
     int conversion_val = -1;
     if (conversion != NULL) {
-        assert(conversion->kind == Name_kind);
-        Py_UCS4 first = PyUnicode_READ_CHAR(conversion->v.Name.id, 0);
+        expr_ty conversion_expr = (expr_ty) conversion->result;
+        assert(conversion_expr->kind == Name_kind);
+        Py_UCS4 first = PyUnicode_READ_CHAR(conversion_expr->v.Name.id, 0);
 
-        if (PyUnicode_GET_LENGTH(conversion->v.Name.id) > 1 ||
+        if (PyUnicode_GET_LENGTH(conversion_expr->v.Name.id) > 1 ||
             !(first == 's' || first == 'r' || first == 'a')) {
-            RAISE_SYNTAX_ERROR_KNOWN_LOCATION(conversion,
+            RAISE_SYNTAX_ERROR_KNOWN_LOCATION(conversion_expr,
                                               "f-string: invalid conversion character %R: expected 's', 'r', or 'a'",
-                                              conversion->v.Name.id);
+                                              conversion_expr->v.Name.id);
             return NULL;
         }
 
@@ -1410,7 +1416,7 @@ expr_ty _PyPegen_formatted_value(Parser *p, expr_ty expression, Token *debug, ex
     }
 
     expr_ty formatted_value = _PyAST_FormattedValue(
-        expression, conversion_val, format,
+        expression, conversion_val, format ? (expr_ty) format->result : NULL,
         lineno, col_offset, end_lineno,
         end_col_offset, arena
     );
@@ -1418,22 +1424,26 @@ expr_ty _PyPegen_formatted_value(Parser *p, expr_ty expression, Token *debug, ex
     if (debug) {
         /* Find the non whitespace token after the "=" */
         int debug_end_line, debug_end_offset;
+        PyObject *debug_metadata;
 
         if (conversion) {
-            debug_end_line = conversion->lineno;
-            debug_end_offset = conversion->col_offset;
+            debug_end_line = ((expr_ty) conversion->result)->lineno;
+            debug_end_offset = ((expr_ty) conversion->result)->col_offset;
+            debug_metadata = conversion->metadata;
         }
         else if (format) {
-            debug_end_line = format->lineno;
-            debug_end_offset = format->col_offset + 1; // HACK: ??
+            debug_end_line = ((expr_ty) format->result)->lineno;
+            debug_end_offset = ((expr_ty) format->result)->col_offset + 1;
+            debug_metadata = format->metadata;
         }
         else {
             debug_end_line = end_lineno;
             debug_end_offset = end_col_offset;
+            debug_metadata = closing_brace->metadata;
         }
 
-        expr_ty debug_text = decode_fstring_buffer(p, lineno, col_offset + 1,
-                                                   debug_end_line, debug_end_offset - 1);
+        expr_ty debug_text = _PyAST_Constant(debug_metadata, NULL, lineno, col_offset + 1, debug_end_line,
+                                             debug_end_offset - 1, p->arena);
         if (!debug_text) {
             return NULL;
         }