cloudiful_redactor/session/
restore.rs

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

cloudiful_redactor/session/
restore.rs
