mirror of
https://github.com/0glabs/0g-storage-node.git
synced 2025-11-04 00:27:39 +00:00
244 lines
6.6 KiB
Python
244 lines
6.6 KiB
Python
import base64
|
|
|
|
from eth_utils import decode_hex
|
|
from math import log2
|
|
from utility.merkle_tree import add_0x_prefix, Leaf, MerkleTree
|
|
from utility.spec import ENTRY_SIZE, PORA_CHUNK_SIZE
|
|
|
|
|
|
def log2_pow2(n):
|
|
return int(log2(((n ^ (n - 1)) >> 1) + 1))
|
|
|
|
|
|
def next_pow2(input):
|
|
x = input
|
|
x -= 1
|
|
x |= x >> 16
|
|
x |= x >> 8
|
|
x |= x >> 4
|
|
x |= x >> 2
|
|
x |= x >> 1
|
|
x += 1
|
|
|
|
return x
|
|
|
|
|
|
def bytes_to_entries(size_bytes):
|
|
if size_bytes % ENTRY_SIZE == 0:
|
|
return size_bytes // ENTRY_SIZE
|
|
else:
|
|
return size_bytes // ENTRY_SIZE + 1
|
|
|
|
|
|
def create_submission(data):
|
|
submission = []
|
|
submission.append(len(data))
|
|
submission.append(b"")
|
|
submission.append([])
|
|
|
|
offset = 0
|
|
nodes = []
|
|
for chunks in split_nodes(len(data)):
|
|
node_hash = create_node(data, offset, chunks)
|
|
nodes.append(node_hash)
|
|
|
|
height = int(log2(chunks))
|
|
submission[2].append([decode_hex(node_hash.decode("utf-8")), height])
|
|
offset += chunks * ENTRY_SIZE
|
|
|
|
root_hash = nodes[-1]
|
|
for i in range(len(nodes) - 2, -1, -1):
|
|
tree = MerkleTree()
|
|
tree.add_leaf(Leaf(nodes[i]))
|
|
tree.add_leaf(Leaf(root_hash))
|
|
root_hash = tree.get_root_hash()
|
|
|
|
return submission, add_0x_prefix(root_hash.decode("utf-8"))
|
|
|
|
|
|
def split_nodes(data_len):
|
|
nodes = []
|
|
|
|
chunks = bytes_to_entries(data_len)
|
|
padded_chunks, chunks_next_pow2 = compute_padded_size(chunks)
|
|
next_chunk_size = chunks_next_pow2
|
|
|
|
while padded_chunks > 0:
|
|
if padded_chunks >= next_chunk_size:
|
|
padded_chunks -= next_chunk_size
|
|
nodes.append(next_chunk_size)
|
|
|
|
next_chunk_size >>= 1
|
|
|
|
return nodes
|
|
|
|
|
|
def compute_padded_size(chunk_len):
|
|
chunks_next_pow2 = next_pow2(chunk_len)
|
|
|
|
if chunks_next_pow2 == chunk_len:
|
|
return chunks_next_pow2, chunks_next_pow2
|
|
|
|
min_chunk = 1 if chunks_next_pow2 < 16 else chunks_next_pow2 // 16
|
|
padded_chunks = ((chunk_len - 1) // min_chunk + 1) * min_chunk
|
|
|
|
return padded_chunks, chunks_next_pow2
|
|
|
|
|
|
def create_node(data, offset, chunks):
|
|
batch = chunks
|
|
if chunks > PORA_CHUNK_SIZE:
|
|
batch = PORA_CHUNK_SIZE
|
|
|
|
return create_segment_node(data, offset, ENTRY_SIZE * batch, ENTRY_SIZE * chunks)
|
|
|
|
|
|
def create_segment_node(data, offset, batch, size):
|
|
tree = MerkleTree()
|
|
n = len(data)
|
|
for i in range(offset, offset + size, batch):
|
|
start = i
|
|
end = min(offset + size, i + batch)
|
|
|
|
if start >= n:
|
|
tree.add_leaf(Leaf(segment_root(b"\x00" * (end - start))))
|
|
elif end > n:
|
|
tree.add_leaf(Leaf(segment_root(data[start:] + b"\x00" * (end - n))))
|
|
else:
|
|
tree.add_leaf(Leaf(segment_root(data[start:end])))
|
|
|
|
return tree.get_root_hash()
|
|
|
|
|
|
segment_root_cached_chunks = None
|
|
segment_root_cached_output = None
|
|
|
|
|
|
def segment_root(chunks):
|
|
global segment_root_cached_chunks, segment_root_cached_output
|
|
|
|
if segment_root_cached_chunks == chunks:
|
|
return segment_root_cached_output
|
|
|
|
data_len = len(chunks)
|
|
if data_len == 0:
|
|
return b"\x00" * 32
|
|
|
|
tree = MerkleTree()
|
|
for i in range(0, data_len, ENTRY_SIZE):
|
|
tree.encrypt(chunks[i : i + ENTRY_SIZE])
|
|
|
|
digest = tree.get_root_hash()
|
|
|
|
segment_root_cached_chunks = chunks
|
|
segment_root_cached_output = digest
|
|
return digest
|
|
|
|
|
|
def generate_merkle_tree(data):
|
|
chunks = bytes_to_entries(len(data))
|
|
padded_chunks, _ = compute_padded_size(chunks)
|
|
|
|
tree = MerkleTree()
|
|
for i in range(padded_chunks):
|
|
if i * ENTRY_SIZE > len(data):
|
|
tree.encrypt(b"\x00" * ENTRY_SIZE)
|
|
elif (i + 1) * ENTRY_SIZE > len(data):
|
|
tree.encrypt(
|
|
data[i * ENTRY_SIZE :] + b"\x00" * ((i + 1) * ENTRY_SIZE - len(data))
|
|
)
|
|
else:
|
|
tree.encrypt(data[i * ENTRY_SIZE : (i + 1) * ENTRY_SIZE])
|
|
|
|
return tree
|
|
|
|
|
|
def generate_merkle_tree_by_batch(data):
|
|
chunks = bytes_to_entries(len(data))
|
|
padded_chunks, _ = compute_padded_size(chunks)
|
|
|
|
tree = MerkleTree()
|
|
for i in range(0, padded_chunks, PORA_CHUNK_SIZE):
|
|
if i * ENTRY_SIZE >= len(data):
|
|
tree.add_leaf(
|
|
Leaf(
|
|
segment_root(
|
|
b"\x00" * ENTRY_SIZE * min(PORA_CHUNK_SIZE, padded_chunks - i)
|
|
)
|
|
)
|
|
)
|
|
elif (i + PORA_CHUNK_SIZE) * ENTRY_SIZE > len(data):
|
|
tree.add_leaf(
|
|
Leaf(
|
|
segment_root(
|
|
data[i * ENTRY_SIZE :]
|
|
+ b"\x00"
|
|
* (
|
|
min(padded_chunks, i + PORA_CHUNK_SIZE) * ENTRY_SIZE
|
|
- len(data)
|
|
)
|
|
)
|
|
)
|
|
)
|
|
else:
|
|
tree.add_leaf(
|
|
Leaf(
|
|
segment_root(
|
|
data[i * ENTRY_SIZE : (i + PORA_CHUNK_SIZE) * ENTRY_SIZE]
|
|
)
|
|
)
|
|
)
|
|
|
|
return tree, add_0x_prefix(tree.decode_value(tree.get_root_hash()))
|
|
|
|
|
|
def submit_data(client, data):
|
|
# NOTE: we assume the data is unique in this function, otherwise zgs_getFileInfo will only get the information of the first data with same root
|
|
shard_config = client.rpc.zgs_getShardConfig()
|
|
shard_id = int(shard_config["shardId"])
|
|
num_shard = int(shard_config["numShard"])
|
|
|
|
segments = data_to_segments(data)
|
|
file_info = client.zgs_get_file_info(segments[0]["root"])
|
|
start_seg_index = file_info["tx"]["startEntryIndex"] // 1024
|
|
for index, segment in enumerate(segments):
|
|
if (start_seg_index + index) % num_shard == shard_id:
|
|
client.zgs_upload_segment(segment)
|
|
return segments
|
|
|
|
|
|
def data_to_segments(data):
|
|
tree, root_hash = generate_merkle_tree_by_batch(data)
|
|
chunks = bytes_to_entries(len(data))
|
|
|
|
segments = []
|
|
idx = 0
|
|
while idx * PORA_CHUNK_SIZE < chunks:
|
|
proof = tree.proof_at(idx)
|
|
|
|
tmp = (
|
|
data[
|
|
idx
|
|
* ENTRY_SIZE
|
|
* PORA_CHUNK_SIZE : (idx + 1)
|
|
* ENTRY_SIZE
|
|
* PORA_CHUNK_SIZE
|
|
]
|
|
if len(data) >= (idx + 1) * PORA_CHUNK_SIZE * ENTRY_SIZE
|
|
else data[idx * ENTRY_SIZE * PORA_CHUNK_SIZE :]
|
|
+ b"\x00" * (chunks * ENTRY_SIZE - len(data))
|
|
)
|
|
|
|
segment = {
|
|
"root": root_hash,
|
|
"data": base64.b64encode(tmp).decode("utf-8"),
|
|
"index": idx,
|
|
"proof": proof,
|
|
"fileSize": len(data),
|
|
}
|
|
|
|
segments.append(segment)
|
|
idx += 1
|
|
|
|
return segments
|