Skip to main content

rust_ipfs/unixfs/
add.rs

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    /// Sets the multihash used for produced blocks. Defaults to sha2-256. A hash other than
118    /// sha2-256 implies CIDv1.
119    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                                            //TODO: Determine best course to prevent additional allocation
293                                            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}