1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
use bytes::Buf;
use h2::RecvStream;
use http::Request;
use crate::server::connection::BodyReadingBuffer;
use crate::server::errors::{ServerError, WaterErrors};
use crate::server::HttpStream;
pub (crate) enum StreamBytesPuller<'a> {
H1(H1BytesPuller<'a>),
H2(H2BytesPuller<'a>)
}
pub (crate) struct H1BytesPuller<'a> {
pub(crate)stream:&'a mut HttpStream,
pub(crate)reading_buffer:&'a mut BodyReadingBuffer,
pub(crate) left_bytes:&'a [u8]
}
#[derive(Debug)]
pub (crate) struct H2BytesPuller<'a>{
pub(crate)batch:&'a mut Request<RecvStream>,
}
/// struct for handling body bytes as chunks with very efficient way
pub struct BytesPuller<'a >{
puller:StreamBytesPuller<'a>,
content_length:usize,
}
impl <'a> BytesPuller <'a> {
pub (crate) fn new(
puller:StreamBytesPuller<'a>,
content_length:usize)->BytesPuller<'a>{
BytesPuller {
puller,
content_length
}
}
/// reading each chunk seperated and parsed to [FnMut] Closure
pub async fn on_chunk(&mut self,mut callback:impl FnMut(&[u8])-> Result<(),()> )
->Result<(),WaterErrors<'_>>{
let content_length = self.content_length;
match &mut self.puller {
StreamBytesPuller::H1(h1) => {
let left_bytes = h1.left_bytes;
let err= Err(WaterErrors::Server(
ServerError::HANDLING_INCOMING_BODY_ERROR
));
#[cfg(feature = "debugging")]
{
tracing::debug!("[BytesPuller]: Content-Length: {content_length} while left bytes {}",left_bytes.len());
}
if content_length <= left_bytes.len() {
let data = &left_bytes[..content_length];
if let Err(_) = callback(data) {
return err
}
return Ok(())
} else {
let mut remaining = content_length;
if let Err(_) = callback(left_bytes) {
return err
}
remaining-=left_bytes.len();
loop {
if remaining == 0 {
return Ok(())
}
if h1.reading_buffer.read_buf(h1.stream).await.is_err() { return err}
let data = h1.reading_buffer.chunk();
let to_index = content_length.min(data.len());
if callback(&data[..to_index]).is_err() { return err}
remaining-=remaining.min(data.len());
continue;
}
}
}
StreamBytesPuller::H2(h2) => {
let mut remaining = self.content_length;
let body_mut = h2.batch.body_mut();
let err = Err(
WaterErrors::Server(
ServerError::HANDLING_INCOMING_BODY_ERROR
)
);
while remaining > 0 {
let data = body_mut.data().await;
match data {
None => {
// Stream ended but we still expected more bytes!
return err;
}
Some(data_result) => {
match data_result {
Ok(data) => {
let chunk_bytes = data.as_ref();
// Protect against client sending more data than Content-Length
let to_consume = remaining.min(chunk_bytes.len());
if to_consume == 0 {
break;
}
if callback(&chunk_bytes[..to_consume]).is_err() {
return err;
}
remaining -= to_consume;
}
Err(_) => {
return err;
}
}
}
}
}
// If remaining is 0, we got everything perfectly.
if remaining == 0 {
return Ok(());
} else {
return err;
}
}
}
}
}