Skip to main content

webview_bundle_protocol/protocol/
bundle.rs

1use crate::mime_type::MimeType;
2use crate::source::Source;
3use crate::uri::{DefaultUriResolver, UriResolver};
4use async_trait::async_trait;
5use dashmap::DashMap;
6use std::borrow::Cow;
7use std::sync::Arc;
8use tokio::sync::OnceCell;
9use webview_bundle::http::{header, HeaderValue, Method, Request, Response};
10use webview_bundle::BundleManifest;
11
12pub struct BundleProtocol<S: Source> {
13  source: Arc<S>,
14  uri_resolver: Box<dyn UriResolver + 'static>,
15  manifests: DashMap<String, Arc<OnceCell<Arc<BundleManifest>>>>,
16}
17
18impl<S: Source> BundleProtocol<S> {
19  pub fn new(source: S) -> Self {
20    Self {
21      source: Arc::new(source),
22      uri_resolver: Box::new(DefaultUriResolver),
23      manifests: DashMap::default(),
24    }
25  }
26}
27
28impl<S: Source> BundleProtocol<S> {
29  pub async fn load_manifest(&self, name: &str) -> crate::Result<Arc<BundleManifest>> {
30    if let Some(entry) = self.manifests.get(name) {
31      if let Some(m) = entry.get() {
32        return Ok(m.clone());
33      }
34    }
35    let cell_arc = {
36      let entry = self.manifests.entry(name.to_string()).or_default();
37      entry.clone()
38    };
39    let m = cell_arc
40      .get_or_try_init(|| async {
41        let m = self.source.fetch_manifest(name).await?;
42        Ok::<Arc<BundleManifest>, crate::Error>(Arc::new(m))
43      })
44      .await?
45      .clone();
46    Ok(m)
47  }
48
49  pub fn unload_manifest(&self, name: &str) -> bool {
50    self.manifests.remove(name).is_some()
51  }
52}
53
54#[async_trait]
55impl<S: Source> super::Protocol for BundleProtocol<S> {
56  async fn handle(&self, request: Request<Vec<u8>>) -> crate::Result<Response<Cow<'static, [u8]>>> {
57    let name = self
58      .uri_resolver
59      .resolve_bundle(request.uri())
60      .ok_or(crate::Error::BundleNotFound)?;
61    let path = self.uri_resolver.resolve_path(request.uri());
62
63    let mut resp = Response::builder();
64
65    if !(request.method() == Method::GET || request.method() == Method::HEAD) {
66      let method_not_allowed = resp.status(405).body(Vec::new().into())?;
67      return Ok(method_not_allowed);
68    }
69
70    let manifest = self.load_manifest(&name).await?;
71    if !manifest.index().contains_path(&path) {
72      let not_found = resp.status(404).body(Vec::new().into())?;
73      return Ok(not_found);
74    }
75
76    let reader = self.source.reader(&name).await?;
77    let data = manifest.async_get_data(reader, &path).await?.unwrap();
78
79    if let Some(headers_mut) = resp.headers_mut() {
80      if let Some(headers) = manifest.index().get_entry(&path).map(|x| x.headers()) {
81        headers_mut.clone_from(headers);
82      }
83      // append if content-type header does not exists
84      if !headers_mut.contains_key("content-type") {
85        let mime = MimeType::parse(&data, &path);
86        headers_mut.insert(header::CONTENT_TYPE, HeaderValue::try_from(&mime)?);
87      }
88      // insert content-length header
89      headers_mut.insert(header::CONTENT_LENGTH, data.len().into());
90    }
91
92    let response = resp.status(200).body(data.into())?;
93    Ok(response)
94  }
95}
96
97#[cfg(test)]
98mod tests {
99  use super::*;
100  use crate::protocol::Protocol;
101  use crate::source::FileSource;
102  use std::path::PathBuf;
103
104  #[tokio::test]
105  async fn load_many_at_once() {
106    let base_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
107      .join("tests")
108      .join("fixtures");
109    let source = FileSource::new(base_dir);
110    let protocol = Arc::new(BundleProtocol::new(source));
111    let mut handles = Vec::new();
112    for _i in 0..10 {
113      let p = protocol.clone();
114      let handle = tokio::spawn(async move {
115        let _ = p.load_manifest("nextjs.wvb").await;
116      });
117      handles.push(handle);
118    }
119    for h in handles {
120      h.await.unwrap();
121    }
122  }
123
124  #[tokio::test]
125  async fn load_unload_sequential() {
126    let base_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
127      .join("tests")
128      .join("fixtures");
129    let source = FileSource::new(base_dir);
130    let protocol = BundleProtocol::new(source);
131
132    let m1 = protocol.load_manifest("nextjs.wvb").await.unwrap();
133    assert!(
134      protocol.unload_manifest("nextjs.wvb"),
135      "unload should remove existing entry"
136    );
137    let m2 = protocol.load_manifest("nextjs.wvb").await.unwrap();
138    assert!(
139      !Arc::ptr_eq(&m1, &m2),
140      "after unload, reloading should produce a new Arc"
141    );
142
143    assert!(protocol.unload_manifest("nextjs.wvb"));
144    let m3 = protocol.load_manifest("nextjs.wvb").await.unwrap();
145    assert!(!Arc::ptr_eq(&m2, &m3));
146
147    assert!(protocol.unload_manifest("nextjs.wvb"));
148    let m4 = protocol.load_manifest("nextjs.wvb").await.unwrap();
149    assert!(!Arc::ptr_eq(&m3, &m4));
150  }
151
152  #[tokio::test]
153  async fn load_unload_concurrently() {
154    use std::sync::Arc;
155    use tokio::sync::Barrier;
156    use tokio::task::JoinSet;
157
158    let base_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
159      .join("tests")
160      .join("fixtures");
161    let source = FileSource::new(base_dir);
162    let protocol = Arc::new(BundleProtocol::new(source));
163
164    // 1) initial loads. test singleflight
165    let n = 5usize;
166    let mut set = JoinSet::new();
167    for _i in 0..n {
168      let p = protocol.clone();
169      set.spawn(async move { p.load_manifest("nextjs.wvb").await });
170    }
171    let mut initials = Vec::with_capacity(n);
172    while let Some(res) = set.join_next().await {
173      let v = res.unwrap().unwrap();
174      initials.push(v);
175    }
176    for m in &initials[1..] {
177      assert!(Arc::ptr_eq(&initials[0], m));
178    }
179
180    // 2) before/after barriers
181    let barrier_before_unload = Arc::new(Barrier::new(n + 1));
182    let barrier_after_unload = Arc::new(Barrier::new(n + 1));
183
184    let mut before_set = JoinSet::new();
185    for _i in 0..n {
186      let p = protocol.clone();
187      let before = barrier_before_unload.clone();
188      before_set.spawn(async move {
189        before.wait().await;
190        p.load_manifest("nextjs.wvb").await
191      });
192    }
193    let mut after_set = JoinSet::new();
194    for _i in 0..n {
195      let p = protocol.clone();
196      let after = barrier_after_unload.clone();
197      after_set.spawn(async move {
198        after.wait().await;
199        p.load_manifest("nextjs.wvb").await
200      });
201    }
202
203    barrier_before_unload.wait().await;
204    assert!(protocol.unload_manifest("nextjs.wvb"));
205    barrier_after_unload.wait().await;
206
207    let mut before_jobs = Vec::with_capacity(n);
208    while let Some(res) = before_set.join_next().await {
209      let v = res.unwrap().unwrap();
210      before_jobs.push(v);
211    }
212    let mut after_jobs = Vec::with_capacity(n);
213    while let Some(res) = after_set.join_next().await {
214      let v = res.unwrap().unwrap();
215      after_jobs.push(v);
216    }
217    // before jobs should be same with initial loads
218    for m in &before_jobs {
219      assert!(Arc::ptr_eq(&initials[0], m));
220    }
221    // after jobs should be not same with initial loads
222    for m in &after_jobs {
223      assert!(!Arc::ptr_eq(&initials[0], m));
224    }
225    for m in &before_jobs[1..] {
226      assert!(Arc::ptr_eq(&before_jobs[0], m));
227    }
228    for m in &after_jobs[1..] {
229      assert!(Arc::ptr_eq(&after_jobs[0], m));
230    }
231  }
232
233  #[tokio::test]
234  async fn smoke() {
235    let base_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
236      .join("tests")
237      .join("fixtures");
238    let source = FileSource::new(base_dir);
239    let protocol = Arc::new(BundleProtocol::new(source));
240    let resp = protocol
241      .handle(
242        Request::builder()
243          .uri("https://nextjs.wvb/index.html")
244          .method("GET")
245          .body(vec![])
246          .unwrap(),
247      )
248      .await
249      .unwrap();
250    assert_eq!(resp.status(), 200);
251    let resp = protocol
252      .handle(
253        Request::builder()
254          .uri("https://nextjs.wvb/not_found.html")
255          .method("GET")
256          .body(vec![])
257          .unwrap(),
258      )
259      .await
260      .unwrap();
261    assert_eq!(resp.status(), 404);
262    let mut handlers = vec![];
263    for _ in 1..100 {
264      let p = protocol.clone();
265      let handle = tokio::spawn(async move {
266        p.handle(
267          Request::builder()
268            .uri("https://nextjs.wvb/index.html")
269            .method("GET")
270            .body(vec![])
271            .unwrap(),
272        )
273        .await
274      });
275      handlers.push(handle);
276    }
277    for h in handlers {
278      let resp = h.await.unwrap().unwrap();
279      assert_eq!(resp.status(), 200);
280    }
281  }
282
283  #[tokio::test]
284  async fn resolve_index_html() {
285    let base_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
286      .join("tests")
287      .join("fixtures");
288    let source = FileSource::new(base_dir);
289    let protocol = Arc::new(BundleProtocol::new(source));
290    let resp = protocol
291      .handle(
292        Request::builder()
293          .uri("https://nextjs.wvb")
294          .method("GET")
295          .body(vec![])
296          .unwrap(),
297      )
298      .await
299      .unwrap();
300    assert_eq!(resp.status(), 200);
301  }
302
303  #[tokio::test]
304  async fn content_type() {
305    let base_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
306      .join("tests")
307      .join("fixtures");
308    let source = FileSource::new(base_dir);
309    let protocol = Arc::new(BundleProtocol::new(source));
310    let resp = protocol
311      .handle(
312        Request::builder()
313          .uri("https://nextjs.wvb/_next/static/chunks/framework-98177fb2e8834792.js")
314          .method("GET")
315          .body(vec![])
316          .unwrap(),
317      )
318      .await
319      .unwrap();
320    assert_eq!(
321      resp.headers().get(header::CONTENT_TYPE).unwrap(),
322      HeaderValue::from_static("text/javascript")
323    );
324    let resp = protocol
325      .handle(
326        Request::builder()
327          .uri("https://nextjs.wvb/_next/static/css/419406682a95b9a9.css")
328          .method("GET")
329          .body(vec![])
330          .unwrap(),
331      )
332      .await
333      .unwrap();
334    assert_eq!(
335      resp.headers().get(header::CONTENT_TYPE).unwrap(),
336      HeaderValue::from_static("text/css")
337    );
338    let resp = protocol
339      .handle(
340        Request::builder()
341          .uri("https://nextjs.wvb/_next/static/media/build.583ad785.png")
342          .method("GET")
343          .body(vec![])
344          .unwrap(),
345      )
346      .await
347      .unwrap();
348    assert_eq!(
349      resp.headers().get(header::CONTENT_TYPE).unwrap(),
350      HeaderValue::from_static("image/png")
351    );
352  }
353
354  #[tokio::test]
355  async fn content_length() {
356    let base_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
357      .join("tests")
358      .join("fixtures");
359    let source = FileSource::new(base_dir);
360    let protocol = Arc::new(BundleProtocol::new(source));
361    let resp = protocol
362      .handle(
363        Request::builder()
364          .uri("https://nextjs.wvb/_next/static/chunks/framework-98177fb2e8834792.js")
365          .method("GET")
366          .body(vec![])
367          .unwrap(),
368      )
369      .await
370      .unwrap();
371    assert_eq!(
372      resp.headers().get(header::CONTENT_LENGTH).unwrap(),
373      HeaderValue::from_static("139833")
374    );
375    let resp = protocol
376      .handle(
377        Request::builder()
378          .uri("https://nextjs.wvb/_next/static/css/419406682a95b9a9.css")
379          .method("GET")
380          .body(vec![])
381          .unwrap(),
382      )
383      .await
384      .unwrap();
385    assert_eq!(
386      resp.headers().get(header::CONTENT_LENGTH).unwrap(),
387      HeaderValue::from_static("13856")
388    );
389    let resp = protocol
390      .handle(
391        Request::builder()
392          .uri("https://nextjs.wvb/_next/static/media/build.583ad785.png")
393          .method("GET")
394          .body(vec![])
395          .unwrap(),
396      )
397      .await
398      .unwrap();
399    assert_eq!(
400      resp.headers().get(header::CONTENT_LENGTH).unwrap(),
401      HeaderValue::from_static("475918")
402    );
403  }
404
405  #[tokio::test]
406  async fn not_found() {
407    let base_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
408      .join("tests")
409      .join("fixtures");
410    let source = FileSource::new(base_dir);
411    let protocol = Arc::new(BundleProtocol::new(source));
412    let resp = protocol
413      .handle(
414        Request::builder()
415          .uri("https://nextjs.wvb/path/not/exists")
416          .method("GET")
417          .body(vec![])
418          .unwrap(),
419      )
420      .await
421      .unwrap();
422    assert_eq!(resp.status(), 404);
423  }
424
425  #[tokio::test]
426  async fn bundle_not_found() {
427    let base_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
428      .join("tests")
429      .join("fixtures");
430    let source = FileSource::new(base_dir);
431    let protocol = Arc::new(BundleProtocol::new(source));
432    let err = protocol
433      .handle(
434        Request::builder()
435          .uri("https://not_exsits_bundle.wvb")
436          .method("GET")
437          .body(vec![])
438          .unwrap(),
439      )
440      .await
441      .unwrap_err();
442    assert!(matches!(err, crate::Error::BundleNotFound));
443  }
444}