Skip to content

Commit

Permalink
Merge pull request TeXitoi#38 from dgriffen/maniac_oom
Browse files Browse the repository at this point in the history 
Maniac oom
  • Loading branch information
@ZoeyR
ZoeyR authored Jun 10, 2018
2 parents 0c19f0c + 99c11f9 commit d42d7bf
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 44 deletions.
105 changes: 62 additions & 43 deletions flif/src/maniac/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,27 +45,7 @@ impl<'a> ManiacTree<'a> {
}

pub fn size(&self) -> usize {
use self::ManiacNode::*;

let mut size = 0;
let mut stack = vec![0];
loop {
let index = match stack.pop() {
Some(index) => index,
None => break size,
};

size += 1;
match self.nodes[index] {
Property { .. } | InactiveProperty { .. } | Inner { .. } => {
stack.push(2 * index + 2);
stack.push(2 * index + 1);
}
_ => {
continue;
}
};
}
return self.nodes.len();
}

pub fn depth(&self) -> usize {
Expand All @@ -83,9 +63,11 @@ impl<'a> ManiacTree<'a> {
largest_depth = ::std::cmp::max(largest_depth, depth);

match self.nodes[index] {
Property { .. } | InactiveProperty { .. } | Inner { .. } => {
stack.push((2 * index + 2, depth + 1));
stack.push((2 * index + 1, depth + 1));
Property { left, right, .. }
| InactiveProperty { left, right, .. }
| Inner { left, right, .. } => {
stack.push((right, depth + 1));
stack.push((left, depth + 1));
}
_ => {
continue;
Expand Down Expand Up @@ -119,36 +101,36 @@ impl<'a> ManiacTree<'a> {
) -> Result<Vec<ManiacNode<'a>>> {
use self::ManiacNode::*;

let mut result_vec = vec![];
let mut node_count = 0;
let mut result_vec = vec![ManiacNode::InactiveLeaf];
let mut process_stack = vec![(0, prange)];
loop {
let (index, prange) = match process_stack.pop() {
Some(process) => process,
_ => break,
};

if node_count > limits.maniac_nodes {
if result_vec.len() > limits.maniac_nodes {
Err(Error::LimitViolation(format!(
"number of maniac nodes exceeds limit"
)))?;
}

node_count += 1;
let child_start = result_vec.len();
let node = if index == 0 {
Self::create_node(rac, context, update_table, &prange)?
Self::create_node(child_start, rac, context, update_table, &prange)?
} else {
Self::create_inner_node(rac, context, &prange)?
Self::create_inner_node(child_start, rac, context, &prange)?
};

if index >= result_vec.len() {
result_vec.resize(index + 1, ManiacNode::InactiveLeaf);
}

let (property, test_value) = match node {
Property { id, value, .. }
| InactiveProperty { id, value, .. }
| Inner { id, value } => (id, value),
| Inner { id, value, .. } => {
if child_start >= result_vec.len() {
result_vec.resize(child_start + 2, ManiacNode::InactiveLeaf);
}
(id, value)
}
_ => {
result_vec[index] = node;
continue;
Expand All @@ -161,15 +143,16 @@ impl<'a> ManiacTree<'a> {
let mut right_prange = prange;
right_prange[property as usize].max = test_value;

process_stack.push((2 * index + 2, right_prange));
process_stack.push((2 * index + 1, left_prange));
process_stack.push((child_start + 1, right_prange));
process_stack.push((child_start, left_prange));
result_vec[index] = node;
}

Ok(result_vec)
}

fn create_node<R: Read>(
child_start: usize,
rac: &mut Rac<R>,
context: &mut [ChanceTable; 3],
update_table: &'a UpdateTable,
Expand All @@ -183,6 +166,10 @@ impl<'a> ManiacTree<'a> {
}
property -= 1;

if prange[property as usize].min >= prange[property as usize].max {
Err(Error::InvalidOperation(format!("Invalid maniac tree")))?
}

let counter = rac.read_near_zero(1 as i32, 512 as i32, &mut context[1])?;
let test_value = rac.read_near_zero(
prange[property as usize].min,
Expand All @@ -195,10 +182,13 @@ impl<'a> ManiacTree<'a> {
table: chance_table,
value: test_value,
counter: counter as u32,
left: child_start,
right: child_start + 1,
})
}

fn create_inner_node<R: Read>(
child_start: usize,
rac: &mut Rac<R>,
context: &mut [ChanceTable; 3],
prange: &[ColorRange],
Expand All @@ -210,6 +200,10 @@ impl<'a> ManiacTree<'a> {
}
property -= 1;

if prange[property as usize].min >= prange[property as usize].max {
Err(Error::InvalidOperation(format!("Invalid maniac tree")))?
}

let counter = rac.read_near_zero(1 as i32, 512 as i32, &mut context[1])?;
let test_value = rac.read_near_zero(
prange[property as usize].min,
Expand All @@ -221,6 +215,8 @@ impl<'a> ManiacTree<'a> {
id: property,
value: test_value,
counter: counter as u32,
left: child_start,
right: child_start + 1,
})
}

Expand All @@ -237,11 +233,16 @@ impl<'a> ManiacTree<'a> {
let (lnodes, rnodes) = &mut self.nodes.split_at_mut(node_index + 1);
let node = &mut lnodes[node_index];
match node {
Inner { id, value } => {
Inner {
id,
value,
left,
right,
} => {
if pvec[*id as usize] > *value {
node_index = 2 * node_index + 1;
node_index = *left;
} else {
node_index = 2 * node_index + 2;
node_index = *right;
}
}
Leaf(table) => {
Expand All @@ -254,6 +255,8 @@ impl<'a> ManiacTree<'a> {
value,
counter: 0,
table,
left,
right,
} => {
let mut left_table = table.clone();
let mut right_table = table.clone();
Expand All @@ -264,13 +267,15 @@ impl<'a> ManiacTree<'a> {
rac.read_near_zero(min, max, &mut right_table)?
};

rnodes[node_index].activate(left_table);
rnodes[node_index + 1].activate(right_table);
rnodes[*left - node_index - 1].activate(left_table);
rnodes[*right - node_index - 1].activate(right_table);
(
val,
Inner {
id: *id,
value: *value,
left: *left,
right: *right,
},
)
}
Expand Down Expand Up @@ -332,16 +337,22 @@ enum ManiacNode<'a> {
value: i16,
table: ChanceTable<'a>,
counter: u32,
left: usize,
right: usize,
},
InactiveProperty {
id: isize,
value: i16,
counter: u32,
left: usize,
right: usize,
},
/// Inner nodes are property nodes whose counters have reached zero. They no longer have a context associated with them.
Inner {
id: isize,
value: i16,
left: usize,
right: usize,
},
/// Leaf nodes are nodes that can never become inner nodes
Leaf(ChanceTable<'a>),
Expand All @@ -354,11 +365,19 @@ impl<'a> ManiacNode<'a> {
use self::ManiacNode::*;
*self = match self {
InactiveLeaf => Leaf(table),
InactiveProperty { id, value, counter } => Property {
InactiveProperty {
id,
value,
counter,
left,
right,
} => Property {
id: *id,
value: *value,
counter: *counter,
table: table,
left: *left,
right: *right,
},
_ => return,
}
Expand Down
22 changes: 21 additions & 1 deletion flif/tests/invalid_inputs.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ fn invalid_bytes_per_channel() {

/// Tests an issue found in [#30](https://github.com/dgriffen/flif.rs/issues/30)
#[test]
fn maniac_stack_overflow() {
fn ycocg_stack_overflow() {
let bytes = b"FLIF41\x02\x01\x00pr@\x015\xc6\xe3d\xbfct\x00i\x005FLI)F\xca\xcdi\x00r\x00\xfft\x11-FLIF12i\x00r\x00\xfft\x11\x00\xfft\x11-FLIF12i\x00r\x00\xfft\x11-le\x00FLI 11\xe3d\xbfct\x00i\xf9\xf9\x07\xff5\xff\x00\x00";
let limits = flif::Limits {
metadata_chunk: 32,
Expand All @@ -29,3 +29,23 @@ fn maniac_stack_overflow() {
};
let _ = Flif::decode_with_limits(bytes.as_ref(), limits).map(|img| img.get_raw_pixels());
}

/// Tests an issue found in [#34](https://github.com/dgriffen/flif.rs/issues/34)
#[test]
fn memory_growth() {
let bytes = b"FLIF11F\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x04\x00\x00\x00FLIF\x00\x00L\xc5XifI\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00FLIF\x00\x00\x00\x00\x00\x00\x00";
let limits = flif::Limits {
metadata_chunk: 32,
metadata_count: 8,
pixels: 1 << 16,
maniac_nodes: 512,
};
match Flif::decode_with_limits(bytes.as_ref(), limits) {
Err(Error::InvalidOperation(ref message)) if message.contains("maniac") => {}
Err(err) => panic!(
"Expected an Error::InvalidOperation indicating the maniac tree was invalid, got {:?}",
err
),
_ => panic!("Expected an Error::InvalidOperation indicating the maniac tree was invalid, got a valid image instead")
}
}

0 comments on commit d42d7bf

Please sign in to comment.