cookie_b/
lib.rs

1#![feature(let_chains)]
2
3use std::{
4  future::Future,
5  pin::Pin,
6  task::{Context, Poll},
7};
8
9use header_host::header_host;
10use tower::{Layer, Service};
11use rand::{RngCore, SeedableRng, rngs::StdRng};
12use axum::{
13  body::Body,
14  http::{Request, Response},
15};
16
17#[derive(Clone)]
18pub struct BrowserIdLayer;
19
20fn gen_browser_bin() -> [u8; 16] {
21  let mut rng = StdRng::from_rng(&mut rand::rng());
22  let mut bytes = [0u8; 16];
23  rng.fill_bytes(&mut bytes);
24  bytes
25}
26
27impl<S> Layer<S> for BrowserIdLayer {
28  type Service = BrowserIdService<S>;
29
30  fn layer(&self, service: S) -> Self::Service {
31    BrowserIdService { inner: service }
32  }
33}
34
35#[derive(Clone)]
36pub struct BrowserIdService<S> {
37  inner: S,
38}
39
40#[derive(Clone)]
41pub struct Browser {
42  pub bin: [u8; 16],
43  pub renew: bool,
44}
45
46const COOKIE_REFRESH: &str = "r=";
47
48impl<S, B> Service<Request<B>> for BrowserIdService<S>
49where
50  S: Service<Request<B>, Response = Response<Body>> + Send + 'static,
51  S::Future: Send + 'static,
52  B: Send + 'static,
53{
54  type Response = S::Response;
55  type Error = S::Error;
56  type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
57
58  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
59    self.inner.poll_ready(cx)
60  }
61
62  fn call(&mut self, mut request: Request<B>) -> Self::Future {
63    let headers = request.headers();
64    let cookie_li = headers
65      .get("cookie")
66      .and_then(|cookie| cookie.to_str().ok())
67      .unwrap_or("")
68      .split(";");
69
70    let mut no_refresh = true;
71    let mut cookie_browser_bin = None;
72
73    for i in cookie_li {
74      let i = i.trim_start();
75      if let Some(b) = i.strip_prefix("b=") {
76        cookie_browser_bin = Some(b.to_owned());
77      } else if i == COOKIE_REFRESH {
78        no_refresh = false;
79        if cookie_browser_bin.is_some() {
80          break;
81        }
82      }
83    }
84
85    let no_browser_bin = cookie_browser_bin.is_none();
86    let browser_bin;
87
88    #[allow(clippy::never_loop)]
89    loop {
90      if let Some(cookie_browser_bin) = cookie_browser_bin
91        && let Ok(id) = ub64::b64d(&cookie_browser_bin)
92        && let Ok::<[u8; 16], _>(id) = id.try_into()
93      {
94        browser_bin = id;
95        break;
96      }
97      browser_bin = gen_browser_bin();
98      break;
99    }
100
101    let host = if (no_browser_bin || no_refresh)
102      && let Ok(host) = xerr::ok!(header_host(headers))
103    {
104      Some(host.to_owned())
105    } else {
106      None
107    };
108
109    request.extensions_mut().insert(Browser {
110      bin: browser_bin,
111      renew: no_refresh && !no_browser_bin,
112    });
113
114    let future = self.inner.call(request);
115
116    if let Some(host) = host {
117      Box::pin(async move {
118        let mut response = future.await?;
119        let mut cookie = cookie_set::new(xtld::host_tld(host), response.headers_mut());
120        // r如果没了就自动给b续期，防止b过期消失(chrome的cookie最长有效期400天)
121        cookie.set_max("b", ub64::b64e(browser_bin));
122        cookie.set("r", "", 999999);
123        Ok(response)
124      })
125    } else {
126      Box::pin(future)
127    }
128  }
129}
cookie_b/lib.rs

cookie_b/
lib.rs
