1  import eth2spec.phase0.spec as spec
  2  from eth2spec.phase0.spec import process_deposit
  3  from eth2spec.test.context import spec_state_test, expect_assertion_error, always_bls
  4  from eth2spec.test.helpers.deposits import prepare_state_and_deposit, sign_deposit_data
  5  from eth2spec.test.helpers.state import get_balance
  6  from eth2spec.test.helpers.keys import privkeys
  7  
  8  
  9  def run_deposit_processing(state, deposit, validator_index, valid=True, effective=True):
 10      """
 11      Run ``process_deposit``, yielding:
 12        - pre-state ('pre')
 13        - deposit ('deposit')
 14        - post-state ('post').
 15      If ``valid == False``, run expecting ``AssertionError``
 16      """
 17      pre_validator_count = len(state.validator_registry)
 18      pre_balance = 0
 19      if validator_index < pre_validator_count:
 20          pre_balance = get_balance(state, validator_index)
 21      else:
 22          # if it is a new validator, it should be right at the end of the current registry.
 23          assert validator_index == pre_validator_count
 24  
 25      yield 'pre', state
 26      yield 'deposit', deposit
 27  
 28      if not valid:
 29          expect_assertion_error(lambda: process_deposit(state, deposit))
 30          yield 'post', None
 31          return
 32  
 33      process_deposit(state, deposit)
 34  
 35      yield 'post', state
 36  
 37      if not effective:
 38          assert len(state.validator_registry) == pre_validator_count
 39          assert len(state.balances) == pre_validator_count
 40          if validator_index < pre_validator_count:
 41              assert get_balance(state, validator_index) == pre_balance
 42      else:
 43          if validator_index < pre_validator_count:
 44              # top-up
 45              assert len(state.validator_registry) == pre_validator_count
 46              assert len(state.balances) == pre_validator_count
 47          else:
 48              # new validator
 49              assert len(state.validator_registry) == pre_validator_count + 1
 50              assert len(state.balances) == pre_validator_count + 1
 51          assert get_balance(state, validator_index) == pre_balance + deposit.data.amount
 52  
 53      assert state.deposit_index == state.latest_eth1_data.deposit_count
 54  
 55  
 56  @spec_state_test
 57  def test_new_deposit(state):
 58      # fresh deposit = next validator index = validator appended to registry
 59      validator_index = len(state.validator_registry)
 60      amount = spec.MAX_EFFECTIVE_BALANCE
 61      deposit = prepare_state_and_deposit(state, validator_index, amount, signed=True)
 62  
 63      yield from run_deposit_processing(state, deposit, validator_index)
 64  
 65  
 66  @always_bls
 67  @spec_state_test
 68  def test_invalid_sig_new_deposit(state):
 69      # fresh deposit = next validator index = validator appended to registry
 70      validator_index = len(state.validator_registry)
 71      amount = spec.MAX_EFFECTIVE_BALANCE
 72      deposit = prepare_state_and_deposit(state, validator_index, amount)
 73      yield from run_deposit_processing(state, deposit, validator_index, valid=True, effective=False)
 74  
 75  
 76  @spec_state_test
 77  def test_success_top_up(state):
 78      validator_index = 0
 79      amount = spec.MAX_EFFECTIVE_BALANCE // 4
 80      deposit = prepare_state_and_deposit(state, validator_index, amount, signed=True)
 81  
 82      yield from run_deposit_processing(state, deposit, validator_index)
 83  
 84  
 85  @always_bls
 86  @spec_state_test
 87  def test_invalid_sig_top_up(state):
 88      validator_index = 0
 89      amount = spec.MAX_EFFECTIVE_BALANCE // 4
 90      deposit = prepare_state_and_deposit(state, validator_index, amount)
 91  
 92      # invalid signatures, in top-ups, are allowed!
 93      yield from run_deposit_processing(state, deposit, validator_index, valid=True, effective=True)
 94  
 95  
 96  @spec_state_test
 97  def test_invalid_withdrawal_credentials_top_up(state):
 98      validator_index = 0
 99      amount = spec.MAX_EFFECTIVE_BALANCE // 4
100      withdrawal_credentials = spec.BLS_WITHDRAWAL_PREFIX_BYTE + spec.hash(b"junk")[1:]
101      deposit = prepare_state_and_deposit(
102          state,
103          validator_index,
104          amount,
105          withdrawal_credentials=withdrawal_credentials,
106      )
107  
108      # inconsistent withdrawal credentials, in top-ups, are allowed!
109      yield from run_deposit_processing(state, deposit, validator_index, valid=True, effective=True)
110  
111  
112  @spec_state_test
113  def test_wrong_index(state):
114      validator_index = len(state.validator_registry)
115      amount = spec.MAX_EFFECTIVE_BALANCE
116      deposit = prepare_state_and_deposit(state, validator_index, amount)
117  
118      # mess up deposit_index
119      deposit.index = state.deposit_index + 1
120  
121      sign_deposit_data(state, deposit.data, privkeys[validator_index])
122  
123      yield from run_deposit_processing(state, deposit, validator_index, valid=False)
124  
125  
126  # TODO: test invalid signature
127  
128  
129  @spec_state_test
130  def test_bad_merkle_proof(state):
131      validator_index = len(state.validator_registry)
132      amount = spec.MAX_EFFECTIVE_BALANCE
133      deposit = prepare_state_and_deposit(state, validator_index, amount)
134  
135      # mess up merkle branch
136      deposit.proof[-1] = spec.ZERO_HASH
137  
138      sign_deposit_data(state, deposit.data, privkeys[validator_index])
139  
140      yield from run_deposit_processing(state, deposit, validator_index, valid=False)