parent
d44fa5c40b
commit
abd9d2c941
@ -0,0 +1,83 @@
|
||||
//! For API calls if enabled
|
||||
use super::*;
|
||||
use std::{
|
||||
fmt,
|
||||
error,
|
||||
iter,
|
||||
convert::Infallible,
|
||||
};
|
||||
use futures::{
|
||||
stream::{
|
||||
self,
|
||||
BoxStream,
|
||||
StreamExt,
|
||||
},
|
||||
};
|
||||
|
||||
fn aggregate(mut body: impl Buf) -> Result<String, std::string::FromUtf8Error>
|
||||
{
|
||||
let mut output = Vec::new();
|
||||
while body.has_remaining() {
|
||||
let bytes = body.bytes();
|
||||
output.extend_from_slice(&bytes[..]);
|
||||
let cnt = bytes.len();
|
||||
body.advance(cnt);
|
||||
}
|
||||
|
||||
String::from_utf8(output)
|
||||
}
|
||||
|
||||
pub async fn single(host: IpAddr, num: Option<usize>, body: impl Buf) -> Result<impl warp::Reply, warp::reject::Rejection>
|
||||
{
|
||||
single_stream(host, num, body).await
|
||||
.map(|rx| Response::new(Body::wrap_stream(rx.map(move |x| {
|
||||
info!("{} <- {:?}", host, x);
|
||||
x
|
||||
}))))
|
||||
.map_err(warp::reject::custom)
|
||||
}
|
||||
|
||||
async fn single_stream(host: IpAddr, num: Option<usize>, body: impl Buf) -> Result<BoxStream<'static, Result<String, Infallible>>, ApiError>
|
||||
{
|
||||
let body = aggregate(body)?;
|
||||
info!("{} <- {:?}", host, &body[..]);
|
||||
|
||||
let mut chain = Chain::new();
|
||||
feed::feed(&mut chain, body);
|
||||
match num {
|
||||
None => Ok(stream::iter(iter::once(Ok(chain.generate_str()))).boxed()),
|
||||
Some(num) => {
|
||||
let (mut tx, rx) = mpsc::channel(num);
|
||||
tokio::spawn(async move {
|
||||
for string in chain.str_iter_for(num) {
|
||||
tx.send(string).await.expect("Failed to send string to body");
|
||||
}
|
||||
});
|
||||
Ok(StreamExt::map(rx, |x| Ok::<_, Infallible>(x)).boxed())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum ApiError {
|
||||
Body,
|
||||
}
|
||||
impl warp::reject::Reject for ApiError{}
|
||||
impl error::Error for ApiError{}
|
||||
impl std::fmt::Display for ApiError
|
||||
{
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
|
||||
{
|
||||
match self {
|
||||
Self::Body => write!(f, "invalid data in request body"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<std::string::FromUtf8Error> for ApiError
|
||||
{
|
||||
fn from(_: std::string::FromUtf8Error) -> Self
|
||||
{
|
||||
Self::Body
|
||||
}
|
||||
}
|
Loading…
Reference in new issue