Segment: bound MobileSAM prompt points, TFLite metadata and folded DEQUANTIZE sizes (#3509)

crates/segment/src/mobilesam.rs

@@ -19,6 +19,9 @@ const IMG: usize = 1024;
 const EMB: usize = 64;
 const DIM: usize = 256;
 const LOW: usize = 256;
+/// Most prompt points the decoder takes (each is an attention token, and self-attention is
+/// quadratic in them); SAM is used with a handful. More are thinned evenly, keeping both ends.
+const MAX_POINTS: usize = 256;
 const MEAN: [f32; 3] = [123.675, 116.28, 103.53];
 const STD: [f32; 3] = [58.395, 57.12, 57.375];
 
@@ -609,7 +612,7 @@ impl MobileSam {
     /// Foreground probability (`w×h`) for a prompt on an embedded image, and the predicted IoU.
     pub fn predict(&self, e: &Embedding, prompt: &Prompt) -> (Vec<f32>, f32) {
         let s = e.scale;
-        let points: Vec<([f32; 2], bool)> = prompt.points.iter().map(|(p, fg)| ([p[0] * s, p[1] * s], *fg)).collect();
+        let points: Vec<([f32; 2], bool)> = thin(&prompt.points).iter().map(|(p, fg)| ([p[0] * s, p[1] * s], *fg)).collect();
         let bbox = prompt.bbox.map(|b| b.map(|v| v * s));
         // A prior mask (probabilities at image size) → logits on the low-res grid.
         let mask: Option<Vec<f32>> = prompt.mask.as_ref().filter(|m| m.len() == e.w * e.h).map(|m| {
@@ -675,3 +678,28 @@ fn sample(rgb: &[[f32; 3]], w: usize, h: usize, x: f32, y: f32, foot: f32) -> [f
     }
     s.map(|v| v / f32::max(n, 1.0))
 }
+
+/// `points`, or `MAX_POINTS` evenly spaced ones (first and last kept) when there are more.
+fn thin(points: &[([f32; 2], bool)]) -> Vec<([f32; 2], bool)> {
+    let n = points.len();
+    if n <= MAX_POINTS {
+        return points.to_vec();
+    }
+    (0..MAX_POINTS).filter_map(|i| points.get(i * (n - 1) / (MAX_POINTS - 1)).copied()).collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn many_prompt_points_are_thinned_to_the_budget() {
+        let few: Vec<_> = (0..MAX_POINTS).map(|i| ([i as f32, 0.0], i % 2 == 0)).collect();
+        assert_eq!(thin(&few), few, "a prompt within the budget is untouched");
+        let many: Vec<_> = (0..1_000_000).map(|i| ([i as f32, 0.0], i % 2 == 0)).collect();
+        let t = thin(&many);
+        assert_eq!(t.len(), MAX_POINTS);
+        assert_eq!((t[0].0, t[MAX_POINTS - 1].0), (many[0].0, many[many.len() - 1].0));
+        assert!(t.windows(2).all(|w| w[0].0[0] < w[1].0[0]), "spread in order");
+    }
+}

crates/segment/src/tflite.rs

@@ -131,6 +131,8 @@ const MAX_VALUES: u64 = 1 << 30;
 const MAX_LIVE: u64 = 1 << 29;
 /// The most constant values decoded and laid out at load (a weight shared by several tensors or
 /// operators counts once per use).
+/// Bytes of tensor names and shapes a model may carry in all (the official ones take well under 1 MB).
+const MAX_META: usize = 1 << 24;
 const MAX_CONSTS: usize = 1 << 27;
 
 /// Tensor types (`TensorType`).
@@ -261,6 +263,7 @@ impl Model {
         let mut ints: Vec<Option<Vec<i32>>> = vec![None; n];
         let mut data: Vec<(usize, u8, usize, &[u8])> = vec![];
         let mut consts = 0usize;
+        let mut meta = 0usize;
         for (i, t) in tensors.iter().enumerate() {
             let shape: Vec<usize> = fb
                 .ints(*t, 0)?
@@ -269,7 +272,13 @@ impl Model {
                 .collect::<std::result::Result<_, _>>()?;
             let ty = fb.byte(*t, 1, FLOAT32)?;
             let buf = fb.int(*t, 2, 0)? as u32 as usize;
-            names.push(fb.string(*t, 3)?);
+            let name = fb.string(*t, 3)?;
+            // Names and shapes are copied per tensor: bound them in all.
+            meta = meta.saturating_add(name.len()).saturating_add(shape.len().saturating_mul(8));
+            if meta > MAX_META {
+                return Err("the model's tensor metadata is too large".into());
+            }
+            names.push(name);
             // Every size computed from shapes later stays in range, even a product of some of the
             // dimensions (a zero dimension does not let the others grow without bound).
             let count = shape.iter().try_fold(1usize, |a, &d| a.checked_mul(d)).ok_or("a tensor is too large")?;
@@ -329,6 +338,11 @@ impl Model {
                 && let Some(i) = ins.first().copied().flatten()
                 && let Some(len) = floats[i].as_ref().map(Vec::len)
             {
+                // The folded constant is what the graph outputs and the run's memory estimate count:
+                // it must be the declared size.
+                if shapes[output].iter().try_fold(1usize, |a, &d| a.checked_mul(d)) != Some(len) {
+                    return Err(format!("{}: DEQUANTIZE output size disagrees with its input", names[output]));
+                }
                 spend(len)?;
                 floats[output] = floats[i].clone();
                 continue;

crates/segment/src/tflite/tests.rs

@@ -334,6 +334,26 @@ fn shapes_computed_at_run_time_must_match_the_declared_ones() {
     refused(&zero, "tensor is too large");
 }
 
+#[test]
+fn a_folded_dequantize_has_its_declared_size() {
+    let half = |n: usize| vec![0u8; 2 * n];
+    let fold = |declared: Vec<i32>| {
+        model(vec![half(4)], vec![(vec![1, 4], FLOAT16, Some(0)), (declared, FLOAT32, None)], vec![(op::DEQUANTIZE, vec![0], vec![1], vec![])], vec![], vec![1])
+    };
+    // A declared output of one value would let the folded four-value constant be copied out per reference.
+    refused(&fold(vec![1, 1]), "DEQUANTIZE output size");
+    assert!(Model::read(&fold(vec![1, 4])).is_ok());
+}
+
+#[test]
+fn tensor_metadata_is_bounded_in_all() {
+    // A rank of 4 million dimensions of 1 is a valid shape by size, but copied per tensor.
+    let m = model(vec![], vec![(vec![1; 1 << 22], FLOAT32, None)], vec![], vec![], vec![]);
+    refused(&m, "metadata is too large");
+    let ok = model(vec![], (0..70).map(|_| (vec![1, 2, 3], FLOAT32, None)).collect(), vec![], vec![], vec![]);
+    assert!(Model::read(&ok).is_ok());
+}
+
 #[test]
 fn shared_weights_count_once_per_use() {
     // One 128 KB buffer decoded for 2100 tensors would be 2100 copies of 2¹⁶ floats.