Repository: vyperlang/vyper

Issue:
Augassign can bypass function parameters mutability check
### What's your issue about?
The function parameter should be immutable and should not be modified during the function. The direct assignment to a function parameter is forbidden.  However, there is no such check in the Augassign. For example, the following program can compile:
```
@public
def foo(x:num):
    x += 5     
```

### How can it be fixed?

In the stmt.py, add mutability check to function parameters.

#### Cute Animal Picture

![image](https://user-images.githubusercontent.com/5641590/33234523-ca146bb2-d1ed-11e7-94cb-9005aa316e05.png)




Discussion:

Resolving patch:
diff --git a/viper/parser/stmt.py b/viper/parser/stmt.py
--- a/viper/parser/stmt.py
+++ b/viper/parser/stmt.py
@@ -80,22 +80,6 @@ def assign(self):
         from .parser import (
             make_setter,
         )
-
-        if isinstance(self.stmt.targets[0], ast.Subscript) and self.context.in_for_loop:  # Check if we are doing assignment of an iteration loop.
-            raise_exception = False
-            if isinstance(self.stmt.targets[0].value, ast.Attribute):
-                list_name = "%s.%s" % (self.stmt.targets[0].value.value.id, self.stmt.targets[0].value.attr)
-                if list_name in self.context.in_for_loop:
-                    raise_exception = True
-
-            if isinstance(self.stmt.targets[0].value, ast.Name) and \
-               self.stmt.targets[0].value.id in self.context.in_for_loop:
-                list_name = self.stmt.targets[0].value.id
-                raise_exception = True
-
-            if raise_exception:
-                raise StructureException("Altering list '%s' which is being iterated!" % list_name, self.stmt)
-
         # Assignment (eg. x[4] = y)
         if len(self.stmt.targets) != 1:
             raise StructureException("Assignment statement must have one target", self.stmt)
@@ -105,11 +89,8 @@ def assign(self):
             variable_loc = LLLnode.from_list(pos, typ=sub.typ, location='memory', pos=getpos(self.stmt), annotation=self.stmt.targets[0].id)
             o = make_setter(variable_loc, sub, 'memory', pos=getpos(self.stmt))
         else:
-            target = Expr.parse_variable_location(self.stmt.targets[0], self.context)
-            if target.location == 'storage' and self.context.is_constant:
-                raise ConstancyViolationException("Cannot modify storage inside a constant function!", self.stmt.targets[0])
-            if not target.mutable:
-                raise ConstancyViolationException("Cannot modify function argument", self.stmt.targets[0])
+            # Checks to see if assignment is valid
+            target = self.get_target(self.stmt.targets[0])
             o = make_setter(target, sub, target.location, pos=getpos(self.stmt))
         o.pos = getpos(self.stmt)
         return o
@@ -170,7 +151,6 @@ def parse_for(self):
         from .parser import (
             parse_body,
         )
-
         # Type 0 for, eg. for i in list(): ...
         if self._is_list_iter():
             return self.parse_for_list()
@@ -204,9 +184,11 @@ def parse_for(self):
             start = Expr.parse_value_expr(self.stmt.iter.args[0], self.context)
             rounds = self.stmt.iter.args[1].right.n
         varname = self.stmt.target.id
-        pos = self.context.vars[varname].pos if varname in self.context.forvars else self.context.new_variable(varname, BaseType('num'))
-        o = LLLnode.from_list(['repeat', pos, start, rounds, parse_body(self.stmt.body, self.context)], typ=None, pos=getpos(self.stmt))
+        pos = self.context.new_variable(varname, BaseType('num'))
         self.context.forvars[varname] = True
+        o = LLLnode.from_list(['repeat', pos, start, rounds, parse_body(self.stmt.body, self.context)], typ=None, pos=getpos(self.stmt))
+        del self.context.vars[varname]
+        del self.context.forvars[varname]
         return o
 
     def _is_list_iter(self):
@@ -241,7 +223,7 @@ def parse_for_list(self):
         varname = self.stmt.target.id
         value_pos = self.context.new_variable(varname, BaseType(subtype))
         i_pos = self.context.new_variable('_index_for_' + varname, BaseType(subtype))
-
+        self.context.forvars[varname] = True
         if iter_var_type:  # Is a list that is already allocated to memory.
             self.context.set_in_for_loop(self.stmt.iter.id)  # make sure list cannot be altered whilst iterating.
             iter_var = self.context.vars.get(self.stmt.iter.id)
@@ -254,7 +236,6 @@ def parse_for_list(self):
                 ['repeat', i_pos, 0, iter_var.size, body], typ=None, pos=getpos(self.stmt)
             )
             self.context.remove_in_for_loop(self.stmt.iter.id)
-            return o
         elif isinstance(self.stmt.iter, ast.List):  # List gets defined in the for statement.
             # Allocate list to memory.
             count = iter_list_node.typ.count
@@ -274,7 +255,6 @@ def parse_for_list(self):
                     setter,
                     ['repeat', i_pos, 0, count, body]], typ=None, pos=getpos(self.stmt)
             )
-            return o
         elif isinstance(self.stmt.iter, ast.Attribute):  # List is contained in storage.
             count = iter_list_node.typ.count
             self.context.set_in_for_loop(iter_list_node.annotation)  # make sure list cannot be altered whilst iterating.
@@ -288,18 +268,19 @@ def parse_for_list(self):
                     ['repeat', i_pos, 0, count, body]], typ=None, pos=getpos(self.stmt)
             )
             self.context.remove_in_for_loop(iter_list_node.annotation)
-            return o
+        del self.context.vars[varname]
+        del self.context.vars['_index_for_' + varname]
+        del self.context.forvars[varname]
+        return o
 
     def aug_assign(self):
-        target = Expr.parse_variable_location(self.stmt.target, self.context)
+        target = self.get_target(self.stmt.target)
         sub = Expr.parse_value_expr(self.stmt.value, self.context)
         if not isinstance(self.stmt.op, (ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Mod)):
             raise Exception("Unsupported operator for augassign")
         if not isinstance(target.typ, BaseType):
             raise TypeMismatchException("Can only use aug-assign operators with simple types!", self.stmt.target)
         if target.location == 'storage':
-            if self.context.is_constant:
-                raise ConstancyViolationException("Cannot modify storage inside a constant function!", self.stmt.target)
             o = Expr.parse_value_expr(ast.BinOp(left=LLLnode.from_list(['sload', '_stloc'], typ=target.typ, pos=target.pos),
                                     right=sub, op=self.stmt.op, lineno=self.stmt.lineno, col_offset=self.stmt.col_offset), self.context)
             return LLLnode.from_list(['with', '_stloc', target, ['sstore', '_stloc', base_type_conversion(o, o.typ, target.typ)]], typ=None, pos=getpos(self.stmt))
@@ -420,3 +401,28 @@ def increment_dynamic_offset(dynamic_spot):
                                         typ=None, pos=getpos(self.stmt))
         else:
             raise TypeMismatchException("Can only return base type!", self.stmt)
+
+    def get_target(self, target):
+        if isinstance(target, ast.Subscript) and self.context.in_for_loop:  # Check if we are doing assignment of an iteration loop.
+            raise_exception = False
+            if isinstance(target.value, ast.Attribute):
+                list_name = "%s.%s" % (target.value.value.id, target.value.attr)
+                if list_name in self.context.in_for_loop:
+                    raise_exception = True
+
+            if isinstance(target.value, ast.Name) and \
+               target.value.id in self.context.in_for_loop:
+                list_name = target.value.id
+                raise_exception = True
+
+            if raise_exception:
+                raise StructureException("Altering list '%s' which is being iterated!" % list_name, self.stmt)
+
+        if isinstance(target, ast.Name) and target.id in self.context.forvars:
+            raise StructureException("Altering iterator '%s' which is in use!" % target.id, self.stmt)
+        target = Expr.parse_variable_location(target, self.context)
+        if target.location == 'storage' and self.context.is_constant:
+            raise ConstancyViolationException("Cannot modify storage inside a constant function: %s" % target.annotation)
+        if not target.mutable:
+            raise ConstancyViolationException("Cannot modify function argument: %s" % target.annotation)
+        return target


Test patch:
diff --git a/tests/parser/features/iteration/test_for_in_list.py b/tests/parser/features/iteration/test_for_in_list.py
--- a/tests/parser/features/iteration/test_for_in_list.py
+++ b/tests/parser/features/iteration/test_for_in_list.py
@@ -1,7 +1,7 @@
 import pytest
 from tests.setup_transaction_tests import chain as s, tester as t, ethereum_utils as u, check_gas, \
     get_contract_with_gas_estimation, get_contract
-from viper.exceptions import StructureException
+from viper.exceptions import StructureException, VariableDeclarationException
 
 
 def test_basic_for_in_list():
@@ -80,6 +80,68 @@ def data() -> address:
     assert c.data() == "0x82a978b3f5962a5b0957d9ee9eef472ee55b42f1"
 
 
+def test_multiple_for_loops_1():
+    code = """
+@public
+def foo(x: num): 
+    p = 0 
+    for i in range(3): 
+        p += i
+    for i in range(4):
+        p += i
+"""
+    get_contract_with_gas_estimation(code)
+
+def test_multiple_for_loops_2():
+    code = """
+@public
+def foo(x: num): 
+    p = 0 
+    for i in range(3): 
+        p += i
+    for i in [1, 2, 3, 4]: 
+        p += i
+"""
+    get_contract_with_gas_estimation(code)
+
+
+def test_multiple_for_loops_3():
+    code = """
+@public
+def foo(x: num): 
+    p = 0 
+    for i in [1, 2, 3, 4]: 
+        p += i
+    for i in [1, 2, 3, 4]: 
+        p += i
+"""
+    get_contract_with_gas_estimation(code)
+
+
+def test_multiple_loops_4():
+    code = """
+@public
+def foo():
+    for i in range(10):
+        pass
+    for i in range(20):
+        pass
+"""
+    get_contract_with_gas_estimation(code)
+
+
+def test_using_index_variable_after_loop():
+    code = """
+@public
+def foo():
+    for i in range(10):
+        pass
+    i = 100  # create new variable i
+    i = 200  # look up the variable i and check whether it is in forvars
+"""
+    get_contract_with_gas_estimation(code)
+
+
 def test_basic_for_list_storage_address():
     code = """
 addresses: address[3]
@@ -142,7 +204,7 @@ def i_return(break_count: num) -> decimal:
     assert c.ret(0) == c.i_return(0) == 0.0001
 
 
-def test_altering_list_within_for_loop(assert_compile_failed):
+def test_altering_list_within_for_loop_1(assert_compile_failed):
     code = """
 @public
 def data() -> num:
@@ -159,6 +221,18 @@ def data() -> num:
     assert_compile_failed(lambda: get_contract_with_gas_estimation(code), StructureException)
 
 
+def test_altering_list_within_for_loop_2(assert_compile_failed):
+    code = """
+@public
+def foo():
+    s = [1, 2, 3, 4, 5, 6]
+    count = 0
+    for i in s:
+        s[count] += 1
+"""
+    assert_compile_failed(lambda: get_contract_with_gas_estimation(code), StructureException)
+
+
 def test_altering_list_within_for_loop_storage(assert_compile_failed):
     code = """
 s: num[6]
@@ -179,3 +253,45 @@ def data() -> num:
     """
 
     assert_compile_failed(lambda: get_contract_with_gas_estimation(code), StructureException)
+
+
+def test_invalid_nested_for_loop_1(assert_compile_failed):
+    code = """
+@public
+def foo(x: num):
+    for i in range(4):
+        for i in range(5):
+            pass
+"""
+    assert_compile_failed(lambda: get_contract_with_gas_estimation(code),VariableDeclarationException)
+
+
+def test_invalid_nested_for_loop_2(assert_compile_failed):
+    code = """
+@public
+def foo(x: num):
+    for i in [1,2]:
+        for i in [1,2]:
+            pass
+"""
+    assert_compile_failed(lambda: get_contract_with_gas_estimation(code),VariableDeclarationException)
+
+
+def test_invalid_iterator_assignment_1(assert_compile_failed):
+    code = """
+@public
+def foo(x: num):
+    for i in [1,2]:
+        i = 2
+"""
+    assert_compile_failed(lambda: get_contract_with_gas_estimation(code), StructureException)
+
+
+def test_invalid_iterator_assignment_2(assert_compile_failed):
+    code = """
+@public
+def foo(x: num):
+    for i in [1,2]:
+        i += 2
+"""
+    assert_compile_failed(lambda: get_contract_with_gas_estimation(code), StructureException)
diff --git a/tests/parser/features/test_assignment.py b/tests/parser/features/test_assignment.py
--- a/tests/parser/features/test_assignment.py
+++ b/tests/parser/features/test_assignment.py
@@ -1,9 +1,10 @@
 import pytest
 from tests.setup_transaction_tests import chain as s, tester as t, ethereum_utils as u, check_gas, \
     get_contract_with_gas_estimation, get_contract
+from viper.exceptions import ConstancyViolationException
 
 
-def test_augassign_test():
+def test_augassign():
     augassign_test = """
 @public
 def augadd(x: num, y: num) -> num:
@@ -44,3 +45,21 @@ def augmod(x: num, y: num) -> num:
     assert c.augdiv(5, 12) == 0
     assert c.augmod(5, 12) == 5
     print('Passed aug-assignment test')
+
+
+def test_invalid_assign(assert_compile_failed):
+    code = """
+@public
+def foo(x:num):
+    x = 5   
+"""
+    assert_compile_failed(lambda: get_contract_with_gas_estimation(code), ConstancyViolationException)
+
+
+def test_invalid_augassign(assert_compile_failed):
+    code = """
+@public
+def foo(x:num):
+    x += 5     
+"""
+    assert_compile_failed(lambda: get_contract_with_gas_estimation(code), ConstancyViolationException)