1use std::path::{Path, PathBuf};
2
3use anyhow::{anyhow, Result};
4use tokio::{
5 fs::File,
6 io::{AsyncRead, AsyncWrite, AsyncWriteExt},
7};
8use tokio_tar::{Builder, Entry, EntryType, Header};
9
10#[derive(Clone, Debug, Eq, PartialEq)]
11pub enum VfsNodeType {
12 Directory,
13 RegularFile,
14 Symlink,
15 Hardlink,
16 Fifo,
17 CharDevice,
18 BlockDevice,
19}
20
21#[derive(Clone, Debug)]
22pub struct VfsNode {
23 pub name: String,
24 pub size: u64,
25 pub children: Vec<VfsNode>,
26 pub typ: VfsNodeType,
27 pub uid: u64,
28 pub gid: u64,
29 pub link_name: Option<String>,
30 pub mode: u32,
31 pub mtime: u64,
32 pub dev_major: Option<u32>,
33 pub dev_minor: Option<u32>,
34 pub disk_path: Option<PathBuf>,
35}
36
37impl VfsNode {
38 pub fn from<X: AsyncRead + Unpin>(entry: &Entry<X>) -> Result<VfsNode> {
39 let header = entry.header();
40 let name = entry
41 .path()?
42 .file_name()
43 .ok_or(anyhow!("unable to get file name for entry"))?
44 .to_string_lossy()
45 .to_string();
46 let typ = header.entry_type();
47 let vtype = if typ.is_symlink() {
48 VfsNodeType::Symlink
49 } else if typ.is_hard_link() {
50 VfsNodeType::Hardlink
51 } else if typ.is_dir() {
52 VfsNodeType::Directory
53 } else if typ.is_fifo() {
54 VfsNodeType::Fifo
55 } else if typ.is_block_special() {
56 VfsNodeType::BlockDevice
57 } else if typ.is_character_special() {
58 VfsNodeType::CharDevice
59 } else if typ.is_file() {
60 VfsNodeType::RegularFile
61 } else {
62 return Err(anyhow!("unable to determine vfs type for entry"));
63 };
64
65 Ok(VfsNode {
66 name,
67 size: header.size()?,
68 children: vec![],
69 typ: vtype,
70 uid: header.uid()?,
71 gid: header.gid()?,
72 link_name: header.link_name()?.map(|x| x.to_string_lossy().to_string()),
73 mode: header.mode()?,
74 mtime: header.mtime()?,
75 dev_major: header.device_major()?,
76 dev_minor: header.device_minor()?,
77 disk_path: None,
78 })
79 }
80
81 pub fn lookup(&self, path: &Path) -> Option<&VfsNode> {
82 let mut node = self;
83 for part in path {
84 node = node
85 .children
86 .iter()
87 .find(|child| child.name == part.to_string_lossy())?;
88 }
89 Some(node)
90 }
91
92 pub fn lookup_mut(&mut self, path: &Path) -> Option<&mut VfsNode> {
93 let mut node = self;
94 for part in path {
95 node = node
96 .children
97 .iter_mut()
98 .find(|child| child.name == part.to_string_lossy())?;
99 }
100 Some(node)
101 }
102
103 pub fn remove(&mut self, path: &Path) -> Option<(&mut VfsNode, VfsNode)> {
104 let parent = path.parent()?;
105 let node = self.lookup_mut(parent)?;
106 let file_name = path.file_name()?;
107 let file_name = file_name.to_string_lossy();
108 let position = node
109 .children
110 .iter()
111 .position(|child| file_name == child.name)?;
112 let removed = node.children.remove(position);
113 Some((node, removed))
114 }
115
116 pub fn create_tar_header(&self) -> Result<Header> {
117 let mut header = Header::new_ustar();
118 header.set_entry_type(match self.typ {
119 VfsNodeType::Directory => EntryType::Directory,
120 VfsNodeType::CharDevice => EntryType::Char,
121 VfsNodeType::BlockDevice => EntryType::Block,
122 VfsNodeType::Fifo => EntryType::Fifo,
123 VfsNodeType::Hardlink => EntryType::Link,
124 VfsNodeType::Symlink => EntryType::Symlink,
125 VfsNodeType::RegularFile => EntryType::Regular,
126 });
127 header.set_uid(self.uid);
128 header.set_gid(self.gid);
129
130 if let Some(device_major) = self.dev_major {
131 header.set_device_major(device_major)?;
132 }
133
134 if let Some(device_minor) = self.dev_minor {
135 header.set_device_minor(device_minor)?;
136 }
137 header.set_mtime(self.mtime);
138 header.set_mode(self.mode);
139
140 if let Some(link_name) = self.link_name.as_ref() {
141 header.set_link_name(PathBuf::from(link_name))?;
142 }
143 header.set_size(self.size);
144 Ok(header)
145 }
146
147 pub async fn write_to_tar<W: AsyncWrite + Unpin + Send>(
148 &self,
149 path: &Path,
150 builder: &mut Builder<W>,
151 ) -> Result<()> {
152 let mut header = self.create_tar_header()?;
153 header.set_path(path)?;
154 header.set_cksum();
155 if let Some(disk_path) = self.disk_path.as_ref() {
156 builder
157 .append(&header, File::open(disk_path).await?)
158 .await?;
159 } else {
160 builder.append(&header, &[] as &[u8]).await?;
161 }
162 Ok(())
163 }
164}
165
166#[derive(Clone, Debug)]
167pub struct VfsTree {
168 pub root: VfsNode,
169}
170
171impl Default for VfsTree {
172 fn default() -> Self {
173 Self::new()
174 }
175}
176
177impl VfsTree {
178 pub fn new() -> VfsTree {
179 VfsTree {
180 root: VfsNode {
181 name: "".to_string(),
182 size: 0,
183 children: vec![],
184 typ: VfsNodeType::Directory,
185 uid: 0,
186 gid: 0,
187 link_name: None,
188 mode: 0,
189 mtime: 0,
190 dev_major: None,
191 dev_minor: None,
192 disk_path: None,
193 },
194 }
195 }
196
197 pub fn insert_tar_entry<X: AsyncRead + Unpin>(&mut self, entry: &Entry<X>) -> Result<&VfsNode> {
198 let mut meta = VfsNode::from(entry)?;
199 let path = entry.path()?.to_path_buf();
200 let parent = if let Some(parent) = path.parent() {
201 self.root.lookup_mut(parent)
202 } else {
203 Some(&mut self.root)
204 };
205
206 let Some(parent) = parent else {
207 return Err(anyhow!("unable to find parent of entry"));
208 };
209
210 let position = parent
211 .children
212 .iter()
213 .position(|child| meta.name == child.name);
214
215 if let Some(position) = position {
216 let old = parent.children.remove(position);
217 if meta.typ == VfsNodeType::Directory {
218 meta.children = old.children;
219 }
220 }
221 parent.children.push(meta.clone());
222 let Some(reference) = parent.children.iter().find(|child| child.name == meta.name) else {
223 return Err(anyhow!("unable to find inserted child in vfs"));
224 };
225 Ok(reference)
226 }
227
228 pub fn set_disk_path(&mut self, path: &Path, disk_path: &Path) -> Result<()> {
229 let Some(node) = self.root.lookup_mut(path) else {
230 return Err(anyhow!(
231 "unable to find node {:?} to set disk path to",
232 path
233 ));
234 };
235 node.disk_path = Some(disk_path.to_path_buf());
236 Ok(())
237 }
238
239 pub async fn write_to_tar<W: AsyncWrite + Unpin + Send + 'static>(
240 &self,
241 write: W,
242 ) -> Result<()> {
243 let mut builder = Builder::new(write);
244 let mut queue = vec![(PathBuf::from(""), &self.root)];
245
246 while !queue.is_empty() {
247 let (mut path, node) = queue.remove(0);
248 if !node.name.is_empty() {
249 path.push(&node.name);
250 }
251 if path.components().count() != 0 {
252 node.write_to_tar(&path, &mut builder).await?;
253 }
254 for child in &node.children {
255 queue.push((path.clone(), child));
256 }
257 }
258
259 let mut write = builder.into_inner().await?;
260 write.flush().await?;
261 drop(write);
262 Ok(())
263 }
264}