use std::collections::hash_map::Entry;
use std::collections::HashMap;
use parking_lot::Mutex;
use crate::backend::{BackendKey, Exclusivity, KeychainBackend};
use crate::error::{KeystoreError, Result};
#[derive(Default)]
pub struct MemoryBackend {
inner: Mutex<HashMap<BackendKey, Vec<u8>>>,
}
impl MemoryBackend {
pub fn new() -> Self {
Self::default()
}
}
impl KeychainBackend for MemoryBackend {
fn read(&self, key: &BackendKey) -> Result<Vec<u8>> {
self.inner.lock().get(key).cloned().ok_or_else(|| {
KeystoreError::from(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("key not found: {key}"),
))
})
}
fn write(&self, key: &BackendKey, data: &[u8]) -> Result<()> {
self.inner.lock().insert(key.clone(), data.to_vec());
Ok(())
}
fn write_new(&self, key: &BackendKey, data: &[u8]) -> Result<()> {
match self.inner.lock().entry(key.clone()) {
Entry::Occupied(_) => Err(KeystoreError::AlreadyExists(key.as_str().to_string())),
Entry::Vacant(slot) => {
slot.insert(data.to_vec());
Ok(())
}
}
}
fn write_new_exclusivity(&self) -> Exclusivity {
Exclusivity::Atomic
}
fn delete(&self, key: &BackendKey) -> Result<()> {
self.inner.lock().remove(key);
Ok(())
}
fn list(&self, prefix: &str) -> Result<Vec<BackendKey>> {
Ok(self
.inner
.lock()
.keys()
.filter(|k| k.as_str().starts_with(prefix))
.cloned()
.collect())
}
fn exists(&self, key: &BackendKey) -> Result<bool> {
Ok(self.inner.lock().contains_key(key))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn write_new_establishes_once_and_then_refuses() {
let be = MemoryBackend::new();
let key = BackendKey::new("coupled");
be.write_new(&key, b"established").unwrap();
let err = be.write_new(&key, b"usurper").unwrap_err();
assert!(
matches!(err, KeystoreError::AlreadyExists(ref k) if k == "coupled"),
"the collision must be adoptable: {err:?}"
);
assert_eq!(be.read(&key).unwrap(), b"established");
be.write(&key, b"replaced").unwrap();
assert_eq!(be.read(&key).unwrap(), b"replaced");
}
#[test]
fn only_one_concurrent_write_new_can_win() {
use std::sync::{Arc, Barrier};
const RACERS: usize = 16;
let be = Arc::new(MemoryBackend::new());
let key = BackendKey::new("contended");
let gate = Arc::new(Barrier::new(RACERS));
let winners = std::thread::scope(|scope| {
let handles: Vec<_> = (0..RACERS)
.map(|i| {
let (be, gate, key) = (Arc::clone(&be), Arc::clone(&gate), key.clone());
scope.spawn(move || {
gate.wait();
be.write_new(&key, &[i as u8; 8]).is_ok()
})
})
.collect();
handles
.into_iter()
.map(|h| h.join().unwrap())
.filter(|won| *won)
.count()
});
assert_eq!(winners, 1, "exactly one racer may establish; {winners} did");
}
#[test]
fn memory_backend_claims_exclusive_creation() {
assert_eq!(
MemoryBackend::new().write_new_exclusivity(),
Exclusivity::Atomic
);
}
#[test]
fn roundtrip() {
let be = MemoryBackend::new();
let k = BackendKey::new("x");
be.write(&k, b"data").unwrap();
assert_eq!(be.read(&k).unwrap(), b"data");
assert!(be.exists(&k).unwrap());
be.delete(&k).unwrap();
assert!(!be.exists(&k).unwrap());
}
#[test]
fn read_missing_key_is_not_found() {
let be = MemoryBackend::new();
let err = be.read(&BackendKey::new("absent")).unwrap_err();
match err {
KeystoreError::Backend(io) => {
assert_eq!(io.kind(), std::io::ErrorKind::NotFound);
}
other => panic!("expected Backend(NotFound), got {other:?}"),
}
}
#[test]
fn write_overwrites_in_place() {
let be = MemoryBackend::new();
let k = BackendKey::new("k");
be.write(&k, b"first").unwrap();
be.write(&k, b"second").unwrap();
assert_eq!(be.read(&k).unwrap(), b"second");
}
#[test]
fn list_filters_by_prefix() {
let be = MemoryBackend::new();
be.write(&BackendKey::new("validator/a"), b"1").unwrap();
be.write(&BackendKey::new("validator/b"), b"2").unwrap();
be.write(&BackendKey::new("wallet/c"), b"3").unwrap();
let mut matched: Vec<String> = be
.list("validator/")
.unwrap()
.into_iter()
.map(|k| k.as_str().to_string())
.collect();
matched.sort();
assert_eq!(matched, vec!["validator/a", "validator/b"]);
assert_eq!(be.list("").unwrap().len(), 3);
assert!(be.list("none/").unwrap().is_empty());
}
#[test]
fn default_is_empty() {
let be = MemoryBackend::default();
assert!(be.list("").unwrap().is_empty());
assert!(!be.exists(&BackendKey::new("anything")).unwrap());
}
}