Skip to main content

submilli_engine/stdlib/code/
mod.rs

1//! Contained workspace navigation and anchor-based editing, gated by existing fs grants.
2use crate::runtime::host::{abi_arg, abi_result};
3mod budget;
4pub mod declaration;
5mod myers;
6mod patch;
7#[cfg(test)]
8mod tests;
9mod text;
10mod walk;
11
12use crate::runtime::fuel;
13use crate::runtime::{
14    StoreData,
15    host::{register_host_fn, write_submilli_string_struct_units},
16    intrinsic_types::build_intrinsic_types,
17    prelude::vtable::read_string_units,
18};
19use crate::stdlib::shared::{
20    atomic_write, check_security, check_security_call, contain_trap, require_writable,
21    resolve_content_or_trap,
22};
23use budget::{Budget, OutputBudget};
24use serde_json::{Value, json};
25use std::io::Read;
26use wasmtime::{Caller, FuncType, HeapType, Linker, RefType, Result, Val, ValType, bail};
27pub const MODULE_NAME: &str = "submilli:code";
28pub use declaration::package_declaration;
29
30pub fn install(linker: &mut Linker<StoreData>) -> Result<()> {
31    let engine = linker.engine().clone();
32    let intr = build_intrinsic_types(&engine)?;
33    let string = ValType::Ref(RefType::new(false, HeapType::ConcreteStruct(intr.string)));
34    let object = ValType::Ref(RefType::new(
35        false,
36        HeapType::ConcreteStruct(intr.object_shape.clone()),
37    ));
38    let options = ValType::Ref(RefType::new(true, HeapType::ConcreteStruct(intr.object)));
39    for (name, params, result) in [
40        (
41            "read",
42            vec![string.clone(), ValType::F64, ValType::F64],
43            object.clone(),
44        ),
45        ("search", vec![string.clone(), options], object.clone()),
46        ("glob", vec![string.clone()], object.clone()),
47        ("tree", vec![string.clone(), ValType::F64], object.clone()),
48        (
49            "edit",
50            vec![
51                string.clone(),
52                string.clone(),
53                string.clone(),
54                ValType::I32,
55                ValType::F64,
56            ],
57            object.clone(),
58        ),
59        (
60            "insertAt",
61            vec![string.clone(), ValType::F64, string.clone()],
62            object.clone(),
63        ),
64        (
65            "diffText",
66            vec![string.clone(), string.clone()],
67            string.clone(),
68        ),
69        (
70            "diffFiles",
71            vec![string.clone(), string.clone()],
72            string.clone(),
73        ),
74        ("applyPatch", vec![string.clone(), string], object),
75    ] {
76        register_host_fn(
77            linker,
78            MODULE_NAME,
79            crate::mangle::package_symbol(MODULE_NAME, name),
80            FuncType::new(&engine, params, [result]),
81            name == "diffText",
82            move |caller, params, results| {
83                let mut budget = Budget::new(caller);
84                *abi_result(results, 0)? = invoke(caller, &mut budget, name, params)?;
85                Ok(())
86            },
87        )?;
88    }
89    Ok(())
90}
91fn invoke(
92    caller: &mut Caller<'_, StoreData>,
93    budget: &mut Budget,
94    op: &str,
95    params: &[Val],
96) -> Result<Val> {
97    if op == "diffText" {
98        let a = units(caller, budget, abi_arg(params, 0)?)?;
99        let b = units(caller, budget, abi_arg(params, 1)?)?;
100        let result = diff(caller, budget, &a, &b)?;
101        return Ok(Val::AnyRef(Some(
102            write_submilli_string_struct_units(caller, &result)?.to_anyref(),
103        )));
104    }
105    let first = argument(caller, budget, abi_arg(params, 0)?)?;
106    let result = match op {
107        "search" => walk::search(caller, budget, &first, abi_arg(params, 1)?)?,
108        "glob" => walk::glob(caller, budget, &first)?,
109        "tree" => walk::tree(
110            caller,
111            budget,
112            &normalize(caller.data().vfs.cwd(), &first)?,
113            integer(abi_arg(params, 1)?, "depth", 0)?,
114        )?,
115        "read" => read_window(
116            caller,
117            budget,
118            &normalize(caller.data().vfs.cwd(), &first)?,
119            params,
120        )?,
121        "diffFiles" => {
122            let second = argument(caller, budget, abi_arg(params, 1)?)?;
123            let a = read_file(
124                caller,
125                budget,
126                &normalize(caller.data().vfs.cwd(), &first)?,
127                op,
128            )?;
129            let b = read_file(
130                caller,
131                budget,
132                &normalize(caller.data().vfs.cwd(), &second)?,
133                op,
134            )?;
135            budget.charge(caller, (a.len() + b.len()).saturating_mul(2))?;
136            let result = diff(
137                caller,
138                budget,
139                &a.encode_utf16().collect::<Vec<_>>(),
140                &b.encode_utf16().collect::<Vec<_>>(),
141            )?;
142            return Ok(Val::AnyRef(Some(
143                write_submilli_string_struct_units(caller, &result)?.to_anyref(),
144            )));
145        }
146        "edit" | "insertAt" | "applyPatch" => {
147            return mutate(
148                caller,
149                budget,
150                &normalize(caller.data().vfs.cwd(), &first)?,
151                op,
152                params,
153            );
154        }
155        _ => {
156            return Err(crate::runtime::host::invariant_trap(
157                "unknown code operation",
158            ));
159        }
160    };
161    encode(caller, budget, &result)
162}
163fn encode(caller: &mut Caller<'_, StoreData>, budget: &mut Budget, result: &Value) -> Result<Val> {
164    let encoded = serde_json::to_string(result)?;
165    budget.check_size(encoded.len())?;
166    budget.charge(caller, encoded.len().saturating_mul(4))?;
167    crate::stdlib::session::value::deserialize(caller, &encoded.encode_utf16().collect::<Vec<_>>())
168}
169fn read_window(
170    caller: &mut Caller<'_, StoreData>,
171    budget: &mut Budget,
172    path: &str,
173    params: &[Val],
174) -> Result<Value> {
175    let start = integer(abi_arg(params, 1)?, "offset", 1)? - 1;
176    let limit = integer(abi_arg(params, 2)?, "limit", 1)?;
177    let source = read_file(caller, budget, path, "read")?;
178    budget.charge(
179        caller,
180        source
181            .bytes()
182            .filter(|&b| b == b'\n')
183            .count()
184            .saturating_add(1)
185            .saturating_mul(32),
186    )?;
187    let lines = text::lines(source.strip_prefix('\u{feff}').unwrap_or(&source));
188    let mut output = OutputBudget::new();
189    let mut values = Vec::new();
190    for (index, line) in lines.iter().enumerate().skip(start).take(limit) {
191        if !output.reserve(caller, budget, OutputBudget::line_bytes(line))? {
192            break;
193        }
194        values.push(json!({"line":index+1,"text":text::line_text(line)}));
195    }
196    Ok(
197        json!({"path":path,"truncated":start.saturating_add(values.len()) < lines.len(),"lines":values}),
198    )
199}
200fn mutate(
201    caller: &mut Caller<'_, StoreData>,
202    budget: &mut Budget,
203    path: &str,
204    op: &str,
205    params: &[Val],
206) -> Result<Val> {
207    let original = read_file(caller, budget, path, op)?;
208    budget.charge(caller, original.len().saturating_mul(64))?;
209    // Splicing the change back into the original rewrites the whole text.
210    fuel::charge(&mut *caller, fuel::SCAN, original.len() as u64)?;
211    let change = match op {
212        "edit" => prepare_edit(caller, budget, &original, params)?,
213        "insertAt" => text::insert(
214            &original,
215            integer(abi_arg(params, 1)?, "line", 1)?,
216            &argument(caller, budget, abi_arg(params, 2)?)?,
217        )?,
218        "applyPatch" => {
219            let patch = argument(caller, budget, abi_arg(params, 1)?)?;
220            budget.charge(caller, patch.len().saturating_mul(32))?;
221            // Each hunk is located by scanning the original.
222            fuel::charge(
223                &mut *caller,
224                fuel::SCAN,
225                (original.len() as u64).saturating_mul(patch.lines().count().max(1) as u64),
226            )?;
227            fuel::charge(&mut *caller, fuel::PARSE, patch.len() as u64)?;
228            patch::apply(&original, &patch)?
229        }
230        _ => {
231            return Err(crate::runtime::host::invariant_trap(
232                "unknown code mutation",
233            ));
234        }
235    };
236    budget.check_size(change.text.len())?;
237    let success = change.diagnostics.is_empty();
238    let changed = success && change.text != original;
239    let patch = if changed {
240        String::from_utf16(&diff(
241            caller,
242            budget,
243            &original.encode_utf16().collect::<Vec<_>>(),
244            &change.text.encode_utf16().collect::<Vec<_>>(),
245        )?)?
246    } else {
247        String::new()
248    };
249    let ticket = check_security_call(
250        &mut *caller,
251        "fs.write",
252        json!({"path":path, "length": change.text.len(), "diff": patch}),
253    )?;
254    // Allocate the return value before commit so a guest allocation failure cannot hide a write.
255    let result =
256        json!({"success":success,"changed":changed,"diff":patch,"diagnostics":change.diagnostics});
257    let result = encode(caller, budget, &result)?;
258    if changed {
259        let resolved = resolve_content_or_trap(caller.data(), path, op)?;
260        require_writable(&*caller, ticket, resolved.placement(), "fs.write", path)?;
261        let permissions = resolved
262            .metadata()
263            .map_err(|e| contain_trap(op, path, &e))?
264            .permissions();
265        let quota = resolved.placement().quota().cloned();
266        atomic_write(
267            &resolved,
268            change.text.as_bytes(),
269            Some(permissions),
270            path,
271            op,
272            quota,
273        )?;
274    }
275    Ok(result)
276}
277fn prepare_edit(
278    caller: &mut Caller<'_, StoreData>,
279    budget: &mut Budget,
280    original: &str,
281    params: &[Val],
282) -> Result<text::Edit> {
283    let old = argument(caller, budget, abi_arg(params, 1)?)?;
284    let new = argument(caller, budget, abi_arg(params, 2)?)?;
285    // Naive substring search: every position against the whole needle.
286    fuel::charge(
287        &mut *caller,
288        fuel::SCAN,
289        (original.len() as u64).saturating_mul(old.len().max(1) as u64),
290    )?;
291    if old.is_empty() {
292        bail!("code.edit: oldString must not be empty; use insertAt");
293    }
294    let matches = original.match_indices(&old).count();
295    let replace_all = match abi_arg(params, 3)? {
296        Val::I32(value) => *value != 0,
297        _ => {
298            return Err(crate::runtime::host::invariant_trap(
299                "code.edit: invalid boolean argument",
300            ));
301        }
302    };
303    let selected = if replace_all {
304        matches
305    } else {
306        usize::from(matches > 0)
307    };
308    budget.charge(
309        caller,
310        text::MAX_DIAGNOSTICS
311            .min(original.len().saturating_add(1))
312            .saturating_mul(1024),
313    )?;
314    let length = original
315        .len()
316        .saturating_sub(selected.saturating_mul(old.len()))
317        .saturating_add(selected.saturating_mul(new.len()));
318    budget.check_size(length)?;
319    budget.charge(caller, length.saturating_mul(2))?;
320    let maximum = if matches == 0 {
321        budget.max_bytes.min(
322            original
323                .len()
324                .saturating_mul(16)
325                .saturating_add(new.len().saturating_mul(4)),
326        )
327    } else {
328        budget.max_bytes
329    };
330    text::replace(
331        original,
332        &old,
333        &new,
334        replace_all,
335        integer(abi_arg(params, 4)?, "nearLine", 0)?,
336        maximum,
337    )
338}
339
340fn normalize(cwd: &str, path: &str) -> Result<String> {
341    crate::runtime::fs::guest_normalize(cwd, path)
342        .map_err(|e| wasmtime::Error::msg(format!("code: {e}")))
343}
344fn gate(caller: &mut Caller<'_, StoreData>, capability: &str, path: &str) -> Result<()> {
345    check_security(caller, capability, json!({"path":path,"recursive":true}))
346}
347fn gate_call(
348    caller: &mut Caller<'_, StoreData>,
349    capability: &str,
350    path: &str,
351) -> Result<Option<crate::runtime::CallTicket>> {
352    check_security_call(caller, capability, json!({"path":path,"recursive":true}))
353}
354fn read_file(
355    caller: &mut Caller<'_, StoreData>,
356    budget: &mut Budget,
357    path: &str,
358    op: &str,
359) -> Result<String> {
360    read_contents(caller, budget, path, op, false)
361}
362fn read_contents(
363    caller: &mut Caller<'_, StoreData>,
364    budget: &mut Budget,
365    path: &str,
366    op: &str,
367    skip_binary: bool,
368) -> Result<String> {
369    let ticket = gate_call(caller, "fs.read", path)?;
370    let resolved = resolve_content_or_trap(caller.data(), path, op)?;
371    let (file, metadata) = resolved
372        .open_regular()
373        .map_err(|e| contain_trap(op, path, &e))?;
374    let len = usize::try_from(metadata.len())?;
375    budget.check_size(len)?;
376    budget.charge(caller, len.saturating_mul(8).saturating_add(1024))?;
377    fuel::charge(&mut *caller, fuel::IO, len as u64)?;
378    fuel::charge(&mut *caller, fuel::SCAN, len as u64)?;
379    let mut bytes = Vec::with_capacity(len);
380    file.take((len as u64).saturating_add(1))
381        .read_to_end(&mut bytes)?;
382    if bytes.len() > len {
383        bail!("code.{op}: file grew during read; retry");
384    }
385    crate::runtime::call_log::record_payload(
386        &*caller,
387        ticket,
388        crate::runtime::call_log::Side::Response,
389        || crate::runtime::call_log::Payload::meta(Value::Null).with_body(&bytes),
390    );
391    if skip_binary && bytes.contains(&0) {
392        return Ok(String::new());
393    }
394    String::from_utf8(bytes)
395        .map_err(|_| wasmtime::Error::msg(format!("code.{op}: {path} is not valid UTF-8")))
396}
397fn argument(caller: &mut Caller<'_, StoreData>, budget: &mut Budget, val: &Val) -> Result<String> {
398    String::from_utf16(&units(caller, budget, val)?).map_err(|_| {
399        wasmtime::Error::msg("code: filesystem text must not contain lone UTF-16 surrogates")
400    })
401}
402fn units(caller: &mut Caller<'_, StoreData>, budget: &mut Budget, val: &Val) -> Result<Vec<u16>> {
403    let value = read_string_units(caller, val, "code argument")?;
404    budget.check_size(value.len().saturating_mul(2))?;
405    budget.charge(caller, value.len().saturating_mul(4))?;
406    Ok(value)
407}
408fn diff(
409    caller: &mut Caller<'_, StoreData>,
410    budget: &mut Budget,
411    a: &[u16],
412    b: &[u16],
413) -> Result<Vec<u16>> {
414    budget.charge(caller, (a.len() + b.len()).saturating_mul(16))?;
415    let na = text::line_count(a);
416    let nb = text::line_count(b);
417    budget.charge(caller, (na + nb).saturating_mul(1024))?;
418    text::check_diff_size(na, nb)?;
419    fuel::charge(&mut *caller, fuel::SCAN, (a.len() + b.len()) as u64)?;
420    let result = text::diff_charged(a, b, |steps| fuel::charge(&mut *caller, fuel::PARSE, steps))?;
421    budget.check_size(result.len().saturating_mul(2))?;
422    Ok(result)
423}
424fn integer(val: &Val, name: &str, min: usize) -> Result<usize> {
425    let n = match val {
426        Val::F64(bits) => f64::from_bits(*bits),
427        _ => {
428            return Err(crate::runtime::host::invariant_trap(
429                "code: invalid numeric argument",
430            ));
431        }
432    };
433    if !n.is_finite() || n.fract() != 0.0 || n < min as f64 || n > u32::MAX as f64 {
434        bail!(
435            "code: {name} must be an integer between {min} and {}",
436            u32::MAX
437        );
438    }
439    Ok(n as usize)
440}