use std::fmt;
use std::io;
use std::sync::Arc;
pub use build::BodyBuilder;
use ureq_proto::http::header;
use ureq_proto::BodyMode;
use crate::http;
use crate::run::BodyHandler;
use crate::Error;
use self::limit::LimitReader;
use self::lossy::LossyUtf8Reader;
mod build;
mod limit;
mod lossy;
#[cfg(feature = "charset")]
mod charset;
#[cfg(feature = "gzip")]
mod gzip;
#[cfg(feature = "brotli")]
mod brotli;
const MAX_BODY_SIZE: u64 = 10 * 1024 * 1024;
pub struct Body {
source: BodyDataSource,
info: Arc<ResponseInfo>,
}
enum BodyDataSource {
Handler(Box<BodyHandler>),
Reader(Box<dyn io::Read + Send + Sync>),
}
#[derive(Clone)]
pub(crate) struct ResponseInfo {
content_encoding: ContentEncoding,
mime_type: Option<String>,
charset: Option<String>,
body_mode: BodyMode,
}
impl Body {
pub fn builder() -> BodyBuilder {
BodyBuilder::new()
}
pub(crate) fn new(handler: BodyHandler, info: ResponseInfo) -> Self {
Body {
source: BodyDataSource::Handler(Box::new(handler)),
info: Arc::new(info),
}
}
pub fn mime_type(&self) -> Option<&str> {
self.info.mime_type.as_deref()
}
pub fn charset(&self) -> Option<&str> {
self.info.charset.as_deref()
}
pub fn content_length(&self) -> Option<u64> {
match self.info.body_mode {
BodyMode::NoBody => None,
BodyMode::LengthDelimited(v) => Some(v),
BodyMode::Chunked => None,
BodyMode::CloseDelimited => None,
}
}
pub fn as_reader(&mut self) -> BodyReader {
self.with_config().reader()
}
pub fn into_reader(self) -> BodyReader<'static> {
self.into_with_config().reader()
}
pub fn read_to_string(&mut self) -> Result<String, Error> {
self.with_config()
.limit(MAX_BODY_SIZE)
.lossy_utf8(true)
.read_to_string()
}
pub fn read_to_vec(&mut self) -> Result<Vec<u8>, Error> {
self.with_config()
.limit(MAX_BODY_SIZE)
.read_to_vec()
}
#[cfg(feature = "json")]
pub fn read_json<T: serde::de::DeserializeOwned>(&mut self) -> Result<T, Error> {
let reader = self.with_config().limit(MAX_BODY_SIZE).reader();
let value: T = serde_json::from_reader(reader)?;
Ok(value)
}
pub fn with_config(&mut self) -> BodyWithConfig {
let handler = (&mut self.source).into();
BodyWithConfig::new(handler, self.info.clone())
}
pub fn into_with_config(self) -> BodyWithConfig<'static> {
let handler = self.source.into();
BodyWithConfig::new(handler, self.info)
}
}
pub struct BodyWithConfig<'a> {
handler: BodySourceRef<'a>,
info: Arc<ResponseInfo>,
limit: u64,
lossy_utf8: bool,
}
impl<'a> BodyWithConfig<'a> {
fn new(handler: BodySourceRef<'a>, info: Arc<ResponseInfo>) -> Self {
BodyWithConfig {
handler,
info,
limit: u64::MAX,
lossy_utf8: false,
}
}
pub fn limit(mut self, value: u64) -> Self {
self.limit = value;
self
}
pub fn lossy_utf8(mut self, value: bool) -> Self {
self.lossy_utf8 = value;
self
}
fn do_build(self) -> BodyReader<'a> {
BodyReader::new(
LimitReader::new(self.handler, self.limit),
&self.info,
self.info.body_mode,
self.lossy_utf8,
)
}
pub fn reader(self) -> BodyReader<'a> {
self.do_build()
}
pub fn read_to_string(self) -> Result<String, Error> {
use std::io::Read;
let mut reader = self.do_build();
let mut buf = String::new();
reader.read_to_string(&mut buf)?;
Ok(buf)
}
pub fn read_to_vec(self) -> Result<Vec<u8>, Error> {
use std::io::Read;
let mut reader = self.do_build();
let mut buf = Vec::new();
reader.read_to_end(&mut buf)?;
Ok(buf)
}
#[cfg(feature = "json")]
pub fn read_json<T: serde::de::DeserializeOwned>(self) -> Result<T, Error> {
let reader = self.do_build();
let value: T = serde_json::from_reader(reader)?;
Ok(value)
}
}
#[derive(Debug, Clone, Copy)]
enum ContentEncoding {
None,
Gzip,
Brotli,
Unknown,
}
impl ResponseInfo {
pub fn new(headers: &http::HeaderMap, body_mode: BodyMode) -> Self {
let content_encoding = headers
.get(header::CONTENT_ENCODING)
.and_then(|v| v.to_str().ok())
.map(ContentEncoding::from)
.unwrap_or(ContentEncoding::None);
let (mime_type, charset) = headers
.get(header::CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.map(split_content_type)
.unwrap_or((None, None));
ResponseInfo {
content_encoding,
mime_type,
charset,
body_mode,
}
}
fn is_text(&self) -> bool {
self.mime_type
.as_deref()
.map(|s| s.starts_with("text/"))
.unwrap_or(false)
}
}
fn split_content_type(content_type: &str) -> (Option<String>, Option<String>) {
let mut split = content_type.split(';');
let Some(mime_type) = split.next() else {
return (None, None);
};
let mut charset = None;
for maybe_charset in split {
let maybe_charset = maybe_charset.trim();
if let Some(s) = maybe_charset.strip_prefix("charset=") {
charset = Some(s.to_string());
}
}
(Some(mime_type.to_string()), charset)
}
pub struct BodyReader<'a> {
reader: MaybeLossyDecoder<CharsetDecoder<ContentDecoder<LimitReader<BodySourceRef<'a>>>>>,
outgoing_body_mode: BodyMode,
}
impl<'a> BodyReader<'a> {
fn new(
reader: LimitReader<BodySourceRef<'a>>,
info: &ResponseInfo,
incoming_body_mode: BodyMode,
lossy_utf8: bool,
) -> BodyReader<'a> {
let mut outgoing_body_mode = incoming_body_mode;
let reader = match info.content_encoding {
ContentEncoding::None | ContentEncoding::Unknown => ContentDecoder::PassThrough(reader),
#[cfg(feature = "gzip")]
ContentEncoding::Gzip => {
debug!("Decoding gzip");
outgoing_body_mode = BodyMode::Chunked;
ContentDecoder::Gzip(Box::new(gzip::GzipDecoder::new(reader)))
}
#[cfg(not(feature = "gzip"))]
ContentEncoding::Gzip => ContentDecoder::PassThrough(reader),
#[cfg(feature = "brotli")]
ContentEncoding::Brotli => {
debug!("Decoding brotli");
outgoing_body_mode = BodyMode::Chunked;
ContentDecoder::Brotli(Box::new(brotli::BrotliDecoder::new(reader)))
}
#[cfg(not(feature = "brotli"))]
ContentEncoding::Brotli => ContentDecoder::PassThrough(reader),
};
let reader = if info.is_text() {
charset_decoder(
reader,
info.mime_type.as_deref(),
info.charset.as_deref(),
&mut outgoing_body_mode,
)
} else {
CharsetDecoder::PassThrough(reader)
};
let reader = if info.is_text() && lossy_utf8 {
MaybeLossyDecoder::Lossy(LossyUtf8Reader::new(reader))
} else {
MaybeLossyDecoder::PassThrough(reader)
};
BodyReader {
outgoing_body_mode,
reader,
}
}
pub(crate) fn body_mode(&self) -> BodyMode {
self.outgoing_body_mode
}
}
#[allow(unused)]
fn charset_decoder<R: io::Read>(
reader: R,
mime_type: Option<&str>,
charset: Option<&str>,
body_mode: &mut BodyMode,
) -> CharsetDecoder<R> {
#[cfg(feature = "charset")]
{
use encoding_rs::{Encoding, UTF_8};
let from = charset
.and_then(|c| Encoding::for_label(c.as_bytes()))
.unwrap_or(UTF_8);
if from == UTF_8 {
CharsetDecoder::PassThrough(reader)
} else {
debug!("Decoding charset {}", from.name());
*body_mode = BodyMode::Chunked;
CharsetDecoder::Decoder(self::charset::CharCodec::new(reader, from, UTF_8))
}
}
#[cfg(not(feature = "charset"))]
{
CharsetDecoder::PassThrough(reader)
}
}
enum MaybeLossyDecoder<R> {
Lossy(LossyUtf8Reader<R>),
PassThrough(R),
}
impl<R: io::Read> io::Read for MaybeLossyDecoder<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self {
MaybeLossyDecoder::Lossy(r) => r.read(buf),
MaybeLossyDecoder::PassThrough(r) => r.read(buf),
}
}
}
impl<'a> io::Read for BodyReader<'a> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.reader.read(buf)
}
}
enum CharsetDecoder<R> {
#[cfg(feature = "charset")]
Decoder(charset::CharCodec<R>),
PassThrough(R),
}
impl<R: io::Read> io::Read for CharsetDecoder<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self {
#[cfg(feature = "charset")]
CharsetDecoder::Decoder(v) => v.read(buf),
CharsetDecoder::PassThrough(v) => v.read(buf),
}
}
}
enum ContentDecoder<R: io::Read> {
#[cfg(feature = "gzip")]
Gzip(Box<gzip::GzipDecoder<R>>),
#[cfg(feature = "brotli")]
Brotli(Box<brotli::BrotliDecoder<R>>),
PassThrough(R),
}
impl<R: io::Read> io::Read for ContentDecoder<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self {
#[cfg(feature = "gzip")]
ContentDecoder::Gzip(v) => v.read(buf),
#[cfg(feature = "brotli")]
ContentDecoder::Brotli(v) => v.read(buf),
ContentDecoder::PassThrough(v) => v.read(buf),
}
}
}
impl fmt::Debug for Body {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Body").finish()
}
}
impl From<&str> for ContentEncoding {
fn from(s: &str) -> Self {
match s {
"gzip" => ContentEncoding::Gzip,
"br" => ContentEncoding::Brotli,
_ => {
debug!("Unknown content-encoding: {}", s);
ContentEncoding::Unknown
}
}
}
}
impl<'a> From<&'a mut BodyDataSource> for BodySourceRef<'a> {
fn from(value: &'a mut BodyDataSource) -> Self {
match value {
BodyDataSource::Handler(v) => Self::HandlerShared(v),
BodyDataSource::Reader(v) => Self::ReaderShared(v),
}
}
}
impl From<BodyDataSource> for BodySourceRef<'static> {
fn from(value: BodyDataSource) -> Self {
match value {
BodyDataSource::Handler(v) => Self::HandlerOwned(v),
BodyDataSource::Reader(v) => Self::ReaderOwned(v),
}
}
}
pub(crate) enum BodySourceRef<'a> {
HandlerShared(&'a mut BodyHandler),
HandlerOwned(Box<BodyHandler>),
ReaderShared(&'a mut (dyn io::Read + Send + Sync)),
ReaderOwned(Box<dyn io::Read + Send + Sync>),
}
impl<'a> io::Read for BodySourceRef<'a> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self {
BodySourceRef::HandlerShared(v) => v.read(buf),
BodySourceRef::HandlerOwned(v) => v.read(buf),
BodySourceRef::ReaderShared(v) => v.read(buf),
BodySourceRef::ReaderOwned(v) => v.read(buf),
}
}
}
#[cfg(all(test, feature = "_test"))]
mod test {
use crate::test::init_test_log;
use crate::transport::set_handler;
use crate::Error;
#[test]
fn content_type_without_charset() {
init_test_log();
set_handler("/get", 200, &[("content-type", "application/json")], b"{}");
let res = crate::get("https://my.test/get").call().unwrap();
assert_eq!(res.body().mime_type(), Some("application/json"));
assert!(res.body().charset().is_none());
}
#[test]
fn content_type_with_charset() {
init_test_log();
set_handler(
"/get",
200,
&[("content-type", "application/json; charset=iso-8859-4")],
b"{}",
);
let res = crate::get("https://my.test/get").call().unwrap();
assert_eq!(res.body().mime_type(), Some("application/json"));
assert_eq!(res.body().charset(), Some("iso-8859-4"));
}
#[test]
fn chunked_transfer() {
init_test_log();
let s = "3\r\n\
hel\r\n\
b\r\n\
lo world!!!\r\n\
0\r\n\
\r\n";
set_handler(
"/get",
200,
&[("transfer-encoding", "chunked")],
s.as_bytes(),
);
let mut res = crate::get("https://my.test/get").call().unwrap();
let b = res.body_mut().read_to_string().unwrap();
assert_eq!(b, "hello world!!!");
}
#[test]
fn large_response_header() {
init_test_log();
set_handler(
"/get",
200,
&[("content-type", &"b".repeat(64 * 1024))],
b"{}",
);
let err = crate::get("https://my.test/get").call().unwrap_err();
assert!(matches!(err, Error::LargeResponseHeader(_, _)));
}
}