vault_audit_tools/utils/
reader.rs1use anyhow::{Context, Result};
29use flate2::read::GzDecoder;
30use std::fs::File;
31use std::io::Read;
32use std::path::Path;
33
34pub fn open_file(path: impl AsRef<Path>) -> Result<Box<dyn Read + Send>> {
60 let path = path.as_ref();
61 let file =
62 File::open(path).with_context(|| format!("Failed to open file: {}", path.display()))?;
63
64 let extension = path.extension().and_then(|e| e.to_str()).unwrap_or("");
65
66 match extension {
67 "gz" => {
68 let decoder = GzDecoder::new(file);
69 Ok(Box::new(decoder))
70 }
71 "zst" => {
72 let decoder = zstd::Decoder::new(file).with_context(|| {
73 format!("Failed to create zstd decoder for: {}", path.display())
74 })?;
75 Ok(Box::new(decoder))
76 }
77 _ => Ok(Box::new(file)),
78 }
79}
80
81#[cfg(test)]
82mod tests {
83 use super::*;
84 use std::io::{BufRead, BufReader, Write};
85 use tempfile::NamedTempFile;
86
87 #[test]
88 fn test_plain_file() {
89 let mut temp = NamedTempFile::new().unwrap();
90 writeln!(temp, "test line 1").unwrap();
91 writeln!(temp, "test line 2").unwrap();
92 temp.flush().unwrap();
93
94 let reader = open_file(temp.path()).unwrap();
95 let buf_reader = BufReader::new(reader);
96 let lines: Vec<String> = buf_reader.lines().collect::<Result<_, _>>().unwrap();
97
98 assert_eq!(lines.len(), 2);
99 assert_eq!(lines[0], "test line 1");
100 assert_eq!(lines[1], "test line 2");
101 }
102
103 #[test]
104 fn test_gzip_file() {
105 use flate2::write::GzEncoder;
106 use flate2::Compression;
107
108 let mut temp = NamedTempFile::with_suffix(".gz").unwrap();
109 {
110 let mut encoder = GzEncoder::new(&mut temp, Compression::default());
111 writeln!(encoder, "compressed line 1").unwrap();
112 writeln!(encoder, "compressed line 2").unwrap();
113 encoder.finish().unwrap();
114 }
115 temp.flush().unwrap();
116
117 let reader = open_file(temp.path()).unwrap();
118 let buf_reader = BufReader::new(reader);
119 let lines: Vec<String> = buf_reader.lines().collect::<Result<_, _>>().unwrap();
120
121 assert_eq!(lines.len(), 2);
122 assert_eq!(lines[0], "compressed line 1");
123 assert_eq!(lines[1], "compressed line 2");
124 }
125
126 #[test]
127 fn test_zstd_file() {
128 let mut temp = NamedTempFile::with_suffix(".zst").unwrap();
129 {
130 let mut encoder = zstd::Encoder::new(&mut temp, 3).unwrap();
131 writeln!(encoder, "zstd line 1").unwrap();
132 writeln!(encoder, "zstd line 2").unwrap();
133 encoder.finish().unwrap();
134 }
135 temp.flush().unwrap();
136
137 let reader = open_file(temp.path()).unwrap();
138 let buf_reader = BufReader::new(reader);
139 let lines: Vec<String> = buf_reader.lines().collect::<Result<_, _>>().unwrap();
140
141 assert_eq!(lines.len(), 2);
142 assert_eq!(lines[0], "zstd line 1");
143 assert_eq!(lines[1], "zstd line 2");
144 }
145}