make protocol handlers registrable (#33104)

Signed-off-by: webbeef <me@webbeef.org>

4 files changed, 416 insertions, 0 deletions

diff --git a/components/net/protocols/blob.rs b/components/net/protocols/blob.rs
new file mode 100644
index 00000000000..c26f95bf8da
--- /dev/null
+++ b/components/net/protocols/blob.rs
@@ -0,0 +1,91 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
+
+use std::future::{ready, Future};
+use std::pin::Pin;
+
+use headers::{HeaderMapExt, Range};
+use http::{Method, StatusCode};
+use log::debug;
+use net_traits::blob_url_store::{parse_blob_url, BlobURLStoreError};
+use net_traits::request::Request;
+use net_traits::response::{Response, ResponseBody};
+use net_traits::{NetworkError, ResourceFetchTiming};
+use tokio::sync::mpsc::unbounded_channel;
+
+use crate::fetch::methods::{Data, DoneChannel, FetchContext};
+use crate::protocols::{
+    get_range_request_bounds, partial_content, range_not_satisfiable_error, ProtocolHandler,
+};
+
+#[derive(Default)]
+pub struct BlobProtocolHander {}
+
+impl ProtocolHandler for BlobProtocolHander {
+    fn load(
+        &self,
+        request: &mut Request,
+        done_chan: &mut DoneChannel,
+        context: &FetchContext,
+    ) -> Pin<Box<dyn Future<Output = Response> + Send>> {
+        let url = request.current_url();
+        debug!("Loading blob {}", url.as_str());
+
+        // Step 2.
+        if request.method != Method::GET {
+            return Box::pin(ready(Response::network_error(NetworkError::Internal(
+                "Unexpected method for blob".into(),
+            ))));
+        }
+
+        let range_header = request.headers.typed_get::<Range>();
+        let is_range_request = range_header.is_some();
+        // We will get a final version of this range once we have
+        // the length of the data backing the blob.
+        let range = get_range_request_bounds(range_header);
+
+        let (id, origin) = match parse_blob_url(&url) {
+            Ok((id, origin)) => (id, origin),
+            Err(error) => {
+                return Box::pin(ready(Response::network_error(NetworkError::Internal(
+                    format!("Invalid blob URL ({error})"),
+                ))));
+            },
+        };
+
+        let mut response = Response::new(url, ResourceFetchTiming::new(request.timing_type()));
+        response.status = Some((StatusCode::OK, "OK".to_string()));
+        response.raw_status = Some((StatusCode::OK.as_u16(), b"OK".to_vec()));
+
+        if is_range_request {
+            partial_content(&mut response);
+        }
+
+        let (mut done_sender, done_receiver) = unbounded_channel();
+        *done_chan = Some((done_sender.clone(), done_receiver));
+        *response.body.lock().unwrap() = ResponseBody::Receiving(vec![]);
+
+        if let Err(err) = context.filemanager.lock().unwrap().fetch_file(
+            &mut done_sender,
+            context.cancellation_listener.clone(),
+            id,
+            &context.file_token,
+            origin,
+            &mut response,
+            range,
+        ) {
+            let _ = done_sender.send(Data::Done);
+            let err = match err {
+                BlobURLStoreError::InvalidRange => {
+                    range_not_satisfiable_error(&mut response);
+                    return Box::pin(ready(response));
+                },
+                _ => format!("{:?}", err),
+            };
+            return Box::pin(ready(Response::network_error(NetworkError::Internal(err))));
+        };
+
+        Box::pin(ready(response))
+    }
+}
diff --git a/components/net/protocols/data.rs b/components/net/protocols/data.rs
new file mode 100644
index 00000000000..3835525d100
--- /dev/null
+++ b/components/net/protocols/data.rs
@@ -0,0 +1,97 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
+
+use std::future::Future;
+use std::pin::Pin;
+
+use data_url::forgiving_base64;
+use headers::{ContentType, HeaderMapExt};
+use http::StatusCode;
+use mime::Mime;
+use net_traits::request::Request;
+use net_traits::response::{Response, ResponseBody};
+use net_traits::{NetworkError, ResourceFetchTiming};
+use percent_encoding::percent_decode;
+use servo_url::ServoUrl;
+use url::Position;
+
+use crate::fetch::methods::{DoneChannel, FetchContext};
+use crate::protocols::ProtocolHandler;
+
+#[derive(Default)]
+pub struct DataProtocolHander {}
+
+enum DecodeError {
+    InvalidDataUri,
+    NonBase64DataUri,
+}
+
+type DecodeData = (Mime, Vec<u8>);
+
+fn decode(url: &ServoUrl) -> Result<DecodeData, DecodeError> {
+    // data_url could do all of this work for us,
+    // except that it currently (Nov 2019) parses mime types into a
+    // different Mime class than other code expects
+
+    assert_eq!(url.scheme(), "data");
+    // Split out content type and data.
+    let parts: Vec<&str> = url[Position::BeforePath..Position::AfterQuery]
+        .splitn(2, ',')
+        .collect();
+    if parts.len() != 2 {
+        return Err(DecodeError::InvalidDataUri);
+    }
+
+    // ";base64" must come at the end of the content type, per RFC 2397.
+    // rust-http will fail to parse it because there's no =value part.
+    let mut ct_str = parts[0];
+    let is_base64 = ct_str.ends_with(";base64");
+    if is_base64 {
+        ct_str = &ct_str[..ct_str.len() - ";base64".len()];
+    }
+    let ct_str = if ct_str.starts_with(";charset=") {
+        format!("text/plain{}", ct_str)
+    } else {
+        ct_str.to_owned()
+    };
+
+    let content_type = ct_str
+        .parse()
+        .unwrap_or_else(|_| "text/plain; charset=US-ASCII".parse().unwrap());
+
+    let mut bytes = percent_decode(parts[1].as_bytes()).collect::<Vec<_>>();
+    if is_base64 {
+        match forgiving_base64::decode_to_vec(&bytes) {
+            Err(..) => return Err(DecodeError::NonBase64DataUri),
+            Ok(data) => bytes = data,
+        }
+    }
+    Ok((content_type, bytes))
+}
+
+impl ProtocolHandler for DataProtocolHander {
+    fn load(
+        &self,
+        request: &mut Request,
+        _done_chan: &mut DoneChannel,
+        _context: &FetchContext,
+    ) -> Pin<Box<dyn Future<Output = Response> + Send>> {
+        let url = request.current_url();
+        let response = match decode(&url) {
+            Ok((mime, bytes)) => {
+                let mut response =
+                    Response::new(url, ResourceFetchTiming::new(request.timing_type()));
+                *response.body.lock().unwrap() = ResponseBody::Done(bytes);
+                response.headers.typed_insert(ContentType::from(mime));
+                response.status = Some((StatusCode::OK, "OK".to_string()));
+                response.raw_status = Some((StatusCode::OK.as_u16(), b"OK".to_vec()));
+                response
+            },
+            Err(_) => {
+                Response::network_error(NetworkError::Internal("Decoding data URL failed".into()))
+            },
+        };
+        Box::pin(std::future::ready(response))
+    }
+}
diff --git a/components/net/protocols/file.rs b/components/net/protocols/file.rs
new file mode 100644
index 00000000000..65e134e4635
--- /dev/null
+++ b/components/net/protocols/file.rs
@@ -0,0 +1,109 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
+
+use std::fs::File;
+use std::future::{ready, Future};
+use std::io::{BufReader, Seek, SeekFrom};
+use std::pin::Pin;
+
+use headers::{ContentType, HeaderMapExt, Range};
+use http::Method;
+use net_traits::request::Request;
+use net_traits::response::{Response, ResponseBody};
+use net_traits::{NetworkError, ResourceFetchTiming};
+use tokio::sync::mpsc::unbounded_channel;
+
+use crate::fetch::methods::{DoneChannel, FetchContext};
+use crate::filemanager_thread::FILE_CHUNK_SIZE;
+use crate::local_directory_listing;
+use crate::protocols::{
+    get_range_request_bounds, partial_content, range_not_satisfiable_error, ProtocolHandler,
+};
+
+#[derive(Default)]
+pub struct FileProtocolHander {}
+
+impl ProtocolHandler for FileProtocolHander {
+    fn load(
+        &self,
+        request: &mut Request,
+        done_chan: &mut DoneChannel,
+        context: &FetchContext,
+    ) -> Pin<Box<dyn Future<Output = Response> + Send>> {
+        let url = request.current_url();
+
+        if request.method != Method::GET {
+            return Box::pin(ready(Response::network_error(NetworkError::Internal(
+                "Unexpected method for file".into(),
+            ))));
+        }
+        let response = if let Ok(file_path) = url.to_file_path() {
+            if file_path.is_dir() {
+                return Box::pin(ready(local_directory_listing::fetch(
+                    request, url, file_path,
+                )));
+            }
+
+            if let Ok(file) = File::open(file_path.clone()) {
+                // Get range bounds (if any) and try to seek to the requested offset.
+                // If seeking fails, bail out with a NetworkError.
+                let file_size = match file.metadata() {
+                    Ok(metadata) => Some(metadata.len()),
+                    Err(_) => None,
+                };
+
+                let mut response =
+                    Response::new(url, ResourceFetchTiming::new(request.timing_type()));
+
+                let range_header = request.headers.typed_get::<Range>();
+                let is_range_request = range_header.is_some();
+                let Ok(range) = get_range_request_bounds(range_header).get_final(file_size) else {
+                    range_not_satisfiable_error(&mut response);
+                    return Box::pin(ready(response));
+                };
+                let mut reader = BufReader::with_capacity(FILE_CHUNK_SIZE, file);
+                if reader.seek(SeekFrom::Start(range.start as u64)).is_err() {
+                    return Box::pin(ready(Response::network_error(NetworkError::Internal(
+                        "Unexpected method for file".into(),
+                    ))));
+                }
+
+                // Set response status to 206 if Range header is present.
+                // At this point we should have already validated the header.
+                if is_range_request {
+                    partial_content(&mut response);
+                }
+
+                // Set Content-Type header.
+                let mime = mime_guess::from_path(file_path).first_or_octet_stream();
+                response.headers.typed_insert(ContentType::from(mime));
+
+                // Setup channel to receive cross-thread messages about the file fetch
+                // operation.
+                let (mut done_sender, done_receiver) = unbounded_channel();
+                *done_chan = Some((done_sender.clone(), done_receiver));
+
+                *response.body.lock().unwrap() = ResponseBody::Receiving(vec![]);
+
+                context.filemanager.lock().unwrap().fetch_file_in_chunks(
+                    &mut done_sender,
+                    reader,
+                    response.body.clone(),
+                    context.cancellation_listener.clone(),
+                    range,
+                );
+
+                response
+            } else {
+                Response::network_error(NetworkError::Internal("Opening file failed".into()))
+            }
+        } else {
+            Response::network_error(NetworkError::Internal(
+                "Constructing file path failed".into(),
+            ))
+        };
+
+        Box::pin(ready(response))
+    }
+}
diff --git a/components/net/protocols/mod.rs b/components/net/protocols/mod.rs
new file mode 100644
index 00000000000..ea64bfcd389
--- /dev/null
+++ b/components/net/protocols/mod.rs
@@ -0,0 +1,119 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
+
+use std::collections::hash_map::Entry;
+use std::collections::HashMap;
+use std::future::Future;
+use std::ops::Bound;
+use std::pin::Pin;
+
+use headers::Range;
+use http::StatusCode;
+use log::error;
+use net_traits::filemanager_thread::RelativePos;
+use net_traits::request::Request;
+use net_traits::response::Response;
+
+use crate::fetch::methods::{DoneChannel, FetchContext, RangeRequestBounds};
+
+mod blob;
+mod data;
+mod file;
+
+use blob::BlobProtocolHander;
+use data::DataProtocolHander;
+use file::FileProtocolHander;
+
+// The set of schemes that can't be registered.
+static FORBIDDEN_SCHEMES: [&str; 4] = ["http", "https", "chrome", "about"];
+
+pub trait ProtocolHandler: Send + Sync {
+    fn load(
+        &self,
+        request: &mut Request,
+        done_chan: &mut DoneChannel,
+        context: &FetchContext,
+    ) -> Pin<Box<dyn Future<Output = Response> + Send>>;
+}
+
+#[derive(Default)]
+pub struct ProtocolRegistry {
+    pub(crate) handlers: HashMap<String, Box<dyn ProtocolHandler>>, // Maps scheme -> handler
+}
+
+impl ProtocolRegistry {
+    pub fn with_internal_protocols() -> Self {
+        let mut registry = Self::default();
+        registry.register("data", DataProtocolHander::default());
+        registry.register("blob", BlobProtocolHander::default());
+        registry.register("file", FileProtocolHander::default());
+        registry
+    }
+
+    pub fn register(&mut self, scheme: &str, handler: impl ProtocolHandler + 'static) -> bool {
+        if FORBIDDEN_SCHEMES.contains(&scheme) {
+            error!("Protocol handler for '{scheme}' is not allowed to be registered.");
+            return false;
+        }
+
+        if let Entry::Vacant(entry) = self.handlers.entry(scheme.into()) {
+            entry.insert(Box::new(handler));
+            true
+        } else {
+            error!("Protocol handler for '{scheme}' is already registered.");
+            false
+        }
+    }
+
+    pub fn get(&self, scheme: &str) -> Option<&dyn ProtocolHandler> {
+        self.handlers.get(scheme).map(|e| e.as_ref())
+    }
+
+    pub fn merge(&mut self, mut other: ProtocolRegistry) {
+        for (scheme, handler) in other.handlers.drain() {
+            if FORBIDDEN_SCHEMES.contains(&scheme.as_str()) {
+                error!("Protocol handler for '{scheme}' is not allowed to be registered.");
+                continue;
+            }
+
+            self.handlers.entry(scheme).or_insert(handler);
+        }
+    }
+}
+
+pub fn range_not_satisfiable_error(response: &mut Response) {
+    let reason = "Range Not Satisfiable".to_owned();
+    response.status = Some((StatusCode::RANGE_NOT_SATISFIABLE, reason.clone()));
+    response.raw_status = Some((StatusCode::RANGE_NOT_SATISFIABLE.as_u16(), reason.into()));
+}
+
+/// Get the range bounds if the `Range` header is present.
+pub fn get_range_request_bounds(range: Option<Range>) -> RangeRequestBounds {
+    if let Some(ref range) = range {
+        let (start, end) = match range
+            .iter()
+            .collect::<Vec<(Bound<u64>, Bound<u64>)>>()
+            .first()
+        {
+            Some(&(Bound::Included(start), Bound::Unbounded)) => (start, None),
+            Some(&(Bound::Included(start), Bound::Included(end))) => {
+                // `end` should be less or equal to `start`.
+                (start, Some(i64::max(start as i64, end as i64)))
+            },
+            Some(&(Bound::Unbounded, Bound::Included(offset))) => {
+                return RangeRequestBounds::Pending(offset);
+            },
+            _ => (0, None),
+        };
+        RangeRequestBounds::Final(RelativePos::from_opts(Some(start as i64), end))
+    } else {
+        RangeRequestBounds::Final(RelativePos::from_opts(Some(0), None))
+    }
+}
+
+pub fn partial_content(response: &mut Response) {
+    let reason = "Partial Content".to_owned();
+    response.status = Some((StatusCode::PARTIAL_CONTENT, reason.clone()));
+    response.raw_status = Some((StatusCode::PARTIAL_CONTENT.as_u16(), reason.into()));
+}