dre_protocol/
util.rs

1//! Small helpers shared by core and first-party plugins, so each rule has one implementation.
2
3/// Parse decimal text (`"-12.30"`) into an integer scaled by `10^scale` for an Arrow decimal.
4/// Negative scales are allowed (`"1200"` at scale -2 is `12`). Digits that the scale would
5/// drop must be zero: a value that doesn't fit the scale is an error, never silently cut.
6pub fn scaled_decimal(text: &str, scale: i8) -> Result<i128, String> {
7    let bad = || format!("`{text}` isn't a decimal that fits scale {scale}");
8    let t = text.trim();
9    let (neg, t) = match t.strip_prefix('-') {
10        Some(r) => (true, r),
11        None => (false, t.strip_prefix('+').unwrap_or(t)),
12    };
13    let (int, frac) = t.split_once('.').unwrap_or((t, ""));
14    if int.is_empty() && frac.is_empty() || !int.bytes().chain(frac.bytes()).all(|b| b.is_ascii_digit()) {
15        return Err(bad());
16    }
17    // All digits, and the power of ten of the last one.
18    let mut digits = format!("{int}{frac}");
19    let mut exp = -(frac.len() as i32);
20    let target = -i32::from(scale);
21    // Drop trailing digits below the scale; they must be zeros.
22    while exp < target {
23        match digits.pop() {
24            Some('0') => exp += 1,
25            Some(_) => return Err(bad()),
26            None => break,
27        }
28    }
29    while exp > target {
30        digits.push('0');
31        exp -= 1;
32    }
33    let digits = digits.trim_start_matches('0');
34    let v: i128 = if digits.is_empty() {
35        0
36    } else {
37        digits.parse().map_err(|_| bad())?
38    };
39    Ok(if neg { -v } else { v })
40}
41
42/// Parse a cell reference (`B12`, `$B$12`) into zero-based `(row, col)`. `$` may only appear
43/// directly before the column letters and before the row number.
44pub fn parse_cell(s: &str) -> Option<(u32, u16)> {
45    let s = s.trim();
46    let s = s.strip_prefix('$').unwrap_or(s);
47    let split = s.find(|c: char| !c.is_ascii_alphabetic())?;
48    let (letters, rest) = s.split_at(split);
49    let digits = rest.strip_prefix('$').unwrap_or(rest);
50    if letters.is_empty()
51        || letters.len() > 3
52        || digits.is_empty()
53        || !digits.bytes().all(|b| b.is_ascii_digit())
54    {
55        return None;
56    }
57    let col = letters
58        .to_ascii_uppercase()
59        .bytes()
60        .fold(0u32, |acc, b| acc * 26 + u32::from(b - b'A' + 1));
61    let row: u32 = digits.parse().ok()?;
62    if col == 0 || col > 16_384 || row == 0 || row > 1_048_576 {
63        return None;
64    }
65    Some((row - 1, (col - 1) as u16))
66}
67
68/// Percent-encode everything except RFC 3986 unreserved characters and, if `keep_slash`, `/`.
69#[must_use]
70pub fn percent_encode(s: &str, keep_slash: bool) -> String {
71    let mut out = String::with_capacity(s.len());
72    for b in s.bytes() {
73        match b {
74            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => out.push(b as char),
75            b'/' if keep_slash => out.push('/'),
76            b => out.push_str(&format!("%{b:02X}")),
77        }
78    }
79    out
80}
81
82/// One line, at most `max` characters, for quoting SQL in messages.
83#[must_use]
84pub fn summarize(text: &str, max: usize) -> String {
85    let one = text.split_whitespace().collect::<Vec<_>>().join(" ");
86    if one.chars().count() > max {
87        format!("{}…", one.chars().take(max).collect::<String>())
88    } else {
89        one
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn scales_decimals_exactly() {
99        assert_eq!(scaled_decimal("12.34", 2), Ok(1234));
100        assert_eq!(scaled_decimal("-0.5", 3), Ok(-500));
101        assert_eq!(scaled_decimal("7", 2), Ok(700));
102        assert_eq!(scaled_decimal("12.3400", 2), Ok(1234));
103        assert_eq!(scaled_decimal("1200", -2), Ok(12));
104        assert_eq!(scaled_decimal("0.000", 0), Ok(0));
105    }
106
107    #[test]
108    fn refuses_to_drop_significant_digits() {
109        assert!(scaled_decimal("12.345", 2).is_err());
110        assert!(scaled_decimal("1250", -2).is_err());
111        assert!(scaled_decimal("abc", 2).is_err());
112        assert!(scaled_decimal("", 2).is_err());
113    }
114
115    #[test]
116    fn parses_cell_references() {
117        assert_eq!(parse_cell("A1"), Some((0, 0)));
118        assert_eq!(parse_cell("$B$12"), Some((11, 1)));
119        assert_eq!(parse_cell("b12"), Some((11, 1)));
120        assert_eq!(parse_cell("XFD1048576"), Some((1_048_575, 16_383)));
121        for bad in ["XFE1", "A0", "1A", "B$1$2", "$$A1", "A", "12", "A1B"] {
122            assert_eq!(parse_cell(bad), None, "{bad}");
123        }
124    }
125
126    #[test]
127    fn percent_encodes() {
128        assert_eq!(percent_encode("a b/c.csv", true), "a%20b/c.csv");
129        assert_eq!(percent_encode("a b/c", false), "a%20b%2Fc");
130    }
131}
132
133/// Write `bytes` to `path` so that a reader, or a crash, sees either the old file or the new one,
134/// never a part: write a temporary file (`.<name>.<pid>.<random>.tmp`) in the same folder, flush
135/// it to disk, rename it over `path`, then flush the folder (Unix; a folder that can't be
136/// flushed, such as some network mounts, is logged at debug and ignored). For DRE's durable
137/// state (`run_results.json`, the drift snapshot, the manifest, `dre.lock`, sessions), not report
138/// outputs.
139pub fn write_atomic(path: &std::path::Path, bytes: &[u8]) -> std::io::Result<()> {
140    write_atomic_mode(path, bytes, None)
141}
142
143/// [`write_atomic`], creating the file with Unix permissions `mode` (e.g. `0o600` for secrets).
144pub fn write_atomic_mode(path: &std::path::Path, bytes: &[u8], mode: Option<u32>) -> std::io::Result<()> {
145    use std::io::Write;
146    let dir = match path.parent() {
147        Some(d) if !d.as_os_str().is_empty() => d.to_path_buf(),
148        _ => std::path::PathBuf::from("."),
149    };
150    let name = path
151        .file_name()
152        .ok_or_else(|| std::io::Error::other(format!("{} has no file name", path.display())))?
153        .to_string_lossy();
154    let tmp = dir.join(format!(".{name}.{}.{}.tmp", std::process::id(), random_suffix()));
155    let write = || -> std::io::Result<()> {
156        let mut opts = std::fs::OpenOptions::new();
157        opts.write(true).create_new(true);
158        #[cfg(unix)]
159        if let Some(m) = mode {
160            use std::os::unix::fs::OpenOptionsExt;
161            opts.mode(m);
162        }
163        #[cfg(not(unix))]
164        let _ = mode;
165        let mut f = opts.open(&tmp)?;
166        f.write_all(bytes)?;
167        f.sync_all()?;
168        drop(f);
169        std::fs::rename(&tmp, path)
170    };
171    if let Err(e) = write() {
172        let _ = std::fs::remove_file(&tmp);
173        return Err(e);
174    }
175    #[cfg(unix)]
176    if let Err(e) = std::fs::File::open(&dir).and_then(|d| d.sync_all()) {
177        crate::log::debug!("can't flush the folder {}: {e}", dir.display());
178    }
179    Ok(())
180}
181
182/// Four characters that differ between calls, for temporary file names.
183fn random_suffix() -> String {
184    use std::sync::atomic::{AtomicU32, Ordering};
185    static COUNTER: AtomicU32 = AtomicU32::new(0);
186    let nanos = std::time::SystemTime::now()
187        .duration_since(std::time::UNIX_EPOCH)
188        .map(|d| d.subsec_nanos())
189        .unwrap_or(0);
190    let mut n = nanos
191        ^ COUNTER
192            .fetch_add(1, Ordering::Relaxed)
193            .wrapping_mul(2_654_435_761);
194    const ALPHABET: &[u8] = b"abcdefghijklmnopqrstuvwxyz234567";
195    (0..4)
196        .map(|_| {
197            let c = ALPHABET[(n % 32) as usize] as char;
198            n /= 32;
199            c
200        })
201        .collect()
202}
203
204#[cfg(test)]
205mod atomic_tests {
206    use super::*;
207
208    #[test]
209    fn writes_replace_whole_files_and_leave_no_temporary_one() {
210        let dir = tempfile::tempdir().unwrap();
211        let p = dir.path().join("run_results.json");
212        write_atomic(&p, b"{\"a\":1}").unwrap();
213        write_atomic(&p, b"{\"a\":2}").unwrap();
214        assert_eq!(std::fs::read_to_string(&p).unwrap(), "{\"a\":2}");
215        let names: Vec<_> = std::fs::read_dir(dir.path())
216            .unwrap()
217            .map(|e| e.unwrap().file_name())
218            .collect();
219        assert_eq!(names, ["run_results.json"]);
220    }
221
222    #[test]
223    fn a_failed_write_leaves_the_old_file() {
224        let dir = tempfile::tempdir().unwrap();
225        let p = dir.path().join("x.json");
226        std::fs::write(&p, "old").unwrap();
227        // The target is a folder: the rename fails, the temporary file goes.
228        let folder = dir.path().join("sub");
229        std::fs::create_dir(&folder).unwrap();
230        std::fs::write(folder.join("keep"), "").unwrap();
231        assert!(write_atomic(&folder, b"new").is_err());
232        assert_eq!(std::fs::read_to_string(&p).unwrap(), "old");
233        let tmp = std::fs::read_dir(dir.path())
234            .unwrap()
235            .filter(|e| {
236                e.as_ref()
237                    .unwrap()
238                    .file_name()
239                    .to_string_lossy()
240                    .ends_with(".tmp")
241            })
242            .count();
243        assert_eq!(tmp, 0);
244    }
245
246    #[cfg(unix)]
247    #[test]
248    fn a_mode_is_applied() {
249        use std::os::unix::fs::PermissionsExt;
250        let dir = tempfile::tempdir().unwrap();
251        let p = dir.path().join("s.json");
252        write_atomic_mode(&p, b"{}", Some(0o600)).unwrap();
253        assert_eq!(std::fs::metadata(&p).unwrap().permissions().mode() & 0o777, 0o600);
254    }
255
256    /// A process killed while writing leaves the old file or a new one, whole.
257    #[test]
258    fn a_killed_writer_never_leaves_a_partial_file() {
259        if let Some(path) = std::env::var_os("DRE_ATOMIC_CHILD") {
260            // The child: write big files forever.
261            let path = std::path::PathBuf::from(path);
262            let mut i = 0u64;
263            loop {
264                let body = format!("{{\"n\":{i},\"pad\":\"{}\"}}", "x".repeat(1 << 20));
265                write_atomic(&path, body.as_bytes()).unwrap();
266                i += 1;
267            }
268        }
269        let dir = tempfile::tempdir().unwrap();
270        let path = dir.path().join("state.json");
271        write_atomic(&path, b"{\"n\":-1}").unwrap();
272        for _ in 0..5 {
273            let mut child = std::process::Command::new(std::env::current_exe().unwrap())
274                .args([
275                    "--exact",
276                    "util::atomic_tests::a_killed_writer_never_leaves_a_partial_file",
277                    "--nocapture",
278                ])
279                .env("DRE_ATOMIC_CHILD", &path)
280                .stdout(std::process::Stdio::null())
281                .stderr(std::process::Stdio::null())
282                .spawn()
283                .unwrap();
284            std::thread::sleep(std::time::Duration::from_millis(150));
285            child.kill().unwrap();
286            child.wait().unwrap();
287            let text = std::fs::read_to_string(&path).unwrap();
288            let v: serde_json::Value = serde_json::from_str(&text).expect("a whole JSON file");
289            assert!(v["n"].is_i64());
290        }
291    }
292
293    #[test]
294    fn suffixes_differ() {
295        assert_ne!(random_suffix(), random_suffix());
296    }
297}
dre_protocol/util.rs

dre_protocol/
util.rs
