Skip to content

Commit

Permalink
also try extracting payload config as Data<T>
Browse files Browse the repository at this point in the history
  • Loading branch information
robjtede committed Jul 19, 2020
1 parent 971ba3e commit a918080
Showing 1 changed file with 111 additions and 22 deletions.
133 changes: 111 additions & 22 deletions src/types/payload.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,10 @@ use futures_util::future::{err, ok, Either, FutureExt, LocalBoxFuture, Ready};
use futures_util::StreamExt;
use mime::Mime;

use crate::dev;
use crate::extract::FromRequest;
use crate::http::header;
use crate::request::HttpRequest;
use crate::{dev, web};

/// Payload extractor returns request 's payload stream.
///
Expand Down Expand Up @@ -142,13 +142,8 @@ impl FromRequest for Bytes {

#[inline]
fn from_request(req: &HttpRequest, payload: &mut dev::Payload) -> Self::Future {
let tmp;
let cfg = if let Some(cfg) = req.app_data::<PayloadConfig>() {
cfg
} else {
tmp = PayloadConfig::default();
&tmp
};
// allow both Config and Data<Config>
let cfg = PayloadConfig::from_req(req);

if let Err(e) = cfg.check_mimetype(req) {
return Either::Right(err(e));
Expand Down Expand Up @@ -197,13 +192,7 @@ impl FromRequest for String {

#[inline]
fn from_request(req: &HttpRequest, payload: &mut dev::Payload) -> Self::Future {
let tmp;
let cfg = if let Some(cfg) = req.app_data::<PayloadConfig>() {
cfg
} else {
tmp = PayloadConfig::default();
&tmp
};
let cfg = PayloadConfig::from_req(req);

// check content-type
if let Err(e) = cfg.check_mimetype(req) {
Expand Down Expand Up @@ -237,7 +226,12 @@ impl FromRequest for String {
)
}
}
/// Payload configuration for request's payload.

/// Configuration for request's payload.
///
/// Applies to the built-in `Bytes` and `String` extractors. Note that the Payload extractor does
/// not automatically check conformance with this configuration to allow more flexibility when
/// building extractors on top of `Payload`.
#[derive(Clone)]
pub struct PayloadConfig {
limit: usize,
Expand Down Expand Up @@ -284,14 +278,28 @@ impl PayloadConfig {
}
Ok(())
}

/// Allow payload config extraction from app data checking both `T` and `Data<T>`, in that
/// order, and falling back to the default payload config.
fn from_req(req: &HttpRequest) -> &PayloadConfig {
req.app_data::<PayloadConfig>()
.or_else(|| {
req.app_data::<web::Data<PayloadConfig>>()
.map(|d| d.as_ref())
})
.unwrap_or_else(|| &DEFAULT_PAYLOAD_CONFIG)
}
}

// Allow shared refs to default.
static DEFAULT_PAYLOAD_CONFIG: PayloadConfig = PayloadConfig {
limit: 262_144, // 2^18 bytes (~256kB)
mimetype: None,
};

impl Default for PayloadConfig {
fn default() -> Self {
PayloadConfig {
limit: 262_144,
mimetype: None,
}
DEFAULT_PAYLOAD_CONFIG.clone()
}
}

Expand Down Expand Up @@ -407,8 +415,9 @@ mod tests {
use bytes::Bytes;

use super::*;
use crate::http::header;
use crate::test::TestRequest;
use crate::http::{header, StatusCode};
use crate::test::{call_service, init_service, TestRequest};
use crate::{web, App, Responder};

#[actix_rt::test]
async fn test_payload_config() {
Expand All @@ -428,6 +437,86 @@ mod tests {
assert!(cfg.check_mimetype(&req).is_ok());
}

#[actix_rt::test]
async fn test_config_recall_locations() {
async fn bytes_handler(_: Bytes) -> impl Responder {
"payload is probably json bytes"
}

async fn string_handler(_: String) -> impl Responder {
"payload is probably json string"
}

let mut srv = init_service(
App::new()
.service(
web::resource("/bytes-app-data")
.app_data(
PayloadConfig::default().mimetype(mime::APPLICATION_JSON),
)
.route(web::get().to(bytes_handler)),
)
.service(
web::resource("/bytes-data")
.data(PayloadConfig::default().mimetype(mime::APPLICATION_JSON))
.route(web::get().to(bytes_handler)),
)
.service(
web::resource("/string-app-data")
.app_data(
PayloadConfig::default().mimetype(mime::APPLICATION_JSON),
)
.route(web::get().to(string_handler)),
)
.service(
web::resource("/string-data")
.data(PayloadConfig::default().mimetype(mime::APPLICATION_JSON))
.route(web::get().to(string_handler)),
),
)
.await;

let req = TestRequest::with_uri("/bytes-app-data").to_request();
let resp = call_service(&mut srv, req).await;
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

let req = TestRequest::with_uri("/bytes-data").to_request();
let resp = call_service(&mut srv, req).await;
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

let req = TestRequest::with_uri("/string-app-data").to_request();
let resp = call_service(&mut srv, req).await;
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

let req = TestRequest::with_uri("/string-data").to_request();
let resp = call_service(&mut srv, req).await;
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

let req = TestRequest::with_uri("/bytes-app-data")
.header(header::CONTENT_TYPE, mime::APPLICATION_JSON)
.to_request();
let resp = call_service(&mut srv, req).await;
assert_eq!(resp.status(), StatusCode::OK);

let req = TestRequest::with_uri("/bytes-data")
.header(header::CONTENT_TYPE, mime::APPLICATION_JSON)
.to_request();
let resp = call_service(&mut srv, req).await;
assert_eq!(resp.status(), StatusCode::OK);

let req = TestRequest::with_uri("/string-app-data")
.header(header::CONTENT_TYPE, mime::APPLICATION_JSON)
.to_request();
let resp = call_service(&mut srv, req).await;
assert_eq!(resp.status(), StatusCode::OK);

let req = TestRequest::with_uri("/string-data")
.header(header::CONTENT_TYPE, mime::APPLICATION_JSON)
.to_request();
let resp = call_service(&mut srv, req).await;
assert_eq!(resp.status(), StatusCode::OK);
}

#[actix_rt::test]
async fn test_bytes() {
let (req, mut pl) = TestRequest::with_header(header::CONTENT_LENGTH, "11")
Expand Down

0 comments on commit a918080

Please sign in to comment.