Tidy up mixnet verification code

eos/psr/election.py

@@ -65,13 +65,91 @@ class MixingTrustee(Trustee):
 		self.mixnets = [] # TODO: Remove this stuff
 	
 	def compute_challenge(self, question_num):
+		if self._instance[1] % 2 == 1:
+			return self.recurse_parents(Election).mixing_trustees[self._instance[1] - 1].compute_challenge(question_num)
+		
 		sha = SHA256()
 		trustees = self.recurse_parents(Election).mixing_trustees
 		for i in range(len(trustees)):
 			sha.update_text(EosObject.to_json(MixingTrustee._fields['mixed_questions'].element_field.serialise(trustees[i].mixed_questions[question_num])))
 		for i in range(self._instance[1]):
 			sha.update_text(EosObject.to_json(MixingTrustee._fields['response'].element_field.serialise(trustees[i].response[question_num])))
-		return sha
+		return sha.hash_as_bigint()
+	
+	def get_input_answers(self, question_num):
+		if self._instance[1] > 0:
+			# Use the previous mixnet's output
+			return self.recurse_parents(Election).mixing_trustees[self._instance[1] - 1].mixed_questions[question_num]
+		else:
+			# Use the raw ballots from voters
+			orig_answers = []
+			for voter in self.recurse_parents(Election).voters:
+				for ballot in voter.ballots:
+					orig_answers.append(ballot.encrypted_answers[question_num])
+			return orig_answers
+	
+	def verify(self, question_num):
+		# Verify challenge
+		challenge = self.compute_challenge(question_num)
+		if challenge != self.challenge[question_num]:
+			raise Exception('Invalid challenge')
+		
+		orig_answers = self.get_input_answers(question_num)
+		
+		# Prepare challenge bits
+		challenge_bs = InfiniteHashBitStream(challenge)
+		
+		# Check each challenge response
+		responses_iter = iter(self.response[question_num])
+		for k in range(len(self.mixed_questions[question_num])):
+			challenge_bit = challenge_bs.read(1)
+			should_reveal = ((self._instance[1] % 2) == (challenge_bit % 2))
+			if should_reveal:
+				response = next(responses_iter)
+				
+				# Check the commitment matches
+				if self.commitments[question_num][k] != SHA256().update_obj(response).hash_as_bigint():
+					raise Exception('Invalid commitment')
+				
+				# Check the correct challenge/response pair
+				if response.challenge_index != k:
+					raise Exception('Invalid response')
+				
+				if self._instance[1] % 2 == 0:
+					idx_left = response.challenge_index
+					idx_right = response.response_index
+				else:
+					idx_right = response.challenge_index
+					idx_left = response.response_index
+				
+				# Check the shuffle
+				claimed_blocks = self.mixed_questions[question_num][idx_right].blocks
+				for k in range(len(orig_answers[idx_left].blocks)):
+					reencrypted_block, _ = orig_answers[idx_left].blocks[k].reencrypt(response.reenc[k])
+					if claimed_blocks[k].gamma != reencrypted_block.gamma:
+						raise Exception('Reencryption not consistent with challenge response')
+					if claimed_blocks[k].delta != reencrypted_block.delta:
+						raise Exception('Reencryption not consistent with challenge response')
+		
+		# Check the responses are consistent with a permutation
+		challenge_indexes = []
+		response_indexes = []
+		for response in self.response[question_num]:
+			if response.challenge_index in challenge_indexes:
+				raise Exception('Response not consistent with a permutation')
+			if response.response_index in response_indexes:
+				raise Exception('Response not consistent with a permutation')
+			challenge_indexes.append(response.challenge_index)
+			response_indexes.append(response.response_index)
+		
+		# Check the outputs are all different
+		blocks = []
+		for output in self.mixed_questions[question_num]:
+			for block in output.blocks:
+				block = (str(block.gamma), str(block.delta))
+				if block in blocks:
+					raise Exception('Duplicate ciphertexts in output')
+				blocks.append(block)
 
 class PSRElection(Election):
 	_db_name = Election._name

eos/psr/tests.py

@@ -261,48 +261,24 @@ class ElectionTestCase(EosTestCase):
 		election.workflow.get_task('eos.psr.workflow.TaskProveMixes').enter()
 		election.save()
 		
-		def verify_shuffle(i, j, idx_left, idx_right, reencs):
-			if j > 0:
-				orig_answers = election.mixing_trustees[j - 1].mixed_questions[i]
-			else:
-				orig_answers = []
-				for voter in election.voters:
-					for ballot in voter.ballots:
-						orig_answers.append(ballot.encrypted_answers[i])
-			
-			claimed_blocks = election.mixing_trustees[j].mixed_questions[i][idx_right].blocks
-			for k in range(len(orig_answers[idx_left].blocks)):
-				reencrypted_block, _ = orig_answers[idx_left].blocks[k].reencrypt(reencs[k])
-				self.assertEqual(claimed_blocks[k].gamma, reencrypted_block.gamma)
-				self.assertEqual(claimed_blocks[k].delta, reencrypted_block.delta)
-		
 		# Record challenge responses
 		for i in range(len(election.questions)):
 			for j in range(len(election.mixing_trustees)):
 				trustee = election.mixing_trustees[j]
-				if j % 2 == 0:
-					trustee.challenge.append(trustee.compute_challenge(i).hash_as_bigint())
-				else:
-					trustee.challenge.append(election.mixing_trustees[j - 1].challenge[i])
+				trustee.challenge.append(trustee.compute_challenge(i))
 				challenge_bs = InfiniteHashBitStream(trustee.challenge[i])
 				
 				trustee.response.append(EosList())
 				
-				nbits = BigInt(len(trustee.mixed_questions[i])).nbits()
 				for k in range(len(trustee.mixed_questions[i])):
 					challenge_bit = challenge_bs.read(1)
 					should_reveal = ((j % 2) == (challenge_bit % 2))
 					if should_reveal:
 						response = trustee.mixnets[i].challenge(k)
 						trustee.response[i].append(response)
-						
-						# Verify proof
-						self.assertEqual(trustee.commitments[i][k], SHA256().update_obj(response).hash_as_bigint())
-						
-						if j % 2 == 0:
-							verify_shuffle(i, j, response.challenge_index, response.response_index, response.reenc)
-						else:
-							verify_shuffle(i, j, response.response_index, response.challenge_index, response.reenc)
+				
+				# Verify challenge response
+				trustee.verify(i)
 		
 		election.workflow.get_task('eos.psr.workflow.TaskProveMixes').exit()
 		election.save()