Added rewrite_all for body

This commit is contained in:
Emii Tatsuo 2020-12-01 21:41:00 -05:00
parent 64f54e8e58
commit 8f86bac608
Signed by: Emi
GPG key ID: 68FAB2E2E6DFC98B
2 changed files with 91 additions and 1 deletions

View file

@ -1,4 +1,4 @@
use tokio::io::AsyncRead;
use tokio::io::{AsyncRead, AsyncReadExt};
#[cfg(feature="serve_dir")]
use tokio::fs::File;
@ -11,6 +11,51 @@ pub enum Body {
Reader(Box<dyn AsyncRead + Send + Sync + Unpin>),
}
impl Body {
/// Called by [`Response::rewrite_all`]
pub (crate) async fn rewrite_all(&mut self, based_on: &crate::Request) -> std::io::Result<bool> {
let bytes = match self {
Self::Bytes(bytes) => {
let mut newbytes = Vec::new();
std::mem::swap(bytes, &mut newbytes);
newbytes
}
Self::Reader(reader) => {
let mut bytes = Vec::new();
reader.read_to_end(&mut bytes).await?;
bytes
}
};
let mut body = String::from_utf8(bytes).expect("text/gemini wasn't UTF8");
let mut maybe_indx = if body.starts_with("=> ") {
Some(3)
} else {
body.find("\n=> ").map(|offset| offset + 4)
};
while let Some(indx) = maybe_indx {
// Find the end of the link part
let end = (&body[indx..]).find(&[' ', '\n', '\r'][..])
.map(|offset| indx + offset )
.unwrap_or(body.len());
// Perform replacement
if let Some(replacement) = based_on.rewrite_path(&body[indx..end]) {
body.replace_range(indx..end, replacement.as_str());
} else {
return Ok(false)
};
// Find next match
maybe_indx = (&body[indx..]).find("\n=> ").map(|offset| offset + 4 + indx);
}
*self = Self::Bytes(body.into_bytes());
Ok(true)
}
}
impl<D: Borrow<Document>> From<D> for Body {
fn from(document: D) -> Self {
Self::from(document.borrow().to_string())

View file

@ -144,6 +144,51 @@ impl Response {
pub const fn is_success(&self) -> bool {
self.status == 10
}
/// Rewrite any links in this response based on the path identified by a request
///
/// For more information about what rewriting a link means, see
/// [`Request::rewrite_path`].
///
/// Currently, this rewrites any links in:
/// * SUCCESS (10) requests with a `text/gemini` MIME
/// * REDIRECT (3X) requests
///
/// For all other responses, and for any responses without links, this method has no
/// effect.
///
/// If this response contains a reader-based body, this **MAY** load the reader's
/// contents into memory if the mime is "text/gemini". If an IO error occurs during
/// this process, this error will be raised
///
/// If the request does not contain enough information to rewrite a link (in other
/// words, if [`Requet::rewrite_path`] returns [`None`]), then false is returned. In
/// all other cases, this method returns true.
///
/// Panics if a "text/gemini" response is not UTF-8 formatted
pub async fn rewrite_all(&mut self, based_on: &crate::Request) -> std::io::Result<bool> {
#[cfg(feature = "scgi_srv")]
match self.status {
20 if self.meta == "text/gemini" => {
if let Some(body) = self.body.as_mut() {
body.rewrite_all(based_on).await
} else {
Ok(false)
}
},
30 | 31 => {
if let Some(path) = based_on.rewrite_path(&self.meta) {
self.meta = path;
Ok(true)
} else {
Ok(false)
}
},
_ => Ok(true),
}
#[cfg(feature = "gemini_srv")]
Ok(true)
}
}
impl AsRef<Option<Body>> for Response {