Skip to content

Commit 5cff1da

Browse files
algestenalgesten
committed
SendBody to track Content-Length
1 parent 2f63e70 commit 5cff1da

4 files changed

Lines changed: 139 additions & 46 deletions

File tree

‎src/agent.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -213,14 +213,14 @@ impl Agent {
213213
/// ```
214214
pub fn run(&self, request: Request<impl AsSendBody>) -> Result<Response<Body>, Error> {
215215
let (parts, mut body) = request.into_parts();
216-
let body = body.as_body();
216+
let mut body = body.as_body();
217217
let mut request = Request::from_parts(parts, ());
218218

219219
// When using the http-crate API we cannot enforce the correctness of
220220
// Method vs Body combos. This also solves a problem where we can't
221221
// determine if a non-standard method is supposed to have a body such
222222
// as for WebDAV PROPFIND.
223-
let has_body = !matches!(body.body_mode(), BodyMode::NoBody);
223+
let has_body = !matches!(body.body_mode(), Ok(BodyMode::NoBody));
224224
if has_body {
225225
request.extensions_mut().insert(ForceSendBody);
226226
}

‎src/lib.rs‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1031,6 +1031,35 @@ pub(crate) mod test {
10311031
.expect("to send correctly");
10321032
}
10331033

1034+
#[test]
1035+
#[cfg(not(feature = "_test"))]
1036+
fn post_array_body_sends_content_length() {
1037+
init_test_log();
1038+
let mut response = post("http://httpbin.org/post")
1039+
.content_type("application/octet-stream")
1040+
.send(vec![42; 123])
1041+
.expect("to send correctly");
1042+
1043+
let ret = response.body_mut().read_to_string().unwrap();
1044+
assert!(ret.contains("\"Content-Length\": \"123\""));
1045+
}
1046+
1047+
#[test]
1048+
#[cfg(not(feature = "_test"))]
1049+
fn post_file_sends_file_length() {
1050+
init_test_log();
1051+
1052+
let file = std::fs::File::open("LICENSE-MIT").unwrap();
1053+
1054+
let mut response = post("http://httpbin.org/post")
1055+
.content_type("application/octet-stream")
1056+
.send(file)
1057+
.expect("to send correctly");
1058+
1059+
let ret = response.body_mut().read_to_string().unwrap();
1060+
assert!(ret.contains("\"Content-Length\": \"1072\""));
1061+
}
1062+
10341063
#[test]
10351064
#[cfg(not(feature = "_test"))]
10361065
fn username_password_from_uri() {

‎src/run.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -251,15 +251,15 @@ fn add_headers(
251251
call: &mut Call<Prepare>,
252252
agent: &Agent,
253253
config: &Config,
254-
body: &SendBody,
254+
body: &mut SendBody,
255255
uri: &Uri,
256256
) -> Result<(), Error> {
257257
let headers = call.headers();
258258

259259
let send_body_mode = if headers.has_send_body_mode() {
260260
None
261261
} else {
262-
Some(body.body_mode())
262+
Some(body.body_mode()?)
263263
};
264264
let has_header_accept_enc = headers.has_accept_encoding();
265265
let has_header_ua = headers.has_user_agent();

‎src/send_body.rs‎

Lines changed: 106 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ use std::io::{self, Read, Stdin};
33
use std::net::TcpStream;
44

55
use crate::body::{Body, BodyReader};
6-
use crate::http;
76
use crate::util::private::Private;
7+
use crate::{http, Error};
88

99
/// Request body for sending data via POST, PUT and PATCH.
1010
///
@@ -18,23 +18,24 @@ use crate::util::private::Private;
1818
///
1919
pub struct SendBody<'a> {
2020
inner: BodyInner<'a>,
21+
size: Option<Result<u64, Error>>,
2122
ended: bool,
2223
}
2324

2425
impl<'a> SendBody<'a> {
2526
/// Creates an empty body.
2627
pub fn none() -> SendBody<'static> {
27-
BodyInner::None.into()
28+
(None, BodyInner::None).into()
2829
}
2930

3031
/// Creates a body from a shared [`Read`] impl.
3132
pub fn from_reader(reader: &'a mut dyn Read) -> SendBody<'a> {
32-
BodyInner::Reader(reader).into()
33+
(None, BodyInner::Reader(reader)).into()
3334
}
3435

3536
/// Creates a body from an owned [`Read`] impl.
3637
pub fn from_owned_reader(reader: impl Read + 'static) -> SendBody<'static> {
37-
BodyInner::OwnedReader(Box::new(reader)).into()
38+
(None, BodyInner::OwnedReader(Box::new(reader))).into()
3839
}
3940

4041
/// Creates a body to send as JSON from any [`Serialize`](serde::ser::Serialize) value.
@@ -43,7 +44,12 @@ impl<'a> SendBody<'a> {
4344
value: &impl serde::ser::Serialize,
4445
) -> Result<SendBody<'static>, crate::Error> {
4546
let json = serde_json::to_vec_pretty(value)?;
46-
Ok(BodyInner::ByteVec(io::Cursor::new(json)).into())
47+
let body = (
48+
Some(json.len() as u64),
49+
BodyInner::ByteVec(io::Cursor::new(json)),
50+
)
51+
.into();
52+
Ok(body)
4753
}
4854

4955
pub(crate) fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
@@ -73,8 +79,38 @@ impl<'a> SendBody<'a> {
7379
Ok(n)
7480
}
7581

76-
pub(crate) fn body_mode(&self) -> BodyMode {
77-
self.inner.body_mode()
82+
pub(crate) fn body_mode(&mut self) -> Result<BodyMode, Error> {
83+
// Lazily surface a potential error now.
84+
let size = match self.size {
85+
None => None,
86+
Some(Ok(v)) => Some(v),
87+
Some(Err(_)) => {
88+
// unwraps here are ok because we matched exactly this
89+
return Err(self.size.take().unwrap().unwrap_err());
90+
}
91+
};
92+
93+
match &self.inner {
94+
BodyInner::None => return Ok(BodyMode::NoBody),
95+
BodyInner::Body(v) => return Ok(v.body_mode()),
96+
97+
// The others fall through
98+
BodyInner::ByteSlice(_) => {}
99+
#[cfg(feature = "json")]
100+
BodyInner::ByteVec(_) => {}
101+
BodyInner::Reader(_) => {}
102+
BodyInner::OwnedReader(_) => {}
103+
};
104+
105+
// Any other body mode could be LengthDelimited depending on whether
106+
// we have got a size set.
107+
let mode = if let Some(size) = size {
108+
BodyMode::LengthDelimited(size)
109+
} else {
110+
BodyMode::Chunked
111+
};
112+
113+
Ok(mode)
78114
}
79115

80116
/// Turn this `SendBody` into a reader.
@@ -190,6 +226,7 @@ impl<'a> AsSendBody for SendBody<'a> {
190226
BodyInner::Body(v) => BodyInner::Reader(v),
191227
BodyInner::OwnedReader(v) => BodyInner::Reader(v),
192228
},
229+
size: self.size.take(),
193230
ended: self.ended,
194231
}
195232
}
@@ -205,94 +242,117 @@ pub(crate) enum BodyInner<'a> {
205242
OwnedReader(Box<dyn Read>),
206243
}
207244

208-
impl<'a> BodyInner<'a> {
209-
pub fn body_mode(&self) -> BodyMode {
210-
match self {
211-
BodyInner::None => BodyMode::NoBody,
212-
BodyInner::ByteSlice(v) => BodyMode::LengthDelimited(v.len() as u64),
213-
#[cfg(feature = "json")]
214-
BodyInner::ByteVec(v) => BodyMode::LengthDelimited(v.get_ref().len() as u64),
215-
BodyInner::Body(v) => v.body_mode(),
216-
BodyInner::Reader(_) => BodyMode::Chunked,
217-
BodyInner::OwnedReader(_) => BodyMode::Chunked,
218-
}
219-
}
220-
}
245+
// impl<'a> BodyInner<'a> {
246+
// pub fn body_mode(&self) -> BodyMode {
247+
// match self {
248+
// BodyInner::None => BodyMode::NoBody,
249+
// BodyInner::ByteSlice(v) => BodyMode::LengthDelimited(v.len() as u64),
250+
// #[cfg(feature = "json")]
251+
// BodyInner::ByteVec(v) => BodyMode::LengthDelimited(v.get_ref().len() as u64),
252+
// BodyInner::Body(v) => v.body_mode(),
253+
// BodyInner::Reader(_) => BodyMode::Chunked,
254+
// BodyInner::OwnedReader(_) => BodyMode::Chunked,
255+
// }
256+
// }
257+
// }
221258

222259
impl Private for &[u8] {}
223260
impl AsSendBody for &[u8] {
224261
fn as_body(&mut self) -> SendBody {
225-
BodyInner::ByteSlice(self).into()
262+
let inner = BodyInner::ByteSlice(self);
263+
(Some(self.len() as u64), inner).into()
226264
}
227265
}
228266

229267
impl Private for &str {}
230268
impl AsSendBody for &str {
231269
fn as_body(&mut self) -> SendBody {
232-
BodyInner::ByteSlice((*self).as_ref()).into()
270+
let inner = BodyInner::ByteSlice((*self).as_ref());
271+
(Some(self.len() as u64), inner).into()
233272
}
234273
}
235274

236275
impl Private for String {}
237276
impl AsSendBody for String {
238277
fn as_body(&mut self) -> SendBody {
239-
BodyInner::ByteSlice((*self).as_ref()).into()
278+
let inner = BodyInner::ByteSlice((*self).as_ref());
279+
(Some(self.len() as u64), inner).into()
240280
}
241281
}
242282

243283
impl Private for Vec<u8> {}
244284
impl AsSendBody for Vec<u8> {
245285
fn as_body(&mut self) -> SendBody {
246-
BodyInner::ByteSlice((*self).as_ref()).into()
286+
let inner = BodyInner::ByteSlice((*self).as_ref());
287+
(Some(self.len() as u64), inner).into()
247288
}
248289
}
249290

250291
impl Private for &String {}
251292
impl AsSendBody for &String {
252293
fn as_body(&mut self) -> SendBody {
253-
BodyInner::ByteSlice((*self).as_ref()).into()
294+
let inner = BodyInner::ByteSlice((*self).as_ref());
295+
(Some(self.len() as u64), inner).into()
254296
}
255297
}
256298

257299
impl Private for &Vec<u8> {}
258300
impl AsSendBody for &Vec<u8> {
259301
fn as_body(&mut self) -> SendBody {
260-
BodyInner::ByteSlice((*self).as_ref()).into()
302+
let inner = BodyInner::ByteSlice((*self).as_ref());
303+
(Some(self.len() as u64), inner).into()
261304
}
262305
}
263306

264307
impl Private for &File {}
265308
impl AsSendBody for &File {
266309
fn as_body(&mut self) -> SendBody {
267-
BodyInner::Reader(self).into()
310+
let size = lazy_file_size(self);
311+
SendBody {
312+
inner: BodyInner::Reader(self),
313+
size: Some(size),
314+
ended: false,
315+
}
268316
}
269317
}
270318

271-
impl Private for &TcpStream {}
272-
impl AsSendBody for &TcpStream {
319+
impl Private for File {}
320+
impl AsSendBody for File {
273321
fn as_body(&mut self) -> SendBody {
274-
BodyInner::Reader(self).into()
322+
let size = lazy_file_size(self);
323+
SendBody {
324+
inner: BodyInner::Reader(self),
325+
size: Some(size),
326+
ended: false,
327+
}
275328
}
276329
}
277330

278-
impl Private for File {}
279-
impl AsSendBody for File {
331+
fn lazy_file_size(file: &File) -> Result<u64, Error> {
332+
match file.metadata() {
333+
Ok(v) => Ok(v.len()),
334+
Err(e) => Err(e.into()),
335+
}
336+
}
337+
338+
impl Private for &TcpStream {}
339+
impl AsSendBody for &TcpStream {
280340
fn as_body(&mut self) -> SendBody {
281-
BodyInner::Reader(self).into()
341+
(None, BodyInner::Reader(self)).into()
282342
}
283343
}
284344

285345
impl Private for TcpStream {}
286346
impl AsSendBody for TcpStream {
287347
fn as_body(&mut self) -> SendBody {
288-
BodyInner::Reader(self).into()
348+
(None, BodyInner::Reader(self)).into()
289349
}
290350
}
291351

292352
impl Private for Stdin {}
293353
impl AsSendBody for Stdin {
294354
fn as_body(&mut self) -> SendBody {
295-
BodyInner::Reader(self).into()
355+
(None, BodyInner::Reader(self)).into()
296356
}
297357
}
298358

@@ -307,14 +367,15 @@ impl Private for UnixStream {}
307367
#[cfg(target_family = "unix")]
308368
impl AsSendBody for UnixStream {
309369
fn as_body(&mut self) -> SendBody {
310-
BodyInner::Reader(self).into()
370+
(None, BodyInner::Reader(self)).into()
311371
}
312372
}
313373

314-
impl<'a> From<BodyInner<'a>> for SendBody<'a> {
315-
fn from(inner: BodyInner<'a>) -> Self {
374+
impl<'a> From<(Option<u64>, BodyInner<'a>)> for SendBody<'a> {
375+
fn from((size, inner): (Option<u64>, BodyInner<'a>)) -> Self {
316376
SendBody {
317377
inner,
378+
size: size.map(Ok),
318379
ended: false,
319380
}
320381
}
@@ -323,27 +384,30 @@ impl<'a> From<BodyInner<'a>> for SendBody<'a> {
323384
impl Private for Body {}
324385
impl AsSendBody for Body {
325386
fn as_body(&mut self) -> SendBody {
326-
BodyInner::Body(Box::new(self.as_reader())).into()
387+
let size = self.content_length();
388+
(size, BodyInner::Body(Box::new(self.as_reader()))).into()
327389
}
328390
}
329391

330392
impl Private for Response<Body> {}
331393
impl AsSendBody for Response<Body> {
332394
fn as_body(&mut self) -> SendBody {
333-
BodyInner::Body(Box::new(self.body_mut().as_reader())).into()
395+
let size = self.body().content_length();
396+
(size, BodyInner::Body(Box::new(self.body_mut().as_reader()))).into()
334397
}
335398
}
336399

337400
impl<const N: usize> Private for &[u8; N] {}
338401
impl<const N: usize> AsSendBody for &[u8; N] {
339402
fn as_body(&mut self) -> SendBody {
340-
BodyInner::ByteSlice(self.as_slice()).into()
403+
let inner = BodyInner::ByteSlice((*self).as_ref());
404+
(Some(self.len() as u64), inner).into()
341405
}
342406
}
343407

344408
impl Private for () {}
345409
impl AsSendBody for () {
346410
fn as_body(&mut self) -> SendBody {
347-
BodyInner::None.into()
411+
(None, BodyInner::None).into()
348412
}
349413
}

0 commit comments

Comments
 (0)