107 lines
3.1 KiB
Python
107 lines
3.1 KiB
Python
|
from hashlib import sha256
|
||
|
|
from random import randbytes
|
||
|
|
from typing import Self
|
||
|
|
|
||
|
|
import dispersal.proto as proto
|
||
|
|
import multiaddr
|
||
|
|
import trio
|
||
|
|
from constants import HASH_LENGTH, PROTOCOL_ID
|
||
|
|
from libp2p import host, new_host
|
||
|
|
from libp2p.network.stream.net_stream_interface import INetStream
|
||
|
|
from libp2p.peer.peerinfo import info_from_p2p_addr
|
||
|
|
|
||
|
|
|
||
|
|
class Executor:
|
||
|
|
"""
|
||
|
|
A class for simulating a simple executor.
|
||
|
|
|
||
|
|
Runs on hardcoded port.
|
||
|
|
Creates random data and disperses it.
|
||
|
|
One packet represents a subnet, and each packet is sent
|
||
|
|
to one DANode.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
listen_addr: multiaddr.Multiaddr
|
||
|
|
host: host
|
||
|
|
port: int
|
||
|
|
num_subnets: int
|
||
|
|
node_list: {}
|
||
|
|
# size of a packet
|
||
|
|
data_size: int
|
||
|
|
# holds random data for dispersal
|
||
|
|
data: []
|
||
|
|
# stores hashes of the data for later verification
|
||
|
|
data_hashes: []
|
||
|
|
blob_id: int
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def new(cls, port, node_list, num_subnets, data_size) -> Self:
|
||
|
|
self = cls()
|
||
|
|
self.listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
|
||
|
|
self.host = new_host()
|
||
|
|
self.port = port
|
||
|
|
self.num_subnets = num_subnets
|
||
|
|
self.data_size = data_size
|
||
|
|
# one packet per subnet
|
||
|
|
self.data = [[] * data_size] * num_subnets
|
||
|
|
# one hash per packet. **assumes 256 hash length**
|
||
|
|
self.data_hashes = [[] * HASH_LENGTH] * num_subnets
|
||
|
|
self.node_list = node_list
|
||
|
|
# create random simulated data right from the beginning
|
||
|
|
self.__create_data()
|
||
|
|
return self
|
||
|
|
|
||
|
|
def get_id(self):
|
||
|
|
return self.host.get_id()
|
||
|
|
|
||
|
|
def net_iface(self):
|
||
|
|
return self.host
|
||
|
|
|
||
|
|
def get_port(self):
|
||
|
|
return self.port
|
||
|
|
|
||
|
|
def get_hash(self, index: int):
|
||
|
|
return self.data_hashes[index]
|
||
|
|
|
||
|
|
def __create_data(self):
|
||
|
|
"""
|
||
|
|
Create random data for dispersal
|
||
|
|
One packet of self.data_size length per subnet
|
||
|
|
"""
|
||
|
|
id = sha256()
|
||
|
|
for i in range(self.num_subnets):
|
||
|
|
self.data[i] = randbytes(self.data_size)
|
||
|
|
self.data_hashes[i] = sha256(self.data[i]).hexdigest()
|
||
|
|
id.update(self.data[i])
|
||
|
|
self.blob_id = id.digest()
|
||
|
|
|
||
|
|
async def disperse(self, nursery):
|
||
|
|
"""
|
||
|
|
Disperse the data to the DA network.
|
||
|
|
Sends one packet of data per network node
|
||
|
|
"""
|
||
|
|
|
||
|
|
async with self.host.run(listen_addrs=[self.listen_addr]):
|
||
|
|
for subnet, nodes in self.node_list.items():
|
||
|
|
# get first node of each subnet
|
||
|
|
n = nodes[0]
|
||
|
|
# connect to it...
|
||
|
|
await self.host.connect(n)
|
||
|
|
|
||
|
|
# ...and send (async)
|
||
|
|
stream = await self.host.new_stream(n.peer_id, [PROTOCOL_ID])
|
||
|
|
nursery.start_soon(self.write_data, stream, subnet)
|
||
|
|
|
||
|
|
async def write_data(self, stream: INetStream, index: int) -> None:
|
||
|
|
"""
|
||
|
|
Send data to peer (async)
|
||
|
|
The index is the subnet number
|
||
|
|
"""
|
||
|
|
|
||
|
|
blob_id = self.blob_id
|
||
|
|
blob_data = self.data[index]
|
||
|
|
|
||
|
|
message = proto.new_dispersal_req_msg(blob_id, blob_data)
|
||
|
|
await stream.write(message)
|