1use crate::TryFromConn;
2use std::{future::Future, marker::PhantomData, sync::Arc};
3use trillium::{Conn, Handler, Info, Status, Upgrade};
4
5pub trait MutBorrowConn<'conn, ReturnType, Additional>: Send + Sync + 'conn {
7 type Fut: Future<Output = ReturnType> + Send + 'conn;
9 fn call(&self, conn: &'conn mut Conn, additional: Additional) -> Self::Fut;
11}
12
13impl<'conn, Fun, Fut, ReturnType, Additional> MutBorrowConn<'conn, ReturnType, Additional> for Fun
14where
15 Fun: Fn(&'conn mut Conn, Additional) -> Fut + Send + Sync + 'conn,
16 Fut: Future<Output = ReturnType> + Send + 'conn,
17{
18 type Fut = Fut;
19
20 fn call(&self, conn: &'conn mut Conn, additional: Additional) -> Fut {
21 self(conn, additional)
22 }
23}
24
25#[derive(Debug)]
35pub struct ApiHandler<F, OutputHandler, TryFromConn>(
36 F,
37 PhantomData<OutputHandler>,
38 PhantomData<TryFromConn>,
39);
40
41impl<TryFromConnHandler, OutputHandler, Extracted>
42 ApiHandler<TryFromConnHandler, OutputHandler, Extracted>
43where
44 TryFromConnHandler: for<'a> MutBorrowConn<'a, OutputHandler, Extracted>,
45 OutputHandler: Handler,
46 Extracted: TryFromConn,
47{
48 pub fn new(api_handler: TryFromConnHandler) -> Self {
51 Self::from(api_handler)
52 }
53}
54
55impl<TryFromConnHandler, OutputHandler, Extracted> From<TryFromConnHandler>
56 for ApiHandler<TryFromConnHandler, OutputHandler, Extracted>
57where
58 TryFromConnHandler: for<'a> MutBorrowConn<'a, OutputHandler, Extracted>,
59 OutputHandler: Handler,
60 Extracted: TryFromConn,
61{
62 fn from(value: TryFromConnHandler) -> Self {
63 Self(value, PhantomData, PhantomData)
64 }
65}
66
67pub fn api<TryFromConnHandler, OutputHandler, Extracted>(
72 api_handler: TryFromConnHandler,
73) -> ApiHandler<TryFromConnHandler, OutputHandler, Extracted>
74where
75 TryFromConnHandler: for<'a> MutBorrowConn<'a, OutputHandler, Extracted>,
76 Extracted: TryFromConn,
77 OutputHandler: Handler,
78{
79 ApiHandler::from(api_handler)
80}
81
82impl<TryFromConnHandler, OutputHandler, Extracted> Handler
83 for ApiHandler<TryFromConnHandler, OutputHandler, Extracted>
84where
85 TryFromConnHandler: for<'a> MutBorrowConn<'a, OutputHandler, Extracted>,
86 Extracted: TryFromConn,
87 Extracted::Error: Handler,
88 OutputHandler: Handler,
89{
90 async fn run(&self, mut conn: Conn) -> Conn {
91 let mut output_handler: Result<OutputHandler, <Extracted as TryFromConn>::Error> =
92 match Extracted::try_from_conn(&mut conn).await {
93 Ok(extracted) => Ok(self.0.call(&mut conn, extracted).await),
94 Err(error_handler) => Err(error_handler),
95 };
96
97 if let Some(info) = conn.state_mut::<Info>() {
98 output_handler.init(info).await;
99 } else {
100 output_handler.init(&mut Info::default()).await;
101 }
102 let mut conn = output_handler.run(conn).await;
103 if conn.status().is_none() && conn.response_body().is_some() {
104 conn.set_status(Status::Ok);
105 }
106 conn.with_state(OutputHandlerWrapper(
107 Arc::new(output_handler),
108 PhantomData::<Self>,
109 ))
110 }
111
112 async fn before_send(&self, conn: Conn) -> Conn {
113 match conn
114 .state::<OutputHandlerWrapper<Self, OutputHandler, <Extracted as TryFromConn>::Error>>()
115 .cloned()
116 {
117 Some(OutputHandlerWrapper(handler, _)) => handler.before_send(conn).await,
118 _ => conn,
119 }
120 }
121
122 fn has_upgrade(&self, upgrade: &Upgrade) -> bool {
123 upgrade
124 .state()
125 .get::<OutputHandlerWrapper<Self, OutputHandler, <Extracted as TryFromConn>::Error>>()
126 .cloned()
127 .is_some_and(|OutputHandlerWrapper(handler, _)| handler.has_upgrade(upgrade))
128 }
129
130 async fn upgrade(&self, upgrade: Upgrade) {
131 if let Some(OutputHandlerWrapper(handler, _)) = upgrade
132 .state()
133 .get::<OutputHandlerWrapper<Self, OutputHandler, <Extracted as TryFromConn>::Error>>()
134 .cloned()
135 {
136 handler.upgrade(upgrade).await
137 }
138 }
139}
140
141struct OutputHandlerWrapper<TFC, OH, EH>(Arc<Result<OH, EH>>, PhantomData<TFC>);
142
143impl<TFC, OH, EH> Clone for OutputHandlerWrapper<TFC, OH, EH> {
144 fn clone(&self) -> Self {
145 Self(Arc::clone(&self.0), self.1)
146 }
147}