ahri_tre_lake/
capture_worker.rs

1//! Fixed source capture protocol. This process can read the approved Lake, but
2//! accepts only adapter-generated relation reads. User SQL runs in a different,
3//! authority-stripped process after this one has exited and released credentials.
4use crate::{DisclosureDatasetInput, DuckLakeOpenedCatalog, LakeError};
5use duckdb::{
6    Connection,
7    arrow::record_batch::RecordBatch,
8    vtab::arrow::{ArrowVTab, arrow_recordbatch_to_query_params},
9};
10use serde::{Deserialize, Serialize};
11use std::{
12    io::{Read, Write},
13    path::{Path, PathBuf},
14    process::{Command, Stdio},
15    time::Duration,
16};
17
18#[derive(Clone)]
19pub struct DisclosureCaptureAuthority {
20    metadata_schema: String,
21    extensions: Vec<String>,
22    setup: Vec<std::sync::Arc<ahri_tre_secrets::SecretMaterial>>,
23    attach: String,
24}
25#[derive(Serialize, Deserialize)]
26struct CaptureSource {
27    metadata_schema: String,
28    extensions: Vec<String>,
29    setup: Vec<String>,
30    attach: String,
31}
32impl DisclosureCaptureAuthority {
33    /// Derived from an already approved Lake connection; neither SQL nor Secret
34    /// material can be supplied through the public disclosure request.
35    pub fn from_opened(opened: &mut DuckLakeOpenedCatalog) -> Result<Self, LakeError> {
36        let mut statement = opened.connection.prepare("SELECT install_path FROM duckdb_extensions() WHERE loaded AND installed AND extension_name IN ('postgres_scanner','ducklake','httpfs','azure') ORDER BY extension_name")?;
37        let extensions = statement
38            .query_map([], |row| row.get::<_, String>(0))?
39            .collect::<Result<Vec<_>, _>>()?;
40        if !extensions
41            .iter()
42            .any(|path| path.ends_with("ducklake.duckdb_extension"))
43            || opened.attach_plan.create_if_not_exists
44        {
45            return Err(LakeError::QueryNotAdmissible);
46        }
47        Ok(Self {
48            metadata_schema: opened.attach_plan.metadata_schema.clone(),
49            extensions,
50            setup: std::mem::take(&mut opened.attach_plan.pre_attach_sql)
51                .into_iter()
52                .map(|sql| {
53                    ahri_tre_secrets::SecretMaterial::try_from(sql.into_bytes())
54                        .map(std::sync::Arc::new)
55                        .map_err(|_| LakeError::QueryNotAdmissible)
56                })
57                .collect::<Result<_, _>>()?,
58            attach: opened.attach_plan.attach_sql.clone(),
59        })
60    }
61}
62impl std::fmt::Debug for DisclosureCaptureAuthority {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        f.write_str("DisclosureCaptureAuthority(<resolved>)")
65    }
66}
67
68#[derive(Serialize, Deserialize)]
69struct CaptureRequest {
70    authority: CaptureSource,
71    inputs: Vec<DisclosureDatasetInput>,
72    output: PathBuf,
73    deadline: chrono::DateTime<chrono::Utc>,
74}
75
76pub(crate) fn capture(
77    authority: &DisclosureCaptureAuthority,
78    inputs: &[DisclosureDatasetInput],
79    directory: &Path,
80    deadline: chrono::DateTime<chrono::Utc>,
81    executable: &Path,
82) -> Result<(), LakeError> {
83    if deadline <= chrono::Utc::now() {
84        return Err(LakeError::QueryNotAdmissible);
85    }
86    let request = CaptureRequest {
87        authority: CaptureSource {
88            metadata_schema: authority.metadata_schema.clone(),
89            extensions: authority.extensions.clone(),
90            attach: authority.attach.clone(),
91            setup: authority
92                .setup
93                .iter()
94                .map(|secret| {
95                    secret.expose(|bytes| {
96                        String::from_utf8(bytes.to_vec()).map_err(|_| LakeError::QueryNotAdmissible)
97                    })
98                })
99                .collect::<Result<_, _>>()?,
100        },
101        inputs: inputs.to_vec(),
102        output: directory.join("inputs.duckdb"),
103        deadline,
104    };
105    let encoded = serde_json::to_vec(&request).map_err(|_| LakeError::QueryNotAdmissible)?;
106    if encoded.len() > 1024 * 1024 {
107        return Err(LakeError::QueryNotAdmissible);
108    }
109    let mut command = Command::new(executable);
110    crate::disclosure_inputs::parent_bound(&mut command)?;
111    let mut child = command
112        .env_clear()
113        .arg("--capture")
114        .stdin(Stdio::piped())
115        .stdout(Stdio::null())
116        .stderr(Stdio::null())
117        .spawn()?;
118    let sent = child
119        .stdin
120        .take()
121        .ok_or(LakeError::QueryNotAdmissible)?
122        .write_all(&encoded);
123    if sent.is_err() {
124        let _ = child.kill();
125        let _ = child.wait();
126        return Err(LakeError::QueryNotAdmissible);
127    }
128    loop {
129        match child.try_wait() {
130            Ok(Some(status)) if status.success() => return Ok(()),
131            Ok(Some(_)) => return Err(LakeError::QueryNotAdmissible),
132            Ok(None) if chrono::Utc::now() < deadline => {
133                std::thread::sleep(Duration::from_millis(10))
134            }
135            _ => {
136                let _ = child.kill();
137                let _ = child.wait();
138                return Err(LakeError::QueryNotAdmissible);
139            }
140        }
141    }
142}
143
144/// Internal executable entry point. Its stdin is a private pipe from the Lake
145/// adapter, never an HTTP request. Errors are deliberately content-free.
146#[doc(hidden)]
147pub fn run_capture_worker() -> Result<(), LakeError> {
148    let mut bytes = Vec::new();
149    std::io::stdin().take(1_048_577).read_to_end(&mut bytes)?;
150    if bytes.len() > 1_048_576 {
151        return Err(LakeError::QueryNotAdmissible);
152    }
153    let request: CaptureRequest =
154        serde_json::from_slice(&bytes).map_err(|_| LakeError::QueryNotAdmissible)?;
155    let seconds = (request.deadline - chrono::Utc::now()).num_milliseconds();
156    if seconds <= 0
157        || seconds > i64::from(u32::MAX) * 1000
158        || request.inputs.is_empty()
159        || request.inputs.len() > 128
160    {
161        return Err(LakeError::QueryNotAdmissible);
162    }
163    #[cfg(target_os = "linux")]
164    // SAFETY: alarm receives a bounded integer; the process has no signal handler.
165    unsafe {
166        libc::alarm(((seconds + 999) / 1000) as u32);
167    }
168    let config = || {
169        duckdb::Config::default()
170            .threads(1)?
171            .max_memory("256MB")?
172            .with("temp_directory", "")
173    };
174    let mut source = Connection::open_in_memory_with_flags(config()?)?;
175    for extension in request.authority.extensions {
176        source.execute_batch(&format!("LOAD '{}'", extension.replace('\'', "''")))?;
177    }
178    for setup in request.authority.setup {
179        source
180            .execute_batch(&setup)
181            .map_err(|_| LakeError::QueryNotAdmissible)?;
182    }
183    source
184        .execute_batch(&request.authority.attach)
185        .map_err(|_| LakeError::QueryNotAdmissible)?;
186    let output = Connection::open_with_flags(&request.output, config()?)?;
187    output.register_table_function::<ArrowVTab>("tre_input_batch")?;
188    let transaction = source.transaction()?;
189    let mut total_bytes = 0usize;
190    let mut total_rows = 0usize;
191    for input in request.inputs {
192        let relation = crate::dataset_loading::dataset_table_relation(
193            input.study,
194            input.name.as_str(),
195            input.major,
196            input.minor,
197            input.patch,
198        );
199        let (identity, _) = crate::dataset_physical_identity(
200            &transaction,
201            &request.authority.metadata_schema,
202            input.study,
203            input.name.as_str(),
204            input.major,
205            input.minor,
206            input.patch,
207        )?;
208        if identity != input.table_uuid || input.lake_snapshot < 0 {
209            return Err(LakeError::QueryNotAdmissible);
210        }
211        let sql = format!(
212            "SELECT * FROM \"{}\".\"{}\".\"{}\" AT (VERSION => {}) LIMIT 10000001",
213            crate::LAKE_ALIAS,
214            relation.schema_name.replace('"', "\"\""),
215            relation.table_name.replace('"', "\"\""),
216            input.lake_snapshot
217        );
218        let table = format!("input_{}", input.version.0.simple());
219        let mut statement = transaction.prepare(&sql)?;
220        let batches = statement.query_arrow([])?;
221        let schema = batches.get_schema();
222        if schema.fields().len() > 512 {
223            return Err(LakeError::QueryNotAdmissible);
224        }
225        output.execute(
226            &format!("CREATE TABLE {table} AS SELECT * FROM tre_input_batch(?,?)"),
227            arrow_recordbatch_to_query_params(RecordBatch::new_empty(schema)),
228        )?;
229        let mut appender = output.appender(&table)?;
230        let mut rows = 0;
231        for batch in batches {
232            rows += batch.num_rows();
233            total_rows += batch.num_rows();
234            total_bytes = total_bytes
235                .checked_add(batch.get_array_memory_size())
236                .ok_or(LakeError::QueryNotAdmissible)?;
237            if total_rows > 10_000_000
238                || total_bytes > 1_000_000_000
239                || batch.get_array_memory_size() > 8 * 1024 * 1024
240                || chrono::Utc::now() >= request.deadline
241            {
242                return Err(LakeError::QueryNotAdmissible);
243            }
244            appender.append_record_batch(batch)?;
245        }
246        appender.flush()?;
247        if rows != statement.row_count() {
248            return Err(LakeError::QueryNotAdmissible);
249        }
250    }
251    transaction.commit()?;
252    output.execute_batch("CHECKPOINT")?;
253    if chrono::Utc::now() >= request.deadline {
254        return Err(LakeError::QueryNotAdmissible);
255    }
256    Ok(())
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262    #[test]
263    fn supervised_capture_and_restricted_query_use_only_private_inputs() {
264        if std::env::var("RUN_LIVE_GOVERNANCE").as_deref() != Ok("true") {
265            return;
266        }
267        use std::os::unix::fs::PermissionsExt;
268        let worker = std::env::var("AHRI_QUERY_WORKER_TEST_BINARY")
269            .expect("live acceptance requires the built worker");
270        let id = uuid::Uuid::new_v4().simple().to_string();
271        let root = std::env::temp_dir().join(format!("capture-{id}"));
272        std::fs::create_dir(&root).unwrap();
273        std::fs::set_permissions(&root, std::fs::Permissions::from_mode(0o700)).unwrap();
274        let wrapper = root.join("worker");
275        let quote = |value: &str| format!("'{}'", value.replace('\'', "'\\''"));
276        std::fs::write(
277            &wrapper,
278            format!(
279                "#!/bin/sh\nexport LD_LIBRARY_PATH={}\nexec {} \"$@\"\n",
280                quote(&std::env::var("LD_LIBRARY_PATH").unwrap_or_default()),
281                quote(&worker)
282            ),
283        )
284        .unwrap();
285        std::fs::set_permissions(&wrapper, std::fs::Permissions::from_mode(0o700)).unwrap();
286        let connection = Connection::open_in_memory().unwrap();
287        let attach = format!(
288            "ATTACH 'ducklake:postgres:host=127.0.0.1 port=55409 user=postgres dbname=ticket09' AS tre_lake (DATA_PATH '{}', METADATA_SCHEMA 'capture_{id}')",
289            root.join("data").display()
290        );
291        connection
292            .execute_batch("INSTALL postgres; LOAD postgres; INSTALL ducklake; LOAD ducklake;")
293            .unwrap();
294        connection.execute_batch(&attach).unwrap();
295        let mut input = DisclosureDatasetInput {
296            table_uuid: uuid::Uuid::nil(),
297            lake_snapshot: 0,
298            study: ahri_tre_types::StudyId(uuid::Uuid::new_v4()),
299            version: ahri_tre_types::VersionId(uuid::Uuid::new_v4()),
300            name: ahri_tre_types::NcName::parse("source").unwrap(),
301            major: 1,
302            minor: 0,
303            patch: 0,
304        };
305        let relation = crate::dataset_loading::dataset_table_relation(
306            input.study,
307            input.name.as_str(),
308            1,
309            0,
310            0,
311        );
312        connection.execute_batch(&format!("CREATE SCHEMA tre_lake.\"{}\"; CREATE TABLE tre_lake.{} AS SELECT * FROM (VALUES (2,'a'),(3,'b')) v(x,label)",relation.schema_name,relation.qualified_name())).unwrap();
313        (input.table_uuid, input.lake_snapshot) = crate::dataset_physical_identity(
314            &connection,
315            &format!("capture_{id}"),
316            input.study,
317            input.name.as_str(),
318            1,
319            0,
320            0,
321        )
322        .unwrap();
323        let extensions = connection.prepare("SELECT install_path FROM duckdb_extensions() WHERE loaded AND installed AND extension_name IN ('postgres_scanner','ducklake') ORDER BY extension_name").unwrap().query_map([], |row| row.get::<_,String>(0)).unwrap().collect::<Result<Vec<_>,_>>().unwrap();
324        let authority = DisclosureCaptureAuthority {
325            metadata_schema: format!("capture_{id}"),
326            extensions,
327            setup: Vec::new(),
328            attach,
329        };
330        let scratch_root = root.join("scratch");
331        std::fs::create_dir(&scratch_root).unwrap();
332        std::fs::set_permissions(&scratch_root, std::fs::Permissions::from_mode(0o700)).unwrap();
333        let scratch =
334            crate::TrustedScratch::open(&scratch_root, std::iter::empty::<&Path>()).unwrap();
335        let attempt = scratch
336            .create_attempt(crate::ScratchAttemptId::new(&id).unwrap())
337            .unwrap();
338        let bindings = std::collections::BTreeMap::from([(
339            "source".into(),
340            crate::QueryBinding::Dataset(input.version.0),
341        )]);
342        let plan = crate::analyze_disclosure_query("SELECT sum(x) FROM source", &bindings).unwrap();
343        let prepared = crate::PreparedDisclosureQuery::capture(
344            &authority,
345            plan,
346            &[input.clone()],
347            &attempt,
348            chrono::Utc::now() + chrono::Duration::seconds(30),
349            &wrapper,
350        )
351        .unwrap();
352        // Subsequent Lake mutation cannot change this captured input version.
353        connection
354            .execute_batch(&format!(
355                "DELETE FROM tre_lake.{}",
356                relation.qualified_name()
357            ))
358            .unwrap();
359        let stream = prepared.execute(&wrapper, 30, 1_000_000, 100).unwrap();
360        let batches = arrow_ipc::reader::StreamReader::try_new(stream, None)
361            .unwrap()
362            .collect::<Result<Vec<_>, _>>()
363            .unwrap();
364        assert_eq!(
365            batches[0]
366                .column(0)
367                .as_any()
368                .downcast_ref::<duckdb::arrow::array::Decimal128Array>()
369                .unwrap()
370                .value(0),
371            5
372        );
373        let plan = crate::analyze_dataset_transform(
374            &format!("SELECT sum(x) AS total FROM {}", relation.qualified_name()),
375            input.study,
376            &input.name,
377            input.version,
378            1,
379            0,
380            0,
381        )
382        .unwrap();
383        let prepared = crate::PreparedDisclosureQuery::capture(
384            &authority,
385            plan,
386            &[input.clone()],
387            &attempt,
388            chrono::Utc::now() + chrono::Duration::seconds(30),
389            &wrapper,
390        )
391        .unwrap();
392        let output = prepared
393            .materialize(
394                &wrapper,
395                chrono::Utc::now() + chrono::Duration::seconds(30),
396                1_000_000,
397                100,
398            )
399            .unwrap();
400        let result: i64 = connection
401            .query_row(
402                "SELECT CAST(total AS BIGINT) FROM read_parquet(?)",
403                [output.path.to_string_lossy().as_ref()],
404                |row| row.get(0),
405            )
406            .unwrap();
407        assert_eq!(result, 5);
408        drop(output);
409        connection
410            .execute_batch(&format!(
411                "INSERT INTO tre_lake.{} VALUES (1,repeat('x',16000000))",
412                relation.qualified_name()
413            ))
414            .unwrap();
415        (_, input.lake_snapshot) = crate::dataset_physical_identity(
416            &connection,
417            &format!("capture_{id}"),
418            input.study,
419            input.name.as_str(),
420            1,
421            0,
422            0,
423        )
424        .unwrap();
425        let plan = crate::analyze_disclosure_query("SELECT * FROM source", &bindings).unwrap();
426        assert!(
427            crate::PreparedDisclosureQuery::capture(
428                &authority,
429                plan,
430                &[input],
431                &attempt,
432                chrono::Utc::now() + chrono::Duration::seconds(30),
433                &wrapper
434            )
435            .is_err()
436        );
437        drop(attempt);
438        drop(scratch);
439        drop(connection);
440        std::fs::remove_dir_all(root).unwrap();
441    }
442}
ahri_tre_lake/capture_worker.rs

ahri_tre_lake/
capture_worker.rs
