surrealml_core/storage/header/
input_dims.rs

1//! InputDims is a struct that holds the dimensions of the input tensors for the model.
2use std::fmt;
3
4use crate::errors::error::{SurrealError, SurrealErrorStatus};
5use crate::safe_eject;
6
7/// InputDims is a struct that holds the dimensions of the input tensors for the model.
8///
9/// # Fields
10/// * `dims` - The dimensions of the input tensors.
11#[derive(Debug, PartialEq)]
12pub struct InputDims {
13	pub dims: [i32; 2],
14}
15
16impl InputDims {
17	/// Creates a new `InputDims` struct with all zeros.
18	///
19	/// # Returns
20	/// A new `InputDims` struct with all zeros.
21	pub fn fresh() -> Self {
22		InputDims {
23			dims: [0, 0],
24		}
25	}
26
27	/// Creates a new `InputDims` struct from a string.
28	///
29	/// # Arguments
30	/// * `data` - The dimensions as a string.
31	///
32	/// # Returns
33	/// A new `InputDims` struct, or a `SurrealError` if the dimensions are malformed.
34	pub fn from_string(data: String) -> Result<InputDims, SurrealError> {
35		if data == *"" {
36			return Ok(InputDims::fresh());
37		}
38		let parts: Vec<&str> = data.split(",").collect();
39		// Reject input that does not contain exactly two dimensions so that the indexing
40		// below cannot panic on attacker-controlled header data.
41		if parts.len() != 2 {
42			return Err(SurrealError::new(
43				format!(
44					"invalid input dimensions '{}': expected 2 comma-separated values, found {}",
45					data,
46					parts.len()
47				),
48				SurrealErrorStatus::BadRequest,
49			));
50		}
51		Ok(InputDims {
52			dims: [
53				safe_eject!(parts[0].parse::<i32>(), SurrealErrorStatus::BadRequest),
54				safe_eject!(parts[1].parse::<i32>(), SurrealErrorStatus::BadRequest),
55			],
56		})
57	}
58}
59
60impl fmt::Display for InputDims {
61	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62		if self.dims == [0, 0] {
63			write!(f, "")
64		} else {
65			write!(f, "{},{}", self.dims[0], self.dims[1])
66		}
67	}
68}
69
70#[cfg(test)]
71pub mod tests {
72
73	use super::*;
74
75	#[test]
76	fn test_fresh() {
77		let input_dims = InputDims::fresh();
78		assert_eq!(input_dims.dims[0], 0);
79		assert_eq!(input_dims.dims[1], 0);
80	}
81
82	#[test]
83	fn test_from_string() {
84		let input_dims = InputDims::from_string("1,2".to_string()).unwrap();
85		assert_eq!(input_dims.dims[0], 1);
86		assert_eq!(input_dims.dims[1], 2);
87	}
88
89	#[test]
90	fn test_to_string() {
91		let input_dims = InputDims::from_string("1,2".to_string()).unwrap();
92		assert_eq!(input_dims.to_string(), "1,2".to_string());
93	}
94
95	#[test]
96	fn test_from_string_empty_is_fresh() {
97		let input_dims = InputDims::from_string("".to_string()).unwrap();
98		assert_eq!(input_dims, InputDims::fresh());
99	}
100
101	// Regression tests for GHSA-jwr6-6444-28xv: malformed dimensions in a `.surml`
102	// header must surface a structured error instead of panicking (the release
103	// profile uses `panic = 'abort'`, so a panic here would crash the server).
104	#[test]
105	fn test_from_string_non_numeric_errors() {
106		// The exact malformed value from the advisory proof-of-concept.
107		let err = InputDims::from_string("bad".to_string()).unwrap_err();
108		assert_eq!(err.status, SurrealErrorStatus::BadRequest);
109	}
110
111	#[test]
112	fn test_from_string_too_few_dims_errors() {
113		let err = InputDims::from_string("1".to_string()).unwrap_err();
114		assert_eq!(err.status, SurrealErrorStatus::BadRequest);
115	}
116
117	#[test]
118	fn test_from_string_too_many_dims_errors() {
119		let err = InputDims::from_string("1,2,3".to_string()).unwrap_err();
120		assert_eq!(err.status, SurrealErrorStatus::BadRequest);
121	}
122}
surrealml_core/storage/header/input_dims.rs

surrealml_core/storage/header/
input_dims.rs
