1use std::task::{Context, Poll};
2
3use crate::{
4 Block,
5 repo::{DefaultStorage, Repo},
6};
7use bytes::Bytes;
8use either::Either;
9#[allow(unused_imports)]
10use futures::{
11 FutureExt, Stream, StreamExt, TryFutureExt,
12 future::BoxFuture,
13 stream::{BoxStream, FusedStream},
14};
15use ipld_core::cid::Version;
16use multihash_codetable::Code;
17use rust_unixfs::Metadata;
18use rust_unixfs::file::adder::{Chunker, FileAdderBuilder};
19#[cfg(not(target_arch = "wasm32"))]
20use std::path::{Path, PathBuf};
21use std::pin::Pin;
22#[cfg(not(target_arch = "wasm32"))]
23use tokio_util::io::ReaderStream;
24use tracing::{Instrument, Span};
25
26use crate::{Ipfs, IpfsPath};
27
28use super::{TraversalFailed, UnixfsStatus};
29
30pub enum AddOpt {
31 #[cfg(not(target_arch = "wasm32"))]
32 File(PathBuf),
33 Stream {
34 name: Option<String>,
35 total: Option<usize>,
36 stream: BoxStream<'static, Result<Bytes, std::io::Error>>,
37 },
38}
39
40#[cfg(not(target_arch = "wasm32"))]
41impl From<PathBuf> for AddOpt {
42 fn from(path: PathBuf) -> Self {
43 AddOpt::File(path)
44 }
45}
46
47#[cfg(not(target_arch = "wasm32"))]
48impl From<&Path> for AddOpt {
49 fn from(path: &Path) -> Self {
50 AddOpt::File(path.to_path_buf())
51 }
52}
53
54#[must_use = "does nothing unless you `.await` or poll the stream"]
55pub struct UnixfsAdd {
56 core: Option<Either<Ipfs, Repo<DefaultStorage>>>,
57 opt: Option<AddOpt>,
58 span: Span,
59 chunk: Chunker,
60 cid_version: Version,
61 raw_leaves: Option<bool>,
62 hasher: Code,
63 metadata: Metadata,
64 pin: bool,
65 provide: bool,
66 wrap: bool,
67 stream: Option<BoxStream<'static, UnixfsStatus>>,
68}
69
70impl UnixfsAdd {
71 pub fn with_ipfs(ipfs: &Ipfs, opt: impl Into<AddOpt>) -> Self {
72 Self::with_either(Either::Left(ipfs.clone()), opt)
73 }
74
75 pub fn with_repo(repo: &Repo<DefaultStorage>, opt: impl Into<AddOpt>) -> Self {
76 Self::with_either(Either::Right(repo.clone()), opt)
77 }
78
79 fn with_either(core: Either<Ipfs, Repo<DefaultStorage>>, opt: impl Into<AddOpt>) -> Self {
80 let opt = opt.into();
81 Self {
82 core: Some(core),
83 opt: Some(opt),
84 span: Span::current(),
85 chunk: Chunker::Size(256 * 1024),
86 cid_version: Version::V1,
87 raw_leaves: None,
88 hasher: Code::Sha2_256,
89 metadata: Metadata::default(),
90 pin: true,
91 provide: false,
92 wrap: false,
93 stream: None,
94 }
95 }
96
97 pub fn span(mut self, span: Span) -> Self {
98 self.span = span;
99 self
100 }
101
102 pub fn chunk(mut self, chunk: Chunker) -> Self {
103 self.chunk = chunk;
104 self
105 }
106
107 pub fn cid_version(mut self, version: Version) -> Self {
108 self.cid_version = version;
109 self
110 }
111
112 pub fn raw_leaves(mut self, raw_leaves: bool) -> Self {
113 self.raw_leaves = Some(raw_leaves);
114 self
115 }
116
117 pub fn hasher(mut self, hasher: Code) -> Self {
120 self.hasher = hasher;
121 self
122 }
123
124 pub fn metadata(mut self, metadata: Metadata) -> Self {
125 self.metadata = metadata;
126 self
127 }
128
129 pub fn pin(mut self, pin: bool) -> Self {
130 self.pin = pin;
131 self
132 }
133
134 pub fn provide(mut self) -> Self {
135 self.provide = true;
136 self
137 }
138
139 pub fn wrap(mut self) -> Self {
140 self.wrap = true;
141 self
142 }
143}
144
145impl Stream for UnixfsAdd {
146 type Item = UnixfsStatus;
147 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
148 if self.core.is_none() && self.stream.is_none() {
149 return Poll::Ready(None);
150 }
151 loop {
152 match &mut self.stream {
153 None => {
154 let (ipfs, repo) = match self.core.take().expect("ipfs or repo is used") {
155 Either::Left(ipfs) => {
156 let repo = ipfs.repo().clone();
157 (Some(ipfs), repo)
158 }
159 Either::Right(repo) => (None, repo),
160 };
161 let option = self.opt.take().expect("option already constructed");
162 let chunk = self.chunk;
163 let cid_version = self.cid_version;
164 let raw_leaves = self.raw_leaves;
165 let hasher = self.hasher;
166 let metadata = self.metadata.clone();
167 let pin = self.pin;
168 let provide = self.provide;
169 let wrap = self.wrap;
170
171 let stream = async_stream::stream! {
172 let _g = repo.gc_guard().await;
173
174 let mut written = 0;
175
176 let (name, total_size, mut stream) = match option {
177 #[cfg(not(target_arch = "wasm32"))]
178 AddOpt::File(path) => match tokio::fs::File::open(path.clone())
179 .and_then(|file| async move {
180 let size = file.metadata().await?.len() as usize;
181
182 let stream = ReaderStream::new(file);
183
184 let name: Option<String> = path.file_name().map(|f| f.to_string_lossy().to_string());
185
186 Ok((name, Some(size), stream.boxed()))
187 }).await {
188 Ok(s) => s,
189 Err(e) => {
190 yield UnixfsStatus::FailedStatus { written, total_size: None, error: e.into() };
191 return;
192 }
193 },
194 AddOpt::Stream { name, total, stream } => (name, total, stream),
195 };
196
197 let mut adder = {
198 let mut builder = FileAdderBuilder::default()
199 .with_chunker(chunk)
200 .with_cid_version(cid_version)
201 .with_hasher(hasher)
202 .with_metadata(metadata);
203 if let Some(raw_leaves) = raw_leaves {
204 builder = builder.with_raw_leaves(raw_leaves);
205 }
206 builder.build()
207 };
208
209 yield UnixfsStatus::ProgressStatus { written, total_size };
210
211 while let Some(buffer) = stream.next().await {
212 let buffer = match buffer {
213 Ok(buf) => buf,
214 Err(e) => {
215 yield UnixfsStatus::FailedStatus { written, total_size, error: e.into() };
216 return;
217 }
218 };
219
220 let mut total = 0;
221 while total < buffer.len() {
222 let (blocks, consumed) = adder.push(&buffer[total..]);
223 let mut to_put = Vec::new();
224 for (cid, block) in blocks {
225 match Block::new(cid, block) {
226 Ok(block) => to_put.push(block),
227 Err(e) => {
228 yield UnixfsStatus::FailedStatus { written, total_size, error: e.into() };
229 return;
230 }
231 }
232 }
233 if let Err(e) = repo.put_blocks(to_put).await {
234 yield UnixfsStatus::FailedStatus { written, total_size, error: e };
235 return;
236 }
237 total += consumed;
238 written += consumed;
239 }
240
241 yield UnixfsStatus::ProgressStatus { written, total_size };
242 }
243
244 let blocks = adder.finish();
245 let mut last_cid = None;
246 let mut to_put = Vec::new();
247
248 for (cid, block) in blocks {
249 match Block::new(cid, block) {
250 Ok(block) => to_put.push(block),
251 Err(e) => {
252 yield UnixfsStatus::FailedStatus { written, total_size, error: e.into() };
253 return;
254 }
255 }
256 last_cid = Some(cid);
257 }
258
259 if let Err(e) = repo.put_blocks(to_put).await {
260 yield UnixfsStatus::FailedStatus { written, total_size, error: e };
261 return;
262 }
263
264 let cid = match last_cid {
265 Some(cid) => cid,
266 None => {
267 yield UnixfsStatus::FailedStatus { written, total_size, error: TraversalFailed::Io(std::io::ErrorKind::InvalidData.into()).into() };
268 return;
269 }
270 };
271
272 let mut path = IpfsPath::from(cid);
273
274 if wrap {
275 if let Some(name) = name {
276 let result = {
277 let repo = repo.clone();
278 async move {
279 let mut opts = rust_unixfs::dir::builder::TreeOptions::default();
280 opts.wrap_with_directory();
281 opts.cid_version(cid_version);
282 opts.hasher(hasher);
283
284 let mut tree = rust_unixfs::dir::builder::BufferingTreeBuilder::new(opts);
285 tree.put_link(&name, cid, written as _)?;
286
287 let mut iter = tree.build();
288 let mut cids = Vec::new();
289 let mut to_put = Vec::new();
290
291 while let Some(node) = iter.next_borrowed() {
292 let node = node?;
294 let cid = node.cid.to_owned();
295 to_put.push(Block::new(cid, node.block.to_vec())?);
296 cids.push(cid);
297 }
298
299 repo.put_blocks(to_put).await?;
300 let cid = cids.last().ok_or(anyhow::anyhow!("no cid available"))?;
301 let path = IpfsPath::from(*cid).sub_path(&name)?;
302
303 Ok::<_, anyhow::Error>(path)
304 }
305 };
306
307 path = match result.await {
308 Ok(path) => path,
309 Err(e) => {
310 yield UnixfsStatus::FailedStatus { written, total_size, error: e };
311 return;
312 }
313 };
314 }
315 }
316
317 let cid = path.root().cid().copied().expect("Cid is apart of the path");
318
319 if pin && !repo.is_pinned(&cid).await.unwrap_or_default() {
320 if let Err(e) = repo.pin(cid).recursive().await {
321 error!("Unable to pin {cid}: {e}");
322 }
323 }
324
325 if provide {
326 if let Some(ipfs) = ipfs {
327 if let Err(e) = ipfs.provide(cid).await {
328 error!("Unable to provide {cid}: {e}");
329 }
330 }
331 }
332
333
334 yield UnixfsStatus::CompletedStatus { path, written, total_size }
335 };
336
337 self.stream = Some(stream.boxed());
338 }
339 Some(stream) => match futures::ready!(stream.poll_next_unpin(cx)) {
340 Some(item) => {
341 if matches!(
342 item,
343 UnixfsStatus::FailedStatus { .. }
344 | UnixfsStatus::CompletedStatus { .. }
345 ) {
346 self.stream.take();
347 }
348 return Poll::Ready(Some(item));
349 }
350 None => {
351 self.stream.take();
352 return Poll::Ready(None);
353 }
354 },
355 }
356 }
357 }
358}
359
360impl std::future::IntoFuture for UnixfsAdd {
361 type Output = Result<IpfsPath, anyhow::Error>;
362
363 type IntoFuture = BoxFuture<'static, Self::Output>;
364
365 fn into_future(mut self) -> Self::IntoFuture {
366 let span = self.span.clone();
367 async move {
368 while let Some(status) = self.next().await {
369 match status {
370 UnixfsStatus::CompletedStatus { path, .. } => return Ok(path),
371 UnixfsStatus::FailedStatus { error, .. } => {
372 return Err(error);
373 }
374 _ => {}
375 }
376 }
377 Err::<_, anyhow::Error>(anyhow::anyhow!("Unable to add file"))
378 }
379 .instrument(span)
380 .boxed()
381 }
382}
383
384impl FusedStream for UnixfsAdd {
385 fn is_terminated(&self) -> bool {
386 self.stream.is_none() && self.core.is_none()
387 }
388}