mirror of
https://github.com/neondatabase/neon.git
synced 2026-10-19 02:02:53 +00:00
Conrad Ludgate
9cfe08e3d9
## Problem Despite making password hashing async, it can still take time away from the network code. ## Summary of changes Introduce a custom threadpool, inspired by rayon. Features: ### Fairness Each task is tagged with it's endpoint ID. The more times we have seen the endpoint, the more likely we are to skip the task if it comes up in the queue. This is using a min-count-sketch estimator for the number of times we have seen the endpoint, resetting it every 1000+ steps. Since tasks are immediately rescheduled if they do not complete, the worker could get stuck in a "always work available loop". To combat this, we check the global queue every 61 steps to ensure all tasks quickly get a worker assigned to them. ### Balanced Using crossbeam_deque, like rayon does, we have workstealing out of the box. I've tested it a fair amount and it seems to balance the workload accordingly
90 lines
2.2 KiB
Rust
90 lines
2.2 KiB
Rust
use hmac::{
|
|
digest::{consts::U32, generic_array::GenericArray},
|
|
Hmac, Mac,
|
|
};
|
|
use sha2::Sha256;
|
|
|
|
pub struct Pbkdf2 {
|
|
hmac: Hmac<Sha256>,
|
|
prev: GenericArray<u8, U32>,
|
|
hi: GenericArray<u8, U32>,
|
|
iterations: u32,
|
|
}
|
|
|
|
// inspired from <https://github.com/neondatabase/rust-postgres/blob/20031d7a9ee1addeae6e0968e3899ae6bf01cee2/postgres-protocol/src/authentication/sasl.rs#L36-L61>
|
|
impl Pbkdf2 {
|
|
pub fn start(str: &[u8], salt: &[u8], iterations: u32) -> Self {
|
|
let hmac =
|
|
Hmac::<Sha256>::new_from_slice(str).expect("HMAC is able to accept all key sizes");
|
|
|
|
let prev = hmac
|
|
.clone()
|
|
.chain_update(salt)
|
|
.chain_update(1u32.to_be_bytes())
|
|
.finalize()
|
|
.into_bytes();
|
|
|
|
Self {
|
|
hmac,
|
|
// one consumed for the hash above
|
|
iterations: iterations - 1,
|
|
hi: prev,
|
|
prev,
|
|
}
|
|
}
|
|
|
|
pub fn cost(&self) -> u32 {
|
|
(self.iterations).clamp(0, 4096)
|
|
}
|
|
|
|
pub fn turn(&mut self) -> std::task::Poll<[u8; 32]> {
|
|
let Self {
|
|
hmac,
|
|
prev,
|
|
hi,
|
|
iterations,
|
|
} = self;
|
|
|
|
// only do 4096 iterations per turn before sharing the thread for fairness
|
|
let n = (*iterations).clamp(0, 4096);
|
|
for _ in 0..n {
|
|
*prev = hmac.clone().chain_update(*prev).finalize().into_bytes();
|
|
|
|
for (hi, prev) in hi.iter_mut().zip(*prev) {
|
|
*hi ^= prev;
|
|
}
|
|
}
|
|
|
|
*iterations -= n;
|
|
if *iterations == 0 {
|
|
std::task::Poll::Ready((*hi).into())
|
|
} else {
|
|
std::task::Poll::Pending
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::Pbkdf2;
|
|
use pbkdf2::pbkdf2_hmac_array;
|
|
use sha2::Sha256;
|
|
|
|
#[test]
|
|
fn works() {
|
|
let salt = b"sodium chloride";
|
|
let pass = b"Ne0n_!5_50_C007";
|
|
|
|
let mut job = Pbkdf2::start(pass, salt, 600000);
|
|
let hash = loop {
|
|
let std::task::Poll::Ready(hash) = job.turn() else {
|
|
continue;
|
|
};
|
|
break hash;
|
|
};
|
|
|
|
let expected = pbkdf2_hmac_array::<Sha256, 32>(pass, salt, 600000);
|
|
assert_eq!(hash, expected)
|
|
}
|
|
}
|