Skip to main content

trillium_api/
api_handler.rs

1use crate::TryFromConn;
2use std::{future::Future, marker::PhantomData, sync::Arc};
3use trillium::{Conn, Handler, Info, Status, Upgrade};
4
5// A trait for `async fn(conn: &mut Conn, additional: Additional) -> ReturnType`
6pub trait MutBorrowConn<'conn, ReturnType, Additional>: Send + Sync + 'conn {
7    /// the returned future
8    type Fut: Future<Output = ReturnType> + Send + 'conn;
9    /// executes this function
10    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/// An interface layer built on trillium
26///
27/// This handler provides the capacity to extract various components of a conn such as deserializing
28/// a body, and supports returning handlers that will be called on the returned conn.
29///
30/// If [`ApiHandler`] encounters an error of any sort before the user-provided logic is executed, it
31/// will put an [`Error`](crate::Error) into the conn's state. A default error handler is provided.
32///
33/// More documentation for this type is needed, hence the -rc semver on this crate
34#[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    /// constructs a new [`ApiHandler`] from the provided
49    /// `async fn(&mut conn, TryFromConn) -> impl Handler`
50    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
67/// constructs a new [`ApiHandler`] from the provided
68/// `async fn(&mut conn, TryFromConn) -> impl Handler`
69///
70/// convenience function for [`ApiHandler::new`]
71pub 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}