viceroy_lib/
linking.rs

1//! Linking and name resolution.
2
3use {
4    crate::{
5        Error, body::Body, config::ExperimentalModule, execute::ExecuteCtx, logging::LogEndpoint,
6        session::Session, wiggle_abi,
7    },
8    hyper::Response,
9    std::collections::HashSet,
10    wasmtime::{GuestProfiler, Linker, Store, StoreLimits, StoreLimitsBuilder, UpdateDeadline},
11    wasmtime_wasi::p1::WasiP1Ctx,
12    wasmtime_wasi_nn::witx::WasiNnCtx,
13};
14
15pub struct Limiter {
16    /// Total memory allocated so far.
17    pub memory_allocated: usize,
18    /// The internal limiter we use to actually answer calls
19    internal: StoreLimits,
20}
21
22impl Limiter {
23    pub fn for_wasip2() -> Self {
24        Self::new(100, 100, 100)
25    }
26
27    pub fn for_wasip1() -> Self {
28        Self::new(1, 1, 1)
29    }
30
31    fn new(max_instances: usize, max_memories: usize, max_tables: usize) -> Self {
32        Limiter {
33            memory_allocated: 0,
34            internal: StoreLimitsBuilder::new()
35                .instances(max_instances)
36                .memories(max_memories)
37                .memory_size(128 * 1024 * 1024)
38                .table_elements(98765)
39                .tables(max_tables)
40                .build(),
41        }
42    }
43}
44
45impl wasmtime::ResourceLimiter for Limiter {
46    fn memory_growing(
47        &mut self,
48        current: usize,
49        desired: usize,
50        maximum: Option<usize>,
51    ) -> anyhow::Result<bool> {
52        // limit the amount of memory that an instance can use to (roughly) 128MB, erring on
53        // the side of letting things run that might get killed on Compute, because we are not
54        // tracking some runtime factors in this count.
55        let result = self.internal.memory_growing(current, desired, maximum);
56
57        if matches!(result, Ok(true)) {
58            // Track the diff in memory allocated over time. As each instance will start with 0 and
59            // gradually resize, this will track the total allocations throughout the lifetime of the
60            // instance.
61            self.memory_allocated += desired - current;
62        }
63
64        result
65    }
66
67    fn table_growing(
68        &mut self,
69        current: usize,
70        desired: usize,
71        maximum: Option<usize>,
72    ) -> anyhow::Result<bool> {
73        self.internal.table_growing(current, desired, maximum)
74    }
75
76    fn memory_grow_failed(&mut self, error: anyhow::Error) -> anyhow::Result<()> {
77        self.internal.memory_grow_failed(error)
78    }
79
80    fn table_grow_failed(&mut self, error: anyhow::Error) -> anyhow::Result<()> {
81        self.internal.table_grow_failed(error)
82    }
83
84    fn instances(&self) -> usize {
85        self.internal.instances()
86    }
87
88    fn tables(&self) -> usize {
89        self.internal.tables()
90    }
91
92    fn memories(&self) -> usize {
93        self.internal.memories()
94    }
95}
96
97#[allow(unused)]
98pub struct ComponentCtx {
99    pub wasi_ctx: wasmtime_wasi::WasiCtx,
100    pub wasi_table: wasmtime_wasi::ResourceTable,
101    pub wasi_random: wasmtime_wasi::random::WasiRandomCtx,
102    pub(crate) session: Session,
103    guest_profiler: Option<Box<GuestProfiler>>,
104}
105
106/// An extension trait for users of `ComponentCtx` to access the session.
107pub trait SessionView {
108    fn session(&self) -> &Session;
109    fn session_mut(&mut self) -> &mut Session;
110}
111
112impl SessionView for ComponentCtx {
113    fn session(&self) -> &Session {
114        &self.session
115    }
116
117    fn session_mut(&mut self) -> &mut Session {
118        &mut self.session
119    }
120}
121
122impl ComponentCtx {
123    pub fn wasi(&mut self) -> &mut wasmtime_wasi::WasiCtx {
124        &mut self.wasi_ctx
125    }
126
127    pub fn session(&mut self) -> &mut Session {
128        &mut self.session
129    }
130
131    pub fn take_guest_profiler(&mut self) -> Option<Box<GuestProfiler>> {
132        self.guest_profiler.take()
133    }
134
135    pub fn limiter(&self) -> &Limiter {
136        self.session.limiter()
137    }
138
139    pub fn close_downstream_response_sender(&mut self, resp: Response<Body>) {
140        self.session.close_downstream_response_sender(resp)
141    }
142
143    /// Initialize a new [`Store`][store], given an [`ExecuteCtx`][ctx].
144    ///
145    /// [ctx]: ../wiggle_abi/struct.ExecuteCtx.html
146    /// [store]: https://docs.rs/wasmtime/latest/wasmtime/struct.Store.html
147    pub(crate) fn create_store(
148        ctx: &ExecuteCtx,
149        session: Session,
150        guest_profiler: Option<GuestProfiler>,
151        extra_init: impl FnOnce(&mut wasmtime_wasi::WasiCtxBuilder),
152    ) -> Result<Store<ComponentCtx>, anyhow::Error> {
153        let mut builder = make_wasi_ctx(ctx, &session);
154
155        extra_init(&mut builder);
156
157        let wasm_ctx = Self {
158            wasi_table: wasmtime_wasi::ResourceTable::new(),
159            wasi_ctx: builder.build(),
160            wasi_random: wasmtime_wasi::random::WasiRandomCtx::default(),
161            session,
162            guest_profiler: guest_profiler.map(Box::new),
163        };
164        let mut store = Store::new(ctx.engine(), wasm_ctx);
165        store.set_epoch_deadline(1);
166
167        // instrument hostcalls to have those show up in profiles
168        store.call_hook(|mut store, kind| {
169            if let Some(mut prof) = store.data_mut().guest_profiler.take() {
170                prof.call_hook(&store, kind);
171                store.data_mut().guest_profiler = Some(prof);
172            }
173            Ok(())
174        });
175
176        // sampling profiler
177        store.epoch_deadline_callback(|mut store| {
178            if let Some(mut prof) = store.data_mut().guest_profiler.take() {
179                prof.sample(&store, std::time::Duration::ZERO);
180                store.data_mut().guest_profiler = Some(prof);
181            }
182            Ok(UpdateDeadline::Yield(1))
183        });
184
185        store.limiter(|ctx| ctx.session.limiter_mut());
186        Ok(store)
187    }
188}
189
190impl wasmtime_wasi::WasiView for ComponentCtx {
191    fn ctx(&mut self) -> wasmtime_wasi::WasiCtxView<'_> {
192        wasmtime_wasi::WasiCtxView {
193            ctx: &mut self.wasi_ctx,
194            table: &mut self.wasi_table,
195        }
196    }
197}
198
199impl wasmtime_wasi_io::IoView for ComponentCtx {
200    fn table(&mut self) -> &mut wasmtime_wasi::ResourceTable {
201        &mut self.wasi_table
202    }
203}
204
205pub struct WasmCtx {
206    wasi: WasiP1Ctx,
207    wasi_nn: WasiNnCtx,
208    session: Session,
209    guest_profiler: Option<Box<GuestProfiler>>,
210}
211
212impl WasmCtx {
213    pub fn wasi(&mut self) -> &mut WasiP1Ctx {
214        &mut self.wasi
215    }
216
217    fn wasi_nn(&mut self) -> &mut WasiNnCtx {
218        &mut self.wasi_nn
219    }
220
221    pub fn session(&mut self) -> &mut Session {
222        &mut self.session
223    }
224
225    pub fn take_guest_profiler(&mut self) -> Option<Box<GuestProfiler>> {
226        self.guest_profiler.take()
227    }
228
229    pub fn limiter(&self) -> &Limiter {
230        self.session.limiter()
231    }
232}
233
234impl WasmCtx {
235    pub fn close_downstream_response_sender(&mut self, resp: Response<Body>) {
236        self.session.close_downstream_response_sender(resp)
237    }
238}
239
240/// Initialize a new [`Store`][store], given an [`ExecuteCtx`][ctx].
241///
242/// [ctx]: ../wiggle_abi/struct.ExecuteCtx.html
243/// [store]: https://docs.rs/wasmtime/latest/wasmtime/struct.Store.html
244pub(crate) fn create_store(
245    ctx: &ExecuteCtx,
246    session: Session,
247    guest_profiler: Option<GuestProfiler>,
248    extra_init: impl FnOnce(&mut wasmtime_wasi::WasiCtxBuilder),
249) -> Result<Store<WasmCtx>, anyhow::Error> {
250    let mut builder = make_wasi_ctx(ctx, &session);
251
252    extra_init(&mut builder);
253
254    let wasi = builder.build_p1();
255    let (backends, registry) = wasmtime_wasi_nn::preload(&[])?;
256    let wasi_nn = WasiNnCtx::new(backends, registry);
257    let wasm_ctx = WasmCtx {
258        wasi,
259        wasi_nn,
260        session,
261        guest_profiler: guest_profiler.map(Box::new),
262    };
263    let mut store = Store::new(ctx.engine(), wasm_ctx);
264    store.set_epoch_deadline(1);
265
266    // instrument hostcalls to have those show up in profiles
267    store.call_hook(|mut store, kind| {
268        if let Some(mut prof) = store.data_mut().guest_profiler.take() {
269            prof.call_hook(&store, kind);
270            store.data_mut().guest_profiler = Some(prof);
271        }
272        Ok(())
273    });
274
275    // sampling profiler
276    store.epoch_deadline_callback(|mut store| {
277        if let Some(mut prof) = store.data_mut().guest_profiler.take() {
278            prof.sample(&store, std::time::Duration::ZERO);
279            store.data_mut().guest_profiler = Some(prof);
280        }
281        Ok(UpdateDeadline::Yield(1))
282    });
283
284    store.limiter(|ctx| ctx.session.limiter_mut());
285    Ok(store)
286}
287
288/// Constructs a `WasiCtxBuilder` for _each_ incoming request.
289fn make_wasi_ctx(ctx: &ExecuteCtx, session: &Session) -> wasmtime_wasi::WasiCtxBuilder {
290    let mut wasi_ctx = wasmtime_wasi::WasiCtxBuilder::new();
291
292    // Viceroy provides the same `FASTLY_*` environment variables that the production
293    // Compute platform provides:
294
295    wasi_ctx
296        // These variables are stubbed out for compatibility
297        .env("FASTLY_CACHE_GENERATION", "0")
298        .env("FASTLY_CUSTOMER_ID", "0000000000000000000000")
299        .env("FASTLY_POP", "XXX")
300        .env("FASTLY_REGION", "Somewhere")
301        .env("FASTLY_SERVICE_ID", "0000000000000000000000")
302        .env("FASTLY_SERVICE_VERSION", "0")
303        // signal that we're in a local testing environment
304        .env("FASTLY_HOSTNAME", "localhost")
305        // ...which is not the staging environment
306        .env("FASTLY_IS_STAGING", "0")
307        // request IDs start at 0 and increment, rather than being UUIDs, for ease of testing
308        .env("FASTLY_TRACE_ID", &format!("{:032x}", session.session_id()));
309
310    if ctx.log_stdout() {
311        wasi_ctx.stdout(LogEndpoint::new(b"stdout", ctx.capture_logs()));
312    } else {
313        wasi_ctx.inherit_stdout();
314    }
315
316    if ctx.log_stderr() {
317        wasi_ctx.stderr(LogEndpoint::new(b"stderr", ctx.capture_logs()));
318    } else {
319        wasi_ctx.inherit_stderr();
320    }
321
322    wasi_ctx
323}
324
325pub fn link_host_functions(
326    linker: &mut Linker<WasmCtx>,
327    experimental_modules: &HashSet<ExperimentalModule>,
328) -> Result<(), Error> {
329    experimental_modules
330        .iter()
331        .try_for_each(|experimental_module| match experimental_module {
332            ExperimentalModule::WasiNn => {
333                wasmtime_wasi_nn::witx::add_to_linker(linker, WasmCtx::wasi_nn)
334            }
335        })?;
336
337    wasmtime_wasi::p1::add_to_linker_async(linker, WasmCtx::wasi)?;
338    wiggle_abi::fastly_abi::add_to_linker(linker, WasmCtx::session)?;
339    wiggle_abi::fastly_acl::add_to_linker(linker, WasmCtx::session)?;
340    wiggle_abi::fastly_async_io::add_to_linker(linker, WasmCtx::session)?;
341    wiggle_abi::fastly_backend::add_to_linker(linker, WasmCtx::session)?;
342    wiggle_abi::fastly_cache::add_to_linker(linker, WasmCtx::session)?;
343    wiggle_abi::fastly_compute_runtime::add_to_linker(linker, WasmCtx::session)?;
344    wiggle_abi::fastly_config_store::add_to_linker(linker, WasmCtx::session)?;
345    wiggle_abi::fastly_device_detection::add_to_linker(linker, WasmCtx::session)?;
346    wiggle_abi::fastly_dictionary::add_to_linker(linker, WasmCtx::session)?;
347    wiggle_abi::fastly_erl::add_to_linker(linker, WasmCtx::session)?;
348    wiggle_abi::fastly_geo::add_to_linker(linker, WasmCtx::session)?;
349    wiggle_abi::fastly_http_body::add_to_linker(linker, WasmCtx::session)?;
350    wiggle_abi::fastly_http_cache::add_to_linker(linker, WasmCtx::session)?;
351    wiggle_abi::fastly_http_downstream::add_to_linker(linker, WasmCtx::session)?;
352    wiggle_abi::fastly_http_req::add_to_linker(linker, WasmCtx::session)?;
353    wiggle_abi::fastly_http_resp::add_to_linker(linker, WasmCtx::session)?;
354    wiggle_abi::fastly_image_optimizer::add_to_linker(linker, WasmCtx::session)?;
355    wiggle_abi::fastly_kv_store::add_to_linker(linker, WasmCtx::session)?;
356    wiggle_abi::fastly_log::add_to_linker(linker, WasmCtx::session)?;
357    wiggle_abi::fastly_object_store::add_to_linker(linker, WasmCtx::session)?;
358    wiggle_abi::fastly_purge::add_to_linker(linker, WasmCtx::session)?;
359    wiggle_abi::fastly_secret_store::add_to_linker(linker, WasmCtx::session)?;
360    wiggle_abi::fastly_shielding::add_to_linker(linker, WasmCtx::session)?;
361    wiggle_abi::fastly_uap::add_to_linker(linker, WasmCtx::session)?;
362    link_legacy_aliases(linker)?;
363    Ok(())
364}
365
366fn link_legacy_aliases(linker: &mut Linker<WasmCtx>) -> Result<(), Error> {
367    linker.alias("fastly_abi", "init", "env", "xqd_init")?;
368
369    let body = "fastly_http_body";
370    linker.alias(body, "append", "env", "xqd_body_append")?;
371    linker.alias(body, "new", "env", "xqd_body_new")?;
372    linker.alias(body, "read", "env", "xqd_body_read")?;
373    linker.alias(body, "write", "env", "xqd_body_write")?;
374    linker.alias(body, "close", "env", "xqd_body_close")?;
375    // `xqd_body_close_downstream` is deprecated since `fastly-sys:0.3.4`, and renamed to
376    // `xqd_body_close`. We include it here under both names for compatibility's sake.
377    linker.alias(body, "close", "env", "xqd_body_close_downstream")?;
378
379    linker.alias("fastly_log", "endpoint_get", "env", "xqd_log_endpoint_get")?;
380    linker.alias("fastly_log", "write", "env", "xqd_log_write")?;
381
382    let req = "fastly_http_req";
383    linker.alias(
384        req,
385        "body_downstream_get",
386        "env",
387        "xqd_req_body_downstream_get",
388    )?;
389    linker.alias(
390        req,
391        "cache_override_set",
392        "env",
393        "xqd_req_cache_override_set",
394    )?;
395    linker.alias(
396        req,
397        "downstream_client_ip_addr",
398        "env",
399        "xqd_req_downstream_client_ip_addr",
400    )?;
401    linker.alias(
402        req,
403        "downstream_client_request_id",
404        "env",
405        "xqd_req_downstream_client_request_id",
406    )?;
407    linker.alias(
408        req,
409        "downstream_tls_cipher_openssl_name",
410        "env",
411        "xqd_req_downstream_tls_cipher_openssl_name",
412    )?;
413    linker.alias(
414        req,
415        "downstream_tls_protocol",
416        "env",
417        "xqd_req_downstream_tls_protocol",
418    )?;
419    linker.alias(
420        req,
421        "downstream_tls_client_hello",
422        "env",
423        "xqd_req_downstream_tls_client_hello",
424    )?;
425    linker.alias(req, "new", "env", "xqd_req_new")?;
426
427    linker.alias(req, "header_names_get", "env", "xqd_req_header_names_get")?;
428    linker.alias(
429        req,
430        "original_header_names_get",
431        "env",
432        "xqd_req_original_header_names_get",
433    )?;
434    linker.alias(
435        req,
436        "original_header_count",
437        "env",
438        "xqd_req_original_header_count",
439    )?;
440    linker.alias(req, "header_value_get", "env", "xqd_req_header_value_get")?;
441    linker.alias(req, "header_values_get", "env", "xqd_req_header_values_get")?;
442    linker.alias(req, "header_values_set", "env", "xqd_req_header_values_set")?;
443    linker.alias(req, "header_insert", "env", "xqd_req_header_insert")?;
444    linker.alias(req, "header_append", "env", "xqd_req_header_append")?;
445    linker.alias(req, "header_remove", "env", "xqd_req_header_remove")?;
446    linker.alias(req, "method_get", "env", "xqd_req_method_get")?;
447    linker.alias(req, "method_set", "env", "xqd_req_method_set")?;
448    linker.alias(req, "uri_get", "env", "xqd_req_uri_get")?;
449    linker.alias(req, "uri_set", "env", "xqd_req_uri_set")?;
450    linker.alias(req, "version_get", "env", "xqd_req_version_get")?;
451    linker.alias(req, "version_set", "env", "xqd_req_version_set")?;
452    linker.alias(req, "send", "env", "xqd_req_send")?;
453    linker.alias(req, "send_async", "env", "xqd_req_send_async")?;
454    linker.alias(
455        req,
456        "send_async_streaming",
457        "env",
458        "xqd_req_send_async_streaming",
459    )?;
460    linker.alias(req, "pending_req_poll", "env", "xqd_pending_req_poll")?;
461    linker.alias(req, "pending_req_wait", "env", "xqd_pending_req_wait")?;
462    linker.alias(req, "pending_req_select", "env", "xqd_pending_req_select")?;
463
464    let resp = "fastly_http_resp";
465    linker.alias(resp, "new", "env", "xqd_resp_new")?;
466
467    linker.alias(resp, "header_names_get", "env", "xqd_resp_header_names_get")?;
468    linker.alias(resp, "header_value_get", "env", "xqd_resp_header_value_get")?;
469    linker.alias(
470        resp,
471        "header_values_get",
472        "env",
473        "xqd_resp_header_values_get",
474    )?;
475    linker.alias(
476        resp,
477        "header_values_set",
478        "env",
479        "xqd_resp_header_values_set",
480    )?;
481    linker.alias(resp, "header_insert", "env", "xqd_resp_header_insert")?;
482    linker.alias(resp, "header_append", "env", "xqd_resp_header_append")?;
483    linker.alias(resp, "header_remove", "env", "xqd_resp_header_remove")?;
484    linker.alias(resp, "version_get", "env", "xqd_resp_version_get")?;
485    linker.alias(resp, "version_set", "env", "xqd_resp_version_set")?;
486    linker.alias(resp, "status_get", "env", "xqd_resp_status_get")?;
487    linker.alias(resp, "status_set", "env", "xqd_resp_status_set")?;
488    Ok(())
489}
viceroy_lib/linking.rs

viceroy_lib/
linking.rs
