webview_bundle_protocol/protocol/
local.rs

1use async_trait::async_trait;
2use dashmap::DashMap;
3use std::borrow::Cow;
4use std::collections::HashMap;
5use webview_bundle::http;
6use webview_bundle::http::Uri;
7
8pub trait LocalUriResolver: Send + Sync {
9  /// Resolve localhost uri from original request uri.
10  fn resolve_localhost(&self, uri: &Uri) -> Option<String>;
11
12  /// Get proxied localhost uri from original request uri.
13  fn get_localhost_uri(&self, uri: &Uri) -> Option<String> {
14    if let Some(localhost) = self.resolve_localhost(uri) {
15      let decoded_path = percent_encoding::percent_decode(uri.path().as_bytes())
16        .decode_utf8_lossy()
17        .to_string();
18      // Append the path of the original URI to the localhost URI.
19      let url = format!(
20        "{}/{}{}",
21        localhost.trim_end_matches('/'),
22        decoded_path.trim_start_matches('/'),
23        match uri.query() {
24          Some(q) => format!("?{q}"),
25          None => "".to_string(),
26        },
27      );
28      return Some(url);
29    }
30    None
31  }
32}
33
34#[derive(Default)]
35pub struct MappingLocalUriResolver {
36  mapping: HashMap<String, String>,
37}
38
39impl MappingLocalUriResolver {
40  pub fn new<T: Into<HashMap<String, String>>>(mapping: T) -> Self {
41    Self {
42      mapping: mapping.into(),
43    }
44  }
45}
46
47impl LocalUriResolver for MappingLocalUriResolver {
48  fn resolve_localhost(&self, uri: &Uri) -> Option<String> {
49    match uri.host() {
50      Some(host) => self.mapping.get(host).cloned(),
51      None => None,
52    }
53  }
54}
55
56#[derive(Clone)]
57struct CachedResponse {
58  status: http::StatusCode,
59  headers: http::HeaderMap,
60  body: bytes::Bytes,
61}
62
63pub struct LocalProtocol {
64  uri_resolver: Box<dyn LocalUriResolver + 'static>,
65  cache: DashMap<String, CachedResponse>,
66}
67
68impl LocalProtocol {
69  pub fn new<T: Into<HashMap<String, String>>>(mapping: T) -> Self {
70    Self {
71      uri_resolver: Box::new(MappingLocalUriResolver::new(mapping)),
72      cache: DashMap::default(),
73    }
74  }
75}
76
77#[async_trait]
78impl super::Protocol for LocalProtocol {
79  async fn handle(
80    &self,
81    request: http::Request<Vec<u8>>,
82  ) -> crate::Result<http::Response<Cow<'static, [u8]>>> {
83    let url = self
84      .uri_resolver
85      .get_localhost_uri(request.uri())
86      .ok_or(crate::Error::LocalNotFound)?;
87
88    let mut builder = http::Response::builder();
89
90    let client = reqwest::ClientBuilder::new();
91    let mut proxy_builder = client.build()?.request(request.method().clone(), &url);
92    proxy_builder = proxy_builder.body(request.body().clone());
93    let r = proxy_builder.send().await?;
94    let mut response = None;
95    if r.status() == http::StatusCode::NOT_MODIFIED {
96      response = self.cache.get(&url)
97    }
98    let response = if let Some(r) = response {
99      r
100    } else {
101      let status = r.status();
102      let headers = r.headers().clone();
103      let body = r.bytes().await?;
104      let response = CachedResponse {
105        status,
106        headers,
107        body,
108      };
109      self.cache.insert(url.to_string(), response);
110      self.cache.get(&url).unwrap()
111    };
112    for (name, value) in &response.headers {
113      builder = builder.header(name, value);
114    }
115    let resp = builder
116      .status(response.status)
117      .body(response.body.to_vec().into())?;
118    Ok(resp)
119  }
120}
121
122#[cfg(test)]
123mod tests {
124  use crate::{LocalProtocol, Protocol};
125  use std::collections::HashMap;
126  use std::iter;
127  use std::net::{SocketAddr, TcpListener};
128  use std::sync::atomic::{AtomicUsize, Ordering};
129  use std::sync::Arc;
130  use tiny_http::{Header as TinyHeader, Method, Response as TinyResponse, Server as TinyServer};
131  use webview_bundle::http;
132
133  fn server() -> (SocketAddr, std::thread::JoinHandle<()>) {
134    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
135    let addr = listener.local_addr().unwrap();
136    let server = TinyServer::from_listener(listener, None).unwrap();
137
138    let counter = Arc::new(AtomicUsize::new(0));
139    let counter_for_thread = counter.clone();
140
141    let handle = std::thread::spawn(move || {
142      for request in server.incoming_requests() {
143        let n = counter_for_thread.fetch_add(1, Ordering::SeqCst) + 1;
144        if request.method() == &Method::Get && request.url().starts_with("/index.html") {
145          if n == 1 {
146            let mut resp = TinyResponse::from_string("Hello World");
147            resp.add_header(TinyHeader::from_bytes("Content-Type", "text/plain").unwrap());
148            resp.add_header(TinyHeader::from_bytes("ETag", "\"v1\"").unwrap());
149            let _ = request.respond(resp);
150          } else {
151            // After first response, server will return 304 because content is not changed.
152            let mut resp = TinyResponse::empty(304);
153            resp.add_header(TinyHeader::from_bytes("ETag", "\"v1\"").unwrap());
154            let _ = request.respond(resp);
155          }
156        } else {
157          let _ = request.respond(TinyResponse::empty(404));
158        }
159      }
160    });
161
162    (addr, handle)
163  }
164
165  #[tokio::test]
166  async fn smoke() {
167    let (addr, _) = server();
168    let mapping =
169      iter::once(("nextjs.wvb".to_string(), format!("http://{addr}"))).collect::<HashMap<_, _>>();
170    let protocol = LocalProtocol::new(mapping);
171
172    let first_req = http::Request::builder()
173      .uri("https://nextjs.wvb/index.html")
174      .method("GET")
175      .body(Vec::new())
176      .unwrap();
177    let first_resp = protocol.handle(first_req).await.unwrap();
178    assert_eq!(first_resp.status(), 200);
179    assert_eq!(
180      first_resp.headers().get("content-type").unwrap(),
181      "text/plain"
182    );
183    assert_eq!(first_resp.body().as_ref(), b"Hello World");
184
185    let second_req = http::Request::builder()
186      .uri("https://nextjs.wvb/index.html")
187      .method("GET")
188      .body(Vec::new())
189      .unwrap();
190    let second_resp = protocol.handle(second_req).await.unwrap();
191    assert_eq!(second_resp.status(), 200);
192    assert_eq!(
193      second_resp.headers().get("content-type").unwrap(),
194      "text/plain"
195    );
196    assert_eq!(second_resp.headers().get("etag").unwrap(), "\"v1\"");
197    assert_eq!(first_resp.body().as_ref(), b"Hello World");
198  }
199}
webview_bundle_protocol/protocol/local.rs

webview_bundle_protocol/protocol/
local.rs
