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, Error, 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 struct StateBacked(String);
226
227 impl<'request> FromRequest<(&'request Request, &'request [u8])> for StateBacked {
228 type Error = Error;
229
230 async fn from_request(
231 input: (&'request Request, &'request [u8]),
232 ) -> Result<Self, Self::Error> {
233 let State(value) = State::<String>::from_request(input).await?;
234 Ok(Self(value.as_str().to_owned()))
235 }
236 }
237
238 impl<'request> FromRequest<(&'request Request, &'request [u8])> for BufferedBytes {
239 type Error = Infallible;
240
241 const BUFFERED: bool = true;
242
243 async fn from_request(
244 input: (&'request Request, &'request [u8]),
245 ) -> Result<Self, Self::Error> {
246 Ok(Self(input.1.to_vec()))
247 }
248 }
249
250 fn block_on<F: Future>(future: F) -> F::Output {
251 let mut future = std::pin::pin!(future);
252 let waker = Waker::noop();
253 let mut context = Context::from_waker(waker);
254
255 loop {
256 match future.as_mut().poll(&mut context) {
257 Poll::Ready(output) => return output,
258 Poll::Pending => std::thread::yield_now(),
259 }
260 }
261 }
262
263 fn request(body: &[u8], polls: Rc<Cell<usize>>) -> Request {
264 Request::from_parts(
265 Method::GET,
266 "/",
267 None,
268 Headers::new(),
269 Box::new(ProbeStream {
270 body: body.to_vec(),
271 sent: false,
272 polls,
273 }),
274 )
275 }
276
277 async fn one(_a0: Method) {}
278
279 async fn two(_a0: Method, _a1: Method) {}
280
281 async fn stream_last(_a0: Method, _a1: Body) {}
282
283 async fn leave_stream_unread(_body: Body) -> &'static str {
284 "unread"
285 }
286
287 async fn buffered_then_stream(
288 first: BufferedBytes,
289 second: BufferedBytes,
290 mut body: Body,
291 ) -> Vec<u8> {
292 assert_eq!(first.0, second.0);
293 body.next().await.unwrap().unwrap().to_vec()
294 }
295
296 async fn buffered_body(Bytes(bytes): Bytes) -> Vec<u8> {
297 bytes
298 }
299
300 async fn streaming_body(mut body: Body) -> Result<Vec<u8>, StreamError> {
301 let mut bytes = Vec::new();
302
303 while let Some(chunk) = body.next().await {
304 bytes.extend_from_slice(chunk?);
305 }
306
307 Ok(bytes)
308 }
309
310 async fn application_state(State(value): State<String>) -> String {
311 value.as_str().to_owned()
312 }
313
314 async fn request_extension(Extension(value): Extension<u64>) -> String {
315 value.to_string()
316 }
317
318 async fn state_backed(StateBacked(value): StateBacked) -> String {
319 value
320 }
321
322 #[allow(clippy::too_many_arguments)]
323 async fn sixteen(
324 _a0: Method,
325 _a1: Method,
326 _a2: Method,
327 _a3: Method,
328 _a4: Method,
329 _a5: Method,
330 _a6: Method,
331 _a7: Method,
332 _a8: Method,
333 _a9: Method,
334 _a10: Method,
335 _a11: Method,
336 _a12: Method,
337 _a13: Method,
338 _a14: Method,
339 _a15: Method,
340 ) {
341 }
342
343 fn assert_handler<Arguments, Input, H: Handler<Arguments, Input>>(_handler: H) {}
344
345 #[test]
346 fn implements_supported_arities() {
347 assert_handler::<(Method,), (), _>(one);
348 assert_handler::<(Method, Method), (), _>(two);
349 assert_handler::<(Method, Body), Request, _>(stream_last);
350 assert_handler::<
351 (
352 Method,
353 Method,
354 Method,
355 Method,
356 Method,
357 Method,
358 Method,
359 Method,
360 Method,
361 Method,
362 Method,
363 Method,
364 Method,
365 Method,
366 Method,
367 Method,
368 ),
369 (),
370 _,
371 >(sixteen);
372 }
373
374 #[test]
375 fn streaming_only_does_not_preconsume_the_body() {
376 let polls = Rc::new(Cell::new(0));
377 let application = Router::new(Config::new(), ("/".GET(leave_stream_unread),));
378 let response = block_on(application.handle(request(b"stream", Rc::clone(&polls))));
379
380 assert_eq!(response.body(), b"unread");
381 assert_eq!(polls.get(), 0);
382 }
383
384 #[test]
385 fn buffered_extractors_share_one_collection_before_streaming() {
386 let polls = Rc::new(Cell::new(0));
387 let application = Router::new(Config::new(), ("/".GET(buffered_then_stream),));
388 let response = block_on(application.handle(request(b"replayed", Rc::clone(&polls))));
389
390 assert_eq!(response.body(), b"replayed");
391 assert_eq!(polls.get(), 2);
392 }
393
394 #[test]
395 fn body_limit_applies_to_buffered_and_streaming_extractors() {
396 let buffered = Router::new(Config::new(), ("/".GET(buffered_body),)).body_limit(3);
397 let response = block_on(buffered.handle(request(b"four", Rc::new(Cell::new(0)))));
398 assert_eq!(response.status(), 413);
399
400 let streaming = Router::new(Config::new(), ("/".GET(streaming_body),)).body_limit(3);
401 let response = block_on(streaming.handle(request(b"four", Rc::new(Cell::new(0)))));
402 assert_eq!(response.status(), 413);
403 }
404
405 #[test]
406 fn extracts_application_state_and_request_extensions() {
407 let application =
408 Router::new(Config::new(), ("/".GET(application_state),)).state("ready".to_owned());
409 let response = block_on(application.handle(request(b"", Rc::new(Cell::new(0)))));
410 assert_eq!(response.body(), b"ready");
411
412 let application = Router::new(Config::new(), ("/".GET(request_extension),));
413 let mut request = request(b"", Rc::new(Cell::new(0)));
414 request.extensions.insert(42_u64);
415 let response = block_on(application.handle(request));
416 assert_eq!(response.body(), b"42");
417 }
418
419 #[test]
420 fn nested_extractors_can_propagate_missing_values_into_error() {
421 let application =
422 Router::new(Config::new(), ("/".GET(state_backed),)).state("ready".to_owned());
423 let response = block_on(application.handle(request(b"", Rc::new(Cell::new(0)))));
424 assert_eq!(response.body(), b"ready");
425
426 let application = Router::new(Config::new(), ("/".GET(state_backed),));
427 let response = block_on(application.handle(request(b"", Rc::new(Cell::new(0)))));
428 assert_eq!(response.status(), 500);
429 assert_eq!(
430 response.body(),
431 br#"{"error":{"code":"application.state.unavailable","message":"application state is unavailable","fields":[]}}"#,
432 );
433 }
434}