gotham/state/
mod.rs

1//! Defines types for passing request state through `Middleware` and `Handler` implementations
2
3pub(crate) mod client_addr;
4mod data;
5mod from_state;
6mod request_id;
7
8use hyper::http::request;
9use hyper::upgrade::OnUpgrade;
10use hyper::{Body, Request};
11use log::{debug, trace};
12use std::any::{Any, TypeId};
13use std::collections::HashMap;
14use std::hash::{BuildHasherDefault, Hasher};
15use std::net::SocketAddr;
16
17pub use crate::state::client_addr::client_addr;
18pub use crate::state::data::StateData;
19pub use crate::state::from_state::FromState;
20pub use crate::state::request_id::request_id;
21
22use crate::helpers::http::request::path::RequestPathSegments;
23use crate::state::client_addr::put_client_addr;
24pub(crate) use crate::state::request_id::set_request_id;
25
26// https://docs.rs/http/0.2.5/src/http/extensions.rs.html#8-28
27// With TypeIds as keys, there's no need to hash them. They are already hashes
28// themselves, coming from the compiler. The IdHasher just holds the u64 of
29// the TypeId, and then returns it, instead of doing any bit fiddling.
30#[derive(Default)]
31struct IdHasher(u64);
32
33impl Hasher for IdHasher {
34    fn write(&mut self, _: &[u8]) {
35        unreachable!("TypeId calls write_u64");
36    }
37
38    #[inline]
39    fn write_u64(&mut self, id: u64) {
40        self.0 = id;
41    }
42
43    #[inline]
44    fn finish(&self) -> u64 {
45        self.0
46    }
47}
48
49/// Provides storage for request state, and stores one item of each type. The types used for
50/// storage must implement the [`StateData`] trait to allow its storage, which is usually done
51/// by adding `#[derive(StateData)]` on the type in question.
52///
53/// # Examples
54///
55/// ```rust
56/// use gotham::state::{State, StateData};
57///
58/// #[derive(StateData)]
59/// struct MyStruct {
60///     value: i32,
61/// }
62/// # fn main() {
63/// #   State::with_new(|state| {
64/// #
65/// state.put(MyStruct { value: 1 });
66/// assert_eq!(state.borrow::<MyStruct>().value, 1);
67/// #
68/// #   });
69/// # }
70/// ```
71pub struct State {
72    data: HashMap<TypeId, Box<dyn Any + Send>, BuildHasherDefault<IdHasher>>,
73}
74
75impl State {
76    /// Creates a new, empty `State` container. This is for internal Gotham use, because the
77    /// ability to create a new `State` container would allow for libraries and applications to
78    /// incorrectly discard important internal data.
79    pub(crate) fn new() -> State {
80        State {
81            data: HashMap::default(),
82        }
83    }
84
85    /// Creates a new, empty `State` and yields it mutably into the provided closure. This is
86    /// intended only for use in the documentation tests for `State`, since the `State` container
87    /// cannot be constructed otherwise.
88    #[doc(hidden)]
89    pub fn with_new<F>(f: F)
90    where
91        F: FnOnce(&mut State),
92    {
93        f(&mut State::new())
94    }
95
96    /// Instantiate a new `State` for a given `Request`. This is primarily useful if you're calling
97    /// Gotham from your own Hyper service.
98    pub fn from_request(req: Request<Body>, client_addr: SocketAddr) -> Self {
99        let mut state = Self::new();
100
101        put_client_addr(&mut state, client_addr);
102
103        let (
104            request::Parts {
105                method,
106                uri,
107                version,
108                headers,
109                mut extensions,
110                ..
111            },
112            body,
113        ) = req.into_parts();
114
115        state.put(RequestPathSegments::new(uri.path()));
116        state.put(method);
117        state.put(uri);
118        state.put(version);
119        state.put(headers);
120        state.put(body);
121
122        if let Some(on_upgrade) = extensions.remove::<OnUpgrade>() {
123            state.put(on_upgrade);
124        }
125
126        {
127            let request_id = set_request_id(&mut state);
128            debug!(
129                "[DEBUG][{}][Thread][{:?}]",
130                request_id,
131                std::thread::current().id(),
132            );
133        };
134
135        state
136    }
137
138    /// Puts a value into the `State` storage. One value of each type is retained. Successive calls
139    /// to `put` will overwrite the existing value of the same type.
140    ///
141    /// # Examples
142    ///
143    /// ```rust
144    /// # extern crate gotham;
145    /// # #[macro_use]
146    /// # extern crate gotham_derive;
147    /// #
148    /// # use gotham::state::State;
149    /// #
150    /// # #[derive(StateData)]
151    /// # struct MyStruct {
152    /// #     value: i32
153    /// # }
154    /// #
155    /// # #[derive(StateData)]
156    /// # struct AnotherStruct {
157    /// #     value: &'static str
158    /// # }
159    /// #
160    /// # fn main() {
161    /// #   State::with_new(|state| {
162    /// #
163    /// state.put(MyStruct { value: 1 });
164    /// assert_eq!(state.borrow::<MyStruct>().value, 1);
165    ///
166    /// state.put(AnotherStruct { value: "a string" });
167    /// state.put(MyStruct { value: 100 });
168    ///
169    /// assert_eq!(state.borrow::<AnotherStruct>().value, "a string");
170    /// assert_eq!(state.borrow::<MyStruct>().value, 100);
171    /// #
172    /// #   });
173    /// # }
174    /// ```
175    pub fn put<T>(&mut self, t: T)
176    where
177        T: StateData,
178    {
179        let type_id = TypeId::of::<T>();
180        trace!(" inserting record to state for type_id `{:?}`", type_id);
181        self.data.insert(type_id, Box::new(t));
182    }
183
184    /// Determines if the current value exists in `State` storage.
185    ///
186    /// # Examples
187    ///
188    /// ```rust
189    /// # extern crate gotham;
190    /// # #[macro_use]
191    /// # extern crate gotham_derive;
192    /// #
193    /// # use gotham::state::State;
194    /// #
195    /// # #[derive(StateData)]
196    /// # struct MyStruct {
197    /// #     value: i32
198    /// # }
199    /// #
200    /// # #[derive(StateData)]
201    /// # struct AnotherStruct {
202    /// # }
203    /// #
204    /// # fn main() {
205    /// #   State::with_new(|state| {
206    /// #
207    /// state.put(MyStruct { value: 1 });
208    /// assert!(state.has::<MyStruct>());
209    /// assert_eq!(state.borrow::<MyStruct>().value, 1);
210    ///
211    /// assert!(!state.has::<AnotherStruct>());
212    /// #
213    /// #   });
214    /// # }
215    /// ```
216    pub fn has<T>(&self) -> bool
217    where
218        T: StateData,
219    {
220        let type_id = TypeId::of::<T>();
221        self.data.get(&type_id).is_some()
222    }
223
224    /// Tries to borrow a value from the `State` storage.
225    ///
226    /// # Examples
227    ///
228    /// ```rust
229    /// # extern crate gotham;
230    /// # #[macro_use]
231    /// # extern crate gotham_derive;
232    /// #
233    /// # use gotham::state::State;
234    /// #
235    /// # #[derive(StateData)]
236    /// # struct MyStruct {
237    /// #     value: i32
238    /// # }
239    /// #
240    /// # #[derive(StateData)]
241    /// # struct AnotherStruct {
242    /// # }
243    /// #
244    /// # fn main() {
245    /// #   State::with_new(|state| {
246    /// #
247    /// state.put(MyStruct { value: 1 });
248    /// assert!(state.try_borrow::<MyStruct>().is_some());
249    /// assert_eq!(state.try_borrow::<MyStruct>().unwrap().value, 1);
250    ///
251    /// assert!(state.try_borrow::<AnotherStruct>().is_none());
252    /// #
253    /// #   });
254    /// # }
255    /// ```
256    pub fn try_borrow<T>(&self) -> Option<&T>
257    where
258        T: StateData,
259    {
260        let type_id = TypeId::of::<T>();
261        trace!(" borrowing state data for type_id `{:?}`", type_id);
262        self.data.get(&type_id).and_then(|b| b.downcast_ref::<T>())
263    }
264
265    /// Borrows a value from the `State` storage.
266    ///
267    /// # Panics
268    ///
269    /// If a value of type `T` is not present in `State`.
270    ///
271    /// # Examples
272    ///
273    /// ```rust
274    /// # extern crate gotham;
275    /// # #[macro_use]
276    /// # extern crate gotham_derive;
277    /// #
278    /// # use gotham::state::State;
279    /// #
280    /// # #[derive(StateData)]
281    /// # struct MyStruct {
282    /// #     value: i32
283    /// # }
284    /// #
285    /// # fn main() {
286    /// #   State::with_new(|state| {
287    /// #
288    /// state.put(MyStruct { value: 1 });
289    /// assert_eq!(state.borrow::<MyStruct>().value, 1);
290    /// #
291    /// #   });
292    /// # }
293    /// ```
294    pub fn borrow<T>(&self) -> &T
295    where
296        T: StateData,
297    {
298        self.try_borrow()
299            .expect("required type is not present in State container")
300    }
301
302    /// Tries to mutably borrow a value from the `State` storage.
303    ///
304    /// # Examples
305    ///
306    /// ```rust
307    /// # extern crate gotham;
308    /// # #[macro_use]
309    /// # extern crate gotham_derive;
310    /// #
311    /// # use gotham::state::State;
312    /// #
313    /// # #[derive(StateData)]
314    /// # struct MyStruct {
315    /// #     value: i32
316    /// # }
317    /// #
318    /// # #[derive(StateData)]
319    /// # struct AnotherStruct {
320    /// # }
321    /// #
322    /// # fn main() {
323    /// #   State::with_new(|state| {
324    /// #
325    /// state.put(MyStruct { value: 100 });
326    ///
327    /// if let Some(a) = state.try_borrow_mut::<MyStruct>() {
328    ///     a.value += 10;
329    /// }
330    ///
331    /// assert_eq!(state.borrow::<MyStruct>().value, 110);
332    ///
333    /// assert!(state.try_borrow_mut::<AnotherStruct>().is_none());
334    /// #   });
335    /// # }
336    pub fn try_borrow_mut<T>(&mut self) -> Option<&mut T>
337    where
338        T: StateData,
339    {
340        let type_id = TypeId::of::<T>();
341        trace!(" mutably borrowing state data for type_id `{:?}`", type_id);
342        self.data
343            .get_mut(&type_id)
344            .and_then(|b| b.downcast_mut::<T>())
345    }
346
347    /// Mutably borrows a value from the `State` storage.
348    ///
349    /// # Panics
350    ///
351    /// If a value of type `T` is not present in `State`.
352    ///
353    /// # Examples
354    ///
355    /// ```rust
356    /// # extern crate gotham;
357    /// # #[macro_use]
358    /// # extern crate gotham_derive;
359    /// #
360    /// # use gotham::state::State;
361    /// #
362    /// # #[derive(StateData)]
363    /// # struct MyStruct {
364    /// #     value: i32
365    /// # }
366    /// #
367    /// # #[derive(StateData)]
368    /// # struct AnotherStruct {
369    /// # }
370    /// #
371    /// # fn main() {
372    /// #   State::with_new(|state| {
373    /// #
374    /// state.put(MyStruct { value: 100 });
375    ///
376    /// {
377    ///     let a = state.borrow_mut::<MyStruct>();
378    ///     a.value += 10;
379    /// }
380    ///
381    /// assert_eq!(state.borrow::<MyStruct>().value, 110);
382    ///
383    /// assert!(state.try_borrow_mut::<AnotherStruct>().is_none());
384    /// #
385    /// #   });
386    /// # }
387    pub fn borrow_mut<T>(&mut self) -> &mut T
388    where
389        T: StateData,
390    {
391        self.try_borrow_mut()
392            .expect("required type is not present in State container")
393    }
394
395    /// Tries to move a value out of the `State` storage and return ownership.
396    ///
397    /// # Examples
398    ///
399    /// ```rust
400    /// # extern crate gotham;
401    /// # #[macro_use]
402    /// # extern crate gotham_derive;
403    /// #
404    /// # use gotham::state::State;
405    /// #
406    /// # #[derive(StateData)]
407    /// # struct MyStruct {
408    /// #     value: i32
409    /// # }
410    /// #
411    /// # #[derive(StateData)]
412    /// # struct AnotherStruct {
413    /// # }
414    /// #
415    /// # fn main() {
416    /// #   State::with_new(|state| {
417    /// #
418    /// state.put(MyStruct { value: 110 });
419    ///
420    /// assert_eq!(state.try_take::<MyStruct>().unwrap().value, 110);
421    ///
422    /// assert!(state.try_take::<MyStruct>().is_none());
423    /// assert!(state.try_borrow_mut::<MyStruct>().is_none());
424    /// assert!(state.try_borrow::<MyStruct>().is_none());
425    ///
426    /// assert!(state.try_take::<AnotherStruct>().is_none());
427    /// #
428    /// #   });
429    /// # }
430    pub fn try_take<T>(&mut self) -> Option<T>
431    where
432        T: StateData,
433    {
434        let type_id = TypeId::of::<T>();
435        trace!(
436            " taking ownership from state data for type_id `{:?}`",
437            type_id
438        );
439        self.data
440            .remove(&type_id)
441            .and_then(|b| b.downcast::<T>().ok())
442            .map(|b| *b)
443    }
444
445    /// Moves a value out of the `State` storage and returns ownership.
446    ///
447    /// # Panics
448    ///
449    /// If a value of type `T` is not present in `State`.
450    ///
451    /// # Examples
452    ///
453    /// ```rust
454    /// # extern crate gotham;
455    /// # #[macro_use]
456    /// # extern crate gotham_derive;
457    /// #
458    /// # use gotham::state::State;
459    /// #
460    /// # #[derive(StateData)]
461    /// # struct MyStruct {
462    /// #     value: i32
463    /// # }
464    /// #
465    /// # fn main() {
466    /// #   State::with_new(|state| {
467    /// #
468    /// state.put(MyStruct { value: 110 });
469    ///
470    /// assert_eq!(state.take::<MyStruct>().value, 110);
471    ///
472    /// assert!(state.try_take::<MyStruct>().is_none());
473    /// assert!(state.try_borrow_mut::<MyStruct>().is_none());
474    /// assert!(state.try_borrow::<MyStruct>().is_none());
475    /// #
476    /// #   });
477    /// # }
478    pub fn take<T>(&mut self) -> T
479    where
480        T: StateData,
481    {
482        self.try_take()
483            .expect("required type is not present in State container")
484    }
485}