From 8361f714dd69e91afd766f0c47f132a2f081ad48 Mon Sep 17 00:00:00 2001
From: crthpl <56052645+crthpl@users.noreply.github.com>
Date: Wed, 12 May 2021 01:44:47 -0700
Subject: [PATCH] cgen: fix deadlock when returning/breaking in `lock` (#10079)

---
 vlib/v/ast/ast.v                   |   1 +
 vlib/v/ast/scope.v                 |   2 +-
 vlib/v/checker/checker.v           |   3 +
 vlib/v/gen/c/cgen.v                | 127 ++++++++++++++++++++++++-----
 vlib/v/gen/c/fn.v                  |   2 +
 vlib/v/parser/lock.v               |   6 +-
 vlib/v/tests/break_in_lock_test.v  |  53 ++++++++++++
 vlib/v/tests/return_in_lock_test.v |  53 ++++++++++++
 8 files changed, 223 insertions(+), 24 deletions(-)
 create mode 100644 vlib/v/tests/break_in_lock_test.v
 create mode 100644 vlib/v/tests/return_in_lock_test.v

diff --git a/vlib/v/ast/ast.v b/vlib/v/ast/ast.v
index d14f995454..55ca2cb43f 100644
--- a/vlib/v/ast/ast.v
+++ b/vlib/v/ast/ast.v
@@ -743,6 +743,7 @@ pub mut:
 	lockeds []Ident // `x`, `y` in `lock x, y {`
 	is_expr bool
 	typ     Type
+	scope   &Scope
 }
 
 pub struct MatchExpr {
diff --git a/vlib/v/ast/scope.v b/vlib/v/ast/scope.v
index fecb8c6f56..a5fd867bc8 100644
--- a/vlib/v/ast/scope.v
+++ b/vlib/v/ast/scope.v
@@ -193,7 +193,7 @@ pub fn (s &Scope) innermost(pos int) &Scope {
 }
 
 [inline]
-fn (s &Scope) contains(pos int) bool {
+pub fn (s &Scope) contains(pos int) bool {
 	return pos >= s.start_pos && pos <= s.end_pos
 }
 
diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index 165d9152b7..35c07cb793 100644
--- a/vlib/v/checker/checker.v
+++ b/vlib/v/checker/checker.v
@@ -3843,6 +3843,9 @@ fn (mut c Checker) stmt(node ast.Stmt) {
 				node.idx_in_fn = c.table.cur_fn.defer_stmts.len
 				c.table.cur_fn.defer_stmts << unsafe { &node }
 			}
+			if c.locked_names.len != 0 || c.rlocked_names.len != 0 {
+				c.error('defers are not allowed in lock statements', node.pos)
+			}
 			c.stmts(node.stmts)
 		}
 		ast.EnumDecl {
diff --git a/vlib/v/gen/c/cgen.v b/vlib/v/gen/c/cgen.v
index 77ec1bd6fd..6be19846aa 100644
--- a/vlib/v/gen/c/cgen.v
+++ b/vlib/v/gen/c/cgen.v
@@ -83,10 +83,14 @@ mut:
 	optionals              []string // to avoid duplicates TODO perf, use map
 	chan_pop_optionals     []string // types for `x := <-ch or {...}`
 	chan_push_optionals    []string // types for `ch <- x or {...}`
-	shareds                []int    // types with hidden mutex for which decl has been emitted
-	inside_ternary         int      // ?: comma separated statements on a single line
-	inside_map_postfix     bool     // inside map++/-- postfix expr
-	inside_map_infix       bool     // inside map<</+=/-= infix expr
+	cur_lock               ast.LockExpr
+	mtxs                   string // array of mutexes if the `lock` has multiple variables
+	labeled_loops          map[string]&ast.Stmt
+	inner_loop             &ast.Stmt
+	shareds                []int // types with hidden mutex for which decl has been emitted
+	inside_ternary         int   // ?: comma separated statements on a single line
+	inside_map_postfix     bool  // inside map++/-- postfix expr
+	inside_map_infix       bool  // inside map<</+=/-= infix expr
 	inside_map_index       bool
 	inside_opt_data        bool
 	inside_if_optional     bool
@@ -217,6 +221,7 @@ pub fn gen(files []ast.File, table &ast.Table, pref &pref.Preferences) string {
 		indent: -1
 		module_built: module_built
 		timers: util.new_timers(timers_should_print)
+		inner_loop: &ast.EmptyStmt{}
 	}
 	g.timers.start('cgen init')
 	for mod in g.table.modules {
@@ -1106,7 +1111,30 @@ fn (mut g Gen) stmt(node ast.Stmt) {
 		}
 		ast.BranchStmt {
 			g.write_v_source_line_info(node.pos)
+
 			if node.label != '' {
+				x := g.labeled_loops[node.label] or {
+					panic('$node.label doesn\'t exist $g.file.path, $node.pos')
+				}
+				match x {
+					ast.ForCStmt {
+						if x.scope.contains(g.cur_lock.pos.pos) {
+							g.unlock_locks()
+						}
+					}
+					ast.ForInStmt {
+						if x.scope.contains(g.cur_lock.pos.pos) {
+							g.unlock_locks()
+						}
+					}
+					ast.ForStmt {
+						if x.scope.contains(g.cur_lock.pos.pos) {
+							g.unlock_locks()
+						}
+					}
+					else {}
+				}
+
 				if node.kind == .key_break {
 					g.writeln('goto ${node.label}__break;')
 				} else {
@@ -1114,6 +1142,25 @@ fn (mut g Gen) stmt(node ast.Stmt) {
 					g.writeln('goto ${node.label}__continue;')
 				}
 			} else {
+				inner_loop := g.inner_loop
+				match inner_loop {
+					ast.ForCStmt {
+						if inner_loop.scope.contains(g.cur_lock.pos.pos) {
+							g.unlock_locks()
+						}
+					}
+					ast.ForInStmt {
+						if inner_loop.scope.contains(g.cur_lock.pos.pos) {
+							g.unlock_locks()
+						}
+					}
+					ast.ForStmt {
+						if inner_loop.scope.contains(g.cur_lock.pos.pos) {
+							g.unlock_locks()
+						}
+					}
+					else {}
+				}
 				// continue or break
 				if g.is_autofree && !g.is_builtin_mod {
 					g.trace_autofree('// free before continue/break')
@@ -1192,23 +1239,44 @@ fn (mut g Gen) stmt(node ast.Stmt) {
 		ast.ForCStmt {
 			prev_branch_parent_pos := g.branch_parent_pos
 			g.branch_parent_pos = node.pos.pos
+			save_inner_loop := g.inner_loop
+			g.inner_loop = unsafe { &node }
+			if node.label != '' {
+				g.labeled_loops[node.label] = unsafe { &node }
+			}
 			g.write_v_source_line_info(node.pos)
 			g.for_c_stmt(node)
 			g.branch_parent_pos = prev_branch_parent_pos
+			g.labeled_loops.delete(node.label)
+			g.inner_loop = save_inner_loop
 		}
 		ast.ForInStmt {
 			prev_branch_parent_pos := g.branch_parent_pos
 			g.branch_parent_pos = node.pos.pos
+			save_inner_loop := g.inner_loop
+			g.inner_loop = unsafe { &node }
+			if node.label != '' {
+				g.labeled_loops[node.label] = unsafe { &node }
+			}
 			g.write_v_source_line_info(node.pos)
 			g.for_in_stmt(node)
 			g.branch_parent_pos = prev_branch_parent_pos
+			g.labeled_loops.delete(node.label)
+			g.inner_loop = save_inner_loop
 		}
 		ast.ForStmt {
 			prev_branch_parent_pos := g.branch_parent_pos
 			g.branch_parent_pos = node.pos.pos
+			save_inner_loop := g.inner_loop
+			g.inner_loop = unsafe { &node }
+			if node.label != '' {
+				g.labeled_loops[node.label] = unsafe { &node }
+			}
 			g.write_v_source_line_info(node.pos)
 			g.for_stmt(node)
 			g.branch_parent_pos = prev_branch_parent_pos
+			g.labeled_loops.delete(node.label)
+			g.inner_loop = save_inner_loop
 		}
 		ast.GlobalDecl {
 			g.global_decl(node)
@@ -3872,6 +3940,12 @@ fn (mut g Gen) infix_expr(node ast.InfixExpr) {
 }
 
 fn (mut g Gen) lock_expr(node ast.LockExpr) {
+	g.cur_lock = unsafe { node } // is ok because it is discarded at end of fn
+	defer {
+		g.cur_lock = ast.LockExpr{
+			scope: 0
+		}
+	}
 	tmp_result := if node.is_expr { g.new_tmp_var() } else { '' }
 	mut cur_line := ''
 	if node.is_expr {
@@ -3920,30 +3994,19 @@ fn (mut g Gen) lock_expr(node ast.LockExpr) {
 		g.writeln('\t\tsync__RwMutex_lock((sync__RwMutex*)_arr_$mtxs[$mtxs]);')
 		g.writeln('}')
 	}
+	println('')
+	g.mtxs = mtxs
+	defer {
+		g.mtxs = ''
+	}
+
 	g.writeln('/*lock*/ {')
 	g.stmts_with_tmp_var(node.stmts, tmp_result)
 	if node.is_expr {
 		g.writeln(';')
 	}
 	g.writeln('}')
-	if node.lockeds.len == 0 {
-		// this should not happen
-	} else if node.lockeds.len == 1 {
-		id := node.lockeds[0]
-		name := id.name
-		deref := if id.is_mut { '->' } else { '.' }
-		lock_prefix := if node.is_rlock[0] { 'r' } else { '' }
-		g.writeln('sync__RwMutex_${lock_prefix}unlock(&$name${deref}mtx);')
-	} else {
-		// unlock in reverse order
-		g.writeln('for (int $mtxs=${node.lockeds.len - 1}; $mtxs>=0; $mtxs--) {')
-		g.writeln('\tif ($mtxs && _arr_$mtxs[$mtxs] == _arr_$mtxs[$mtxs-1]) continue;')
-		g.writeln('\tif (_isrlck_$mtxs[$mtxs])')
-		g.writeln('\t\tsync__RwMutex_runlock((sync__RwMutex*)_arr_$mtxs[$mtxs]);')
-		g.writeln('\telse')
-		g.writeln('\t\tsync__RwMutex_unlock((sync__RwMutex*)_arr_$mtxs[$mtxs]);')
-		g.write('}')
-	}
+	g.unlock_locks()
 	if node.is_expr {
 		g.writeln('')
 		g.write(cur_line)
@@ -3951,6 +4014,25 @@ fn (mut g Gen) lock_expr(node ast.LockExpr) {
 	}
 }
 
+fn (mut g Gen) unlock_locks() {
+	if g.cur_lock.lockeds.len == 0 {
+	} else if g.cur_lock.lockeds.len == 1 {
+		id := g.cur_lock.lockeds[0]
+		name := id.name
+		deref := if id.is_mut { '->' } else { '.' }
+		lock_prefix := if g.cur_lock.is_rlock[0] { 'r' } else { '' }
+		g.writeln('sync__RwMutex_${lock_prefix}unlock(&$name${deref}mtx);')
+	} else {
+		g.writeln('for (int $g.mtxs=${g.cur_lock.lockeds.len - 1}; $g.mtxs>=0; $g.mtxs--) {')
+		g.writeln('\tif ($g.mtxs && _arr_$g.mtxs[$g.mtxs] == _arr_$g.mtxs[$g.mtxs-1]) continue;')
+		g.writeln('\tif (_isrlck_$g.mtxs[$g.mtxs])')
+		g.writeln('\t\tsync__RwMutex_runlock((sync__RwMutex*)_arr_$g.mtxs[$g.mtxs]);')
+		g.writeln('\telse')
+		g.writeln('\t\tsync__RwMutex_unlock((sync__RwMutex*)_arr_$g.mtxs[$g.mtxs]);')
+		g.write('}')
+	}
+}
+
 fn (mut g Gen) need_tmp_var_in_match(node ast.MatchExpr) bool {
 	if node.is_expr && node.return_type != ast.void_type && node.return_type != 0 {
 		sym := g.table.get_type_symbol(node.return_type)
@@ -4760,6 +4842,7 @@ fn (mut g Gen) return_stmt(node ast.Return) {
 			return
 		}
 	}
+
 	g.inside_return = true
 	defer {
 		g.inside_return = false
diff --git a/vlib/v/gen/c/fn.v b/vlib/v/gen/c/fn.v
index fcbdf871d5..6c1bb41b38 100644
--- a/vlib/v/gen/c/fn.v
+++ b/vlib/v/gen/c/fn.v
@@ -347,6 +347,8 @@ fn (g &Gen) defer_flag_var(stmt &ast.DeferStmt) string {
 }
 
 fn (mut g Gen) write_defer_stmts_when_needed() {
+	// unlock all mutexes, in case we are in a lock statement. defers are not allowed in lock statements
+	g.unlock_locks()
 	if g.defer_stmts.len > 0 {
 		g.write_defer_stmts()
 	}
diff --git a/vlib/v/parser/lock.v b/vlib/v/parser/lock.v
index b2e0bceac7..b8fe74a99d 100644
--- a/vlib/v/parser/lock.v
+++ b/vlib/v/parser/lock.v
@@ -5,6 +5,7 @@ import v.ast
 fn (mut p Parser) lock_expr() ast.LockExpr {
 	// TODO Handle aliasing sync
 	p.register_auto_import('sync')
+	p.open_scope()
 	mut pos := p.tok.position()
 	mut lockeds := []ast.Ident{}
 	mut is_rlocked := []bool{}
@@ -39,12 +40,15 @@ fn (mut p Parser) lock_expr() ast.LockExpr {
 			p.check(.comma)
 		}
 	}
-	stmts := p.parse_block()
+	stmts := p.parse_block_no_scope(false)
+	scope := p.scope
+	p.close_scope()
 	pos.update_last_line(p.prev_tok.line_nr)
 	return ast.LockExpr{
 		lockeds: lockeds
 		stmts: stmts
 		is_rlock: is_rlocked
 		pos: pos
+		scope: scope
 	}
 }
diff --git a/vlib/v/tests/break_in_lock_test.v b/vlib/v/tests/break_in_lock_test.v
new file mode 100644
index 0000000000..02f09a79e4
--- /dev/null
+++ b/vlib/v/tests/break_in_lock_test.v
@@ -0,0 +1,53 @@
+import time
+
+struct AA {
+mut:
+	b string
+}
+
+const (
+	sleep_time = time.millisecond * 50
+)
+
+fn test_return_lock() {
+	shared s := AA{'3'}
+	go printer(shared s)
+	go fn (shared s AA) {
+		start := time.now()
+		for {
+			reader(shared s)
+			if time.now() - start > sleep_time {
+				exit(0)
+			}
+		}
+	}(shared s)
+	time.sleep(sleep_time * 2)
+	assert false
+}
+
+fn printer(shared s AA) {
+	start := time.now()
+	for {
+		lock s {
+			assert s.b in ['0', '1', '2', '3', '4', '5']
+		}
+		if time.now() - start > time.millisecond * 50 {
+			exit(0)
+		}
+	}
+}
+
+fn reader(shared s AA) {
+	mut i := 0
+	for {
+		i++
+		x := i.str()
+		lock s {
+			s.b = x
+			if s.b == '5' {
+				// this test checks if cgen unlocks the mutex here
+				break
+			}
+		}
+	}
+}
diff --git a/vlib/v/tests/return_in_lock_test.v b/vlib/v/tests/return_in_lock_test.v
new file mode 100644
index 0000000000..2d7997dddc
--- /dev/null
+++ b/vlib/v/tests/return_in_lock_test.v
@@ -0,0 +1,53 @@
+import time
+
+struct AA {
+mut:
+	b string
+}
+
+const (
+	sleep_time = time.millisecond * 50
+)
+
+fn test_return_lock() {
+	shared s := AA{'3'}
+	go printer(shared s)
+	go fn (shared s AA) {
+		start := time.now()
+		for {
+			reader(shared s)
+			if time.now() - start > sleep_time {
+				exit(0)
+			}
+		}
+	}(shared s)
+	time.sleep(sleep_time * 2)
+	assert false
+}
+
+fn printer(shared s AA) {
+	start := time.now()
+	for {
+		lock s {
+			assert s.b in ['0', '1', '2', '3', '4', '5']
+		}
+		if time.now() - start > time.millisecond * 50 {
+			exit(0)
+		}
+	}
+}
+
+fn reader(shared s AA) {
+	mut i := 0
+	for {
+		i++
+		x := i.str()
+		lock s {
+			s.b = x
+			if s.b == '5' {
+				// this test checks if cgen unlocks the mutex here
+				return
+			}
+		}
+	}
+}