Skip to main content

trailbase_wasm/
lib.rs

1#![forbid(unsafe_code, clippy::unwrap_used)]
2#![allow(clippy::needless_return)]
3#![warn(clippy::await_holding_lock, clippy::inefficient_to_string)]
4
5pub mod wit {
6  wit_bindgen::generate!({
7      world: "trailbase:component/interfaces",
8      path: [
9          // Order-sensitive: will import *.wit from the folder.
10          "wit/deps-0.2.6/random",
11          "wit/deps-0.2.6/io",
12          "wit/deps-0.2.6/clocks",
13          "wit/deps-0.2.6/filesystem",
14          "wit/deps-0.2.6/sockets",
15          "wit/deps-0.2.6/cli",
16          "wit/deps-0.2.6/http",
17          "wit/keyvalue-0.2.0-draft",
18          // Ours:
19          "wit/trailbase/database",
20          "wit/trailbase/component",
21      ],
22      pub_export_macro: true,
23      default_bindings_module: "trailbase_wasm::wit",
24      // additional_derives: [PartialEq, Eq, Hash, Clone],
25      generate_all,
26  });
27}
28
29pub mod db;
30pub mod fetch;
31pub mod fs;
32pub mod http;
33pub mod job;
34pub mod kv;
35pub mod time;
36
37use std::sync::OnceLock;
38use trailbase_wasm_common::{HttpContext, HttpContextKind};
39use wstd::http::Request;
40use wstd::http::body::IncomingBody;
41use wstd::http::server::{Finished, Responder};
42
43use crate::http::{HttpRoute, Method, StatusCode, empty_error_response};
44use crate::job::Job;
45
46// Needed for export macro
47pub use static_assertions::assert_impl_all;
48pub use wstd::wasip2 as __wasi;
49
50pub use crate::wit::exports::trailbase::component::init_endpoint::Arguments;
51
52pub mod sqlite {
53  pub use crate::wit::exports::trailbase::component::init_endpoint::SqliteFunctionFlags;
54  pub use crate::wit::exports::trailbase::component::sqlite_function_endpoint::{Error, Value};
55}
56
57pub mod rand {
58  pub use wstd::rand::{get_insecure_random_bytes, get_random_bytes};
59}
60
61#[macro_export]
62macro_rules! export {
63    ($impl:ident) => {
64        ::trailbase_wasm::assert_impl_all!($impl: ::trailbase_wasm::Guest);
65        // Register InitEndpoint.
66        type _TrailbaseHandlerIdent = ::trailbase_wasm::TrailbaseHandler<$impl>;
67        ::trailbase_wasm::wit::export!(_TrailbaseHandlerIdent);
68        // Register Incoming HTTP handler.
69        type _HttpHandlerIdent = ::trailbase_wasm::HttpIncomingHandler<$impl>;
70        ::trailbase_wasm::__wasi::http::proxy::export!(
71            _HttpHandlerIdent with_types_in ::trailbase_wasm::__wasi);
72    };
73}
74
75#[derive(Debug)]
76pub struct Args {
77  pub version: Option<String>,
78}
79
80type SqliteFunctionHandler =
81  Box<dyn Fn(Vec<sqlite::Value>) -> Result<sqlite::Value, sqlite::Error> + Send + Sync>;
82
83pub struct SqliteFunction {
84  name: String,
85  num_args: u32,
86  flags: Vec<sqlite::SqliteFunctionFlags>,
87  handler: SqliteFunctionHandler,
88}
89
90impl SqliteFunction {
91  pub fn new<const N: usize>(
92    name: impl std::string::ToString,
93    f: impl Fn([sqlite::Value; N]) -> Result<sqlite::Value, sqlite::Error> + Send + Sync + 'static,
94    flags: &[sqlite::SqliteFunctionFlags],
95  ) -> Self {
96    return Self {
97      name: name.to_string(),
98      num_args: N as u32,
99      flags: flags.into(),
100      handler: Box::new(move |args| {
101        return f(args.try_into().expect("wrong number of arguments"));
102      }),
103    };
104  }
105}
106
107pub trait Guest {
108  fn init(_: Args) {}
109
110  fn http_handlers() -> Vec<HttpRoute> {
111    return vec![];
112  }
113
114  fn job_handlers() -> Vec<Job> {
115    return vec![];
116  }
117
118  fn sqlite_scalar_functions() -> Vec<SqliteFunction> {
119    return vec![];
120  }
121}
122
123pub struct TrailbaseHandler<T: Guest> {
124  phantom: std::marker::PhantomData<T>,
125}
126
127impl<T: Guest> TrailbaseHandler<T> {
128  fn call_init_once(args: Args) -> &'static () {
129    static S: OnceLock<()> = OnceLock::new();
130    return S.get_or_init(|| T::init(args));
131  }
132
133  fn get_sqlite_scalar_functions() -> &'static [SqliteFunction] {
134    // NOTE: This assumes that there's only one `T`, since static are shared across generics.
135    static FUNCS: OnceLock<Vec<SqliteFunction>> = OnceLock::new();
136    return FUNCS.get_or_init(T::sqlite_scalar_functions);
137  }
138}
139
140impl<T: Guest> crate::wit::exports::trailbase::component::init_endpoint::Guest
141  for TrailbaseHandler<T>
142{
143  fn init_http_handlers(
144    args: Arguments,
145  ) -> wit::exports::trailbase::component::init_endpoint::HttpHandlers {
146    Self::call_init_once(Args {
147      version: args.version,
148    });
149
150    return wit::exports::trailbase::component::init_endpoint::HttpHandlers {
151      handlers: T::http_handlers()
152        .into_iter()
153        .map(|route| (to_method_type(route.method), route.path))
154        .collect(),
155    };
156  }
157
158  fn init_job_handlers(
159    args: Arguments,
160  ) -> wit::exports::trailbase::component::init_endpoint::JobHandlers {
161    Self::call_init_once(Args {
162      version: args.version,
163    });
164
165    return wit::exports::trailbase::component::init_endpoint::JobHandlers {
166      handlers: T::job_handlers()
167        .into_iter()
168        .map(|config| (config.name, config.spec))
169        .collect(),
170    };
171  }
172
173  fn init_sqlite_functions(
174    args: Arguments,
175  ) -> wit::exports::trailbase::component::init_endpoint::SqliteFunctions {
176    use wit::exports::trailbase::component::init_endpoint::{
177      SqliteFunctions, SqliteScalarFunction,
178    };
179
180    // QUESTION: Should we ensure that init is called only once?
181    Self::call_init_once(Args {
182      version: args.version,
183    });
184
185    return SqliteFunctions {
186      scalar_functions: Self::get_sqlite_scalar_functions()
187        .iter()
188        .map(|f| SqliteScalarFunction {
189          name: f.name.clone(),
190          num_args: f.num_args,
191          function_flags: f.flags.clone(),
192        })
193        .collect(),
194    };
195  }
196}
197
198impl<T: Guest> crate::wit::exports::trailbase::component::sqlite_function_endpoint::Guest
199  for TrailbaseHandler<T>
200{
201  fn dispatch_scalar_function(
202    args: crate::wit::exports::trailbase::component::sqlite_function_endpoint::Arguments,
203  ) -> Result<
204    crate::wit::exports::trailbase::component::sqlite_function_endpoint::Value,
205    crate::wit::exports::trailbase::component::sqlite_function_endpoint::Error,
206  > {
207    use crate::wit::exports::trailbase::component::sqlite_function_endpoint::Error;
208    let f = Self::get_sqlite_scalar_functions()
209      .iter()
210      .find(|f| f.name == args.function_name)
211      .ok_or_else(|| Error::Other("Missing function".to_string()))?;
212
213    return (f.handler)(args.arguments);
214  }
215}
216
217pub struct HttpIncomingHandler<T: Guest> {
218  phantom: std::marker::PhantomData<T>,
219}
220
221impl<T: Guest> HttpIncomingHandler<T> {
222  async fn handle(request: Request<IncomingBody>, responder: Responder) -> Finished {
223    let path = request.uri().path();
224    let method = request.method();
225
226    let Some(context) = request
227      .headers()
228      .get("__context")
229      .and_then(|h| serde_json::from_slice::<HttpContext>(h.as_bytes()).ok())
230    else {
231      return responder
232        .respond(empty_error_response(StatusCode::INTERNAL_SERVER_ERROR))
233        .await;
234    };
235
236    log::debug!("WASM guest received HTTP request {path}: {context:?}");
237
238    match context.kind {
239      HttpContextKind::Http => {
240        if let Some(HttpRoute { handler, .. }) = T::http_handlers()
241          .into_iter()
242          .find(|route| route.method == method && route.path == context.registered_path)
243        {
244          return handler(context, request, responder).await;
245        }
246      }
247      HttpContextKind::Job => {
248        if let Some(Job { handler, .. }) = T::job_handlers()
249          .into_iter()
250          .find(|config| method == Method::GET && config.name == context.registered_path)
251        {
252          return handler(responder).await;
253        }
254      }
255    }
256
257    return responder
258      .respond(empty_error_response(StatusCode::NOT_FOUND))
259      .await;
260  }
261}
262
263impl<T: Guest> ::wstd::wasip2::exports::http::incoming_handler::Guest for HttpIncomingHandler<T> {
264  fn handle(
265    request: ::wstd::wasip2::http::types::IncomingRequest,
266    response_out: ::wstd::wasip2::http::types::ResponseOutparam,
267  ) {
268    let responder = Responder::new(response_out);
269
270    let _finished: Finished = match ::wstd::http::request::try_from_incoming(request) {
271      Ok(request) => ::wstd::runtime::block_on(async { Self::handle(request, responder).await }),
272      Err(err) => responder.fail(err),
273    };
274  }
275}
276
277fn to_method_type(
278  m: Method,
279) -> crate::wit::exports::trailbase::component::init_endpoint::HttpMethodType {
280  use crate::wit::exports::trailbase::component::init_endpoint::HttpMethodType;
281
282  return match m {
283    Method::GET => HttpMethodType::Get,
284    Method::POST => HttpMethodType::Post,
285    Method::HEAD => HttpMethodType::Head,
286    Method::OPTIONS => HttpMethodType::Options,
287    Method::PATCH => HttpMethodType::Patch,
288    Method::DELETE => HttpMethodType::Delete,
289    Method::PUT => HttpMethodType::Put,
290    Method::TRACE => HttpMethodType::Trace,
291    Method::CONNECT => HttpMethodType::Connect,
292    _ => panic!("unknown http method type: {m}"),
293  };
294}