1use osproxy_core::json::scalar_at_path;
13use serde_json::Value;
14
15use crate::error::RewriteError;
16use crate::extract::extract_scalar;
17
18pub fn construct_id(template: &str, partition: &str, doc: &Value) -> Result<String, RewriteError> {
37 construct_id_with(template, partition, |path| {
38 extract_scalar(doc, path.split('.'))
39 })
40}
41
42pub fn construct_id_bytes(
64 template: &str,
65 partition: &str,
66 body: &[u8],
67) -> Result<String, RewriteError> {
68 construct_id_with(template, partition, |path| {
69 scalar_at_path(body, path.split('.')).map_err(RewriteError::from)
70 })
71}
72
73fn construct_id_with<F>(
77 template: &str,
78 partition: &str,
79 resolve_body: F,
80) -> Result<String, RewriteError>
81where
82 F: Fn(&str) -> Result<String, RewriteError>,
83{
84 let mut out = String::with_capacity(template.len());
85 let mut rest = template;
86 while let Some(open) = rest.find('{') {
87 out.push_str(&rest[..open]);
88 let after = &rest[open + 1..];
89 let close = after
90 .find('}')
91 .ok_or_else(|| RewriteError::UnsupportedPlaceholder {
92 placeholder: after.to_owned(),
93 })?;
94 let placeholder = &after[..close];
95 if placeholder == "partition" {
96 out.push_str(partition);
97 } else if let Some(path) = placeholder.strip_prefix("body.") {
98 out.push_str(&resolve_body(path)?);
99 } else {
100 return Err(RewriteError::UnsupportedPlaceholder {
101 placeholder: placeholder.to_owned(),
102 });
103 }
104 rest = &after[close + 1..];
105 }
106 out.push_str(rest);
107 Ok(out)
108}
109
110pub fn map_logical_to_physical(
133 template: &str,
134 partition: &str,
135 logical_id: &str,
136) -> Result<String, RewriteError> {
137 let (prefix, suffix) = id_frame(template, partition)?;
138 Ok(format!("{prefix}{logical_id}{suffix}"))
139}
140
141pub fn map_physical_to_logical(
163 template: &str,
164 partition: &str,
165 physical_id: &str,
166) -> Result<Option<String>, RewriteError> {
167 let (prefix, suffix) = id_frame(template, partition)?;
168 Ok(physical_id
169 .strip_prefix(&prefix)
170 .and_then(|rest| rest.strip_suffix(&suffix))
171 .map(str::to_owned))
172}
173
174fn id_frame(template: &str, partition: &str) -> Result<(String, String), RewriteError> {
180 let mut prefix = String::new();
181 let mut suffix = String::new();
182 let mut seen_body = false;
183 let mut rest = template;
184 while let Some(open) = rest.find('{') {
185 let literal = &rest[..open];
186 let after = &rest[open + 1..];
187 let close = after
188 .find('}')
189 .ok_or_else(|| RewriteError::UnsupportedPlaceholder {
190 placeholder: after.to_owned(),
191 })?;
192 let placeholder = &after[..close];
193 let frame = if seen_body { &mut suffix } else { &mut prefix };
194 frame.push_str(literal);
195 if placeholder == "partition" {
196 frame.push_str(partition);
197 } else if placeholder.strip_prefix("body.").is_some() {
198 if seen_body {
199 return Err(RewriteError::IrreversibleIdTemplate);
200 }
201 seen_body = true;
202 } else {
203 return Err(RewriteError::UnsupportedPlaceholder {
204 placeholder: placeholder.to_owned(),
205 });
206 }
207 rest = &after[close + 1..];
208 }
209 if seen_body {
210 suffix.push_str(rest);
211 } else {
212 return Err(RewriteError::IrreversibleIdTemplate);
213 }
214 Ok((prefix, suffix))
215}
216
217#[cfg(test)]
218mod tests {
219 use super::*;
220 use serde_json::json;
221
222 #[test]
223 fn expands_partition_and_body_placeholders() {
224 let doc = json!({ "k": "natural", "nested": { "v": 9 } });
225 assert_eq!(
226 construct_id("{partition}:{body.k}", "p1", &doc).unwrap(),
227 "p1:natural"
228 );
229 assert_eq!(
230 construct_id("{body.nested.v}-{partition}", "p1", &doc).unwrap(),
231 "9-p1"
232 );
233 }
234
235 #[test]
236 fn literal_only_template_is_copied() {
237 let doc = json!({});
238 assert_eq!(construct_id("static-id", "p", &doc).unwrap(), "static-id");
239 }
240
241 #[test]
242 fn unknown_placeholder_is_rejected() {
243 let doc = json!({});
244 assert_eq!(
245 construct_id("{principal}", "p", &doc).unwrap_err(),
246 RewriteError::UnsupportedPlaceholder {
247 placeholder: "principal".to_owned()
248 }
249 );
250 }
251
252 #[test]
253 fn unterminated_placeholder_is_rejected() {
254 let doc = json!({});
255 assert!(construct_id("{partition", "p", &doc).is_err());
256 }
257
258 #[test]
259 fn missing_body_path_propagates_error() {
260 let doc = json!({ "a": 1 });
261 assert!(construct_id("{body.missing}", "p", &doc).is_err());
262 }
263
264 #[test]
265 fn logical_to_physical_substitutes_natural_key() {
266 assert_eq!(
267 map_logical_to_physical("{partition}:{body.id}", "acme", "7").unwrap(),
268 "acme:7"
269 );
270 assert_eq!(
271 map_logical_to_physical("doc-{body.k}@{partition}", "p1", "abc").unwrap(),
272 "doc-abc@p1"
273 );
274 }
275
276 #[test]
277 fn physical_to_logical_strips_the_frame() {
278 assert_eq!(
279 map_physical_to_logical("{partition}:{body.id}", "acme", "acme:7").unwrap(),
280 Some("7".to_owned())
281 );
282 assert_eq!(
284 map_physical_to_logical("{partition}:{body.id}", "acme", "other:7").unwrap(),
285 None
286 );
287 }
288
289 #[test]
290 fn mapping_round_trips_for_arbitrary_natural_keys() {
291 let template = "{partition}:{body.natural}";
292 for key in ["1001", "a-b", "", "x:y"] {
293 let physical = map_logical_to_physical(template, "acme", key).unwrap();
294 assert_eq!(
295 map_physical_to_logical(template, "acme", &physical).unwrap(),
296 Some(key.to_owned())
297 );
298 }
299 }
300
301 #[test]
302 fn templates_without_exactly_one_body_placeholder_are_irreversible() {
303 assert_eq!(
304 map_logical_to_physical("{partition}:static", "p", "x").unwrap_err(),
305 RewriteError::IrreversibleIdTemplate
306 );
307 assert_eq!(
308 map_logical_to_physical("{body.a}-{body.b}", "p", "x").unwrap_err(),
309 RewriteError::IrreversibleIdTemplate
310 );
311 }
312}