1use crate::primitives::{hardlink, rw};
2use crate::{Error, Result};
3use std::path::{Component, Path, PathBuf};
4
5pub const DEFAULT_COPY_ONLY_THRESHOLD_BYTES: u64 = 4 * 1024 * 1024;
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct WorkspaceReport {
9 pub staging_root: PathBuf,
10 pub destination_root: PathBuf,
11 pub file_count: usize,
12 pub directory_count: usize,
13 pub total_bytes: u64,
14}
15
16pub struct Workspace {
17 staging: PathBuf,
18 dest: PathBuf,
19 committed: bool,
20}
21
22impl Workspace {
23 pub fn new(staging_dir: impl AsRef<Path>, dest_dir: impl AsRef<Path>) -> Result<Self> {
24 let staging_path = staging_dir.as_ref().to_path_buf();
25 let destination_path = dest_dir.as_ref().to_path_buf();
26
27 if !staging_path.exists() {
28 std::fs::create_dir_all(&staging_path).map_err(|e| Error::Write {
29 path: staging_path.clone(),
30 source: e,
31 })?;
32 }
33
34 Ok(Self {
35 staging: staging_path,
36 dest: destination_path,
37 committed: false,
38 })
39 }
40
41 pub fn path(&self) -> &Path {
42 &self.staging
43 }
44
45 pub fn staging_path(&self) -> &Path {
46 &self.staging
47 }
48
49 pub fn destination_path(&self) -> &Path {
50 &self.dest
51 }
52
53 pub fn exists(&self, relative_path: impl AsRef<Path>) -> Result<bool> {
54 Ok(self.resolve(relative_path)?.exists())
55 }
56
57 pub fn create_dir(&self, relative_path: impl AsRef<Path>) -> Result<()> {
58 let path = self.resolve(relative_path)?;
59 if let Some(parent) = path.parent() {
60 std::fs::create_dir_all(parent).map_err(|e| Error::Write {
61 path: parent.to_path_buf(),
62 source: e,
63 })?;
64 }
65 std::fs::create_dir(&path).map_err(|e| Error::Write { path, source: e })
66 }
67
68 pub fn create_dir_all(&self, relative_path: impl AsRef<Path>) -> Result<()> {
69 let path = self.resolve(relative_path)?;
70 std::fs::create_dir_all(&path).map_err(|e| Error::Write { path, source: e })
71 }
72
73 pub fn write(&self, relative_path: impl AsRef<Path>, content: &[u8]) -> Result<()> {
74 self.write_with_options(relative_path, content, rw::Options::default())
75 }
76
77 pub fn write_with_options(
78 &self,
79 relative_path: impl AsRef<Path>,
80 content: &[u8],
81 options: rw::Options,
82 ) -> Result<()> {
83 let path = self.resolve(relative_path)?;
84 if let Some(parent) = path.parent() {
85 std::fs::create_dir_all(parent).map_err(|e| Error::Write {
86 path: parent.to_path_buf(),
87 source: e,
88 })?;
89 }
90 rw::atomic_write(path, content, options)
91 }
92
93 pub fn copy_file(
94 &self,
95 source: impl AsRef<Path>,
96 relative_path: impl AsRef<Path>,
97 ) -> Result<u64> {
98 let source = source.as_ref();
99 let path = self.resolve(relative_path)?;
100 if let Some(parent) = path.parent() {
101 std::fs::create_dir_all(parent).map_err(|e| Error::Write {
102 path: parent.to_path_buf(),
103 source: e,
104 })?;
105 }
106 std::fs::copy(source, &path).map_err(|e| Error::Write { path, source: e })
107 }
108
109 pub fn link_or_copy_file(
110 &self,
111 source: impl AsRef<Path>,
112 relative_path: impl AsRef<Path>,
113 options: hardlink::Options,
114 ) -> Result<()> {
115 let source = source.as_ref();
116 let path = self.resolve(relative_path)?;
117 if let Some(parent) = path.parent() {
118 std::fs::create_dir_all(parent).map_err(|e| Error::Write {
119 path: parent.to_path_buf(),
120 source: e,
121 })?;
122 }
123 hardlink::hardlink_or_copy(source, &path, options)
124 }
125
126 pub fn stage_file_by_size(
127 &self,
128 source: impl AsRef<Path>,
129 relative_path: impl AsRef<Path>,
130 threshold_bytes: u64,
131 options: hardlink::Options,
132 ) -> Result<()> {
133 let source = source.as_ref();
134 if should_copy_only(source, threshold_bytes)? {
135 let _ = self.copy_file(source, relative_path)?;
136 } else {
137 self.link_or_copy_file(source, relative_path, options)?;
138 }
139 Ok(())
140 }
141
142 pub fn read(&self, relative_path: impl AsRef<Path>) -> Result<Vec<u8>> {
143 let path = self.resolve(relative_path)?;
144 rw::atomic_read(path)
145 }
146
147 pub fn report(&self) -> Result<WorkspaceReport> {
148 let mut report = WorkspaceReport {
149 staging_root: self.staging.clone(),
150 destination_root: self.dest.clone(),
151 file_count: 0,
152 directory_count: 0,
153 total_bytes: 0,
154 };
155
156 if self.staging.exists() {
157 self.walk(&self.staging, &mut report)?;
158 }
159
160 Ok(report)
161 }
162
163 pub fn commit(mut self) -> Result<()> {
164 crate::primitives::replace_dir::replace_dir(&self.staging, &self.dest, Default::default())?;
165 self.committed = true;
166 Ok(())
167 }
168
169 fn resolve(&self, relative_path: impl AsRef<Path>) -> Result<PathBuf> {
170 let relative_path = relative_path.as_ref();
171
172 if relative_path.as_os_str().is_empty() {
173 return Err(Error::InvalidInput(
174 "workspace path must not be empty".to_string(),
175 ));
176 }
177
178 let mut sanitized = PathBuf::new();
179 for component in relative_path.components() {
180 match component {
181 Component::Normal(part) => sanitized.push(part),
182 Component::CurDir => {}
183 Component::ParentDir => {
184 return Err(Error::InvalidInput(format!(
185 "workspace path escapes staging root: {}",
186 relative_path.display()
187 )));
188 }
189 Component::RootDir | Component::Prefix(_) => {
190 return Err(Error::InvalidInput(format!(
191 "workspace path must be relative: {}",
192 relative_path.display()
193 )));
194 }
195 }
196 }
197
198 if sanitized.as_os_str().is_empty() {
199 return Err(Error::InvalidInput(format!(
200 "workspace path must contain a normal component: {}",
201 relative_path.display()
202 )));
203 }
204
205 Ok(self.staging.join(sanitized))
206 }
207
208 fn walk(&self, path: &Path, report: &mut WorkspaceReport) -> Result<()> {
209 for entry in std::fs::read_dir(path).map_err(|e| Error::Read {
210 path: path.to_path_buf(),
211 source: e,
212 })? {
213 let entry = entry.map_err(|e| Error::Read {
214 path: path.to_path_buf(),
215 source: e,
216 })?;
217 let entry_path = entry.path();
218 let metadata = entry.metadata().map_err(|e| Error::Read {
219 path: entry_path.clone(),
220 source: e,
221 })?;
222
223 if metadata.is_dir() {
224 report.directory_count += 1;
225 self.walk(&entry_path, report)?;
226 } else {
227 report.file_count += 1;
228 report.total_bytes += metadata.len();
229 }
230 }
231
232 Ok(())
233 }
234}
235
236impl Drop for Workspace {
237 fn drop(&mut self) {
238 if !self.committed {
239 let _ = std::fs::remove_dir_all(&self.staging);
240 }
241 }
242}
243
244pub fn should_copy_only(source: &Path, threshold_bytes: u64) -> Result<bool> {
245 Ok(std::fs::metadata(source)
246 .map_err(|source_error| Error::Read {
247 path: source.to_path_buf(),
248 source: source_error,
249 })?
250 .len()
251 < threshold_bytes)
252}
253
254#[cfg(test)]
255mod tests {
256 use super::*;
257 use tempfile::tempdir;
258
259 #[test]
260 fn test_workspace() {
261 let dir = tempdir().unwrap();
262 let staging = dir.path().join("staging");
263 let dest = dir.path().join("dest");
264 let workspace = Workspace::new(&staging, &dest).unwrap();
265 workspace.write("file.txt", b"data").unwrap();
266 workspace.commit().unwrap();
267 assert!(dest.join("file.txt").exists());
268 }
269
270 #[test]
271 fn test_workspace_cleanup_on_drop() {
272 let dir = tempdir().unwrap();
273 let staging = dir.path().join("staging");
274 {
275 let workspace = Workspace::new(&staging, dir.path().join("dest")).unwrap();
276 workspace.write("file.txt", b"data").unwrap();
277 assert!(staging.exists());
278 }
279 assert!(!staging.exists());
280 }
281
282 #[test]
283 fn test_workspace_create_dirs_and_report() {
284 let dir = tempdir().unwrap();
285 let workspace =
286 Workspace::new(dir.path().join("staging"), dir.path().join("dest")).unwrap();
287
288 workspace.create_dir("bin").unwrap();
289 workspace.create_dir_all("lib/nested").unwrap();
290 workspace.write("bin/tool", b"hello").unwrap();
291 workspace.write("lib/nested/config.toml", b"abc").unwrap();
292
293 let report = workspace.report().unwrap();
294 assert_eq!(report.file_count, 2);
295 assert_eq!(report.directory_count, 3);
296 assert_eq!(report.total_bytes, 8);
297 }
298
299 #[test]
300 fn test_workspace_rejects_escaping_path() {
301 let dir = tempdir().unwrap();
302 let workspace =
303 Workspace::new(dir.path().join("staging"), dir.path().join("dest")).unwrap();
304
305 let result = workspace.write("../escape.txt", b"data");
306 assert!(matches!(result, Err(Error::InvalidInput(_))));
307 }
308
309 #[test]
310 fn test_workspace_link_or_copy_file() {
311 let dir = tempdir().unwrap();
312 let source = dir.path().join("source.txt");
313 std::fs::write(&source, b"data").unwrap();
314
315 let workspace =
316 Workspace::new(dir.path().join("staging"), dir.path().join("dest")).unwrap();
317 workspace
318 .link_or_copy_file(&source, "bin/tool.txt", hardlink::Options::new())
319 .unwrap();
320
321 assert_eq!(workspace.read("bin/tool.txt").unwrap(), b"data");
322 }
323
324 #[test]
325 fn test_workspace_stage_file_by_size_prefers_copy_under_threshold() {
326 let dir = tempdir().unwrap();
327 let source = dir.path().join("source.bin");
328 let destination = dir.path().join("dest");
329 std::fs::write(&source, b"data").unwrap();
330
331 let workspace = Workspace::new(dir.path().join("staging"), &destination).unwrap();
332 workspace
333 .stage_file_by_size(&source, "bin/tool.bin", 1024, hardlink::Options::new())
334 .unwrap();
335 workspace.commit().unwrap();
336
337 assert_eq!(
338 std::fs::read(destination.join("bin/tool.bin")).unwrap(),
339 b"data"
340 );
341 }
342}