trillium_conn_id/
lib.rs

1/*!
2Trillium crate to add identifiers to conns.
3
4This crate provides the following utilities:
5* [`ConnId`] a handler which must be called for the rest of this crate to function
6* [`log_formatter::conn_id`] a formatter to use with trillium_logger
7  (note that this does not depend on the trillium_logger crate and is very lightweight
8  if you do not use that crate)
9* [`ConnIdExt`] an extension trait for retrieving the id from a conn
10
11*/
12#![forbid(unsafe_code)]
13#![deny(
14    missing_copy_implementations,
15    rustdoc::missing_crate_level_docs,
16    missing_debug_implementations,
17    missing_docs,
18    nonstandard_style,
19    unused_qualifications
20)]
21
22use fastrand::Rng;
23use std::{
24    fmt::{Debug, Formatter, Result},
25    iter::repeat_with,
26    ops::Deref,
27    sync::{Arc, Mutex},
28};
29use trillium::{async_trait, Conn, Handler, HeaderName, KnownHeaderName, StateSet};
30
31#[derive(Default)]
32enum IdGenerator {
33    #[default]
34    Default,
35    SeededFastrand(Arc<Mutex<Rng>>),
36    Fn(Box<dyn Fn() -> String + Send + Sync + 'static>),
37}
38
39impl IdGenerator {
40    fn generate(&self) -> Id {
41        match self {
42            IdGenerator::Default => Id::default(),
43            IdGenerator::SeededFastrand(rng) => Id::with_rng(&mut rng.lock().unwrap()),
44            IdGenerator::Fn(gen_fun) => Id(gen_fun()),
45        }
46    }
47}
48
49/**
50Trillium handler to set a identifier for every Conn.
51
52By default, it will use an inbound `x-request-id` request header or if
53that is missing, populate a ten character random id. This handler will
54set an outbound `x-request-id` header as well by default. All of this
55behavior can be customized with [`ConnId::with_request_header`],
56[`ConnId::with_response_header`] and [`ConnId::with_id_generator`]
57*/
58pub struct ConnId {
59    request_header: Option<HeaderName<'static>>,
60    response_header: Option<HeaderName<'static>>,
61    id_generator: IdGenerator,
62}
63
64impl Debug for ConnId {
65    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
66        f.debug_struct("ConnId")
67            .field("request_header", &self.request_header)
68            .field("response_header", &self.response_header)
69            .field("id_generator", &self.id_generator)
70            .finish()
71    }
72}
73
74impl Default for ConnId {
75    fn default() -> Self {
76        Self {
77            request_header: Some(KnownHeaderName::XrequestId.into()),
78            response_header: Some(KnownHeaderName::XrequestId.into()),
79            id_generator: Default::default(),
80        }
81    }
82}
83
84impl ConnId {
85    /**
86    Constructs a new ConnId handler
87    ```
88    # use trillium_testing::prelude::*;
89    # use trillium_conn_id::ConnId;
90    let app = (ConnId::new().with_seed(1000), "ok"); // seeded for testing
91    assert_ok!(
92        get("/").on(&app),
93        "ok",
94        "x-request-id" => "J4lzoPXcT5"
95    );
96
97    assert_headers!(
98        get("/")
99            .with_request_header("x-request-id", "inbound")
100            .on(&app),
101        "x-request-id" => "inbound"
102    );
103    ```
104    */
105    pub fn new() -> Self {
106        Self::default()
107    }
108
109    /**
110    Specifies a request header to use. If this header is provided on
111    the inbound request, the id will be used unmodified. To disable
112    this behavior, see [`ConnId::without_request_header`]
113
114    ```
115    # use trillium_testing::prelude::*;
116    # use trillium_conn_id::ConnId;
117
118    let app = (
119        ConnId::new().with_request_header("x-custom-id"),
120        "ok"
121    );
122
123    assert_headers!(
124        get("/")
125            .with_request_header("x-custom-id", "inbound")
126            .on(&app),
127        "x-request-id" => "inbound"
128    );
129    ```
130
131    */
132    pub fn with_request_header(mut self, request_header: impl Into<HeaderName<'static>>) -> Self {
133        self.request_header = Some(request_header.into());
134        self
135    }
136
137    /**
138    disables the default behavior of reusing an inbound header for
139    the request id. If a ConnId is configured
140    `without_request_header`, a new id will always be generated
141    */
142    pub fn without_request_header(mut self) -> Self {
143        self.request_header = None;
144        self
145    }
146
147    /**
148    Specifies a response header to set. To disable this behavior, see
149    [`ConnId::without_response_header`]
150
151    ```
152    # use trillium_testing::prelude::*;
153    # use trillium_conn_id::ConnId;
154    let app = (
155        ConnId::new()
156            .with_seed(1000) // for testing
157            .with_response_header("x-custom-header"),
158        "ok"
159    );
160
161    assert_headers!(
162        get("/").on(&app),
163        "x-custom-header" => "J4lzoPXcT5"
164    );
165    ```
166    */
167    pub fn with_response_header(mut self, response_header: impl Into<HeaderName<'static>>) -> Self {
168        self.response_header = Some(response_header.into());
169        self
170    }
171
172    /**
173    Disables the default behavior of sending the conn id as a response
174    header. A request id will be available within the application
175    through use of [`ConnIdExt`] but will not be sent as part of the
176    response.
177    */
178    pub fn without_response_header(mut self) -> Self {
179        self.response_header = None;
180        self
181    }
182
183    /**
184    Provide an alternative generator function for ids. The default
185    is a ten-character alphanumeric random sequence.
186
187    ```
188    # use trillium_testing::prelude::*;
189    # use trillium_conn_id::ConnId;
190    # use uuid::Uuid;
191    let app = (
192        ConnId::new().with_id_generator(|| Uuid::new_v4().to_string()),
193        "ok"
194    );
195
196    // assert that the id is a valid uuid, even if we can't assert a specific value
197    assert!(Uuid::parse_str(get("/").on(&app).response_headers().get_str("x-request-id").unwrap()).is_ok());
198    ```
199    */
200    pub fn with_id_generator<F>(mut self, id_generator: F) -> Self
201    where
202        F: Fn() -> String + Send + Sync + 'static,
203    {
204        self.id_generator = IdGenerator::Fn(Box::new(id_generator));
205        self
206    }
207
208    /// seed a shared rng
209    ///
210    /// this is primarily useful for tests
211    pub fn with_seed(mut self, seed: u64) -> Self {
212        self.id_generator = IdGenerator::SeededFastrand(Arc::new(Mutex::new(Rng::with_seed(seed))));
213        self
214    }
215
216    fn generate_id(&self) -> Id {
217        self.id_generator.generate()
218    }
219}
220
221#[derive(Clone, Debug)]
222struct Id(String);
223
224impl Deref for Id {
225    type Target = str;
226
227    fn deref(&self) -> &Self::Target {
228        &self.0
229    }
230}
231
232impl std::fmt::Display for Id {
233    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
234        f.write_str(self)
235    }
236}
237
238impl Default for Id {
239    fn default() -> Self {
240        Self(repeat_with(fastrand::alphanumeric).take(10).collect())
241    }
242}
243
244impl Id {
245    fn with_rng(rng: &mut Rng) -> Self {
246        Self(repeat_with(|| rng.alphanumeric()).take(10).collect())
247    }
248}
249
250#[async_trait]
251impl Handler for ConnId {
252    async fn run(&self, mut conn: Conn) -> Conn {
253        let id = self
254            .request_header
255            .as_ref()
256            .and_then(|request_header| conn.request_headers().get_str(request_header.clone()))
257            .map(|request_header| Id(request_header.to_string()))
258            .unwrap_or_else(|| self.generate_id());
259
260        if let Some(ref response_header) = self.response_header {
261            conn.response_headers_mut()
262                .insert(response_header.clone(), id.to_string());
263        }
264
265        conn.with_state(id)
266    }
267}
268
269/// Extension trait to retrieve an id generated by the [`ConnId`] handler
270pub trait ConnIdExt {
271    /// Retrieves the id for this conn. This method will panic if it
272    /// is run before the [`ConnId`] handler.
273    fn id(&self) -> &str;
274}
275
276impl<ConnLike> ConnIdExt for ConnLike
277where
278    ConnLike: AsRef<StateSet>,
279{
280    fn id(&self) -> &str {
281        self.as_ref()
282            .get::<Id>()
283            .expect("ConnId handler must be run before calling IdConnExt::id")
284    }
285}
286
287/// Formatter for the trillium_log crate
288pub mod log_formatter {
289    use std::borrow::Cow;
290
291    use super::*;
292    /// Formatter for the trillium_log crate. This will be `-` if
293    /// there is no id on the conn.
294    pub fn conn_id(conn: &Conn, _color: bool) -> Cow<'static, str> {
295        conn.state::<Id>()
296            .map(|id| Cow::Owned(id.0.clone()))
297            .unwrap_or_else(|| Cow::Borrowed("-"))
298    }
299}
300
301/// Alias for ConnId::new()
302pub fn conn_id() -> ConnId {
303    ConnId::new()
304}
305
306impl Debug for IdGenerator {
307    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
308        f.write_str(match self {
309            IdGenerator::Default => "IdGenerator::Default",
310            IdGenerator::SeededFastrand(_) => "IdGenerator::SeededFastrand",
311            IdGenerator::Fn(_) => "IdGenerator::Fn",
312        })
313    }
314}
trillium_conn_id/lib.rs

trillium_conn_id/
lib.rs
