rskit_dataset/
payload.rs

1use std::path::{Path, PathBuf};
2
3use rskit_errors::{AppError, AppResult, ErrorCode};
4
5use crate::DatasetLimits;
6
7/// Payload for a single dataset item.
8///
9/// In-memory payloads can only be constructed through checked constructors so
10/// large datasets do not accidentally bypass [`DatasetLimits`].
11#[derive(Debug, Clone)]
12pub struct DataPayload {
13    kind: DataPayloadKind,
14}
15
16#[derive(Debug, Clone)]
17enum DataPayloadKind {
18    Bytes(Vec<u8>),
19    File(PathBuf),
20}
21
22impl DataPayload {
23    /// Create a bounded in-memory payload with explicit limits.
24    pub fn bytes(bytes: Vec<u8>, limits: &DatasetLimits) -> AppResult<Self> {
25        if bytes.len() > limits.max_in_memory_bytes {
26            return Err(AppError::new(
27                ErrorCode::InvalidInput,
28                format!(
29                    "in-memory dataset payload is {} bytes, exceeding max_in_memory_bytes={}",
30                    bytes.len(),
31                    limits.max_in_memory_bytes
32                ),
33            ));
34        }
35        Ok(Self {
36            kind: DataPayloadKind::Bytes(bytes),
37        })
38    }
39
40    /// Create a bounded in-memory payload using [`DatasetLimits::default`].
41    pub fn bytes_default(bytes: Vec<u8>) -> AppResult<Self> {
42        Self::bytes(bytes, &DatasetLimits::default())
43    }
44
45    /// Create a file-backed payload for streaming large data.
46    #[must_use]
47    pub fn file(path: impl Into<PathBuf>) -> Self {
48        Self {
49            kind: DataPayloadKind::File(path.into()),
50        }
51    }
52
53    /// Returns true when this payload is memory-backed.
54    #[must_use]
55    pub fn is_bytes(&self) -> bool {
56        matches!(self.kind, DataPayloadKind::Bytes(_))
57    }
58
59    /// Borrow the file path when this is a file-backed payload.
60    #[must_use]
61    pub fn as_file(&self) -> Option<&Path> {
62        match &self.kind {
63            DataPayloadKind::File(path) => Some(path),
64            DataPayloadKind::Bytes(_) => None,
65        }
66    }
67
68    /// Return the payload size in bytes when it can be determined.
69    pub fn len(&self) -> AppResult<u64> {
70        match &self.kind {
71            DataPayloadKind::Bytes(bytes) => Ok(bytes.len() as u64),
72            DataPayloadKind::File(path) => std::fs::metadata(path)
73                .map(|metadata| metadata.len())
74                .map_err(|error| {
75                    AppError::new(
76                        ErrorCode::Internal,
77                        format!("failed to stat payload file {}: {error}", path.display()),
78                    )
79                }),
80        }
81    }
82
83    /// Return true when the payload has zero bytes.
84    pub fn is_empty(&self) -> AppResult<bool> {
85        self.len().map(|len| len == 0)
86    }
87
88    /// Read a payload into memory only if it is within configured bounds.
89    pub fn read_bytes_bounded(&self, limits: &DatasetLimits) -> AppResult<Vec<u8>> {
90        match &self.kind {
91            DataPayloadKind::Bytes(bytes) => {
92                if bytes.len() > limits.max_in_memory_bytes {
93                    return Err(AppError::new(
94                        ErrorCode::InvalidInput,
95                        format!(
96                            "dataset payload is {} bytes, exceeding max_in_memory_bytes={}",
97                            bytes.len(),
98                            limits.max_in_memory_bytes
99                        ),
100                    ));
101                }
102                Ok(bytes.clone())
103            }
104            DataPayloadKind::File(path) => read_file_bounded(path, limits.max_in_memory_bytes),
105        }
106    }
107
108    /// Write the payload to `path`, streaming file payloads without materializing them.
109    pub fn write_to_path(&self, path: &Path, limits: &DatasetLimits) -> AppResult<u64> {
110        match &self.kind {
111            DataPayloadKind::Bytes(bytes) => {
112                if bytes.len() > limits.max_in_memory_bytes {
113                    return Err(AppError::new(
114                        ErrorCode::InvalidInput,
115                        format!(
116                            "in-memory dataset payload is {} bytes, exceeding max_in_memory_bytes={}",
117                            bytes.len(),
118                            limits.max_in_memory_bytes
119                        ),
120                    ));
121                }
122                std::fs::write(path, bytes).map_err(|error| {
123                    AppError::new(
124                        ErrorCode::Internal,
125                        format!("failed to write dataset item {}: {error}", path.display()),
126                    )
127                })?;
128                Ok(bytes.len() as u64)
129            }
130            DataPayloadKind::File(source) => {
131                if is_same_file(source, path)? {
132                    return self.len();
133                }
134                let mut input = std::fs::File::open(source).map_err(|error| {
135                    AppError::new(
136                        ErrorCode::Internal,
137                        format!("failed to open payload file {}: {error}", source.display()),
138                    )
139                })?;
140                let mut output = std::fs::File::create(path).map_err(|error| {
141                    AppError::new(
142                        ErrorCode::Internal,
143                        format!("failed to create dataset item {}: {error}", path.display()),
144                    )
145                })?;
146                std::io::copy(&mut input, &mut output).map_err(|error| {
147                    AppError::new(
148                        ErrorCode::Internal,
149                        format!(
150                            "failed to stream payload {} to {}: {error}",
151                            source.display(),
152                            path.display()
153                        ),
154                    )
155                })
156            }
157        }
158    }
159}
160
161fn is_same_file(source: &Path, destination: &Path) -> AppResult<bool> {
162    let source_metadata = std::fs::metadata(source).map_err(|error| {
163        AppError::new(
164            ErrorCode::Internal,
165            format!("failed to stat payload file {}: {error}", source.display()),
166        )
167    })?;
168    let destination_metadata = match std::fs::metadata(destination) {
169        Ok(metadata) => metadata,
170        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false),
171        Err(error) => {
172            return Err(AppError::new(
173                ErrorCode::Internal,
174                format!(
175                    "failed to stat dataset item {}: {error}",
176                    destination.display()
177                ),
178            ));
179        }
180    };
181
182    if same_file_metadata(&source_metadata, &destination_metadata) {
183        return Ok(true);
184    }
185
186    let source = std::fs::canonicalize(source).map_err(|error| {
187        AppError::new(
188            ErrorCode::Internal,
189            format!(
190                "failed to canonicalize payload file {}: {error}",
191                source.display()
192            ),
193        )
194    })?;
195    let destination = std::fs::canonicalize(destination).map_err(|error| {
196        AppError::new(
197            ErrorCode::Internal,
198            format!(
199                "failed to canonicalize dataset item {}: {error}",
200                destination.display()
201            ),
202        )
203    })?;
204    Ok(source == destination)
205}
206
207#[cfg(unix)]
208fn same_file_metadata(left: &std::fs::Metadata, right: &std::fs::Metadata) -> bool {
209    use std::os::unix::fs::MetadataExt as _;
210    left.dev() == right.dev() && left.ino() == right.ino()
211}
212
213#[cfg(windows)]
214fn same_file_metadata(left: &std::fs::Metadata, right: &std::fs::Metadata) -> bool {
215    use std::os::windows::fs::MetadataExt as _;
216    matches!(
217        (
218            left.volume_serial_number(),
219            left.file_index(),
220            right.volume_serial_number(),
221            right.file_index()
222        ),
223        (Some(left_volume), Some(left_index), Some(right_volume), Some(right_index))
224            if left_volume == right_volume && left_index == right_index
225    )
226}
227
228#[cfg(not(any(unix, windows)))]
229fn same_file_metadata(_left: &std::fs::Metadata, _right: &std::fs::Metadata) -> bool {
230    false
231}
232fn read_file_bounded(path: &Path, max_bytes: usize) -> AppResult<Vec<u8>> {
233    use std::io::Read as _;
234
235    let mut file = std::fs::File::open(path).map_err(|error| {
236        AppError::new(
237            ErrorCode::Internal,
238            format!("failed to open payload file {}: {error}", path.display()),
239        )
240    })?;
241    let mut bytes = Vec::new();
242    file.by_ref()
243        .take(max_bytes as u64 + 1)
244        .read_to_end(&mut bytes)
245        .map_err(|error| {
246            AppError::new(
247                ErrorCode::Internal,
248                format!("failed to read payload file {}: {error}", path.display()),
249            )
250        })?;
251    if bytes.len() > max_bytes {
252        return Err(AppError::new(
253            ErrorCode::InvalidInput,
254            format!(
255                "dataset payload exceeded max_in_memory_bytes={max_bytes} while reading {}",
256                path.display()
257            ),
258        ));
259    }
260    Ok(bytes)
261}
rskit_dataset/payload.rs

rskit_dataset/
payload.rs
