Skip to main content

rns_net/
resource.rs

1//! Standard-library Resource streaming and disk-delivery types.
2
3use std::fs::{self, File, OpenOptions};
4#[cfg(test)]
5use std::io::Read;
6use std::io::{self, Write};
7use std::path::{Path, PathBuf};
8use std::sync::atomic::{AtomicU64, Ordering};
9
10use rns_core::constants::RESOURCE_AUTO_COMPRESS_MAX_SIZE;
11
12static NEXT_PERSIST_TEMP: AtomicU64 = AtomicU64::new(1);
13static NEXT_RECEIVE_TEMP: AtomicU64 = AtomicU64::new(1);
14
15/// A node-local identifier for an outgoing streaming Resource transfer.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
17pub struct ResourceTransferId(pub u64);
18
19/// Failure reported for the new streaming Resource API.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum ResourceTransferError {
22    Source(String),
23    Storage(String),
24    Protocol(String),
25    Cancelled,
26}
27
28/// Options for an outgoing Resource.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub struct ResourceSendOptions {
31    /// Attempt compression and use the result only when smaller.
32    pub auto_compress: bool,
33    /// Application-defined flag carried in the advertisement (bit 6 of its
34    /// flags, Python `Resource(app_flag=...)`) and reported to the receiver.
35    pub app_flag: bool,
36}
37
38impl Default for ResourceSendOptions {
39    fn default() -> Self {
40        Self {
41            auto_compress: true,
42            app_flag: false,
43        }
44    }
45}
46
47/// How independent incoming Resources are delivered to the application.
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub enum ResourceReceiveMode {
50    /// Assemble the complete logical Resource in memory.
51    ///
52    /// Nodes can verify compressed, single-segment application Resources of
53    /// 16 KiB through the efficient segment limit on a lazy receive worker.
54    /// Authentication remains on the driver. Worker admission reserves owned
55    /// input and worst-case output/metadata capacity: four jobs and 256 MiB,
56    /// allowing one job at the default 64 MiB decoder limit. This reservation
57    /// is not preallocated memory or a whole-node RSS cap; native codec workspace
58    /// and unrelated transfers are additional. Small, split, request/response
59    /// and saturated work stays synchronous. Earlier same-link assemblies finish
60    /// first, so overload can still stall the driver. Shutdown joins active work.
61    Memory { max_bytes: u64 },
62    /// Write verified segments to a managed temporary file.
63    TemporaryFile {
64        directory: PathBuf,
65        max_bytes: Option<u64>,
66    },
67    /// Hand each verified segment to the application as it arrives
68    /// ([`Callbacks::on_resource_segment`](crate::Callbacks::on_resource_segment)),
69    /// in order, then report the whole Resource received
70    /// ([`Callbacks::on_resource_stream_received`](crate::Callbacks::on_resource_stream_received)).
71    ///
72    /// Memory stays bounded by about one segment (at most the efficient
73    /// segment size, 1 MiB) whatever the Resource's size. A Resource of one
74    /// segment arrives the same way, as one segment at offset 0. A transfer
75    /// that fails after some segments were delivered is reported through
76    /// [`Callbacks::on_resource_failed`](crate::Callbacks::on_resource_failed);
77    /// the application discards what it received of it. Requests and
78    /// responses are delivered as before.
79    Stream { max_bytes: Option<u64> },
80}
81
82impl Default for ResourceReceiveMode {
83    fn default() -> Self {
84        Self::Memory {
85            max_bytes: RESOURCE_AUTO_COMPRESS_MAX_SIZE as u64,
86        }
87    }
88}
89
90/// A streamed incoming Resource that has been received in full: every
91/// segment was delivered to [`Callbacks::on_resource_segment`](crate::Callbacks::on_resource_segment).
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct ReceivedResourceStream {
94    /// The Resource's original hash, which its segments were delivered under.
95    pub original_hash: [u8; 32],
96    /// Bytes delivered in all.
97    pub size: u64,
98    pub metadata: Option<Vec<u8>>,
99    /// The application flag the sender set in the advertisement.
100    pub app_flag: bool,
101}
102
103/// A received disk-backed Resource.
104///
105/// The temporary file is removed when this value is dropped unless ownership
106/// is transferred with [`ReceivedResourceFile::persist`].
107#[derive(Debug)]
108pub struct ReceivedResourceFile {
109    path: Option<PathBuf>,
110    pub original_hash: [u8; 32],
111    pub size: u64,
112    pub metadata: Option<Vec<u8>>,
113    /// The application flag the sender set in the advertisement.
114    pub app_flag: bool,
115}
116
117impl ReceivedResourceFile {
118    pub(crate) fn new(
119        path: PathBuf,
120        original_hash: [u8; 32],
121        size: u64,
122        metadata: Option<Vec<u8>>,
123        app_flag: bool,
124    ) -> Self {
125        Self {
126            path: Some(path),
127            original_hash,
128            size,
129            metadata,
130            app_flag,
131        }
132    }
133
134    /// Path of the managed temporary file.
135    pub fn path(&self) -> &Path {
136        self.path
137            .as_deref()
138            .expect("persisted resource has no path")
139    }
140
141    /// Persist this Resource at `destination`.
142    ///
143    /// The final name appears atomically. Cross-filesystem moves are staged in
144    /// the destination directory before the final rename.
145    pub fn persist(
146        mut self,
147        destination: impl AsRef<Path>,
148        overwrite: bool,
149    ) -> io::Result<PathBuf> {
150        let destination = destination.as_ref();
151        let source = self.path.as_ref().expect("persisted resource has no path");
152        match rename_final(source, destination, overwrite) {
153            Ok(()) => {
154                self.path = None;
155                sync_parent(destination)?;
156                return Ok(destination.to_path_buf());
157            }
158            Err(error) if error.raw_os_error() != Some(libc::EXDEV) => return Err(error),
159            Err(_) => {}
160        }
161
162        let parent = destination.parent().unwrap_or_else(|| Path::new("."));
163        let sequence = NEXT_PERSIST_TEMP.fetch_add(1, Ordering::Relaxed);
164        let stage = parent.join(format!(
165            ".rns-resource-{}-{}.tmp",
166            std::process::id(),
167            sequence
168        ));
169        let mut input = File::open(source)?;
170        let mut output = OpenOptions::new()
171            .write(true)
172            .create_new(true)
173            .open(&stage)?;
174        let copied = io::copy(&mut input, &mut output)?;
175        if copied != self.size {
176            let _ = fs::remove_file(&stage);
177            return Err(io::Error::new(
178                io::ErrorKind::UnexpectedEof,
179                "resource size changed",
180            ));
181        }
182        output.flush()?;
183        output.sync_all()?;
184        drop(output);
185        rename_final(&stage, destination, overwrite).inspect_err(|_| {
186            let _ = fs::remove_file(&stage);
187        })?;
188        sync_parent(destination)?;
189        fs::remove_file(source)?;
190        self.path = None;
191        Ok(destination.to_path_buf())
192    }
193}
194
195#[cfg(target_os = "linux")]
196fn rename_final(source: &Path, destination: &Path, overwrite: bool) -> io::Result<()> {
197    if overwrite {
198        return fs::rename(source, destination);
199    }
200    use std::ffi::CString;
201    use std::os::unix::ffi::OsStrExt;
202    let source = CString::new(source.as_os_str().as_bytes())
203        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "source path contains NUL"))?;
204    let destination = CString::new(destination.as_os_str().as_bytes()).map_err(|_| {
205        io::Error::new(io::ErrorKind::InvalidInput, "destination path contains NUL")
206    })?;
207    let result = unsafe {
208        libc::renameat2(
209            libc::AT_FDCWD,
210            source.as_ptr(),
211            libc::AT_FDCWD,
212            destination.as_ptr(),
213            libc::RENAME_NOREPLACE,
214        )
215    };
216    if result == 0 {
217        Ok(())
218    } else {
219        Err(io::Error::last_os_error())
220    }
221}
222
223#[cfg(not(target_os = "linux"))]
224fn rename_final(source: &Path, destination: &Path, overwrite: bool) -> io::Result<()> {
225    if !overwrite && destination.exists() {
226        return Err(io::Error::new(
227            io::ErrorKind::AlreadyExists,
228            "resource destination already exists",
229        ));
230    }
231    fs::rename(source, destination)
232}
233
234impl Drop for ReceivedResourceFile {
235    fn drop(&mut self) {
236        if let Some(path) = self.path.take() {
237            let _ = fs::remove_file(path);
238        }
239    }
240}
241
242fn sync_parent(path: &Path) -> io::Result<()> {
243    if let Some(parent) = path.parent() {
244        File::open(parent)?.sync_all()?;
245    }
246    Ok(())
247}
248
249pub(crate) fn create_receive_file(
250    directory: &Path,
251    original_hash: &[u8; 32],
252) -> io::Result<(File, PathBuf)> {
253    fs::create_dir_all(directory)?;
254    for _ in 0..100 {
255        let sequence = NEXT_RECEIVE_TEMP.fetch_add(1, Ordering::Relaxed);
256        let path = directory.join(format!(
257            "rns-resource-{:02x}{:02x}{:02x}{:02x}-{}-{}.part",
258            original_hash[0],
259            original_hash[1],
260            original_hash[2],
261            original_hash[3],
262            std::process::id(),
263            sequence
264        ));
265        match OpenOptions::new().write(true).create_new(true).open(&path) {
266            Ok(file) => return Ok((file, path)),
267            Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
268            Err(error) => return Err(error),
269        }
270    }
271    Err(io::Error::new(
272        io::ErrorKind::AlreadyExists,
273        "could not allocate unique Resource temporary file",
274    ))
275}
276
277/// Read exactly `length` bytes and reject both early EOF and trailing data.
278#[cfg(test)]
279fn read_declared(mut reader: impl Read, length: u64) -> io::Result<Vec<u8>> {
280    let length = usize::try_from(length)
281        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "length does not fit usize"))?;
282    let mut data = vec![0; length];
283    reader.read_exact(&mut data)?;
284    let mut trailing = [0u8; 1];
285    if reader.read(&mut trailing)? != 0 {
286        return Err(io::Error::new(
287            io::ErrorKind::InvalidData,
288            "resource source exceeds declared length",
289        ));
290    }
291    Ok(data)
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    #[test]
299    fn managed_file_is_deleted_on_drop_and_can_be_persisted() {
300        let root = tempfile::tempdir().unwrap();
301        let temporary = root.path().join("temporary");
302        fs::write(&temporary, b"resource").unwrap();
303        {
304            let artifact = ReceivedResourceFile::new(temporary.clone(), [7; 32], 8, None, false);
305            assert_eq!(artifact.path(), temporary);
306        }
307        assert!(!temporary.exists());
308
309        fs::write(&temporary, b"resource").unwrap();
310        let destination = root.path().join("saved");
311        let artifact = ReceivedResourceFile::new(temporary.clone(), [7; 32], 8, None, false);
312        artifact.persist(&destination, false).unwrap();
313        assert_eq!(fs::read(destination).unwrap(), b"resource");
314        assert!(!temporary.exists());
315    }
316
317    #[test]
318    fn declared_reader_rejects_short_and_long_sources() {
319        assert_eq!(read_declared(&b"abc"[..], 3).unwrap(), b"abc");
320        assert_eq!(
321            read_declared(&b"ab"[..], 3).unwrap_err().kind(),
322            io::ErrorKind::UnexpectedEof
323        );
324        assert_eq!(
325            read_declared(&b"abcd"[..], 3).unwrap_err().kind(),
326            io::ErrorKind::InvalidData
327        );
328    }
329}