Skip to content

Commit d1debb0

Browse files
demosdemondemosdemon
committed
feat: verify node hashes on read
1 parent 2ba9fb8 commit d1debb0

File tree

9 files changed

+395
-247
lines changed

9 files changed

+395
-247
lines changed

‎firewood/src/iter.rs‎

Lines changed: 41 additions & 108 deletions
Large diffs are not rendered by default.

‎firewood/src/merkle/mod.rs‎

Lines changed: 89 additions & 50 deletions
Large diffs are not rendered by default.

‎firewood/src/merkle/parallel.rs‎

Lines changed: 4 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -141,13 +141,8 @@ impl ParallelMerkle {
141141
}
142142

143143
// Return the child as the new root. Update its partial path to include the index value.
144-
let mut child = match child {
145-
Child::Node(child_node) => std::mem::take(child_node),
146-
Child::AddressWithHash(addr, _) => nodestore.read_for_update((*addr).into())?,
147-
Child::MaybePersisted(maybe_persisted, _) => {
148-
nodestore.read_for_update(maybe_persisted.clone())?
149-
}
150-
};
144+
let mut child =
145+
child.read_for_update(nodestore, &Path::from(&[child_index.as_u8()]))?;
151146

152147
// The child's partial path is the concatenation of its (now removed) parent, which
153148
// should always be empty because of our prepare step, its (former) child index, and
@@ -238,17 +233,8 @@ impl ParallelMerkle {
238233
.children
239234
.get_mut(first_path_component)
240235
.take()
241-
.map(|child| -> Result<_, FileIoError> {
242-
match child {
243-
Child::Node(node) => Ok(node),
244-
Child::AddressWithHash(address, _) => {
245-
// Track deletion of the removed child from the root (if it was persisted).
246-
Ok(proposal.read_for_update(address.into())?)
247-
}
248-
Child::MaybePersisted(maybe_persisted, _) => {
249-
Ok(proposal.read_for_update(maybe_persisted)?)
250-
}
251-
}
236+
.map(|child| {
237+
child.read_for_update(proposal, &Path::from(&[first_path_component.as_u8()]))
252238
})
253239
.transpose()?;
254240

‎firewood/src/merkle/tests/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -281,7 +281,7 @@ fn remove_root() {
281281
assert!(merkle.get_value(&key3).unwrap().is_none());
282282
assert!(merkle.remove(&key3).unwrap().is_none());
283283

284-
assert!(merkle.nodestore.root_node().is_none());
284+
assert!(merkle.nodestore.read_root().unwrap().is_none());
285285
}
286286

287287
#[test]
@@ -357,7 +357,7 @@ fn remove_many() {
357357
let got = merkle.get_value(&key).unwrap();
358358
assert!(got.is_none());
359359
}
360-
assert!(merkle.nodestore.root_node().is_none());
360+
assert!(merkle.nodestore.read_root().unwrap().is_none());
361361
}
362362

363363
#[test]

‎storage/src/node/branch.rs‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,58 @@ impl Child {
141141
Child::MaybePersisted(maybe_persisted, _) => maybe_persisted.clone(),
142142
}
143143
}
144+
145+
/// Read the node from storage for update
146+
pub fn read_for_update<S: crate::ReadableStorage>(
147+
&self,
148+
storage: &mut crate::NodeStore<crate::MutableProposal, S>,
149+
leading_path: &Path,
150+
) -> Result<Node, crate::FileIoError> {
151+
match self {
152+
Child::Node(node) => Ok(node.clone()),
153+
Child::AddressWithHash(addr, hash) => {
154+
storage.read_for_update((*addr).into(), hash, leading_path)
155+
}
156+
Child::MaybePersisted(node, hash) => {
157+
storage.read_for_update(node.clone(), hash, leading_path)
158+
}
159+
}
160+
}
161+
162+
/// Read the node from storage
163+
pub fn read_node_from_storage<S: crate::NodeReader>(
164+
&self,
165+
storage: &S,
166+
leading_path: &Path,
167+
) -> Result<SharedNode, crate::FileIoError> {
168+
match self {
169+
Child::Node(node) => Ok(node.clone().into()),
170+
Child::AddressWithHash(addr, hash) => storage.read_node(*addr, hash, leading_path),
171+
Child::MaybePersisted(maybe_persisted, hash) => {
172+
maybe_persisted.as_shared_node(storage, hash, leading_path)
173+
}
174+
}
175+
}
176+
177+
pub(crate) fn read_root<S: crate::NodeReader>(
178+
&self,
179+
storage: &S,
180+
) -> Result<(MaybePersistedNode, Node), crate::FileIoError> {
181+
match self {
182+
Child::Node(node) => Ok((
183+
MaybePersistedNode::from(SharedNode::from(node.clone())),
184+
node.clone(),
185+
)),
186+
Child::AddressWithHash(addr, hash) => {
187+
let node = storage.read_node(*addr, hash, &Path::new())?;
188+
Ok((MaybePersistedNode::from(*addr), (*node).clone()))
189+
}
190+
Child::MaybePersisted(maybe_persisted, hash) => {
191+
let node = maybe_persisted.as_shared_node(storage, hash, &Path::new())?;
192+
Ok((maybe_persisted.clone(), (*node).clone()))
193+
}
194+
}
195+
}
144196
}
145197

146198
#[derive(PartialEq, Eq, Clone)]

‎storage/src/node/persist.rs‎

Lines changed: 43 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
use parking_lot::Mutex;
55
use std::{fmt::Display, sync::Arc};
66

7-
use crate::{FileIoError, LinearAddress, NodeReader, SharedNode};
7+
use crate::{FileIoError, HashType, LinearAddress, NodeReader, Path, SharedNode};
88

99
/// A node that is either in memory or on disk.
1010
///
@@ -86,12 +86,19 @@ impl MaybePersistedNode {
8686
/// Returns a `Result<SharedNode, FileIoError>` where:
8787
/// - `Ok(SharedNode)` contains the node if successfully retrieved
8888
/// - `Err(FileIoError)` if there was an error reading from storage
89-
pub fn as_shared_node<S: NodeReader>(&self, storage: &S) -> Result<SharedNode, FileIoError> {
89+
pub fn as_shared_node<S: NodeReader>(
90+
&self,
91+
storage: &S,
92+
expected_hash: &HashType,
93+
leading_path: &Path,
94+
) -> Result<SharedNode, FileIoError> {
9095
match &*self.0.lock() {
9196
MaybePersisted::Allocated(_, node) | MaybePersisted::Unpersisted(node) => {
9297
Ok(node.clone())
9398
}
94-
MaybePersisted::Persisted(address) => storage.read_node(*address),
99+
MaybePersisted::Persisted(address) => {
100+
storage.read_node(*address, expected_hash, leading_path)
101+
}
95102
}
96103
}
97104

@@ -234,15 +241,24 @@ mod test {
234241
// create as unpersisted
235242
let maybe_persisted_node = MaybePersistedNode::from(node.clone());
236243
let addr = nonzero!(2048u64);
237-
assert_eq!(maybe_persisted_node.as_shared_node(&store).unwrap(), node);
244+
assert_eq!(
245+
maybe_persisted_node
246+
.as_shared_node(&store, &HashType::empty(), &Path::new())
247+
.unwrap(),
248+
node
249+
);
238250
assert_eq!(
239251
Option::<LinearAddress>::None,
240252
Option::from(&maybe_persisted_node)
241253
);
242254

243255
let addr = LinearAddress::new(addr.get()).unwrap();
244256
maybe_persisted_node.persist_at(addr);
245-
assert!(maybe_persisted_node.as_shared_node(&store).is_err());
257+
assert!(
258+
maybe_persisted_node
259+
.as_shared_node(&store, &HashType::empty(), &Path::new())
260+
.is_err()
261+
);
246262
assert_eq!(Some(addr), Option::from(&maybe_persisted_node));
247263
Ok(())
248264
}
@@ -267,8 +283,18 @@ mod test {
267283
let cloned = original.clone();
268284

269285
// Both should be unpersisted initially
270-
assert_eq!(original.as_shared_node(&store).unwrap(), node);
271-
assert_eq!(cloned.as_shared_node(&store).unwrap(), node);
286+
assert_eq!(
287+
original
288+
.as_shared_node(&store, &HashType::empty(), &Path::new())
289+
.unwrap(),
290+
node
291+
);
292+
assert_eq!(
293+
cloned
294+
.as_shared_node(&store, &HashType::empty(), &Path::new())
295+
.unwrap(),
296+
node
297+
);
272298
assert_eq!(Option::<LinearAddress>::None, Option::from(&original));
273299
assert_eq!(Option::<LinearAddress>::None, Option::from(&cloned));
274300

@@ -281,8 +307,16 @@ mod test {
281307

282308
// Both original and clone should now be persisted since they share the same
283309
// mutex-protected pointer
284-
assert!(original.as_shared_node(&store).is_err());
285-
assert!(cloned.as_shared_node(&store).is_err());
310+
assert!(
311+
original
312+
.as_shared_node(&store, &HashType::empty(), &Path::new())
313+
.is_err()
314+
);
315+
assert!(
316+
cloned
317+
.as_shared_node(&store, &HashType::empty(), &Path::new())
318+
.is_err()
319+
);
286320
assert_eq!(Some(addr), Option::from(&original));
287321
assert_eq!(Some(addr), Option::from(&cloned));
288322

‎storage/src/nodestore/hash.rs‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ where
102102
} else {
103103
// If not persisted, we need to get the node to hash it
104104
let node = maybe_persisted
105-
.as_shared_node(&self)
105+
.as_shared_node(&self, &HashType::empty(), &Path::new())
106106
.expect("will never fail for unpersisted nodes");
107107
acc.unhashed.push((idx, node.deref().clone()));
108108
}
@@ -158,14 +158,17 @@ where
158158
trace!("hashed {hashed:?} unhashed {unhashed:?}");
159159
// we were left with one hashed node that must be rehashed
160160
if let [(child_idx, (child_node, child_hash))] = &mut hashed[..] {
161+
let mut path_guard = PathGuard::new(&mut path_prefix);
162+
path_guard.0.extend(b.partial_path.0.iter().copied());
161163
// Extract the address from the MaybePersistedNode
162164
let addr: crate::LinearAddress = child_node
163165
.as_linear_address()
164166
.expect("hashed node should be persisted");
165-
let mut hashable_node = self.read_node(addr)?.deref().clone();
167+
let mut hashable_node = self
168+
.read_node(addr, child_hash, &path_guard)?
169+
.deref()
170+
.clone();
166171
let hash = {
167-
let mut path_guard = PathGuard::new(&mut path_prefix);
168-
path_guard.0.extend(b.partial_path.0.iter().copied());
169172
if unhashed.is_empty() {
170173
hashable_node.update_partial_path(Path::from_nibbles_iterator(
171174
std::iter::once(child_idx.as_u8())

0 commit comments

Comments
 (0)