235 lines
7.8 KiB
Rust
235 lines
7.8 KiB
Rust
//! Middleware plugs for the bench.
|
|
//!
|
|
//! Per the spec §3:
|
|
//! - Logger writes to a bounded ring (no stdout — that would dominate).
|
|
//! - Request ID is generated from a thread-local fast RNG (SplitMix64).
|
|
//! - Auth is constant-time bearer-token compare.
|
|
//!
|
|
//! Each plug is a `Fn(Conn, Next) -> Conn` closure; urus's blanket impl
|
|
//! turns them into Plug values when handed to Pipeline::plug.
|
|
|
|
use std::cell::Cell;
|
|
use std::sync::Arc;
|
|
use std::time::Instant;
|
|
|
|
use crossbeam_queue::ArrayQueue;
|
|
use subtle::ConstantTimeEq;
|
|
use urus::{Conn, Next};
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Log ring
|
|
// ---------------------------------------------------------------------------
|
|
//
|
|
// Bounded MPSC over a fixed-capacity array. We push from many connection
|
|
// actors and never read during the bench (drain on shutdown only). If the
|
|
// ring fills, new entries are dropped — the bench is the priority, not
|
|
// the log. This matches what we said in spec §1.5: warm-up + 60s measured
|
|
// at high RPS would overflow any unbounded buffer.
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct LogEntry {
|
|
pub method: String,
|
|
pub path: String,
|
|
pub status: u16,
|
|
pub elapsed_us: u64,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct LogRing {
|
|
inner: Arc<ArrayQueue<LogEntry>>,
|
|
}
|
|
|
|
impl LogRing {
|
|
pub fn with_capacity(n: usize) -> Self {
|
|
Self { inner: Arc::new(ArrayQueue::new(n)) }
|
|
}
|
|
|
|
/// Push an entry. Drops the entry on overflow (returns Err in
|
|
/// crossbeam's API — we discard).
|
|
pub fn push(&self, e: LogEntry) {
|
|
let _ = self.inner.push(e);
|
|
}
|
|
|
|
/// For inspection after a run.
|
|
#[allow(dead_code)]
|
|
pub fn drain(&self) -> Vec<LogEntry> {
|
|
let mut out = Vec::with_capacity(self.inner.len());
|
|
while let Some(e) = self.inner.pop() { out.push(e); }
|
|
out
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Logger plug
|
|
// ---------------------------------------------------------------------------
|
|
//
|
|
// Wraps `next`. Captures method/path *before* the inner plugs run (the
|
|
// router may move them), then captures status + elapsed after.
|
|
|
|
pub fn logger_plug(conn: Conn, next: Next, ring: &LogRing) -> Conn {
|
|
let start = Instant::now();
|
|
let method = conn.method.as_str().to_string();
|
|
let path = conn.path.clone();
|
|
let conn = next.run(conn);
|
|
let elapsed_us = start.elapsed().as_micros() as u64;
|
|
ring.push(LogEntry {
|
|
method,
|
|
path,
|
|
status: conn.status.unwrap_or(0),
|
|
elapsed_us,
|
|
});
|
|
conn
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Request ID — SplitMix64 in a thread-local, base32-encoded.
|
|
// ---------------------------------------------------------------------------
|
|
//
|
|
// We avoid global RNG / OS getrandom on the hot path. Each scheduler OS
|
|
// thread gets its own seed (Cell + thread_local) and advances it with the
|
|
// SplitMix64 step. 12 bytes of output is encoded with the RFC 4648 base32
|
|
// alphabet (no padding, no lowercase — the header is opaque to the load
|
|
// generator).
|
|
//
|
|
// Crockford-style alphabet would be friendlier to humans but base32-hex
|
|
// is fine here.
|
|
|
|
thread_local! {
|
|
static RNG_STATE: Cell<u64> = const { Cell::new(0) };
|
|
}
|
|
|
|
fn splitmix64(s: u64) -> (u64, u64) {
|
|
// Standard SplitMix64. One 64-bit step gives a uniformly distributed
|
|
// output and the next state.
|
|
let new_state = s.wrapping_add(0x9E3779B97F4A7C15);
|
|
let mut z = new_state;
|
|
z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
|
|
z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
|
|
(new_state, z ^ (z >> 31))
|
|
}
|
|
|
|
fn next_random_bytes(out: &mut [u8; 12]) {
|
|
RNG_STATE.with(|cell| {
|
|
let mut s = cell.get();
|
|
if s == 0 {
|
|
// Lazy first-time seed. Mixing pointer + thread id + clock
|
|
// is sufficient for a bench-local RNG; this isn't crypto.
|
|
let p = cell as *const _ as usize as u64;
|
|
let t = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.map(|d| d.as_nanos() as u64)
|
|
.unwrap_or(0);
|
|
s = p ^ t ^ 0xA5A5_A5A5_A5A5_A5A5;
|
|
if s == 0 { s = 1; }
|
|
}
|
|
let (s1, r1) = splitmix64(s);
|
|
let (s2, r2) = splitmix64(s1);
|
|
cell.set(s2);
|
|
out[0..8].copy_from_slice(&r1.to_le_bytes());
|
|
out[8..12].copy_from_slice(&r2.to_le_bytes()[..4]);
|
|
});
|
|
}
|
|
|
|
// Base32 (RFC 4648) without padding. 12 bytes = 96 bits -> ceil(96/5) = 20
|
|
// output chars exactly. No padding needed.
|
|
fn b32_encode(input: &[u8; 12]) -> String {
|
|
static ALPHABET: &[u8; 32] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
|
|
let mut out = String::with_capacity(20);
|
|
let mut buf: u64 = 0;
|
|
let mut bits: u32 = 0;
|
|
for &b in input {
|
|
buf = (buf << 8) | (b as u64);
|
|
bits += 8;
|
|
while bits >= 5 {
|
|
bits -= 5;
|
|
let idx = ((buf >> bits) & 0x1F) as usize;
|
|
out.push(ALPHABET[idx] as char);
|
|
}
|
|
}
|
|
if bits > 0 {
|
|
let idx = ((buf << (5 - bits)) & 0x1F) as usize;
|
|
out.push(ALPHABET[idx] as char);
|
|
}
|
|
out
|
|
}
|
|
|
|
pub fn request_id_plug(conn: Conn, next: Next) -> Conn {
|
|
let mut raw = [0u8; 12];
|
|
next_random_bytes(&mut raw);
|
|
let id = b32_encode(&raw);
|
|
// Set on the response. We could also stash it in `conn.assigns` for
|
|
// handlers to read, but the bench doesn't need that.
|
|
let conn = next.run(conn);
|
|
conn.put_header("x-request-id", id)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Auth plug
|
|
// ---------------------------------------------------------------------------
|
|
//
|
|
// Read `Authorization: Bearer <token>`; constant-time compare against the
|
|
// configured token. Mismatch -> 401 + halt.
|
|
//
|
|
// The token is borrowed (&[u8]) so the closure that wraps this can hold
|
|
// it via Arc<[u8]> and pass a slice in. No alloc per request.
|
|
|
|
pub fn auth_plug(conn: Conn, next: Next, expected: &[u8]) -> Conn {
|
|
// /ping is the S1 "naked" endpoint per the bench spec — it must not
|
|
// pay the auth cost. We exempt it by path. This is a deliberate
|
|
// exception, not a generic feature; the rest of the pipeline (logger,
|
|
// request_id) still applies to /ping, which is what the spec asks for
|
|
// (S1 measures the framework floor, not "no middleware whatsoever").
|
|
if conn.path == "/ping" {
|
|
return next.run(conn);
|
|
}
|
|
|
|
let header = conn.headers.get("authorization");
|
|
let presented: Option<&[u8]> = header.and_then(|h| {
|
|
// strip the "Bearer " prefix (case-insensitive per RFC 6750, but
|
|
// wrk2 sends the canonical form so an exact prefix check is fine)
|
|
h.strip_prefix("Bearer ").map(|s| s.as_bytes())
|
|
});
|
|
|
|
let ok = match presented {
|
|
Some(p) if p.len() == expected.len() =>
|
|
bool::from(p.ct_eq(expected)),
|
|
_ => false,
|
|
};
|
|
|
|
if ok {
|
|
next.run(conn)
|
|
} else {
|
|
conn.put_status(401)
|
|
.put_header("content-type", "application/json")
|
|
.put_body(&b"{\"error\":\"unauthorized\"}"[..])
|
|
.halt()
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn b32_round_trip_lengths() {
|
|
let mut raw = [0u8; 12];
|
|
for i in 0..12u8 { raw[i as usize] = i; }
|
|
let s = b32_encode(&raw);
|
|
assert_eq!(s.len(), 20);
|
|
assert!(s.chars().all(|c| c.is_ascii_uppercase() || c.is_ascii_digit()));
|
|
}
|
|
|
|
#[test]
|
|
fn rng_advances() {
|
|
let mut a = [0u8; 12];
|
|
let mut b = [0u8; 12];
|
|
next_random_bytes(&mut a);
|
|
next_random_bytes(&mut b);
|
|
assert_ne!(a, b, "consecutive calls must produce different bytes");
|
|
}
|
|
}
|