use core::fmt;
use std::collections::BTreeMap;
use std::path::Path;
use serde::{Deserialize, Serialize};
use crate::file_lock::FileLock;
pub const KV_VERSION: u32 = 1;
pub const MAX_KEY_BYTES: usize = 128;
pub const MAX_VALUE_BYTES: usize = 4096;
#[derive(Debug, Default, Serialize, Deserialize)]
struct KvFile {
version: u32,
entries: BTreeMap<String, String>,
}
#[non_exhaustive]
#[derive(Debug)]
pub enum KvError {
Io(std::io::Error),
Decode(serde_json::Error),
InvalidKey(String),
ValueTooLong {
key: String,
len: usize,
},
FutureVersion(u32),
}
impl fmt::Display for KvError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Io(err) => write!(f, "kv store I/O failed: {err}"),
Self::Decode(err) => write!(f, "kv store failed to parse: {err}"),
Self::InvalidKey(key) => write!(f, "`{key}` is not a valid kv key"),
Self::ValueTooLong { key, len } => write!(
f,
"value for `{key}` is {len} bytes, over the {MAX_VALUE_BYTES}-byte limit"
),
Self::FutureVersion(version) => {
write!(
f,
"kv store is version {version}, newer than this build understands"
)
}
}
}
}
impl core::error::Error for KvError {
fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
match self {
Self::Io(err) => Some(err),
Self::Decode(err) => Some(err),
Self::InvalidKey(_) | Self::ValueTooLong { .. } | Self::FutureVersion(_) => None,
}
}
}
impl From<std::io::Error> for KvError {
fn from(source: std::io::Error) -> Self {
Self::Io(source)
}
}
impl From<serde_json::Error> for KvError {
fn from(source: serde_json::Error) -> Self {
Self::Decode(source)
}
}
fn check_key(key: &str) -> Result<(), KvError> {
let ok = !key.is_empty()
&& key.len() <= MAX_KEY_BYTES
&& !key.starts_with('.')
&& key
.bytes()
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-'));
if ok {
Ok(())
} else {
Err(KvError::InvalidKey(key.to_string()))
}
}
fn read_file(path: &Path) -> Result<KvFile, KvError> {
let raw = match std::fs::read_to_string(path) {
Ok(raw) => raw,
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(KvFile::default()),
Err(err) => return Err(KvError::Io(err)),
};
let file: KvFile = serde_json::from_str(&raw)?;
if file.version > KV_VERSION {
return Err(KvError::FutureVersion(file.version));
}
Ok(file)
}
fn write_file(path: &Path, file: &KvFile) -> Result<(), KvError> {
crate::atomic_file::write_json(path, "kv", file).map_err(KvError::Io)
}
pub fn all(path: &Path) -> Result<BTreeMap<String, String>, KvError> {
let _lock = FileLock::acquire(path)?;
Ok(read_file(path)?.entries)
}
pub fn get(path: &Path, key: &str) -> Result<Option<String>, KvError> {
check_key(key)?;
Ok(all(path)?.remove(key))
}
pub fn set(path: &Path, key: &str, value: &str) -> Result<(), KvError> {
check_key(key)?;
if value.len() > MAX_VALUE_BYTES {
return Err(KvError::ValueTooLong {
key: key.to_string(),
len: value.len(),
});
}
let _lock = FileLock::acquire(path)?;
let mut file = read_file(path)?;
file.version = KV_VERSION;
file.entries.insert(key.to_string(), value.to_string());
write_file(path, &file)
}
pub fn unset(path: &Path, key: &str) -> Result<bool, KvError> {
check_key(key)?;
let _lock = FileLock::acquire(path)?;
let mut file = read_file(path)?;
let was_present = file.entries.remove(key).is_some();
if was_present {
file.version = KV_VERSION;
write_file(path, &file)?;
}
Ok(was_present)
}
pub fn clear(path: &Path) -> Result<u32, KvError> {
let _lock = FileLock::acquire(path)?;
let file = read_file(path)?;
let count = u32::try_from(file.entries.len()).unwrap_or(u32::MAX);
if count > 0 {
write_file(
path,
&KvFile {
version: KV_VERSION,
entries: BTreeMap::new(),
},
)?;
}
Ok(count)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_value_survives_a_write_and_a_read() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("kv.json");
set(&path, "bark.cooldown", "30s").unwrap();
assert_eq!(
get(&path, "bark.cooldown").unwrap(),
Some("30s".to_string())
);
}
#[test]
fn a_store_that_does_not_exist_reads_as_empty() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("kv.json");
assert!(all(&path).unwrap().is_empty());
assert_eq!(get(&path, "anything").unwrap(), None);
}
#[test]
fn unset_reports_whether_the_key_was_there() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("kv.json");
set(&path, "a", "1").unwrap();
assert!(unset(&path, "a").unwrap());
assert!(!unset(&path, "a").unwrap());
}
#[test]
fn clear_empties_the_store_and_counts_what_it_took() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("kv.json");
set(&path, "a", "1").unwrap();
set(&path, "b", "2").unwrap();
assert_eq!(clear(&path).unwrap(), 2);
assert!(all(&path).unwrap().is_empty());
assert_eq!(clear(&path).unwrap(), 0);
}
#[test]
fn the_key_grammar_refuses_what_it_says_it_refuses() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("kv.json");
for bad in [
"", " ", "a b", "a\nb", "a/b", "a:b", ".hidden", "a\"b", "$HOME",
] {
assert!(
matches!(set(&path, bad, "1"), Err(KvError::InvalidKey(_))),
"`{bad}` was accepted as a key"
);
}
for good in ["a", "bark.cooldown", "metrics_port", "a-b", "A1.b-c_d"] {
assert!(set(&path, good, "1").is_ok(), "`{good}` was refused");
}
}
#[test]
fn a_dotted_key_is_one_flat_key_and_not_a_path() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("kv.json");
set(&path, "bark.cooldown", "30s").unwrap();
set(&path, "bark.sink", "discord").unwrap();
let stored = all(&path).unwrap();
assert_eq!(stored.len(), 2);
assert!(stored.contains_key("bark.cooldown"));
assert_eq!(get(&path, "bark").unwrap(), None);
let raw = std::fs::read_to_string(&path).unwrap();
assert!(raw.contains(r#""bark.cooldown""#), "{raw}");
}
#[test]
fn an_oversized_value_is_refused_by_name_and_length() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("kv.json");
let big = "x".repeat(MAX_VALUE_BYTES + 1);
let err = set(&path, "a", &big).unwrap_err();
let KvError::ValueTooLong { key, len } = err else {
panic!("expected ValueTooLong, got {err:?}");
};
assert_eq!(key, "a");
assert_eq!(len, MAX_VALUE_BYTES + 1);
}
#[test]
fn a_store_from_a_future_shep_is_refused_rather_than_replaced() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("kv.json");
std::fs::write(&path, r#"{"version":99,"entries":{"a":"1"}}"#).unwrap();
assert!(matches!(all(&path), Err(KvError::FutureVersion(99))));
assert!(matches!(
set(&path, "b", "2"),
Err(KvError::FutureVersion(99))
));
let raw = std::fs::read_to_string(&path).unwrap();
assert!(raw.contains(r#""a":"1""#), "{raw}");
}
#[cfg(unix)]
#[test]
fn the_store_is_owner_only() {
use std::os::unix::fs::PermissionsExt as _;
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("kv.json");
set(&path, "a", "1").unwrap();
let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777;
assert_eq!(mode, 0o600, "{mode:o}");
}
#[test]
fn two_concurrent_writers_lose_nothing() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("kv.json");
const PER_WRITER: usize = 100;
let (done_tx, done_rx) = std::sync::mpsc::channel();
for writer in 0..2 {
let path = path.clone();
let done_tx = done_tx.clone();
std::thread::spawn(move || {
for n in 0..PER_WRITER {
set(&path, &format!("w{writer}.k{n}"), "v").unwrap();
}
done_tx.send(()).unwrap();
});
}
drop(done_tx);
for _ in 0..2 {
done_rx
.recv_timeout(std::time::Duration::from_secs(60))
.expect("a writer did not finish within 60s");
}
assert_eq!(all(&path).unwrap().len(), PER_WRITER * 2);
}
}