ahri_tre_lake/
disclosure_datafile.rs

1//! Fixed, bounded Datafile capture. Descriptor attestations contain hashes,
2//! while resolved key material exists only in the private capture protocol.
3use crate::{LakeError, ScratchAttempt, ScratchAttemptId};
4use ahri_tre_types::{
5    AdmittedDatafileInput, AssetId, DataFileRecord, GovernedDatafileBinding, StudyId,
6};
7use sha2::{Digest, Sha256};
8use std::{
9    fs::File,
10    io::{Read, Write},
11    path::{Path, PathBuf},
12    process::{Command, Stdio},
13};
14
15fn physical_path(root: &str, datafile: &DataFileRecord) -> Result<PathBuf, LakeError> {
16    let root = std::fs::canonicalize(root)?;
17    let path = std::fs::canonicalize(crate::file_staging::datafile_storage_path(
18        root.to_str().ok_or(LakeError::QueryNotAdmissible)?,
19        &datafile.storage_uri,
20    )?)?;
21    if !path.starts_with(&root) || !path.is_file() {
22        return Err(LakeError::QueryNotAdmissible);
23    }
24    Ok(path)
25}
26
27/// Called by a trusted writer before publishing metadata, or by the operator
28/// during legacy conversion. Ordinary metadata UPDATE cannot create this proof.
29pub fn attest_datafile(
30    root: &str,
31    study: StudyId,
32    asset: AssetId,
33    datafile: &DataFileRecord,
34) -> Result<GovernedDatafileBinding, LakeError> {
35    let path = physical_path(root, datafile)?;
36    let encoded = serde_json::to_vec(datafile).map_err(|_| LakeError::QueryNotAdmissible)?;
37    if encoded.len() > 16 * 1024 {
38        return Err(LakeError::QueryNotAdmissible);
39    }
40    Ok(GovernedDatafileBinding {
41        version_id: datafile.datafile_id.0,
42        asset_id: asset.0,
43        study_id: study.0,
44        storage_identity: crate::file_staging::hex_digest(Sha256::digest(
45            path.as_os_str().as_encoded_bytes(),
46        )),
47        descriptor_fingerprint: crate::file_staging::hex_digest(Sha256::digest(encoded)),
48    })
49}
50
51#[derive(serde::Serialize, serde::Deserialize)]
52struct Capture {
53    representation: (bool, bool),
54    root: String,
55    input: AdmittedDatafileInput,
56    directory: PathBuf,
57    deadline: chrono::DateTime<chrono::Utc>,
58    max_decoded: u64,
59}
60
61pub struct PreparedDisclosureDatafile {
62    scratch: ScratchAttempt,
63}
64impl PreparedDisclosureDatafile {
65    pub fn capture(
66        root: &str,
67        input: AdmittedDatafileInput,
68        scratch: &ScratchAttempt,
69        deadline: chrono::DateTime<chrono::Utc>,
70        max_decoded: u64,
71        executable: &Path,
72    ) -> Result<Self, LakeError> {
73        Self::capture_representation(
74            root,
75            input,
76            scratch,
77            deadline,
78            max_decoded,
79            executable,
80            (true, true),
81        )
82    }
83    pub fn capture_representation(
84        root: &str,
85        input: AdmittedDatafileInput,
86        scratch: &ScratchAttempt,
87        deadline: chrono::DateTime<chrono::Utc>,
88        max_decoded: u64,
89        executable: &Path,
90        representation: (bool, bool),
91    ) -> Result<Self, LakeError> {
92        if input.datafile.encrypted.unwrap_or(false) && !representation.0 && representation.1 {
93            return Err(LakeError::QueryNotAdmissible);
94        }
95        if max_decoded == 0 || max_decoded > i64::MAX as u64 || deadline <= chrono::Utc::now() {
96            return Err(LakeError::QueryNotAdmissible);
97        }
98        let scratch = scratch
99            .create_child(
100                ScratchAttemptId::new(&uuid::Uuid::new_v4().simple().to_string())
101                    .map_err(|_| LakeError::QueryNotAdmissible)?,
102            )
103            .map_err(|_| LakeError::QueryNotAdmissible)?;
104        let request = Capture {
105            representation,
106            root: root.into(),
107            input,
108            directory: scratch.path().to_owned(),
109            deadline,
110            max_decoded,
111        };
112        let encoded = serde_json::to_vec(&request).map_err(|_| LakeError::QueryNotAdmissible)?;
113        if encoded.len() > 16 * 1024 {
114            return Err(LakeError::QueryNotAdmissible);
115        }
116        let mut command = Command::new(executable);
117        crate::disclosure_inputs::parent_bound(&mut command)?;
118        let mut child = command
119            .env_clear()
120            .arg("--capture-datafile")
121            .stdin(Stdio::piped())
122            .stdout(Stdio::null())
123            .stderr(Stdio::null())
124            .spawn()?;
125        if child
126            .stdin
127            .take()
128            .ok_or(LakeError::QueryNotAdmissible)?
129            .write_all(&encoded)
130            .is_err()
131        {
132            let _ = child.kill();
133            let _ = child.wait();
134            return Err(LakeError::QueryNotAdmissible);
135        }
136        loop {
137            match child.try_wait() {
138                Ok(Some(status)) if status.success() && chrono::Utc::now() < deadline => break,
139                Ok(Some(_)) => return Err(LakeError::QueryNotAdmissible),
140                Ok(None) if chrono::Utc::now() < deadline => {
141                    std::thread::sleep(std::time::Duration::from_millis(10))
142                }
143                _ => {
144                    let _ = child.kill();
145                    let _ = child.wait();
146                    return Err(LakeError::QueryNotAdmissible);
147                }
148            }
149        }
150        Ok(Self { scratch })
151    }
152    pub fn open_representation(self, compress: bool) -> Result<Box<dyn Read + Send>, LakeError> {
153        let stream = self.open()?;
154        if compress {
155            Ok(Box::new(zstd::stream::read::Encoder::new(stream, 3)?))
156        } else {
157            Ok(Box::new(stream))
158        }
159    }
160    pub fn open(self) -> Result<DisclosureDatafileStream, LakeError> {
161        let file = File::open(self.scratch.path().join("payload"))?;
162        Ok(DisclosureDatafileStream { file, _input: self })
163    }
164}
165pub struct DisclosureDatafileStream {
166    file: File,
167    _input: PreparedDisclosureDatafile,
168}
169impl Read for DisclosureDatafileStream {
170    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
171        self.file.read(buffer)
172    }
173}
174
175/// Private worker entry point; accepts no user program, URL or destination.
176pub fn run_datafile_capture_worker() -> Result<(), LakeError> {
177    let mut encoded = Vec::new();
178    std::io::stdin()
179        .take(16 * 1024 + 1)
180        .read_to_end(&mut encoded)?;
181    if encoded.len() > 16 * 1024 {
182        return Err(LakeError::QueryNotAdmissible);
183    }
184    let request: Capture =
185        serde_json::from_slice(&encoded).map_err(|_| LakeError::QueryNotAdmissible)?;
186    let seconds = (request.deadline - chrono::Utc::now()).num_seconds();
187    if seconds < 1 || seconds > i64::from(u32::MAX) {
188        return Err(LakeError::QueryNotAdmissible);
189    }
190    // SAFETY: alarm only installs this private worker's integer deadline.
191    unsafe {
192        libc::alarm(seconds as u32);
193    }
194    let expected = attest_datafile(
195        &request.root,
196        StudyId(request.input.binding.study_id),
197        AssetId(request.input.binding.asset_id),
198        &request.input.datafile,
199    )?;
200    if expected != request.input.binding {
201        return Err(LakeError::QueryNotAdmissible);
202    }
203    let source = physical_path(&request.root, &request.input.datafile)?;
204    let mut source = File::open(source)?;
205    let copied = request.directory.join("stored");
206    let mut target = std::fs::OpenOptions::new()
207        .create_new(true)
208        .write(true)
209        .open(&copied)?;
210    // This engineering workspace bound is independent of admitted request budgets.
211    let workspace_bytes = request.max_decoded.min(1_000_000_000);
212    let mut hash = Sha256::new();
213    let mut buffer = [0u8; 64 * 1024];
214    let mut bytes = 0u64;
215    loop {
216        if chrono::Utc::now() >= request.deadline {
217            return Err(LakeError::QueryNotAdmissible);
218        }
219        let read = source.read(&mut buffer)?;
220        if read == 0 {
221            break;
222        }
223        bytes += read as u64;
224        // Independent worker workspace bound; disclosure policy never silently clamps.
225        if bytes > workspace_bytes.saturating_add(64 * 1024) {
226            return Err(LakeError::QueryNotAdmissible);
227        }
228        hash.update(&buffer[..read]);
229        target.write_all(&buffer[..read])?;
230    }
231    target.flush()?;
232    if crate::file_staging::hex_digest(hash.finalize()) != request.input.datafile.digest.as_str()
233        || request
234            .input
235            .datafile
236            .size_bytes
237            .is_some_and(|size| size != bytes)
238    {
239        return Err(LakeError::QueryNotAdmissible);
240    }
241    crate::file_staging::decode_datafile_bounded(
242        &copied,
243        &request.directory.join("payload"),
244        &request.input.datafile,
245        workspace_bytes,
246        request.deadline,
247    )?;
248    let payload = request.directory.join("payload");
249    if !request.representation.0 {
250        std::fs::rename(copied, payload)?;
251    } else if !request.representation.1 {
252        let source = if request.input.datafile.encrypted.unwrap_or(false) {
253            payload.with_extension("decrypted")
254        } else {
255            copied
256        };
257        std::fs::rename(source, payload)?;
258    }
259    Ok(())
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use std::os::unix::fs::PermissionsExt;
266
267    #[test]
268    fn capture_is_immutable_checks_attestation_and_bounds_decoded_content() {
269        let Ok(worker) = std::env::var("AHRI_QUERY_WORKER_TEST_BINARY") else {
270            return;
271        };
272        let id = uuid::Uuid::new_v4().simple().to_string();
273        let root = std::env::temp_dir().join(format!("datafile_capture_{id}"));
274        std::fs::create_dir(&root).unwrap();
275        let wrapper = root.join("worker");
276        let quote = |text: &str| format!("'{}'", text.replace('\'', "'\\''"));
277        std::fs::write(
278            &wrapper,
279            format!(
280                "#!/bin/sh\nexport LD_LIBRARY_PATH={}\nexec {} \"$@\"\n",
281                quote(&std::env::var("LD_LIBRARY_PATH").unwrap_or_default()),
282                quote(&worker)
283            ),
284        )
285        .unwrap();
286        std::fs::set_permissions(&wrapper, std::fs::Permissions::from_mode(0o700)).unwrap();
287        let source = root.join("source.csv");
288        let content = b"participant,value\n1,secret\n";
289        std::fs::write(&source, content).unwrap();
290        let study = StudyId(uuid::Uuid::new_v4());
291        let asset = AssetId(uuid::Uuid::new_v4());
292        let staged = crate::DuckLakeAdapter::new(root.to_str().unwrap())
293            .stage_datafile(&crate::StageDataFileRequest {
294                source_path: source,
295                study_id: study,
296                asset_id: asset,
297                version_id: ahri_tre_types::VersionId(uuid::Uuid::new_v4()),
298                asset_name: ahri_tre_types::NcName::parse("content").unwrap(),
299                major: 1,
300                minor: 0,
301                patch: 0,
302                edam_format: "text/csv".into(),
303                compression: crate::DataFileCompression::Zstd,
304                encryption: crate::DataFileEncryption::Aes256Cbc,
305            })
306            .unwrap();
307        let binding =
308            attest_datafile(root.to_str().unwrap(), study, asset, &staged.datafile).unwrap();
309        let input = AdmittedDatafileInput {
310            binding,
311            datafile: staged.datafile.clone(),
312        };
313        let scratch_root = root.join("scratch");
314        std::fs::create_dir(&scratch_root).unwrap();
315        std::fs::set_permissions(&scratch_root, std::fs::Permissions::from_mode(0o700)).unwrap();
316        let scratch =
317            crate::TrustedScratch::open(&scratch_root, std::iter::empty::<&Path>()).unwrap();
318        let attempt = scratch
319            .create_attempt(ScratchAttemptId::new(&id).unwrap())
320            .unwrap();
321        let deadline = || chrono::Utc::now() + chrono::Duration::seconds(30);
322        assert!(
323            PreparedDisclosureDatafile::capture(
324                root.to_str().unwrap(),
325                input.clone(),
326                &attempt,
327                deadline(),
328                8,
329                &wrapper
330            )
331            .is_err()
332        );
333        let mut changed = input.clone();
334        changed.datafile.encryption_key = Some(vec![0; 32]);
335        assert!(
336            PreparedDisclosureDatafile::capture(
337                root.to_str().unwrap(),
338                changed,
339                &attempt,
340                deadline(),
341                1024,
342                &wrapper
343            )
344            .is_err()
345        );
346        let captured = PreparedDisclosureDatafile::capture(
347            root.to_str().unwrap(),
348            input.clone(),
349            &attempt,
350            deadline(),
351            1024,
352            &wrapper,
353        )
354        .unwrap();
355        std::fs::write(root.join(staged.storage_relative_path), b"replacement").unwrap();
356        let mut bytes = Vec::new();
357        captured.open().unwrap().read_to_end(&mut bytes).unwrap();
358        assert_eq!(bytes, content);
359        assert!(
360            PreparedDisclosureDatafile::capture(
361                root.to_str().unwrap(),
362                input,
363                &attempt,
364                deadline(),
365                1024,
366                &wrapper
367            )
368            .is_err()
369        );
370        attempt.cleanup().unwrap();
371        std::fs::remove_dir_all(root).unwrap();
372    }
373}
ahri_tre_lake/disclosure_datafile.rs

ahri_tre_lake/
disclosure_datafile.rs
