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`, `/`.
69pub fn percent_encode(s: &str, keep_slash: bool) -> String {
70    let mut out = String::with_capacity(s.len());
71    for b in s.bytes() {
72        match b {
73            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => out.push(b as char),
74            b'/' if keep_slash => out.push('/'),
75            b => out.push_str(&format!("%{b:02X}")),
76        }
77    }
78    out
79}
80
81/// One line, at most `max` characters, for quoting SQL in messages.
82pub fn summarize(text: &str, max: usize) -> String {
83    let one = text.split_whitespace().collect::<Vec<_>>().join(" ");
84    if one.chars().count() > max {
85        format!("{}…", one.chars().take(max).collect::<String>())
86    } else {
87        one
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn scales_decimals_exactly() {
97        assert_eq!(scaled_decimal("12.34", 2), Ok(1234));
98        assert_eq!(scaled_decimal("-0.5", 3), Ok(-500));
99        assert_eq!(scaled_decimal("7", 2), Ok(700));
100        assert_eq!(scaled_decimal("12.3400", 2), Ok(1234));
101        assert_eq!(scaled_decimal("1200", -2), Ok(12));
102        assert_eq!(scaled_decimal("0.000", 0), Ok(0));
103    }
104
105    #[test]
106    fn refuses_to_drop_significant_digits() {
107        assert!(scaled_decimal("12.345", 2).is_err());
108        assert!(scaled_decimal("1250", -2).is_err());
109        assert!(scaled_decimal("abc", 2).is_err());
110        assert!(scaled_decimal("", 2).is_err());
111    }
112
113    #[test]
114    fn parses_cell_references() {
115        assert_eq!(parse_cell("A1"), Some((0, 0)));
116        assert_eq!(parse_cell("$B$12"), Some((11, 1)));
117        assert_eq!(parse_cell("b12"), Some((11, 1)));
118        assert_eq!(parse_cell("XFD1048576"), Some((1_048_575, 16_383)));
119        for bad in ["XFE1", "A0", "1A", "B$1$2", "$$A1", "A", "12", "A1B"] {
120            assert_eq!(parse_cell(bad), None, "{bad}");
121        }
122    }
123
124    #[test]
125    fn percent_encodes() {
126        assert_eq!(percent_encode("a b/c.csv", true), "a%20b/c.csv");
127        assert_eq!(percent_encode("a b/c", false), "a%20b%2Fc");
128    }
129}
130
131/// Write `bytes` to `path` so that a reader, or a crash, sees either the old file or the new one,
132/// never a part: write a temporary file (`.<name>.<pid>.<random>.tmp`) in the same folder, flush
133/// it to disk, rename it over `path`, then flush the folder (Unix; a folder that can't be
134/// flushed, such as some network mounts, is logged at debug and ignored). For DRE's durable
135/// state (`run_results.json`, the drift snapshot, the manifest, `dre.lock`, sessions), not report
136/// outputs.
137pub fn write_atomic(path: &std::path::Path, bytes: &[u8]) -> std::io::Result<()> {
138    write_atomic_mode(path, bytes, None)
139}
140
141/// [`write_atomic`], creating the file with Unix permissions `mode` (e.g. `0o600` for secrets).
142pub fn write_atomic_mode(path: &std::path::Path, bytes: &[u8], mode: Option<u32>) -> std::io::Result<()> {
143    use std::io::Write;
144    let dir = match path.parent() {
145        Some(d) if !d.as_os_str().is_empty() => d.to_path_buf(),
146        _ => std::path::PathBuf::from("."),
147    };
148    let name = path
149        .file_name()
150        .ok_or_else(|| std::io::Error::other(format!("{} has no file name", path.display())))?
151        .to_string_lossy();
152    let tmp = dir.join(format!(".{name}.{}.{}.tmp", std::process::id(), random_suffix()));
153    let write = || -> std::io::Result<()> {
154        let mut opts = std::fs::OpenOptions::new();
155        opts.write(true).create_new(true);
156        #[cfg(unix)]
157        if let Some(m) = mode {
158            use std::os::unix::fs::OpenOptionsExt;
159            opts.mode(m);
160        }
161        #[cfg(not(unix))]
162        let _ = mode;
163        let mut f = opts.open(&tmp)?;
164        f.write_all(bytes)?;
165        f.sync_all()?;
166        drop(f);
167        std::fs::rename(&tmp, path)
168    };
169    if let Err(e) = write() {
170        let _ = std::fs::remove_file(&tmp);
171        return Err(e);
172    }
173    #[cfg(unix)]
174    if let Err(e) = std::fs::File::open(&dir).and_then(|d| d.sync_all()) {
175        crate::log::debug!("can't flush the folder {}: {e}", dir.display());
176    }
177    Ok(())
178}
179
180/// Four characters that differ between calls, for temporary file names.
181fn random_suffix() -> String {
182    use std::sync::atomic::{AtomicU32, Ordering};
183    static COUNTER: AtomicU32 = AtomicU32::new(0);
184    let nanos = std::time::SystemTime::now()
185        .duration_since(std::time::UNIX_EPOCH)
186        .map(|d| d.subsec_nanos())
187        .unwrap_or(0);
188    let mut n = nanos
189        ^ COUNTER
190            .fetch_add(1, Ordering::Relaxed)
191            .wrapping_mul(2_654_435_761);
192    const ALPHABET: &[u8] = b"abcdefghijklmnopqrstuvwxyz234567";
193    (0..4)
194        .map(|_| {
195            let c = ALPHABET[(n % 32) as usize] as char;
196            n /= 32;
197            c
198        })
199        .collect()
200}
201
202#[cfg(test)]
203mod atomic_tests {
204    use super::*;
205
206    #[test]
207    fn writes_replace_whole_files_and_leave_no_temporary_one() {
208        let dir = tempfile::tempdir().unwrap();
209        let p = dir.path().join("run_results.json");
210        write_atomic(&p, b"{\"a\":1}").unwrap();
211        write_atomic(&p, b"{\"a\":2}").unwrap();
212        assert_eq!(std::fs::read_to_string(&p).unwrap(), "{\"a\":2}");
213        let names: Vec<_> = std::fs::read_dir(dir.path())
214            .unwrap()
215            .map(|e| e.unwrap().file_name())
216            .collect();
217        assert_eq!(names, ["run_results.json"]);
218    }
219
220    #[test]
221    fn a_failed_write_leaves_the_old_file() {
222        let dir = tempfile::tempdir().unwrap();
223        let p = dir.path().join("x.json");
224        std::fs::write(&p, "old").unwrap();
225        // The target is a folder: the rename fails, the temporary file goes.
226        let folder = dir.path().join("sub");
227        std::fs::create_dir(&folder).unwrap();
228        std::fs::write(folder.join("keep"), "").unwrap();
229        assert!(write_atomic(&folder, b"new").is_err());
230        assert_eq!(std::fs::read_to_string(&p).unwrap(), "old");
231        let tmp = std::fs::read_dir(dir.path())
232            .unwrap()
233            .filter(|e| {
234                e.as_ref()
235                    .unwrap()
236                    .file_name()
237                    .to_string_lossy()
238                    .ends_with(".tmp")
239            })
240            .count();
241        assert_eq!(tmp, 0);
242    }
243
244    #[cfg(unix)]
245    #[test]
246    fn a_mode_is_applied() {
247        use std::os::unix::fs::PermissionsExt;
248        let dir = tempfile::tempdir().unwrap();
249        let p = dir.path().join("s.json");
250        write_atomic_mode(&p, b"{}", Some(0o600)).unwrap();
251        assert_eq!(std::fs::metadata(&p).unwrap().permissions().mode() & 0o777, 0o600);
252    }
253
254    /// A process killed while writing leaves the old file or a new one, whole.
255    #[test]
256    fn a_killed_writer_never_leaves_a_partial_file() {
257        if let Some(path) = std::env::var_os("DRE_ATOMIC_CHILD") {
258            // The child: write big files forever.
259            let path = std::path::PathBuf::from(path);
260            let mut i = 0u64;
261            loop {
262                let body = format!("{{\"n\":{i},\"pad\":\"{}\"}}", "x".repeat(1 << 20));
263                write_atomic(&path, body.as_bytes()).unwrap();
264                i += 1;
265            }
266        }
267        let dir = tempfile::tempdir().unwrap();
268        let path = dir.path().join("state.json");
269        write_atomic(&path, b"{\"n\":-1}").unwrap();
270        for _ in 0..5 {
271            let mut child = std::process::Command::new(std::env::current_exe().unwrap())
272                .args([
273                    "--exact",
274                    "util::atomic_tests::a_killed_writer_never_leaves_a_partial_file",
275                    "--nocapture",
276                ])
277                .env("DRE_ATOMIC_CHILD", &path)
278                .stdout(std::process::Stdio::null())
279                .stderr(std::process::Stdio::null())
280                .spawn()
281                .unwrap();
282            std::thread::sleep(std::time::Duration::from_millis(150));
283            child.kill().unwrap();
284            child.wait().unwrap();
285            let text = std::fs::read_to_string(&path).unwrap();
286            let v: serde_json::Value = serde_json::from_str(&text).expect("a whole JSON file");
287            assert!(v["n"].is_i64());
288        }
289    }
290
291    #[test]
292    fn suffixes_differ() {
293        assert_ne!(random_suffix(), random_suffix());
294    }
295}
dre_protocol/util.rs

dre_protocol/
util.rs
