135 lines
5.4 KiB
Python
135 lines
5.4 KiB
Python
|
from itertools import chain
|
||
|
|
from unittest import TestCase
|
||
|
|
from typing import List, Optional
|
||
|
|
|
||
|
|
from da.common import NodeId, build_attestation_message, BLSPublicKey, NomosDaG2ProofOfPossession as bls_pop
|
||
|
|
from da.api.common import DAApi, VID, Metadata
|
||
|
|
from da.verifier import DAVerifier, DABlob
|
||
|
|
from da.api.test_flow import MockStore
|
||
|
|
from da.dispersal import Dispersal, DispersalSettings
|
||
|
|
from da.test_encoder import TestEncoder
|
||
|
|
from da.encoder import DAEncoderParams, DAEncoder
|
||
|
|
|
||
|
|
|
||
|
|
class DAVerifierWApi:
|
||
|
|
def __init__(self, sk: int, public_keys: List[BLSPublicKey]):
|
||
|
|
self.store = MockStore()
|
||
|
|
self.api = DAApi(self.store)
|
||
|
|
self.verifier = DAVerifier(sk, public_keys)
|
||
|
|
|
||
|
|
def receive_blob(self, blob: DABlob):
|
||
|
|
if attestation := self.verifier.verify(blob):
|
||
|
|
# Warning: If aggregated col commitment and row commitment are the same,
|
||
|
|
# the build_attestation_message method will produce the same output.
|
||
|
|
cert_id = build_attestation_message(blob.aggregated_column_commitment, blob.rows_commitments)
|
||
|
|
self.store.populate(blob, cert_id)
|
||
|
|
return attestation
|
||
|
|
|
||
|
|
def receive_cert(self, vid: VID):
|
||
|
|
# Usually the certificate would be verifier here,
|
||
|
|
# but we are assuming that this it is already coming from the verified block,
|
||
|
|
# in which case all certificates had been already verified by the DA Node.
|
||
|
|
self.api.write(vid.cert_id, vid.metadata)
|
||
|
|
|
||
|
|
def read(self, app_id, indexes) -> List[Optional[DABlob]]:
|
||
|
|
return self.api.read(app_id, indexes)
|
||
|
|
|
||
|
|
|
||
|
|
class TestFullFlow(TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.n_nodes = 16
|
||
|
|
self.nodes_ids = [NodeId(x.to_bytes(length=32, byteorder='big')) for x in range(self.n_nodes)]
|
||
|
|
self.secret_keys = list(range(1, self.n_nodes+1))
|
||
|
|
self.public_keys = [bls_pop.SkToPk(sk) for sk in self.secret_keys]
|
||
|
|
# sort by pk as we do in dispersal
|
||
|
|
self.secret_keys, self.public_keys = zip(
|
||
|
|
*sorted(zip(self.secret_keys, self.public_keys), key=lambda x: x[1])
|
||
|
|
)
|
||
|
|
dispersal_settings = DispersalSettings(
|
||
|
|
self.nodes_ids,
|
||
|
|
self.public_keys,
|
||
|
|
self.n_nodes
|
||
|
|
)
|
||
|
|
self.dispersal = Dispersal(dispersal_settings)
|
||
|
|
self.encoder_test = TestEncoder()
|
||
|
|
self.encoder_test.setUp()
|
||
|
|
|
||
|
|
self.api_nodes = [DAVerifierWApi(k, self.public_keys) for k in self.secret_keys]
|
||
|
|
|
||
|
|
def test_full_flow(self):
|
||
|
|
app_id = int.to_bytes(1)
|
||
|
|
index = 1
|
||
|
|
|
||
|
|
# encoder
|
||
|
|
data = self.encoder_test.data
|
||
|
|
encoding_params = DAEncoderParams(column_count=self.n_nodes // 2, bytes_per_chunk=31)
|
||
|
|
encoded_data = DAEncoder(encoding_params).encode(data)
|
||
|
|
|
||
|
|
# mock send and await method with local verifiers
|
||
|
|
def __send_and_await_response(node: int, blob: DABlob):
|
||
|
|
node = self.api_nodes[int.from_bytes(node)]
|
||
|
|
return node.receive_blob(blob)
|
||
|
|
|
||
|
|
# inject mock send and await method
|
||
|
|
self.dispersal._send_and_await_response = __send_and_await_response
|
||
|
|
certificate = self.dispersal.disperse(encoded_data)
|
||
|
|
|
||
|
|
vid = VID(
|
||
|
|
certificate.id(),
|
||
|
|
Metadata(app_id, index)
|
||
|
|
)
|
||
|
|
|
||
|
|
# verifier
|
||
|
|
for node in self.api_nodes:
|
||
|
|
node.receive_cert(vid)
|
||
|
|
|
||
|
|
# read from api and confirm its working
|
||
|
|
# notice that we need to sort the api_nodes by their public key to have the blobs sorted in the same fashion
|
||
|
|
# we do actually do dispersal.
|
||
|
|
blobs = list(chain.from_iterable(
|
||
|
|
node.read(app_id, [index])
|
||
|
|
for node in sorted(self.api_nodes, key=lambda n: bls_pop.SkToPk(n.verifier.sk))
|
||
|
|
))
|
||
|
|
original_blobs = list(self.dispersal._prepare_data(encoded_data))
|
||
|
|
self.assertEqual(blobs, original_blobs)
|
||
|
|
|
||
|
|
def test_same_blob_multiple_indexes(self):
|
||
|
|
app_id = int.to_bytes(1)
|
||
|
|
indexes = [1, 2, 3] # Different indexes to test with the same blob
|
||
|
|
|
||
|
|
# encoder
|
||
|
|
data = self.encoder_test.data
|
||
|
|
encoding_params = DAEncoderParams(column_count=self.n_nodes // 2, bytes_per_chunk=31)
|
||
|
|
encoded_data = DAEncoder(encoding_params).encode(data)
|
||
|
|
|
||
|
|
# mock send and await method with local verifiers
|
||
|
|
def __send_and_await_response(node: int, blob: DABlob):
|
||
|
|
node = self.api_nodes[int.from_bytes(node)]
|
||
|
|
return node.receive_blob(blob)
|
||
|
|
|
||
|
|
# inject mock send and await method
|
||
|
|
self.dispersal._send_and_await_response = __send_and_await_response
|
||
|
|
certificate = self.dispersal.disperse(encoded_data)
|
||
|
|
|
||
|
|
# Loop through each index and simulate dispersal with the same cert_id but different metadata
|
||
|
|
for index in indexes:
|
||
|
|
vid = VID(
|
||
|
|
certificate.id(),
|
||
|
|
Metadata(app_id, index)
|
||
|
|
)
|
||
|
|
|
||
|
|
# verifier
|
||
|
|
for node in self.api_nodes:
|
||
|
|
node.receive_cert(vid)
|
||
|
|
|
||
|
|
# Verify retrieval for each index
|
||
|
|
for index in indexes:
|
||
|
|
# Notice that we need to sort the api_nodes by their public key to have the blobs sorted in the same fashion
|
||
|
|
# as we do actually do dispersal.
|
||
|
|
blobs = list(chain.from_iterable(
|
||
|
|
node.read(app_id, [index])
|
||
|
|
for node in sorted(self.api_nodes, key=lambda n: bls_pop.SkToPk(n.verifier.sk))
|
||
|
|
))
|
||
|
|
original_blobs = list(self.dispersal._prepare_data(encoded_data))
|
||
|
|
self.assertEqual(blobs, original_blobs, f"Failed at index {index}")
|