Updated code according to review

src/lib.rs

@@ -6,7 +6,6 @@ pub mod cost;
 pub mod dialect;
 pub mod err_utils;
 pub mod f_table;
-pub mod flags;
 pub mod more_ops;
 pub mod number;
 pub mod op_utils;
@@ -26,8 +25,6 @@ pub use chia_dialect::{
     ENABLE_BLS_OPS_OUTSIDE_GUARD, ENABLE_FIXED_DIV, LIMIT_HEAP, MEMPOOL_MODE, NO_UNKNOWN_OPS,
 };
 
-pub use flags::ALLOW_BACKREFS;
-
 #[cfg(feature = "counters")]
 pub use run_program::run_program_with_counters;
 

src/serde/ser_br.rs

@@ -7,6 +7,7 @@ use super::object_cache::{serialized_length, treehash, ObjectCache};
 use super::read_cache_lookup::ReadCacheLookup;
 use super::write_atom::write_atom;
 use crate::allocator::{Allocator, NodePtr, SExp};
+use crate::serde::ser::LimitedWriter;
 
 const BACK_REFERENCE: u8 = 0xfe;
 const CONS_BOX_MARKER: u8 = 0xff;
@@ -77,10 +78,45 @@ pub fn node_to_stream_backrefs<W: io::Write>(
     Ok(())
 }
 
+pub fn node_to_bytes_backrefs_limit(
+    a: &Allocator,
+    node: NodePtr,
+    limit: usize,
+) -> io::Result<Vec<u8>> {
+    let buffer = Cursor::new(Vec::new());
+    let mut writer = LimitedWriter::new(buffer, limit);
+    node_to_stream_backrefs(a, node, &mut writer)?;
+    let vec = writer.into_inner().into_inner();
+    Ok(vec)
+}
+
 pub fn node_to_bytes_backrefs(a: &Allocator, node: NodePtr) -> io::Result<Vec<u8>> {
-    let mut buffer = Cursor::new(Vec::new());
+    node_to_bytes_backrefs_limit(a, node, 2000000)
+}
 
-    node_to_stream_backrefs(a, node, &mut buffer)?;
-    let vec = buffer.into_inner();
-    Ok(vec)
+#[test]
+fn test_serialize_limit() {
+    let mut a = Allocator::new();
+
+    let leaf = a.new_atom(&[1, 2, 3, 4, 5]).unwrap();
+    let l1 = a.new_pair(leaf, leaf).unwrap();
+    let l2 = a.new_pair(l1, l1).unwrap();
+    let l3 = a.new_pair(l2, l2).unwrap();
+
+    let expected = &[255, 255, 255, 133, 1, 2, 3, 4, 5, 254, 2, 254, 2, 254, 2];
+
+    {
+        assert_eq!(node_to_bytes_backrefs(&a, l3).unwrap(), expected);
+    }
+
+    {
+        assert_eq!(node_to_bytes_backrefs_limit(&a, l3, 15).unwrap(), expected);
+    }
+
+    {
+        assert_eq!(
+            node_to_bytes_backrefs_limit(&a, l3, 14).unwrap_err().kind(),
+            io::ErrorKind::OutOfMemory
+        );
+    }
 }

wasm/src/lazy_node.rs

@@ -1,10 +1,10 @@
-use clvmr::allocator::{Allocator, NodePtr, SExp};
+use js_sys::Array;
 use std::rc::Rc;
+use wasm_bindgen::prelude::*;
 
+use crate::flags::ALLOW_BACKREFS;
+use clvmr::allocator::{Allocator, NodePtr, SExp};
 use clvmr::serde::{node_to_bytes, node_to_bytes_backrefs};
-use clvmr::ALLOW_BACKREFS;
-use js_sys::Array;
-use wasm_bindgen::prelude::*;
 
 #[wasm_bindgen]
 #[derive(Clone)]
@@ -45,7 +45,7 @@ impl LazyNode {
         } else {
             node_to_bytes
         };
-serializer(&self.allocator, self.node).ok()
+        serializer(&self.allocator, self.node).ok()
     }
 }
 

wasm/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod flags;
 pub mod lazy_node;
 pub mod run_program;
 pub mod serialize;

wasm/src/run_program.rs

@@ -2,14 +2,14 @@ use js_sys::Array;
 use std::rc::Rc;
 use wasm_bindgen::prelude::*;
 
+use crate::flags::ALLOW_BACKREFS;
 use crate::lazy_node::LazyNode;
 use clvmr::allocator::Allocator;
 use clvmr::chia_dialect::ChiaDialect;
 use clvmr::chia_dialect::NO_UNKNOWN_OPS as _no_unknown_ops;
 use clvmr::cost::Cost;
 use clvmr::run_program::run_program;
 use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs, node_to_bytes};
-use clvmr::ALLOW_BACKREFS;
 
 // When the `wee_alloc` feature is enabled, use `wee_alloc` as the global
 // allocator.

wasm/src/serialize.rs

@@ -1,15 +1,14 @@
-use crate::lazy_node::LazyNode;
-use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs, serialized_length_from_bytes};
-use clvmr::{Allocator, ALLOW_BACKREFS};
 use std::rc::Rc;
 use wasm_bindgen::prelude::wasm_bindgen;
 
+use crate::flags::ALLOW_BACKREFS;
+use crate::lazy_node::LazyNode;
+use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs, serialized_length_from_bytes};
+use clvmr::Allocator;
+
 #[wasm_bindgen]
 pub fn serialized_length(program: &[u8]) -> Result<u64, String> {
-    match serialized_length_from_bytes(program) {
-        Ok(length) => Ok(length),
-        Err(err) => Err(err.to_string()),
-    }
+    serialized_length_from_bytes(program).map_err(|x| x.to_string())
 }
 
 #[wasm_bindgen]
@@ -20,8 +19,6 @@ pub fn sexp_from_bytes(b: &[u8], flag: u32) -> Result<LazyNode, String> {
     } else {
         node_from_bytes
     };
-    match deserializer(&mut allocator, b) {
-        Ok(node) => Ok(LazyNode::new(Rc::new(allocator), node)),
-        Err(err) => Err(err.to_string()),
-    }
+    let node = deserializer(&mut allocator, b).map_err(|e| e.to_string())?;
+    Ok(LazyNode::new(Rc::new(allocator), node))
 }