BREAKING: return only exits current function/modifier

4 files changed, 212 insertions, 44 deletions

diff --git a/docs/contracts.rst b/docs/contracts.rst
index 3d592ecf..5f370951 100644
--- a/docs/contracts.rst
+++ b/docs/contracts.rst
@@ -314,14 +314,40 @@ inheritable properties of contracts and may be overridden by derived contracts.
         }
     }
 
+    contract Mutex {
+        bool locked;
+        modifier noReentrancy() {
+            if (locked) throw;
+            locked = true;
+            _
+            locked = false;
+        }
+
+        /// This function is protected by a mutex, which means that
+        /// reentrant calls from within msg.sender.call cannot call f again.
+        /// The `return 7` statement assigns 7 to the return value but still
+        /// executes the statement `locked = false` in the modifier.
+        function f() noReentrancy returns (uint) {
+            if (!msg.sender.call()) throw;
+            return 7;
+        }
+    }
+
 Multiple modifiers can be applied to a function by specifying them in a
-whitespace-separated list and will be evaluated in order. Explicit returns from
-a modifier or function body immediately leave the whole function, while control
-flow reaching the end of a function or modifier body continues after the "_" in
-the preceding modifier. Arbitrary expressions are allowed for modifier
-arguments and in this context, all symbols visible from the function are
-visible in the modifier. Symbols introduced in the modifier are not visible in
-the function (as they might change by overriding).
+whitespace-separated list and will be evaluated in order.
+
+.. warning::
+    In an earlier version of Solidity, ``return`` statements in functions
+    having modifiers behaved differently.
+
+Explicit returns from a modifier or function body only leave the current
+modifier or function body. Return variables are assigned and
+control flow continues after the "_" in the preceding modifier.
+
+Arbitrary expressions are allowed for modifier arguments and in this context,
+all symbols visible from the function are visible in the modifier. Symbols
+introduced in the modifier are not visible in the function (as they might
+change by overriding).
 
 .. index:: ! constant
 
diff --git a/libsolidity/codegen/ContractCompiler.cpp b/libsolidity/codegen/ContractCompiler.cpp
index bcfd33f2..715852be 100644
--- a/libsolidity/codegen/ContractCompiler.cpp
+++ b/libsolidity/codegen/ContractCompiler.cpp
@@ -431,16 +431,16 @@ bool ContractCompiler::visit(FunctionDefinition const& _function)
 		if (auto c = m_context.nextConstructor(dynamic_cast<ContractDefinition const&>(*_function.scope())))
 			appendBaseConstructor(*c);
 
-	m_returnTag = m_context.newTag();
+	solAssert(m_returnTags.empty(), "");
 	m_breakTags.clear();
 	m_continueTags.clear();
 	m_stackCleanupForReturn = 0;
 	m_currentFunction = &_function;
-	m_modifierDepth = 0;
+	m_modifierDepth = -1;
 
 	appendModifierOrFunctionCode();
 
-	m_context << m_returnTag;
+	solAssert(m_returnTags.empty(), "");
 
 	// Now we need to re-shuffle the stack. For this we keep a record of the stack layout
 	// that shows the target positions of the elements, where "-1" denotes that this element needs
@@ -695,7 +695,7 @@ bool ContractCompiler::visit(Return const& _return)
 	}
 	for (unsigned i = 0; i < m_stackCleanupForReturn; ++i)
 		m_context << Instruction::POP;
-	m_context.appendJumpTo(m_returnTag);
+	m_context.appendJumpTo(m_returnTags.back());
 	m_context.adjustStackOffset(m_stackCleanupForReturn);
 	return false;
 }
@@ -755,9 +755,7 @@ bool ContractCompiler::visit(PlaceholderStatement const& _placeholderStatement)
 {
 	StackHeightChecker checker(m_context);
 	CompilerContext::LocationSetter locationSetter(m_context, _placeholderStatement);
-	++m_modifierDepth;
 	appendModifierOrFunctionCode();
-	--m_modifierDepth;
 	checker.check();
 	return true;
 }
@@ -775,10 +773,15 @@ void ContractCompiler::appendMissingFunctions()
 void ContractCompiler::appendModifierOrFunctionCode()
 {
 	solAssert(m_currentFunction, "");
+	unsigned stackSurplus = 0;
+	Block const* codeBlock = nullptr;
+
+	m_modifierDepth++;
+
 	if (m_modifierDepth >= m_currentFunction->modifiers().size())
 	{
 		solAssert(m_currentFunction->isImplemented(), "");
-		m_currentFunction->body().accept(*this);
+		codeBlock = &m_currentFunction->body();
 	}
 	else
 	{
@@ -786,37 +789,45 @@ void ContractCompiler::appendModifierOrFunctionCode()
 
 		// constructor call should be excluded
 		if (dynamic_cast<ContractDefinition const*>(modifierInvocation->name()->annotation().referencedDeclaration))
-		{
-			++m_modifierDepth;
 			appendModifierOrFunctionCode();
-			--m_modifierDepth;
-			return;
-		}
-
-		ModifierDefinition const& modifier = m_context.functionModifier(modifierInvocation->name()->name());
-		CompilerContext::LocationSetter locationSetter(m_context, modifier);
-		solAssert(modifier.parameters().size() == modifierInvocation->arguments().size(), "");
-		for (unsigned i = 0; i < modifier.parameters().size(); ++i)
+		else
 		{
-			m_context.addVariable(*modifier.parameters()[i]);
-			compileExpression(
-				*modifierInvocation->arguments()[i],
-				modifier.parameters()[i]->annotation().type
-			);
+			ModifierDefinition const& modifier = m_context.functionModifier(modifierInvocation->name()->name());
+			CompilerContext::LocationSetter locationSetter(m_context, modifier);
+			solAssert(modifier.parameters().size() == modifierInvocation->arguments().size(), "");
+			for (unsigned i = 0; i < modifier.parameters().size(); ++i)
+			{
+				m_context.addVariable(*modifier.parameters()[i]);
+				compileExpression(
+					*modifierInvocation->arguments()[i],
+					modifier.parameters()[i]->annotation().type
+				);
+			}
+			for (VariableDeclaration const* localVariable: modifier.localVariables())
+				appendStackVariableInitialisation(*localVariable);
+
+			stackSurplus =
+				CompilerUtils::sizeOnStack(modifier.parameters()) +
+				CompilerUtils::sizeOnStack(modifier.localVariables());
+			codeBlock = &modifier.body();
+
+			codeBlock = &modifier.body();
 		}
-		for (VariableDeclaration const* localVariable: modifier.localVariables())
-			appendStackVariableInitialisation(*localVariable);
+	}
+
+	if (codeBlock)
+	{
+		m_returnTags.push_back(m_context.newTag());
 
-		unsigned const c_stackSurplus = CompilerUtils::sizeOnStack(modifier.parameters()) +
-										CompilerUtils::sizeOnStack(modifier.localVariables());
-		m_stackCleanupForReturn += c_stackSurplus;
+		codeBlock->accept(*this);
 
-		modifier.body().accept(*this);
+		solAssert(!m_returnTags.empty(), "");
+		m_context << m_returnTags.back();
+		m_returnTags.pop_back();
 
-		for (unsigned i = 0; i < c_stackSurplus; ++i)
-			m_context << Instruction::POP;
-		m_stackCleanupForReturn -= c_stackSurplus;
+		CompilerUtils(m_context).popStackSlots(stackSurplus);
 	}
+	m_modifierDepth--;
 }
 
 void ContractCompiler::appendStackVariableInitialisation(VariableDeclaration const& _variable)
diff --git a/libsolidity/codegen/ContractCompiler.h b/libsolidity/codegen/ContractCompiler.h
index d1517e88..0799a543 100644
--- a/libsolidity/codegen/ContractCompiler.h
+++ b/libsolidity/codegen/ContractCompiler.h
@@ -40,11 +40,9 @@ class ContractCompiler: private ASTConstVisitor
 public:
 	explicit ContractCompiler(CompilerContext& _context, bool _optimise):
 		m_optimise(_optimise),
-		m_context(_context),
-		m_returnTag(eth::Tag, u256(-1))
+		m_context(_context)
 	{
 		m_context = CompilerContext();
-		m_returnTag = m_context.newTag();
 	}
 
 	void compileContract(
@@ -122,7 +120,8 @@ private:
 	CompilerContext& m_context;
 	std::vector<eth::AssemblyItem> m_breakTags; ///< tag to jump to for a "break" statement
 	std::vector<eth::AssemblyItem> m_continueTags; ///< tag to jump to for a "continue" statement
-	eth::AssemblyItem m_returnTag; ///< tag to jump to for a "return" statement
+	/// Tag to jump to for a "return" statement, needs to be stacked because of modifiers.
+	std::vector<eth::AssemblyItem> m_returnTags;
 	unsigned m_modifierDepth = 0;
 	FunctionDefinition const* m_currentFunction = nullptr;
 	unsigned m_stackCleanupForReturn = 0; ///< this number of stack elements need to be removed before jump to m_returnTag
diff --git a/test/libsolidity/SolidityEndToEndTest.cpp b/test/libsolidity/SolidityEndToEndTest.cpp
index a1ab7700..48cb29a1 100644
--- a/test/libsolidity/SolidityEndToEndTest.cpp
+++ b/test/libsolidity/SolidityEndToEndTest.cpp
@@ -2390,7 +2390,8 @@ BOOST_AUTO_TEST_CASE(function_modifier_multi_invocation)
 
 BOOST_AUTO_TEST_CASE(function_modifier_multi_with_return)
 {
-	// Here, the explicit return prevents the second execution
+	// Note that return sets the return variable and jumps to the end of the current function or
+	// modifier code block.
 	char const* sourceCode = R"(
 		contract C {
 			modifier repeat(bool twice) { if (twice) _ _ }
@@ -2399,7 +2400,7 @@ BOOST_AUTO_TEST_CASE(function_modifier_multi_with_return)
 	)";
 	compileAndRun(sourceCode);
 	BOOST_CHECK(callContractFunction("f(bool)", false) == encodeArgs(1));
-	BOOST_CHECK(callContractFunction("f(bool)", true) == encodeArgs(1));
+	BOOST_CHECK(callContractFunction("f(bool)", true) == encodeArgs(2));
 }
 
 BOOST_AUTO_TEST_CASE(function_modifier_overriding)
@@ -6880,6 +6881,137 @@ BOOST_AUTO_TEST_CASE(create_dynamic_array_with_zero_length)
 	BOOST_CHECK(callContractFunction("f()") == encodeArgs(u256(7)));
 }
 
+BOOST_AUTO_TEST_CASE(return_does_not_skip_modifier)
+{
+	char const* sourceCode = R"(
+		contract C {
+			uint public x;
+			modifier setsx {
+				_
+				x = 9;
+			}
+			function f() setsx returns (uint) {
+				return 2;
+			}
+		}
+	)";
+	compileAndRun(sourceCode, 0, "C");
+	BOOST_CHECK(callContractFunction("x()") == encodeArgs(u256(0)));
+	BOOST_CHECK(callContractFunction("f()") == encodeArgs(u256(2)));
+	BOOST_CHECK(callContractFunction("x()") == encodeArgs(u256(9)));
+}
+
+BOOST_AUTO_TEST_CASE(break_in_modifier)
+{
+	char const* sourceCode = R"(
+		contract C {
+			uint public x;
+			modifier run() {
+				for (uint i = 0; i < 10; i++) {
+					_
+					break;
+				}
+			}
+			function f() run {
+				x++;
+			}
+		}
+	)";
+	compileAndRun(sourceCode, 0, "C");
+	BOOST_CHECK(callContractFunction("x()") == encodeArgs(u256(0)));
+	BOOST_CHECK(callContractFunction("f()") == encodeArgs());
+	BOOST_CHECK(callContractFunction("x()") == encodeArgs(u256(1)));
+}
+
+BOOST_AUTO_TEST_CASE(stacked_return_with_modifiers)
+{
+	char const* sourceCode = R"(
+		contract C {
+			uint public x;
+			modifier run() {
+				for (uint i = 0; i < 10; i++) {
+					_
+					break;
+				}
+			}
+			function f() run {
+				x++;
+			}
+		}
+	)";
+	compileAndRun(sourceCode, 0, "C");
+	BOOST_CHECK(callContractFunction("x()") == encodeArgs(u256(0)));
+	BOOST_CHECK(callContractFunction("f()") == encodeArgs());
+	BOOST_CHECK(callContractFunction("x()") == encodeArgs(u256(1)));
+}
+
+BOOST_AUTO_TEST_CASE(mutex)
+{
+	char const* sourceCode = R"(
+		contract mutexed {
+			bool locked;
+			modifier protected {
+				if (locked) throw;
+				locked = true;
+				_
+				locked = false;
+			}
+		}
+		contract Fund is mutexed {
+			uint shares;
+			function Fund() { shares = msg.value; }
+			function withdraw(uint amount) protected returns (uint) {
+				// NOTE: It is very bad practice to write this function this way.
+				// Please refer to the documentation of how to do this properly.
+				if (amount > shares) throw;
+				if (!msg.sender.call.value(amount)()) throw;
+				shares -= amount;
+				return shares;
+			}
+			function withdrawUnprotected(uint amount) returns (uint) {
+				// NOTE: It is very bad practice to write this function this way.
+				// Please refer to the documentation of how to do this properly.
+				if (amount > shares) throw;
+				if (!msg.sender.call.value(amount)()) throw;
+				shares -= amount;
+				return shares;
+			}
+		}
+		contract Attacker {
+			Fund public fund;
+			uint callDepth;
+			bool protected;
+			function setProtected(bool _protected) { protected = _protected; }
+			function Attacker(Fund _fund) { fund = _fund; }
+			function attack() returns (uint) {
+				callDepth = 0;
+				return attackInternal();
+			}
+			function attackInternal() internal returns (uint) {
+				if (protected)
+					return fund.withdraw(10);
+				else
+					return fund.withdrawUnprotected(10);
+			}
+			function() {
+				callDepth++;
+				if (callDepth < 4)
+					attackInternal();
+			}
+		}
+	)";
+	compileAndRun(sourceCode, 500, "Fund");
+	auto fund = m_contractAddress;
+	BOOST_CHECK_EQUAL(balanceAt(fund), 500);
+	compileAndRun(sourceCode, 0, "Attacker", encodeArgs(u160(fund)));
+	BOOST_CHECK(callContractFunction("setProtected(bool)", true) == encodeArgs());
+	BOOST_CHECK(callContractFunction("attack()") == encodeArgs());
+	BOOST_CHECK_EQUAL(balanceAt(fund), 500);
+	BOOST_CHECK(callContractFunction("setProtected(bool)", false) == encodeArgs());
+	BOOST_CHECK(callContractFunction("attack()") == encodeArgs(u256(460)));
+	BOOST_CHECK_EQUAL(balanceAt(fund), 460);
+}
+
 BOOST_AUTO_TEST_CASE(failing_ecrecover_invalid_input)
 {
 	// ecrecover should return zero for malformed input