From fcdfce6820373422bcdb5630f9eb63df14fd0764 Mon Sep 17 00:00:00 2001
From: Tom Lane <tgl@sss.pgh.pa.us>
Date: Fri, 21 Aug 2015 20:17:19 -0400
Subject: [PATCH] Detect mismatched CONTINUE and EXIT statements at plpgsql
 compile time.

With a bit of tweaking of the compile namestack data structure, we can
verify at compile time whether a CONTINUE or EXIT is legal.  This is
surely better than leaving it to runtime, both because earlier is better
and because we can issue a proper error pointer.  Also, we can get rid
of the ad-hoc old way of detecting the problem, which only took care of
CONTINUE not EXIT.

Jim Nasby, adjusted a bit by me
---
 src/pl/plpgsql/src/pl_comp.c          |  4 +-
 src/pl/plpgsql/src/pl_exec.c          | 48 +++------------
 src/pl/plpgsql/src/pl_funcs.c         | 23 +++++++-
 src/pl/plpgsql/src/pl_gram.y          | 70 ++++++++++++++++++----
 src/pl/plpgsql/src/plpgsql.h          | 17 +++++-
 src/test/regress/expected/plpgsql.out | 85 +++++++++++++++++++++------
 src/test/regress/sql/plpgsql.sql      | 64 ++++++++++++++++----
 7 files changed, 226 insertions(+), 85 deletions(-)

diff --git a/src/pl/plpgsql/src/pl_comp.c b/src/pl/plpgsql/src/pl_comp.c
index 05268e3d34..1ae4bb7e7c 100644
--- a/src/pl/plpgsql/src/pl_comp.c
+++ b/src/pl/plpgsql/src/pl_comp.c
@@ -371,7 +371,7 @@ do_compile(FunctionCallInfo fcinfo,
 	 * variables (such as FOUND), and is named after the function itself.
 	 */
 	plpgsql_ns_init();
-	plpgsql_ns_push(NameStr(procStruct->proname));
+	plpgsql_ns_push(NameStr(procStruct->proname), PLPGSQL_LABEL_BLOCK);
 	plpgsql_DumpExecTree = false;
 	plpgsql_start_datums();
 
@@ -851,7 +851,7 @@ plpgsql_compile_inline(char *proc_source)
 	function->extra_errors = 0;
 
 	plpgsql_ns_init();
-	plpgsql_ns_push(func_name);
+	plpgsql_ns_push(func_name, PLPGSQL_LABEL_BLOCK);
 	plpgsql_DumpExecTree = false;
 	plpgsql_start_datums();
 
diff --git a/src/pl/plpgsql/src/pl_exec.c b/src/pl/plpgsql/src/pl_exec.c
index fb9333606d..935fa62851 100644
--- a/src/pl/plpgsql/src/pl_exec.c
+++ b/src/pl/plpgsql/src/pl_exec.c
@@ -437,19 +437,9 @@ plpgsql_exec_function(PLpgSQL_function *func, FunctionCallInfo fcinfo,
 	{
 		estate.err_stmt = NULL;
 		estate.err_text = NULL;
-
-		/*
-		 * Provide a more helpful message if a CONTINUE has been used outside
-		 * the context it can work in.
-		 */
-		if (rc == PLPGSQL_RC_CONTINUE)
-			ereport(ERROR,
-					(errcode(ERRCODE_SYNTAX_ERROR),
-					 errmsg("CONTINUE cannot be used outside a loop")));
-		else
-			ereport(ERROR,
-			   (errcode(ERRCODE_S_R_E_FUNCTION_EXECUTED_NO_RETURN_STATEMENT),
-				errmsg("control reached end of function without RETURN")));
+		ereport(ERROR,
+				(errcode(ERRCODE_S_R_E_FUNCTION_EXECUTED_NO_RETURN_STATEMENT),
+				 errmsg("control reached end of function without RETURN")));
 	}
 
 	/*
@@ -791,19 +781,9 @@ plpgsql_exec_trigger(PLpgSQL_function *func,
 	{
 		estate.err_stmt = NULL;
 		estate.err_text = NULL;
-
-		/*
-		 * Provide a more helpful message if a CONTINUE has been used outside
-		 * the context it can work in.
-		 */
-		if (rc == PLPGSQL_RC_CONTINUE)
-			ereport(ERROR,
-					(errcode(ERRCODE_SYNTAX_ERROR),
-					 errmsg("CONTINUE cannot be used outside a loop")));
-		else
-			ereport(ERROR,
-			   (errcode(ERRCODE_S_R_E_FUNCTION_EXECUTED_NO_RETURN_STATEMENT),
-				errmsg("control reached end of trigger procedure without RETURN")));
+		ereport(ERROR,
+				(errcode(ERRCODE_S_R_E_FUNCTION_EXECUTED_NO_RETURN_STATEMENT),
+		 errmsg("control reached end of trigger procedure without RETURN")));
 	}
 
 	estate.err_stmt = NULL;
@@ -919,19 +899,9 @@ plpgsql_exec_event_trigger(PLpgSQL_function *func, EventTriggerData *trigdata)
 	{
 		estate.err_stmt = NULL;
 		estate.err_text = NULL;
-
-		/*
-		 * Provide a more helpful message if a CONTINUE has been used outside
-		 * the context it can work in.
-		 */
-		if (rc == PLPGSQL_RC_CONTINUE)
-			ereport(ERROR,
-					(errcode(ERRCODE_SYNTAX_ERROR),
-					 errmsg("CONTINUE cannot be used outside a loop")));
-		else
-			ereport(ERROR,
-			   (errcode(ERRCODE_S_R_E_FUNCTION_EXECUTED_NO_RETURN_STATEMENT),
-				errmsg("control reached end of trigger procedure without RETURN")));
+		ereport(ERROR,
+				(errcode(ERRCODE_S_R_E_FUNCTION_EXECUTED_NO_RETURN_STATEMENT),
+		 errmsg("control reached end of trigger procedure without RETURN")));
 	}
 
 	estate.err_stmt = NULL;
diff --git a/src/pl/plpgsql/src/pl_funcs.c b/src/pl/plpgsql/src/pl_funcs.c
index 900ba3c13b..1b29024832 100644
--- a/src/pl/plpgsql/src/pl_funcs.c
+++ b/src/pl/plpgsql/src/pl_funcs.c
@@ -51,11 +51,11 @@ plpgsql_ns_init(void)
  * ----------
  */
 void
-plpgsql_ns_push(const char *label)
+plpgsql_ns_push(const char *label, enum PLpgSQL_label_types label_type)
 {
 	if (label == NULL)
 		label = "";
-	plpgsql_ns_additem(PLPGSQL_NSTYPE_LABEL, 0, label);
+	plpgsql_ns_additem(PLPGSQL_NSTYPE_LABEL, (int) label_type, label);
 }
 
 
@@ -206,6 +206,25 @@ plpgsql_ns_lookup_label(PLpgSQL_nsitem *ns_cur, const char *name)
 }
 
 
+/* ----------
+ * plpgsql_ns_find_nearest_loop		Find innermost loop label in namespace chain
+ * ----------
+ */
+PLpgSQL_nsitem *
+plpgsql_ns_find_nearest_loop(PLpgSQL_nsitem *ns_cur)
+{
+	while (ns_cur != NULL)
+	{
+		if (ns_cur->itemtype == PLPGSQL_NSTYPE_LABEL &&
+			ns_cur->itemno == PLPGSQL_LABEL_LOOP)
+			return ns_cur;
+		ns_cur = ns_cur->prev;
+	}
+
+	return NULL;				/* no loop found */
+}
+
+
 /*
  * Statement type as a string, for use in error messages etc.
  */
diff --git a/src/pl/plpgsql/src/pl_gram.y b/src/pl/plpgsql/src/pl_gram.y
index 00978909a3..225c62ab36 100644
--- a/src/pl/plpgsql/src/pl_gram.y
+++ b/src/pl/plpgsql/src/pl_gram.y
@@ -186,7 +186,8 @@ static	void			check_raise_parameters(PLpgSQL_stmt_raise *stmt);
 %type <forvariable>	for_variable
 %type <stmt>	for_control
 
-%type <str>		any_identifier opt_block_label opt_label option_value
+%type <str>		any_identifier opt_block_label opt_loop_label opt_label
+%type <str>		option_value
 
 %type <list>	proc_sect stmt_elsifs stmt_else
 %type <loop_body>	loop_body
@@ -535,7 +536,7 @@ decl_statement	: decl_varname decl_const decl_datatype decl_collate decl_notnull
 										   $4->itemno, $1.name);
 					}
 				| decl_varname opt_scrollable K_CURSOR
-					{ plpgsql_ns_push($1.name); }
+					{ plpgsql_ns_push($1.name, PLPGSQL_LABEL_OTHER); }
 				  decl_cursor_args decl_is_for decl_cursor_query
 					{
 						PLpgSQL_var *new;
@@ -1218,7 +1219,7 @@ opt_case_else	:
 					}
 				;
 
-stmt_loop		: opt_block_label K_LOOP loop_body
+stmt_loop		: opt_loop_label K_LOOP loop_body
 					{
 						PLpgSQL_stmt_loop *new;
 
@@ -1235,7 +1236,7 @@ stmt_loop		: opt_block_label K_LOOP loop_body
 					}
 				;
 
-stmt_while		: opt_block_label K_WHILE expr_until_loop loop_body
+stmt_while		: opt_loop_label K_WHILE expr_until_loop loop_body
 					{
 						PLpgSQL_stmt_while *new;
 
@@ -1253,7 +1254,7 @@ stmt_while		: opt_block_label K_WHILE expr_until_loop loop_body
 					}
 				;
 
-stmt_for		: opt_block_label K_FOR for_control loop_body
+stmt_for		: opt_loop_label K_FOR for_control loop_body
 					{
 						/* This runs after we've scanned the loop body */
 						if ($3->cmd_type == PLPGSQL_STMT_FORI)
@@ -1282,7 +1283,7 @@ stmt_for		: opt_block_label K_FOR for_control loop_body
 						}
 
 						check_labels($1, $4.end_label, $4.end_label_location);
-						/* close namespace started in opt_block_label */
+						/* close namespace started in opt_loop_label */
 						plpgsql_ns_pop();
 					}
 				;
@@ -1603,7 +1604,7 @@ for_variable	: T_DATUM
 					}
 				;
 
-stmt_foreach_a	: opt_block_label K_FOREACH for_variable foreach_slice K_IN K_ARRAY expr_until_loop loop_body
+stmt_foreach_a	: opt_loop_label K_FOREACH for_variable foreach_slice K_IN K_ARRAY expr_until_loop loop_body
 					{
 						PLpgSQL_stmt_foreach_a *new;
 
@@ -1666,6 +1667,42 @@ stmt_exit		: exit_type opt_label opt_exitcond
 						new->label	  = $2;
 						new->cond	  = $3;
 
+						if ($2)
+						{
+							/* We have a label, so verify it exists */
+							PLpgSQL_nsitem *label;
+
+							label = plpgsql_ns_lookup_label(plpgsql_ns_top(), $2);
+							if (label == NULL)
+								ereport(ERROR,
+										(errcode(ERRCODE_SYNTAX_ERROR),
+										 errmsg("label \"%s\" does not exist",
+												$2),
+										 parser_errposition(@2)));
+							/* CONTINUE only allows loop labels */
+							if (label->itemno != PLPGSQL_LABEL_LOOP && !$1)
+								ereport(ERROR,
+										(errcode(ERRCODE_SYNTAX_ERROR),
+										 errmsg("block label \"%s\" cannot be used in CONTINUE",
+												$2),
+										 parser_errposition(@2)));
+						}
+						else
+						{
+							/*
+							 * No label, so make sure there is some loop (an
+							 * unlabelled EXIT does not match a block, so this
+							 * is the same test for both EXIT and CONTINUE)
+							 */
+							if (plpgsql_ns_find_nearest_loop(plpgsql_ns_top()) == NULL)
+								ereport(ERROR,
+										(errcode(ERRCODE_SYNTAX_ERROR),
+										 /* translator: %s is EXIT or CONTINUE */
+										 errmsg("%s cannot be used outside a loop",
+												plpgsql_stmt_typename((PLpgSQL_stmt *) new)),
+										 parser_errposition(@1)));
+						}
+
 						$$ = (PLpgSQL_stmt *)new;
 					}
 				;
@@ -2290,12 +2327,24 @@ expr_until_loop :
 
 opt_block_label	:
 					{
-						plpgsql_ns_push(NULL);
+						plpgsql_ns_push(NULL, PLPGSQL_LABEL_BLOCK);
+						$$ = NULL;
+					}
+				| LESS_LESS any_identifier GREATER_GREATER
+					{
+						plpgsql_ns_push($2, PLPGSQL_LABEL_BLOCK);
+						$$ = $2;
+					}
+				;
+
+opt_loop_label	:
+					{
+						plpgsql_ns_push(NULL, PLPGSQL_LABEL_LOOP);
 						$$ = NULL;
 					}
 				| LESS_LESS any_identifier GREATER_GREATER
 					{
-						plpgsql_ns_push($2);
+						plpgsql_ns_push($2, PLPGSQL_LABEL_LOOP);
 						$$ = $2;
 					}
 				;
@@ -2306,8 +2355,7 @@ opt_label	:
 					}
 				| any_identifier
 					{
-						if (plpgsql_ns_lookup_label(plpgsql_ns_top(), $1) == NULL)
-							yyerror("label does not exist");
+						/* label validity will be checked by outer production */
 						$$ = $1;
 					}
 				;
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index fd5679c23e..696fb6155d 100644
--- a/src/pl/plpgsql/src/plpgsql.h
+++ b/src/pl/plpgsql/src/plpgsql.h
@@ -46,6 +46,17 @@ enum
 	PLPGSQL_NSTYPE_REC
 };
 
+/* ----------
+ * A PLPGSQL_NSTYPE_LABEL stack entry must be one of these types
+ * ----------
+ */
+enum PLpgSQL_label_types
+{
+	PLPGSQL_LABEL_BLOCK,		/* DECLARE/BEGIN block */
+	PLPGSQL_LABEL_LOOP,			/* looping construct */
+	PLPGSQL_LABEL_OTHER			/* anything else */
+};
+
 /* ----------
  * Datum array node types
  * ----------
@@ -331,6 +342,8 @@ typedef struct PLpgSQL_nsitem
 {								/* Item in the compilers namespace tree */
 	int			itemtype;
 	int			itemno;
+	/* For labels, itemno is a value of enum PLpgSQL_label_types. */
+	/* For other itemtypes, itemno is the associated PLpgSQL_datum's dno. */
 	struct PLpgSQL_nsitem *prev;
 	char		name[FLEXIBLE_ARRAY_MEMBER];	/* nul-terminated string */
 } PLpgSQL_nsitem;
@@ -997,7 +1010,8 @@ extern void exec_get_datum_type_info(PLpgSQL_execstate *estate,
  * ----------
  */
 extern void plpgsql_ns_init(void);
-extern void plpgsql_ns_push(const char *label);
+extern void plpgsql_ns_push(const char *label,
+				enum PLpgSQL_label_types label_type);
 extern void plpgsql_ns_pop(void);
 extern PLpgSQL_nsitem *plpgsql_ns_top(void);
 extern void plpgsql_ns_additem(int itemtype, int itemno, const char *name);
@@ -1006,6 +1020,7 @@ extern PLpgSQL_nsitem *plpgsql_ns_lookup(PLpgSQL_nsitem *ns_cur, bool localmode,
 				  const char *name3, int *names_used);
 extern PLpgSQL_nsitem *plpgsql_ns_lookup_label(PLpgSQL_nsitem *ns_cur,
 						const char *name);
+extern PLpgSQL_nsitem *plpgsql_ns_find_nearest_loop(PLpgSQL_nsitem *ns_cur);
 
 /* ----------
  * Other functions in pl_funcs.c
diff --git a/src/test/regress/expected/plpgsql.out b/src/test/regress/expected/plpgsql.out
index 452ef9d836..b10157bb17 100644
--- a/src/test/regress/expected/plpgsql.out
+++ b/src/test/regress/expected/plpgsql.out
@@ -2830,21 +2830,58 @@ NOTICE:  10
  
 (1 row)
 
--- CONTINUE is only legal inside a loop
-create function continue_test2() returns void as $$
+drop function continue_test1();
+drop table conttesttbl;
+-- should fail: CONTINUE is only legal inside a loop
+create function continue_error1() returns void as $$
 begin
     begin
         continue;
     end;
-    return;
 end;
 $$ language plpgsql;
--- should fail
-select continue_test2();
 ERROR:  CONTINUE cannot be used outside a loop
-CONTEXT:  PL/pgSQL function continue_test2()
--- CONTINUE can't reference the label of a named block
-create function continue_test3() returns void as $$
+LINE 4:         continue;
+                ^
+-- should fail: EXIT is only legal inside a loop
+create function exit_error1() returns void as $$
+begin
+    begin
+        exit;
+    end;
+end;
+$$ language plpgsql;
+ERROR:  EXIT cannot be used outside a loop
+LINE 4:         exit;
+                ^
+-- should fail: no such label
+create function continue_error2() returns void as $$
+begin
+    begin
+        loop
+            continue no_such_label;
+        end loop;
+    end;
+end;
+$$ language plpgsql;
+ERROR:  label "no_such_label" does not exist
+LINE 5:             continue no_such_label;
+                             ^
+-- should fail: no such label
+create function exit_error2() returns void as $$
+begin
+    begin
+        loop
+            exit no_such_label;
+        end loop;
+    end;
+end;
+$$ language plpgsql;
+ERROR:  label "no_such_label" does not exist
+LINE 5:             exit no_such_label;
+                         ^
+-- should fail: CONTINUE can't reference the label of a named block
+create function continue_error3() returns void as $$
 begin
     <<begin_block1>>
     begin
@@ -2854,14 +2891,28 @@ begin
     end;
 end;
 $$ language plpgsql;
--- should fail
-select continue_test3();
-ERROR:  CONTINUE cannot be used outside a loop
-CONTEXT:  PL/pgSQL function continue_test3()
-drop function continue_test1();
-drop function continue_test2();
-drop function continue_test3();
-drop table conttesttbl;
+ERROR:  block label "begin_block1" cannot be used in CONTINUE
+LINE 6:             continue begin_block1;
+                             ^
+-- On the other hand, EXIT *can* reference the label of a named block
+create function exit_block1() returns void as $$
+begin
+    <<begin_block1>>
+    begin
+        loop
+            exit begin_block1;
+            raise exception 'should not get here';
+        end loop;
+    end;
+end;
+$$ language plpgsql;
+select exit_block1();
+ exit_block1 
+-------------
+ 
+(1 row)
+
+drop function exit_block1();
 -- verbose end block and end loop
 create function end_label1() returns void as $$
 <<blbl>>
@@ -2891,7 +2942,7 @@ begin
   end loop flbl1;
 end;
 $$ language plpgsql;
-ERROR:  label does not exist at or near "flbl1"
+ERROR:  end label "flbl1" specified for unlabelled block
 LINE 5:   end loop flbl1;
                    ^
 -- should fail: end label does not match start label
diff --git a/src/test/regress/sql/plpgsql.sql b/src/test/regress/sql/plpgsql.sql
index b46439ee70..7b4191ecf9 100644
--- a/src/test/regress/sql/plpgsql.sql
+++ b/src/test/regress/sql/plpgsql.sql
@@ -2361,21 +2361,51 @@ end; $$ language plpgsql;
 
 select continue_test1();
 
--- CONTINUE is only legal inside a loop
-create function continue_test2() returns void as $$
+drop function continue_test1();
+drop table conttesttbl;
+
+-- should fail: CONTINUE is only legal inside a loop
+create function continue_error1() returns void as $$
 begin
     begin
         continue;
     end;
-    return;
 end;
 $$ language plpgsql;
 
--- should fail
-select continue_test2();
+-- should fail: EXIT is only legal inside a loop
+create function exit_error1() returns void as $$
+begin
+    begin
+        exit;
+    end;
+end;
+$$ language plpgsql;
+
+-- should fail: no such label
+create function continue_error2() returns void as $$
+begin
+    begin
+        loop
+            continue no_such_label;
+        end loop;
+    end;
+end;
+$$ language plpgsql;
+
+-- should fail: no such label
+create function exit_error2() returns void as $$
+begin
+    begin
+        loop
+            exit no_such_label;
+        end loop;
+    end;
+end;
+$$ language plpgsql;
 
--- CONTINUE can't reference the label of a named block
-create function continue_test3() returns void as $$
+-- should fail: CONTINUE can't reference the label of a named block
+create function continue_error3() returns void as $$
 begin
     <<begin_block1>>
     begin
@@ -2386,13 +2416,21 @@ begin
 end;
 $$ language plpgsql;
 
--- should fail
-select continue_test3();
+-- On the other hand, EXIT *can* reference the label of a named block
+create function exit_block1() returns void as $$
+begin
+    <<begin_block1>>
+    begin
+        loop
+            exit begin_block1;
+            raise exception 'should not get here';
+        end loop;
+    end;
+end;
+$$ language plpgsql;
 
-drop function continue_test1();
-drop function continue_test2();
-drop function continue_test3();
-drop table conttesttbl;
+select exit_block1();
+drop function exit_block1();
 
 -- verbose end block and end loop
 create function end_label1() returns void as $$
-- 
2.49.0