surrealml_core/storage/header/
origin.rs

1//! Defines the origin of the model in the file.
2use std::fmt;
3
4use super::string_value::StringValue;
5use crate::errors::error::{SurrealError, SurrealErrorStatus};
6use crate::safe_eject_option;
7
8const LOCAL: &str = "local";
9const SURREAL_DB: &str = "surreal_db";
10const NONE: &str = "";
11
12/// Defines the types of origin that are supported.
13///
14/// # Fields
15/// * `Local` - The model was created locally.
16/// * `SurrealDb` - The model was created in the surreal database.
17/// * `None` - The model has no origin
18#[derive(Debug, PartialEq)]
19pub enum OriginValue {
20	Local(StringValue),
21	SurrealDb(StringValue),
22	None(StringValue),
23}
24
25impl OriginValue {
26	/// Creates a new `OriginValue` with no value.
27	///
28	/// # Returns
29	/// A new `OriginValue` with no value.
30	pub fn fresh() -> Self {
31		OriginValue::None(StringValue::fresh())
32	}
33
34	/// Create a `OriginValue` from a string.
35	///
36	/// # Arguments
37	/// * `origin` - The origin as a string.
38	///
39	/// # Returns
40	/// A new `OriginValue`.
41	pub fn from_string(origin: String) -> Result<Self, SurrealError> {
42		match origin.to_lowercase().as_str() {
43			LOCAL => Ok(OriginValue::Local(StringValue::from_string(origin))),
44			SURREAL_DB => Ok(OriginValue::SurrealDb(StringValue::from_string(origin))),
45			NONE => Ok(OriginValue::None(StringValue::from_string(origin))),
46			_ => Err(SurrealError::new(
47				format!("invalid origin: {}", origin),
48				SurrealErrorStatus::BadRequest,
49			)),
50		}
51	}
52}
53
54impl fmt::Display for OriginValue {
55	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56		match self {
57			OriginValue::Local(val) | OriginValue::SurrealDb(val) | OriginValue::None(val) => {
58				write!(f, "{}", val)
59			}
60		}
61	}
62}
63
64/// Defines the origin of the model in the file header.
65///
66/// # Fields
67/// * `origin` - The origin of the model.
68/// * `author` - The author of the model.
69#[derive(Debug, PartialEq)]
70pub struct Origin {
71	pub origin: OriginValue,
72	pub author: StringValue,
73}
74
75impl Origin {
76	/// Creates a new origin with no values.
77	///
78	/// # Returns
79	/// A new origin with no values.
80	pub fn fresh() -> Self {
81		Origin {
82			origin: OriginValue::fresh(),
83			author: StringValue::fresh(),
84		}
85	}
86
87	/// Adds an author to the origin struct.
88	///
89	/// # Arguments
90	/// * `origin` - The origin to be added.
91	pub fn add_author(&mut self, author: String) {
92		self.author = StringValue::from_string(author);
93	}
94
95	/// Adds an origin to the origin struct.
96	///
97	/// # Arguments
98	pub fn add_origin(&mut self, origin: String) -> Result<(), SurrealError> {
99		self.origin = OriginValue::from_string(origin)?;
100		Ok(())
101	}
102
103	/// Creates a new origin from a string.
104	///
105	/// # Arguments
106	/// * `origin` - The origin as a string.
107	///
108	/// # Returns
109	/// A new origin.
110	pub fn from_string(origin: String) -> Result<Self, SurrealError> {
111		if origin == *"" {
112			return Ok(Origin::fresh());
113		}
114		let mut split = origin.split("=>");
115		// Avoid unchecked `unwrap()` on attacker-controlled header data: a malformed
116		// origin field (e.g. one missing the `=>` delimiter) must error, not panic.
117		let author = safe_eject_option!(split.next()).to_string();
118		let origin = safe_eject_option!(split.next()).to_string();
119		Ok(Origin {
120			origin: OriginValue::from_string(origin)?,
121			author: StringValue::from_string(author),
122		})
123	}
124}
125
126impl fmt::Display for Origin {
127	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
128		let is_empty_author = self.author.value.is_none();
129		let is_empty_origin = matches!(self.origin, OriginValue::None(ref s) if s.value.is_none());
130
131		if is_empty_author && is_empty_origin {
132			write!(f, "")
133		} else {
134			write!(f, "{}=>{}", self.author, self.origin)
135		}
136	}
137}
138
139#[cfg(test)]
140mod tests {
141
142	use super::*;
143
144	#[test]
145	fn test_fresh() {
146		let origin = Origin::fresh();
147		assert_eq!(
148			origin,
149			Origin {
150				origin: OriginValue::fresh(),
151				author: StringValue::fresh(),
152			}
153		);
154	}
155
156	#[test]
157	fn test_to_string() {
158		let origin = Origin {
159			origin: OriginValue::from_string("local".to_string()).unwrap(),
160			author: StringValue::from_string("author".to_string()),
161		};
162		assert_eq!(origin.to_string(), "author=>local".to_string());
163
164		let origin = Origin::fresh();
165		assert_eq!(origin.to_string(), "".to_string());
166	}
167
168	#[test]
169	fn test_from_string() {
170		let origin = Origin::from_string("author=>local".to_string()).unwrap();
171		assert_eq!(
172			origin,
173			Origin {
174				origin: OriginValue::from_string("local".to_string()).unwrap(),
175				author: StringValue::from_string("author".to_string()),
176			}
177		);
178
179		let origin = Origin::from_string("=>local".to_string()).unwrap();
180
181		assert_eq!(None, origin.author.value);
182		assert_eq!("local".to_string(), origin.origin.to_string());
183	}
184
185	// Regression test for GHSA-jwr6-6444-28xv: a non-empty origin field without the
186	// `=>` delimiter must error rather than panic on `split.next().unwrap()`.
187	#[test]
188	fn test_from_string_missing_delimiter_errors() {
189		assert!(Origin::from_string("no-delimiter".to_string()).is_err());
190	}
191}
surrealml_core/storage/header/origin.rs

surrealml_core/storage/header/
origin.rs
