Skip to content

Commit

Permalink
Update /proxy/pixiv
Browse files Browse the repository at this point in the history
  • Loading branch information
lifegpc authored Oct 1, 2023
1 parent 8ef659b commit 65a2065
Show file tree
Hide file tree
Showing 11 changed files with 122 additions and 9 deletions.
2 changes: 1 addition & 1 deletion proc_macros/proc_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -938,7 +938,7 @@ pub fn http_error(item: TokenStream) -> TokenStream {
Ok(re) => re,
Err(e) => {
builder = builder.status(#code);
return Ok(builder.body::<Pin<Box<HttpBodyType>>>(Box::pin(Body::from(format!("{}", e))))?);
return Ok(builder.body::<Pin<Box<HttpBodyType>>>(Box::pin(HyperBody::from(format!("{}", e))))?);
}
}
);
Expand Down
3 changes: 3 additions & 0 deletions src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,11 @@ pub enum PixivDownloaderError {
#[cfg(feature = "openssl")]
OpenSSLError(openssl::error::ErrorStack),
ParseIntError(std::num::ParseIntError),
ReqwestError(reqwest::Error),
}

impl std::error::Error for PixivDownloaderError {}

impl From<&str> for PixivDownloaderError {
fn from(p: &str) -> Self {
Self::String(String::from(p))
Expand Down
51 changes: 51 additions & 0 deletions src/server/body/hyper.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
use crate::server::preclude::PixivDownloaderError;
use hyper::body::HttpBody;
use hyper::Body as _Body;
use hyper::HeaderMap;
use std::pin::Pin;
use std::task::{Context, Poll};

pub struct HyperBody {
_body: _Body,
}

impl HyperBody {
pub fn empty() -> Self {
Self {
_body: _Body::empty(),
}
}
}

impl<T: Into<_Body>> From<T> for HyperBody {
fn from(body: T) -> Self {
Self { _body: body.into() }
}
}

impl HttpBody for HyperBody {
type Data = bytes::Bytes;
type Error = PixivDownloaderError;

fn poll_data(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Self::Data, Self::Error>>> {
Pin::new(&mut self._body)
.poll_data(cx)
.map_err(PixivDownloaderError::from)
}

fn poll_trailers(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<Option<HeaderMap>, Self::Error>> {
Pin::new(&mut self._body)
.poll_trailers(cx)
.map_err(PixivDownloaderError::from)
}

fn is_end_stream(&self) -> bool {
self._body.is_end_stream()
}
}
2 changes: 2 additions & 0 deletions src/server/body/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
pub mod hyper;
pub mod response;
42 changes: 42 additions & 0 deletions src/server/body/response.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
use crate::error::PixivDownloaderError;
use hyper::body::HttpBody;
use reqwest::Response;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

pub struct ResponseBody {
res: Response,
}

impl ResponseBody {
pub fn new(res: Response) -> Self {
Self { res }
}
}

impl HttpBody for ResponseBody {
type Data = hyper::body::Bytes;
type Error = PixivDownloaderError;

fn poll_data(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Self::Data, Self::Error>>> {
match Pin::new(&mut Box::pin(self.res.chunk())).poll(cx) {
Poll::Ready(f) => match f {
Ok(Some(data)) => Poll::Ready(Some(Ok(data))),
Ok(None) => Poll::Ready(None),
Err(e) => Poll::Ready(Some(Err(PixivDownloaderError::from(e)))),
},
Poll::Pending => Poll::Pending,
}
}

fn poll_trailers(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<Option<hyper::HeaderMap>, Self::Error>> {
Poll::Ready(Ok(None))
}
}
2 changes: 2 additions & 0 deletions src/server/mod.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
/// Routes about authentication
pub mod auth;
/// Body types which implements [hyper::body::HttpBody]
pub mod body;
pub mod context;
/// CORS Handle
pub mod cors;
Expand Down
5 changes: 4 additions & 1 deletion src/server/preclude.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
pub use super::body::hyper::HyperBody;
pub use super::body::response::ResponseBody;
pub use super::context::ServerContext;
pub use super::params::RequestParams;
pub use super::result::JSONResult;
Expand All @@ -15,4 +17,5 @@ pub use regex::Regex;
pub use std::pin::Pin;
pub use std::sync::Arc;

pub type HttpBodyType = dyn HttpBody<Data = hyper::body::Bytes, Error = hyper::Error> + Send;
pub type HttpBodyType =
dyn HttpBody<Data = hyper::body::Bytes, Error = PixivDownloaderError> + Send;
13 changes: 11 additions & 2 deletions src/server/proxy/pixiv.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
use super::super::preclude::*;
use crate::webclient::WebClient;
use http::Uri;
use std::collections::HashMap;

pub struct ProxyPixivContext {
ctx: Arc<ServerContext>,
Expand All @@ -19,7 +21,7 @@ impl ResponseFor<Body, Pin<Box<HttpBodyType>>> for ProxyPixivContext {
) -> Result<Response<Pin<Box<HttpBodyType>>>, PixivDownloaderError> {
filter_http_methods!(
req,
Box::pin(Body::empty()),
Box::pin(HyperBody::empty()),
true,
self.ctx,
allow_headers = [X_SIGN, X_TOKEN_ID],
Expand All @@ -35,7 +37,14 @@ impl ResponseFor<Body, Pin<Box<HttpBodyType>>> for ProxyPixivContext {
if !host.ends_with(".pximg.net") {
http_error!(403, Err("Host is not allowed."));
}
return Ok(builder.body::<Pin<Box<HttpBodyType>>>(Box::pin(Body::empty()))?);
let client = WebClient::default();
client.set_header("referer", "https://www.pixiv.net/");
let headers = HashMap::new();
let re = http_error!(
502,
client.get(url, headers).await.ok_or("Failed to get image.")
);
return Ok(builder.body::<Pin<Box<HttpBodyType>>>(Box::pin(ResponseBody::new(re)))?);
}
}

Expand Down
6 changes: 3 additions & 3 deletions src/server/service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ impl PixivDownloaderSvc {

impl Service<Request<Body>> for PixivDownloaderSvc {
type Response = Response<Pin<Box<HttpBodyType>>>;
type Error = hyper::Error;
type Error = PixivDownloaderError;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

fn poll_ready(&mut self, _: &mut Context) -> Poll<Result<(), Self::Error>> {
Expand All @@ -43,7 +43,7 @@ impl Service<Request<Body>> for PixivDownloaderSvc {
println!("{}", e);
Ok(Response::builder()
.status(500)
.body::<Pin<Box<HttpBodyType>>>(Box::pin(Body::from(
.body::<Pin<Box<HttpBodyType>>>(Box::pin(HyperBody::from(
"Internal server error",
)))
.unwrap())
Expand All @@ -53,7 +53,7 @@ impl Service<Request<Body>> for PixivDownloaderSvc {
None => Box::pin(async {
Ok(Response::builder()
.status(404)
.body::<Pin<Box<HttpBodyType>>>(Box::pin(Body::from("404 Not Found")))
.body::<Pin<Box<HttpBodyType>>>(Box::pin(HyperBody::from("404 Not Found")))
.unwrap())
}),
}
Expand Down
4 changes: 2 additions & 2 deletions src/server/traits.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use super::body::hyper::HyperBody;
use super::context::ServerContext;
use super::params::RequestParams;
use super::preclude::HttpBodyType;
use crate::error::PixivDownloaderError;
use hyper::Body;
use hyper::Request;
use hyper::Response;
use json::JsonValue;
Expand Down Expand Up @@ -55,7 +55,7 @@ where
);
Ok(Response::from_parts(
parts,
Box::pin(Body::from(body.to_string())),
Box::pin(HyperBody::from(body.to_string())),
))
}
}
1 change: 1 addition & 0 deletions src/ugoira.rs
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,7 @@ impl PartialEq for UgoiraZipError2 {
}

unsafe impl Send for UgoiraZipError2 {}
unsafe impl Sync for UgoiraZipError2 {}

impl ToRawHandle<_ugoira::zip_error_t> for UgoiraZipError2 {
unsafe fn to_raw_handle(&self) -> *mut _ugoira::zip_error_t {
Expand Down

0 comments on commit 65a2065

Please sign in to comment.