spacegate_kernel/utils/
x_request_id.rs

1use std::{
2    hash::Hasher,
3    sync::{atomic::AtomicU64, OnceLock},
4};
5
6use hyper::{http::HeaderValue, Request, Response};
7
8use crate::{helper_layers::function::Inner, SgBody};
9
10/// Generate a `x-request-id` header for the request and response.
11pub trait XRequestIdAlgo {
12    fn generate() -> HeaderValue;
13}
14/// The header name for `x-request-id`.
15pub const X_REQUEST_ID_HEADER_NAME: &str = "x-request-id";
16/// Add a `x-request-id` header to the request and then response.
17///
18/// If the request already has a `x-request-id` header, it will be used.
19pub async fn x_request_id<A: XRequestIdAlgo>(mut request: Request<SgBody>, inner: Inner) -> Response<SgBody> {
20    let id = if let Some(id) = request.headers().get(X_REQUEST_ID_HEADER_NAME) {
21        id.clone()
22    } else {
23        let id = A::generate();
24        request.headers_mut().insert(X_REQUEST_ID_HEADER_NAME, id.clone());
25        id
26    };
27    let mut resp = inner.call(request).await;
28    resp.headers_mut().insert(X_REQUEST_ID_HEADER_NAME, id);
29    resp
30}
31/// # Reference
32/// - discord: https://discord.com/developers/docs/reference#snowflakes
33/// - instagram: https://instagram-engineering.com/sharding-ids-at-instagram-1cf5a71e5a5c
34/// # Bits
35/// - 42: timestamp
36/// - 10: machine id
37/// - 12: increment id
38#[derive(Debug, Default, Clone)]
39pub struct Snowflake;
40
41impl XRequestIdAlgo for Snowflake {
42    fn generate() -> HeaderValue {
43        static INC: AtomicU64 = AtomicU64::new(0);
44        let ts_id = unsafe { std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_unchecked().as_millis() as u64 } << 22;
45        let mach_id = machine_id() << 12;
46        let inc = INC.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
47        let id = ts_id | mach_id | inc;
48        unsafe { HeaderValue::from_str(&format!("{:016x}", id)).unwrap_unchecked() }
49    }
50}
51
52fn machine_id() -> u64 {
53    static MACHINE_ID: OnceLock<u64> = OnceLock::new();
54    *MACHINE_ID.get_or_init(|| {
55        let mid = std::env::var("MACHINE_ID");
56        let mut hasher = std::hash::DefaultHasher::new();
57        if let Ok(mid) = mid {
58            if let Ok(mid) = mid.parse::<u64>() {
59                mid
60            } else {
61                hasher.write(mid.as_bytes());
62                hasher.finish()
63            }
64        } else {
65            #[cfg(target_os = "linux")]
66            {
67                // let's try to read system mid
68                let mid = std::fs::read_to_string("/var/lib/dbus/machine-id").unwrap_or(rand::random::<u64>().to_string());
69                hasher.write(mid.as_bytes());
70                hasher.finish()
71            }
72            #[cfg(not(target_os = "linux"))]
73            {
74                // let's generate random one
75                rand::random::<u64>()
76            }
77        }
78    })
79}
spacegate_kernel/utils/x_request_id.rs

spacegate_kernel/utils/
x_request_id.rs
