cloudiful_redactor/session/
restore.rs

1use std::collections::{BTreeSet, HashMap};
2
3use anyhow::{Result, anyhow};
4
5use crate::types::{RedactionSession, RestoreResult};
6
7pub fn restore_text_with_session(input: &str, session: &RedactionSession) -> RestoreResult {
8    let known_tokens = session
9        .entries
10        .iter()
11        .map(|entry| entry.token.clone())
12        .collect::<BTreeSet<_>>();
13    let token_map = session
14        .entries
15        .iter()
16        .map(|entry| (entry.token.as_str(), entry.original.as_str()))
17        .collect::<HashMap<_, _>>();
18
19    let mut restored_text = String::with_capacity(input.len());
20    let mut restored_count = 0;
21    let mut validation_errors = Vec::new();
22    let mut cursor = 0;
23
24    for token_range in token_like_ranges(input) {
25        restored_text.push_str(&input[cursor..token_range.start]);
26        let candidate = &input[token_range.clone()];
27        if !known_tokens.contains(candidate) {
28            validation_errors.push(format!("unknown or malformed token `{candidate}`"));
29            restored_text.push_str(candidate);
30        } else if let Some(original) = token_map.get(candidate) {
31            restored_text.push_str(original);
32            restored_count += 1;
33        } else {
34            restored_text.push_str(candidate);
35        }
36        cursor = token_range.end;
37    }
38    restored_text.push_str(&input[cursor..]);
39
40    let unresolved_tokens = token_like_ranges(&restored_text)
41        .into_iter()
42        .map(|range| restored_text[range].to_string())
43        .filter(|candidate| candidate.starts_with("__R_"))
44        .collect::<Vec<_>>();
45
46    if !unresolved_tokens.is_empty() {
47        validation_errors.extend(
48            unresolved_tokens
49                .iter()
50                .map(|candidate| format!("unresolved token remained after restore: `{candidate}`")),
51        );
52    }
53
54    RestoreResult {
55        restored_text,
56        restored_count,
57        unresolved_tokens,
58        validation_errors,
59    }
60}
61
62pub fn restore_patch_with_session(patch: &str, session: &RedactionSession) -> RestoreResult {
63    restore_text_with_session(patch, session)
64}
65
66pub fn ensure_restore_valid(result: &RestoreResult) -> Result<()> {
67    if result.is_valid() {
68        return Ok(());
69    }
70
71    let mut messages = Vec::new();
72    if !result.validation_errors.is_empty() {
73        messages.extend(result.validation_errors.clone());
74    }
75    if !result.unresolved_tokens.is_empty() {
76        messages.push(format!(
77            "unresolved tokens: {}",
78            result.unresolved_tokens.join(", ")
79        ));
80    }
81    Err(anyhow!(messages.join("; ")))
82}
83
84fn token_like_ranges(text: &str) -> Vec<std::ops::Range<usize>> {
85    let mut ranges = Vec::new();
86    let bytes = text.as_bytes();
87    let mut index = 0;
88
89    while index + 4 <= bytes.len() {
90        if &bytes[index..index + 4] != b"__R_" {
91            index += 1;
92            continue;
93        }
94
95        let mut end = index + 4;
96        while end < bytes.len() {
97            let byte = bytes[end];
98            if byte.is_ascii_uppercase() || byte.is_ascii_digit() || byte == b'_' {
99                end += 1;
100                continue;
101            }
102            break;
103        }
104
105        if end + 1 < bytes.len() && bytes[end] == b'_' && bytes[end + 1] == b'_' {
106            end += 2;
107        }
108
109        ranges.push(index..end);
110        index = end.max(index + 1);
111    }
112
113    ranges
114}
115
116#[cfg(test)]
117mod tests {
118    use super::restore_text_with_session;
119    use crate::RedactorBuilder;
120
121    #[test]
122    fn restore_streams_multiple_tokens_and_repetitions() {
123        let redactor = RedactorBuilder::new().build();
124        let text = "host=service.example.com alt=service.example.com";
125        let session = redactor.redact_with_session(text).expect("session");
126
127        let restored = restore_text_with_session(&session.redacted_text, &session);
128
129        assert!(restored.is_valid());
130        assert_eq!(restored.restored_text, text);
131        assert_eq!(restored.restored_count, 2);
132    }
133
134    #[test]
135    fn restore_preserves_unknown_token_validation() {
136        let redactor = RedactorBuilder::new().build();
137        let session = redactor
138            .redact_with_session("host=service.example.com")
139            .expect("session");
140        let restored = restore_text_with_session("__R_DOMAIN_001__ __R_DOMAIN_999__", &session);
141
142        assert!(
143            restored
144                .validation_errors
145                .iter()
146                .any(|message| message.contains("unknown or malformed token `__R_DOMAIN_999__`"))
147        );
148        assert_eq!(restored.unresolved_tokens, vec!["__R_DOMAIN_999__"]);
149    }
150}
cloudiful_redactor/session/restore.rs

cloudiful_redactor/session/
restore.rs
