Save/restore more lexer state when skipping text due to \if.

When we implemented \if ... \endif in psql, we arranged to
save/restore the lexer's parenthesis depth counter across any chunk
of input that we're ignoring.  At the time, that was sufficient,
because no other part of PsqlScanState could need to be restored to
its prior value.  However, commit e717a9a18 and follow-ons added
more state fields that ought to be restored to their prior values.
A problem would only be observed if someone tries to \if out a
portion of a CREATE FUNCTION/PROCEDURE command that is relevant to
BEGIN/END matching, which seems like a pretty unusual usage, so the
lack of field reports isn't surprising.  Nonetheless it's a bug.

To fix, replace the simple counter field in ConditionalStack
entries with a pointer to a struct defined by psqlscan_int.h.
(In the back branches, keep the old field and associated functions
to minimize the risk of API/ABI breakage, even though it seems
unlikely that any third-party code is using this.  Making the
new struct private to psqlscan-related code should prevent API/ABI
issues for future additions of this type.)

In itself this is only a minor bug fix, but it's prerequisite
infrastructure for the fix for CVE-2026-6464, which will add
another such field.

Author: Tom Lane <tgl@sss.pgh.pa.us>
Reviewed-by: Noah Misch <noah@leadboat.com>
Backpatch-through: 14
Security: CVE-2026-6464

src/bin/psql/command.c

@@ -3066,8 +3066,8 @@ is_branching_command(const char *cmd)
  * Prepare to possibly restore query buffer to its current state
  * (cf. discard_query_text).
  *
- * We need to remember the length of the query buffer, and the lexer's
- * notion of the parenthesis nesting depth.
+ * We need to remember the length of the query buffer, and assorted
+ * lexer internal state such as parenthesis nesting depth.
  */
 static void
 save_query_text_state(PsqlScanState scan_state, ConditionalStack cstack,
@@ -3075,8 +3075,8 @@ save_query_text_state(PsqlScanState scan_state, ConditionalStack cstack,
 {
 	if (query_buf)
 		conditional_stack_set_query_len(cstack, query_buf->len);
-	conditional_stack_set_paren_depth(cstack,
-									  psql_scan_get_paren_depth(scan_state));
+	conditional_stack_set_lex_state(cstack,
+									psql_scan_get_lex_state(scan_state));
 }
 
 /*
@@ -3085,9 +3085,7 @@ save_query_text_state(PsqlScanState scan_state, ConditionalStack cstack,
  * We must discard data that was appended to query_buf during an inactive
  * \if branch.  We don't have to do anything there if there's no query_buf.
  *
- * Also, reset the lexer state to the same paren depth there was before.
- * (The rest of its state doesn't need attention, since we could not be
- * inside a comment or literal or partial token.)
+ * Also, reset the lexer's state to what it was before.
  */
 static void
 discard_query_text(PsqlScanState scan_state, ConditionalStack cstack,
@@ -3101,8 +3099,8 @@ discard_query_text(PsqlScanState scan_state, ConditionalStack cstack,
 		query_buf->len = new_len;
 		query_buf->data[new_len] = '\0';
 	}
-	psql_scan_set_paren_depth(scan_state,
-							  conditional_stack_get_paren_depth(cstack));
+	psql_scan_set_lex_state(scan_state,
+							conditional_stack_get_lex_state(cstack));
 }
 
 /*

src/bin/psql/psqlscanslash.h

@@ -31,6 +31,11 @@ extern char *psql_scan_slash_option(PsqlScanState state,
 
 extern void psql_scan_slash_command_end(PsqlScanState state);
 
+extern PsqlScanStateSave *psql_scan_get_lex_state(PsqlScanState state);
+
+extern void psql_scan_set_lex_state(PsqlScanState state,
+									const PsqlScanStateSave *lex_state);
+
 extern int	psql_scan_get_paren_depth(PsqlScanState state);
 
 extern void psql_scan_set_paren_depth(PsqlScanState state, int depth);

src/bin/psql/psqlscanslash.l

@@ -693,8 +693,46 @@ psql_scan_slash_command_end(PsqlScanState state)
 	psql_scan_reselect_sql_lexer(state);
 }
 
+/*
+ * Save current lexer state
+ *
+ * Relevant parts of the state are returned in a pg_malloc'd struct.
+ * It is caller's responsibility to free the struct eventually.
+ */
+PsqlScanStateSave *
+psql_scan_get_lex_state(PsqlScanState state)
+{
+	PsqlScanStateSave *lex_state = pg_malloc_object(PsqlScanStateSave);
+	StaticAssertStmt(sizeof(lex_state->identifiers) == sizeof(state->identifiers),
+					 "identifiers array lengths must match");
+
+	lex_state->paren_depth = state->paren_depth;
+	lex_state->begin_depth = state->begin_depth;
+	lex_state->identifier_count = state->identifier_count;
+	memcpy(lex_state->identifiers, state->identifiers,
+		   sizeof(lex_state->identifiers));
+	return lex_state;
+}
+
+/*
+ * Restore lexer state to what it was when saved
+ */
+void
+psql_scan_set_lex_state(PsqlScanState state,
+						const PsqlScanStateSave *lex_state)
+{
+	state->paren_depth = lex_state->paren_depth;
+	state->begin_depth = lex_state->begin_depth;
+	state->identifier_count = lex_state->identifier_count;
+	memcpy(state->identifiers, lex_state->identifiers,
+		   sizeof(state->identifiers));
+}
+
 /*
  * Fetch current paren nesting depth
+ *
+ * (These functions are obsolete, and kept around only to avoid API/ABI
+ * breakage in the back branches.)
  */
 int
 psql_scan_get_paren_depth(PsqlScanState state)

src/fe_utils/conditional.c

@@ -45,6 +45,7 @@ conditional_stack_push(ConditionalStack cstack, ifState new_state)
 	p->if_state = new_state;
 	p->query_len = -1;
 	p->paren_depth = -1;
+	p->lex_state = NULL;
 	p->next = cstack->head;
 	cstack->head = p;
 }
@@ -61,6 +62,8 @@ conditional_stack_pop(ConditionalStack cstack)
 	if (!p)
 		return false;
 	cstack->head = cstack->head->next;
+	if (p->lex_state)
+		free(p->lex_state);
 	free(p);
 	return true;
 }
@@ -154,8 +157,39 @@ conditional_stack_get_query_len(ConditionalStack cstack)
 	return cstack->head->query_len;
 }
 
+/*
+ * Save current lexer state in topmost stack entry.
+ *
+ * The lexer state is presumed to be a single pg_malloc'd chunk.
+ * It will be freed automatically when the stack entry is popped.
+ */
+void
+conditional_stack_set_lex_state(ConditionalStack cstack,
+								struct PsqlScanStateSave *lex_state)
+{
+	Assert(!conditional_stack_empty(cstack));
+	if (cstack->head->lex_state)	/* free old state, if any */
+		free(cstack->head->lex_state);
+	cstack->head->lex_state = lex_state;
+}
+
+/*
+ * Fetch last-recorded lexer state from topmost stack entry.
+ * Will return NULL if no stack or it was never saved.
+ */
+struct PsqlScanStateSave *
+conditional_stack_get_lex_state(ConditionalStack cstack)
+{
+	if (conditional_stack_empty(cstack))
+		return NULL;
+	return cstack->head->lex_state;
+}
+
 /*
  * Save current parenthesis nesting depth in topmost stack entry.
+ *
+ * (These functions are obsolete, and kept around only to avoid API/ABI
+ * breakage in the back branches.)
  */
 void
 conditional_stack_set_paren_depth(ConditionalStack cstack, int depth)

src/include/fe_utils/conditional.h

@@ -49,18 +49,18 @@ typedef enum ifState
  * query_len is used to determine what accumulated text to throw away at the
  * end of an inactive branch.  (We could, perhaps, teach the lexer to not add
  * stuff to the query buffer in the first place when inside an inactive branch;
- * but that would be very invasive.)  We also need to save and restore the
- * lexer's parenthesis nesting depth when throwing away text.  (We don't need
- * to save and restore any of its other state, such as comment nesting depth,
- * because a backslash command could never appear inside a comment or SQL
- * literal.)
+ * but that would be very invasive.)  We also need to save and restore some
+ * lexer state, such as parenthesis nesting depth, when throwing away text.
  */
+struct PsqlScanStateSave;		/* opaque outside lexer */
+
 typedef struct IfStackElem
 {
 	ifState		if_state;		/* current state, see enum above */
 	int			query_len;		/* length of query_buf at last branch start */
-	int			paren_depth;	/* parenthesis depth at last branch start */
+	int			paren_depth;	/* (obsolete, not used anymore) */
 	struct IfStackElem *next;	/* next surrounding \if, if any */
+	struct PsqlScanStateSave *lex_state;	/* lexer state at last branch start */
 } IfStackElem;
 
 typedef struct ConditionalStackData
@@ -93,6 +93,11 @@ extern void conditional_stack_set_query_len(ConditionalStack cstack, int len);
 
 extern int	conditional_stack_get_query_len(ConditionalStack cstack);
 
+extern void conditional_stack_set_lex_state(ConditionalStack cstack,
+											struct PsqlScanStateSave *lex_state);
+
+extern struct PsqlScanStateSave *conditional_stack_get_lex_state(ConditionalStack cstack);
+
 extern void conditional_stack_set_paren_depth(ConditionalStack cstack, int depth);
 
 extern int	conditional_stack_get_paren_depth(ConditionalStack cstack);

src/include/fe_utils/psqlscan.h

@@ -26,6 +26,9 @@
 /* Abstract type for lexer's internal state */
 typedef struct PsqlScanStateData *PsqlScanState;
 
+/* Abstract type for state save/restore */
+typedef struct PsqlScanStateSave PsqlScanStateSave;
+
 /* Termination states for psql_scan() */
 typedef enum
 {

src/include/fe_utils/psqlscan_int.h

@@ -131,6 +131,23 @@ typedef struct PsqlScanStateData
 	void	   *cb_passthrough;
 } PsqlScanStateData;
 
+/*
+ * Conditional scanning (\if ... \endif) needs to be able to reset the
+ * lexer's state to what it was at the beginning of a chunk of text that
+ * we choose to ignore.  PsqlScanStateSave holds the values that need
+ * to be saved and restored.  We assume that saving/restoring happens only
+ * while processing a backslash command, so we needn't save state that is
+ * concerned with comment or SQL literal processing: we won't be inside
+ * one of those.
+ */
+struct PsqlScanStateSave
+{
+	int			paren_depth;	/* depth of nesting in parentheses */
+	int			begin_depth;	/* depth of begin/end pairs */
+	int			identifier_count;	/* identifiers since start of statement */
+	char		identifiers[4]; /* records the first few identifiers */
+};
+
 
 /*
  * Functions exported by psqlscan.l, but only meant for use within

src/test/regress/expected/psql.out

@@ -4542,6 +4542,22 @@ invalid command \lo
 	\echo 'should print #8-1'
 should print #8-1
 \endif
+-- test that begin/end matching ignores to-be-ignored text
+create function silly_function(int) returns int
+begin atomic select $1;
+\if false
+end
+\endif
+;
+end;
+\sf silly_function(int)
+CREATE OR REPLACE FUNCTION public.silly_function(integer)
+ RETURNS integer
+ LANGUAGE sql
+BEGIN ATOMIC
+ SELECT $1;
+END
+drop function silly_function(int);
 -- :{?...} defined variable test
 \set i 1
 \if :{?i}

src/test/regress/sql/psql.sql

@@ -1028,6 +1028,17 @@ select \if false \\ (bogus \else \\ 42 \endif \\ forty_two;
 	\echo 'should print #8-1'
 \endif
 
+-- test that begin/end matching ignores to-be-ignored text
+create function silly_function(int) returns int
+begin atomic select $1;
+\if false
+end
+\endif
+;
+end;
+\sf silly_function(int)
+drop function silly_function(int);
+
 -- :{?...} defined variable test
 \set i 1
 \if :{?i}

src/tools/pgindent/typedefs.list

@@ -2041,6 +2041,7 @@ PsqlScanQuoteType
 PsqlScanResult
 PsqlScanState
 PsqlScanStateData
+PsqlScanStateSave
 PsqlSettings
 Publication
 PublicationActions