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