1use std::future::Future;
2
3use crate::{
4 FromRequest, IntoResponse, Request, Response,
5 openapi::Operation,
6 stream::{BufferedRequestStream, collect_stream},
7};
8
9pub trait Handler<Arguments, Input> {
10 async fn call(&self, request: Request) -> Response;
11
12 #[doc(hidden)]
13 fn openapi() -> Operation;
14}
15
16impl<Output: IntoResponse, Fut: Future<Output = Output>, F: Fn() -> Fut> Handler<(), ()> for F {
17 async fn call(&self, request: Request) -> Response {
18 drop(request);
19 self().await.into_response()
20 }
21
22 fn openapi() -> Operation {
23 let mut operation = Operation::default();
24 Output::openapi(&mut operation);
25 operation.ensure_response();
26 operation
27 }
28}
29
30macro_rules! impl_handler {
31 ([$(($argument:ident, $value:ident)),*]; ($last_argument:ident, $last_value:ident)) => {
32 impl<
33 $(
34 $argument: for<'request> FromRequest<(
35 &'request Request,
36 &'request [u8],
37 )>,
38 )*
39 $last_argument: for<'request> FromRequest<(
40 &'request Request,
41 &'request [u8],
42 )>,
43 Output: IntoResponse,
44 Fut: Future<Output = Output>,
45 F: Fn($($argument,)* $last_argument) -> Fut,
46 > Handler<($($argument,)* $last_argument,), ()> for F
47 {
48 async fn call(&self, mut request: Request) -> Response {
49 let has_buffered = false
50 $(|| <$argument as FromRequest<(&Request, &[u8])>>::BUFFERED)*
51 || <$last_argument as FromRequest<(&Request, &[u8])>>::BUFFERED;
52
53 let buffered = if has_buffered {
54 let body_limit = request.body_limit();
55
56 match collect_stream(request.body.as_mut(), body_limit).await {
57 Ok(buffered) => buffered,
58 Err(error) => return error.into_response(),
59 }
60 } else {
61 Vec::new()
62 };
63
64 $(
65 let $value = match <$argument as FromRequest<(
66 &Request,
67 &[u8],
68 )>>::from_request((&request, buffered.as_slice()))
69 .await
70 {
71 Ok(value) => value,
72 Err(error) => return error.into_response(),
73 };
74 )*
75
76 let $last_value = match <$last_argument as FromRequest<(
77 &Request,
78 &[u8],
79 )>>::from_request((&request, buffered.as_slice()))
80 .await
81 {
82 Ok(value) => value,
83 Err(error) => return error.into_response(),
84 };
85
86 self($($value,)* $last_value).await.into_response()
87 }
88
89 fn openapi() -> Operation {
90 let mut operation = Operation::default();
91 $(
92 <$argument as FromRequest<(
93 &Request,
94 &[u8],
95 )>>::openapi(&mut operation);
96 )*
97 <$last_argument as FromRequest<(
98 &Request,
99 &[u8],
100 )>>::openapi(&mut operation);
101 Output::openapi(&mut operation);
102 operation.ensure_response();
103 operation
104 }
105 }
106
107 impl<
108 $(
109 $argument: for<'request> FromRequest<(
110 &'request Request,
111 &'request [u8],
112 )>,
113 )*
114 $last_argument: FromRequest<Request>,
115 Output: IntoResponse,
116 Fut: Future<Output = Output>,
117 F: Fn($($argument,)* $last_argument) -> Fut,
118 > Handler<($($argument,)* $last_argument,), Request> for F
119 {
120 async fn call(&self, mut request: Request) -> Response {
121 let has_buffered = false
122 $(|| <$argument as FromRequest<(&Request, &[u8])>>::BUFFERED)*;
123
124 let buffered = if has_buffered {
125 let body_limit = request.body_limit();
126
127 match collect_stream(request.body.as_mut(), body_limit).await {
128 Ok(buffered) => buffered,
129 Err(error) => return error.into_response(),
130 }
131 } else {
132 Vec::new()
133 };
134
135 $(
136 let $value = match <$argument as FromRequest<(
137 &Request,
138 &[u8],
139 )>>::from_request((&request, buffered.as_slice()))
140 .await
141 {
142 Ok(value) => value,
143 Err(error) => return error.into_response(),
144 };
145 )*
146
147 if has_buffered {
148 request.body = Box::new(BufferedRequestStream::new(buffered));
149 }
150
151 let $last_value = match <$last_argument as FromRequest<Request>>::from_request(
152 request,
153 )
154 .await
155 {
156 Ok(value) => value,
157 Err(error) => return error.into_response(),
158 };
159
160 self($($value,)* $last_value).await.into_response()
161 }
162
163 fn openapi() -> Operation {
164 let mut operation = Operation::default();
165 $(
166 <$argument as FromRequest<(
167 &Request,
168 &[u8],
169 )>>::openapi(&mut operation);
170 )*
171 <$last_argument as FromRequest<Request>>::openapi(&mut operation);
172 Output::openapi(&mut operation);
173 operation.ensure_response();
174 operation
175 }
176 }
177 };
178}
179
180serverkit_macros::impl_handlers!(16);
181
182#[cfg(test)]
183mod tests {
184 use std::{
185 cell::Cell,
186 convert::Infallible,
187 future::Future,
188 rc::Rc,
189 task::{Context, Poll, Waker},
190 };
191
192 use crate::{
193 Body, Bytes, Config, Extension, FromRequest, Handler, Headers, Method, Request,
194 RequestStream, RouteMethods, Router, State, StreamError,
195 };
196
197 struct ProbeStream {
198 body: Vec<u8>,
199 sent: bool,
200 polls: Rc<Cell<usize>>,
201 }
202
203 impl RequestStream for ProbeStream {
204 fn poll_next(
205 &mut self,
206 _context: &mut Context<'_>,
207 ) -> Poll<Option<Result<(), StreamError>>> {
208 self.polls.set(self.polls.get() + 1);
209
210 if self.sent {
211 Poll::Ready(None)
212 } else {
213 self.sent = true;
214 Poll::Ready(Some(Ok(())))
215 }
216 }
217
218 fn chunk(&self) -> &[u8] {
219 &self.body
220 }
221 }
222
223 struct BufferedBytes(Vec<u8>);
224
225 impl<'request> FromRequest<(&'request Request, &'request [u8])> for BufferedBytes {
226 type Error = Infallible;
227
228 const BUFFERED: bool = true;
229
230 async fn from_request(
231 input: (&'request Request, &'request [u8]),
232 ) -> Result<Self, Self::Error> {
233 Ok(Self(input.1.to_vec()))
234 }
235 }
236
237 fn block_on<F: Future>(future: F) -> F::Output {
238 let mut future = std::pin::pin!(future);
239 let waker = Waker::noop();
240 let mut context = Context::from_waker(waker);
241
242 loop {
243 match future.as_mut().poll(&mut context) {
244 Poll::Ready(output) => return output,
245 Poll::Pending => std::thread::yield_now(),
246 }
247 }
248 }
249
250 fn request(body: &[u8], polls: Rc<Cell<usize>>) -> Request {
251 Request::from_parts(
252 Method::GET,
253 "/",
254 None,
255 Headers::new(),
256 Box::new(ProbeStream {
257 body: body.to_vec(),
258 sent: false,
259 polls,
260 }),
261 )
262 }
263
264 async fn one(_a0: Method) {}
265
266 async fn two(_a0: Method, _a1: Method) {}
267
268 async fn stream_last(_a0: Method, _a1: Body) {}
269
270 async fn leave_stream_unread(_body: Body) -> &'static str {
271 "unread"
272 }
273
274 async fn buffered_then_stream(
275 first: BufferedBytes,
276 second: BufferedBytes,
277 mut body: Body,
278 ) -> Vec<u8> {
279 assert_eq!(first.0, second.0);
280 body.next().await.unwrap().unwrap().to_vec()
281 }
282
283 async fn buffered_body(Bytes(bytes): Bytes) -> Vec<u8> {
284 bytes
285 }
286
287 async fn streaming_body(mut body: Body) -> Result<Vec<u8>, StreamError> {
288 let mut bytes = Vec::new();
289
290 while let Some(chunk) = body.next().await {
291 bytes.extend_from_slice(chunk?);
292 }
293
294 Ok(bytes)
295 }
296
297 async fn application_state(State(value): State<String>) -> String {
298 value.as_str().to_owned()
299 }
300
301 async fn request_extension(Extension(value): Extension<u64>) -> String {
302 value.to_string()
303 }
304
305 #[allow(clippy::too_many_arguments)]
306 async fn sixteen(
307 _a0: Method,
308 _a1: Method,
309 _a2: Method,
310 _a3: Method,
311 _a4: Method,
312 _a5: Method,
313 _a6: Method,
314 _a7: Method,
315 _a8: Method,
316 _a9: Method,
317 _a10: Method,
318 _a11: Method,
319 _a12: Method,
320 _a13: Method,
321 _a14: Method,
322 _a15: Method,
323 ) {
324 }
325
326 fn assert_handler<Arguments, Input, H: Handler<Arguments, Input>>(_handler: H) {}
327
328 #[test]
329 fn implements_supported_arities() {
330 assert_handler::<(Method,), (), _>(one);
331 assert_handler::<(Method, Method), (), _>(two);
332 assert_handler::<(Method, Body), Request, _>(stream_last);
333 assert_handler::<
334 (
335 Method,
336 Method,
337 Method,
338 Method,
339 Method,
340 Method,
341 Method,
342 Method,
343 Method,
344 Method,
345 Method,
346 Method,
347 Method,
348 Method,
349 Method,
350 Method,
351 ),
352 (),
353 _,
354 >(sixteen);
355 }
356
357 #[test]
358 fn streaming_only_does_not_preconsume_the_body() {
359 let polls = Rc::new(Cell::new(0));
360 let application = Router::new(Config::new(), ("/".GET(leave_stream_unread),));
361 let response = block_on(application.handle(request(b"stream", Rc::clone(&polls))));
362
363 assert_eq!(response.body(), b"unread");
364 assert_eq!(polls.get(), 0);
365 }
366
367 #[test]
368 fn buffered_extractors_share_one_collection_before_streaming() {
369 let polls = Rc::new(Cell::new(0));
370 let application = Router::new(Config::new(), ("/".GET(buffered_then_stream),));
371 let response = block_on(application.handle(request(b"replayed", Rc::clone(&polls))));
372
373 assert_eq!(response.body(), b"replayed");
374 assert_eq!(polls.get(), 2);
375 }
376
377 #[test]
378 fn body_limit_applies_to_buffered_and_streaming_extractors() {
379 let buffered = Router::new(Config::new(), ("/".GET(buffered_body),)).body_limit(3);
380 let response = block_on(buffered.handle(request(b"four", Rc::new(Cell::new(0)))));
381 assert_eq!(response.status(), 413);
382
383 let streaming = Router::new(Config::new(), ("/".GET(streaming_body),)).body_limit(3);
384 let response = block_on(streaming.handle(request(b"four", Rc::new(Cell::new(0)))));
385 assert_eq!(response.status(), 413);
386 }
387
388 #[test]
389 fn extracts_application_state_and_request_extensions() {
390 let application =
391 Router::new(Config::new(), ("/".GET(application_state),)).state("ready".to_owned());
392 let response = block_on(application.handle(request(b"", Rc::new(Cell::new(0)))));
393 assert_eq!(response.body(), b"ready");
394
395 let application = Router::new(Config::new(), ("/".GET(request_extension),));
396 let mut request = request(b"", Rc::new(Cell::new(0)));
397 request.insert_extension(42_u64);
398 let response = block_on(application.handle(request));
399 assert_eq!(response.body(), b"42");
400 }
401}