use core::fmt;
use std::collections::BTreeMap;
use std::io::Write as _;
use std::path::Path;
#[cfg(any(unix, windows))]
use std::path::PathBuf;
use serde::{Deserialize, Serialize};
pub const KV_VERSION: u32 = 1;
pub const MAX_KEY_BYTES: usize = 128;
pub const MAX_VALUE_BYTES: usize = 4096;
#[cfg(unix)]
const KV_FILE_MODE: u32 = 0o600;
#[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()))
}
}
#[cfg(any(unix, windows))]
fn lock_path(path: &Path) -> PathBuf {
let mut name = path
.file_name()
.map(std::ffi::OsStr::to_os_string)
.unwrap_or_default();
name.push(".lock");
path.parent().unwrap_or_else(|| Path::new(".")).join(name)
}
struct KvLock {
#[cfg(unix)]
_flock: nix::fcntl::Flock<std::fs::File>,
#[cfg(windows)]
_handle: std::fs::File,
}
impl KvLock {
#[cfg(unix)]
fn acquire(path: &Path) -> std::io::Result<Self> {
use nix::fcntl::{Flock, FlockArg};
use std::os::unix::fs::OpenOptionsExt as _;
let file = std::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(false)
.mode(KV_FILE_MODE)
.open(lock_path(path))?;
Flock::lock(file, FlockArg::LockExclusive)
.map(|flock| Self { _flock: flock })
.map_err(|(_file, errno)| std::io::Error::from(errno))
}
#[cfg(windows)]
fn acquire(path: &Path) -> std::io::Result<Self> {
use std::os::windows::fs::OpenOptionsExt as _;
const ERROR_SHARING_VIOLATION: i32 = 32;
const RETRY_INTERVAL: std::time::Duration = std::time::Duration::from_millis(2);
let lock_path = lock_path(path);
loop {
match std::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(false)
.share_mode(0)
.open(&lock_path)
{
Ok(handle) => return Ok(Self { _handle: handle }),
Err(error) if error.raw_os_error() == Some(ERROR_SHARING_VIOLATION) => {
std::thread::sleep(RETRY_INTERVAL);
}
Err(error) => return Err(error),
}
}
}
}
fn create_kv_file(parent: &Path) -> std::io::Result<tempfile::NamedTempFile> {
let mut builder = tempfile::Builder::new();
builder.prefix("kv").suffix(".tmp");
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt as _;
builder.permissions(std::fs::Permissions::from_mode(KV_FILE_MODE));
}
builder.tempfile_in(parent)
}
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> {
let parent = path.parent().unwrap_or_else(|| Path::new("."));
let mut tmp = create_kv_file(parent)?;
let json = serde_json::to_string_pretty(file)?;
tmp.write_all(json.as_bytes())?;
tmp.write_all(b"\n")?;
tmp.as_file().sync_all()?;
tmp.persist(path).map_err(|err| KvError::Io(err.error))?;
crate::atomic_file::sync_dir(parent)?;
Ok(())
}
pub fn all(path: &Path) -> Result<BTreeMap<String, String>, KvError> {
let _lock = KvLock::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 = KvLock::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 = KvLock::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 = KvLock::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);
}
}