use crate::crypto::WrappedOrgKey;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fmt;
use std::fs::{self, File, OpenOptions};
use std::io::Write;
use std::path::{Path, PathBuf};
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct MemberPublicKey {
pub account_id: String,
pub public_hex: String,
}
#[derive(Debug)]
pub enum OrgKeyDirectoryError {
Io(std::io::Error),
Unauthorized(String),
}
impl fmt::Display for OrgKeyDirectoryError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
OrgKeyDirectoryError::Io(e) => write!(f, "org-key directory io error: {e}"),
OrgKeyDirectoryError::Unauthorized(m) => {
write!(f, "org-key directory unauthorized: {m}")
}
}
}
}
impl std::error::Error for OrgKeyDirectoryError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
OrgKeyDirectoryError::Io(e) => Some(e),
OrgKeyDirectoryError::Unauthorized(_) => None,
}
}
}
pub trait OrgKeyDirectory {
fn publish_wrapped(&mut self, wrapped: &WrappedOrgKey) -> Result<(), OrgKeyDirectoryError>;
fn fetch_wrapped(
&self,
epoch: u64,
recipient_user_id: &str,
) -> Result<Option<WrappedOrgKey>, OrgKeyDirectoryError>;
fn fetch_wrapped_for(
&self,
recipient_user_id: &str,
) -> Result<Vec<WrappedOrgKey>, OrgKeyDirectoryError>;
fn publish_pubkey(
&mut self,
account_id: &str,
public_hex: &str,
) -> Result<(), OrgKeyDirectoryError>;
fn fetch_pubkeys(&self) -> Result<Vec<MemberPublicKey>, OrgKeyDirectoryError>;
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct OrgKeyDirectoryState {
#[serde(default)]
wrapped: BTreeMap<String, BTreeMap<u64, WrappedOrgKey>>,
#[serde(default)]
pubkeys: BTreeMap<String, String>,
}
impl OrgKeyDirectoryState {
fn publish_wrapped(&mut self, wrapped: &WrappedOrgKey) {
self.wrapped
.entry(wrapped.recipient.clone())
.or_default()
.insert(wrapped.epoch, wrapped.clone());
}
fn fetch_wrapped(&self, epoch: u64, recipient_user_id: &str) -> Option<WrappedOrgKey> {
self.wrapped
.get(recipient_user_id)
.and_then(|by_epoch| by_epoch.get(&epoch))
.cloned()
}
fn fetch_wrapped_for(&self, recipient_user_id: &str) -> Vec<WrappedOrgKey> {
self.wrapped
.get(recipient_user_id)
.map(|by_epoch| by_epoch.values().rev().cloned().collect())
.unwrap_or_default()
}
fn publish_pubkey(&mut self, account_id: &str, public_hex: &str) {
self.pubkeys
.insert(account_id.to_string(), public_hex.to_string());
}
fn fetch_pubkeys(&self) -> Vec<MemberPublicKey> {
self.pubkeys
.iter()
.map(|(account_id, public_hex)| MemberPublicKey {
account_id: account_id.clone(),
public_hex: public_hex.clone(),
})
.collect()
}
}
#[derive(Clone, Debug, Default)]
pub struct InMemoryOrgKeyDirectory {
state: OrgKeyDirectoryState,
}
impl InMemoryOrgKeyDirectory {
pub fn new() -> Self {
Self::default()
}
}
impl OrgKeyDirectory for InMemoryOrgKeyDirectory {
fn publish_wrapped(&mut self, wrapped: &WrappedOrgKey) -> Result<(), OrgKeyDirectoryError> {
self.state.publish_wrapped(wrapped);
Ok(())
}
fn fetch_wrapped(
&self,
epoch: u64,
recipient_user_id: &str,
) -> Result<Option<WrappedOrgKey>, OrgKeyDirectoryError> {
Ok(self.state.fetch_wrapped(epoch, recipient_user_id))
}
fn fetch_wrapped_for(
&self,
recipient_user_id: &str,
) -> Result<Vec<WrappedOrgKey>, OrgKeyDirectoryError> {
Ok(self.state.fetch_wrapped_for(recipient_user_id))
}
fn publish_pubkey(
&mut self,
account_id: &str,
public_hex: &str,
) -> Result<(), OrgKeyDirectoryError> {
self.state.publish_pubkey(account_id, public_hex);
Ok(())
}
fn fetch_pubkeys(&self) -> Result<Vec<MemberPublicKey>, OrgKeyDirectoryError> {
Ok(self.state.fetch_pubkeys())
}
}
#[derive(Debug)]
pub struct FsOrgKeyDirectory {
dir: PathBuf,
}
impl FsOrgKeyDirectory {
pub fn open(dir: &Path) -> std::io::Result<Self> {
fs::create_dir_all(dir)?;
Ok(Self {
dir: dir.to_path_buf(),
})
}
fn state_path(&self) -> PathBuf {
self.dir.join("org-keys.json")
}
fn lock_path(&self) -> PathBuf {
self.dir.join("org-keys.lock")
}
fn read_state(path: &Path) -> Result<OrgKeyDirectoryState, OrgKeyDirectoryError> {
match fs::read_to_string(path) {
Ok(raw) => serde_json::from_str(&raw)
.map_err(|e| OrgKeyDirectoryError::Io(std::io::Error::other(e))),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
Ok(OrgKeyDirectoryState::default())
}
Err(e) => Err(OrgKeyDirectoryError::Io(e)),
}
}
fn load(&self) -> Result<OrgKeyDirectoryState, OrgKeyDirectoryError> {
let lock = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(self.lock_path())
.map_err(OrgKeyDirectoryError::Io)?;
lock.lock_shared().map_err(OrgKeyDirectoryError::Io)?; Self::read_state(&self.state_path())
}
fn with_state<T>(
&mut self,
f: impl FnOnce(&mut OrgKeyDirectoryState) -> T,
) -> Result<T, OrgKeyDirectoryError> {
let lock = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(self.lock_path())
.map_err(OrgKeyDirectoryError::Io)?;
lock.lock().map_err(OrgKeyDirectoryError::Io)?;
let mut state = Self::read_state(&self.state_path())?;
let result = f(&mut state);
let tmp = self.dir.join("org-keys.json.tmp");
{
let mut file = File::create(&tmp).map_err(OrgKeyDirectoryError::Io)?;
file.write_all(
serde_json::to_string(&state)
.map_err(|e| OrgKeyDirectoryError::Io(std::io::Error::other(e)))?
.as_bytes(),
)
.map_err(OrgKeyDirectoryError::Io)?;
file.sync_all().map_err(OrgKeyDirectoryError::Io)?;
}
fs::rename(&tmp, self.state_path()).map_err(OrgKeyDirectoryError::Io)?;
Ok(result)
}
}
impl OrgKeyDirectory for FsOrgKeyDirectory {
fn publish_wrapped(&mut self, wrapped: &WrappedOrgKey) -> Result<(), OrgKeyDirectoryError> {
self.with_state(|state| state.publish_wrapped(wrapped))
}
fn fetch_wrapped(
&self,
epoch: u64,
recipient_user_id: &str,
) -> Result<Option<WrappedOrgKey>, OrgKeyDirectoryError> {
Ok(self.load()?.fetch_wrapped(epoch, recipient_user_id))
}
fn fetch_wrapped_for(
&self,
recipient_user_id: &str,
) -> Result<Vec<WrappedOrgKey>, OrgKeyDirectoryError> {
Ok(self.load()?.fetch_wrapped_for(recipient_user_id))
}
fn publish_pubkey(
&mut self,
account_id: &str,
public_hex: &str,
) -> Result<(), OrgKeyDirectoryError> {
self.with_state(|state| state.publish_pubkey(account_id, public_hex))
}
fn fetch_pubkeys(&self) -> Result<Vec<MemberPublicKey>, OrgKeyDirectoryError> {
Ok(self.load()?.fetch_pubkeys())
}
}
#[cfg(test)]
pub(crate) mod conformance {
use super::*;
use crate::crypto::{
derive_ed25519_identity, derive_x25519_identity, wrap_org_key, x25519_public,
StretchedMaster,
};
pub(crate) fn make_wrap(org: &str, epoch: u64, recipient: &str) -> WrappedOrgKey {
let k_org = [7u8; 32];
let sk = derive_x25519_identity(
&StretchedMaster::from_issued_high_entropy(b"login-secret-for-tests", recipient),
recipient,
);
let recipient_pub = x25519_public(&sk);
let signer = derive_ed25519_identity(
&StretchedMaster::from_issued_high_entropy(b"granter-login-for-tests", "acc_granter"),
"acc_granter",
);
wrap_org_key(
&k_org,
org,
epoch,
recipient,
&recipient_pub,
"acc_granter",
&signer,
)
.expect("wrap succeeds")
}
pub(crate) fn round_trip_suite(dir: &mut dyn OrgKeyDirectory) {
let w = make_wrap("acme", 1, "alice");
dir.publish_wrapped(&w).unwrap();
assert_eq!(dir.fetch_wrapped(1, "alice").unwrap().as_ref(), Some(&w));
let got = dir.fetch_wrapped(1, "alice").unwrap().unwrap();
assert_eq!(got.envelope, w.envelope);
assert_eq!(got.ephemeral_pub, w.ephemeral_pub);
assert_eq!(got.car_wrap, w.car_wrap);
assert!(dir.fetch_wrapped(2, "alice").unwrap().is_none());
assert!(dir.fetch_wrapped(1, "bob").unwrap().is_none());
let w2 = make_wrap("acme", 5, "alice");
let w3 = make_wrap("acme", 3, "alice");
dir.publish_wrapped(&w2).unwrap();
dir.publish_wrapped(&w3).unwrap();
let epochs: Vec<u64> = dir
.fetch_wrapped_for("alice")
.unwrap()
.iter()
.map(|w| w.epoch)
.collect();
assert_eq!(epochs, vec![5, 3, 1], "newest epoch first");
assert!(dir.fetch_wrapped_for("nobody").unwrap().is_empty());
dir.publish_pubkey("bob", "beef").unwrap();
dir.publish_pubkey("alice", "cafe").unwrap();
let keys = dir.fetch_pubkeys().unwrap();
assert_eq!(
keys,
vec![
MemberPublicKey {
account_id: "alice".into(),
public_hex: "cafe".into()
},
MemberPublicKey {
account_id: "bob".into(),
public_hex: "beef".into()
},
],
"ordered by account_id"
);
}
}
#[cfg(test)]
mod tests {
use super::conformance::{make_wrap, round_trip_suite};
use super::*;
use serde_json::json;
#[test]
fn in_memory_round_trip() {
let mut dir = InMemoryOrgKeyDirectory::new();
round_trip_suite(&mut dir);
}
#[test]
fn fs_round_trip() {
let tmp = tempfile::tempdir().unwrap();
let mut dir = FsOrgKeyDirectory::open(tmp.path()).unwrap();
round_trip_suite(&mut dir);
}
#[test]
fn publish_wrapped_is_last_write_wins() {
let mut dir = InMemoryOrgKeyDirectory::new();
let mut w = make_wrap("acme", 1, "alice");
dir.publish_wrapped(&w).unwrap();
w.envelope = json!({"car_enc": "org-key-wrap/v1", "nonce": "00", "ct": "ff"});
dir.publish_wrapped(&w).unwrap();
assert_eq!(dir.fetch_wrapped_for("alice").unwrap().len(), 1);
assert_eq!(
dir.fetch_wrapped(1, "alice").unwrap().unwrap().envelope,
w.envelope
);
}
#[test]
fn publish_pubkey_is_last_write_wins() {
let mut dir = InMemoryOrgKeyDirectory::new();
dir.publish_pubkey("alice", "aaaa").unwrap();
dir.publish_pubkey("alice", "bbbb").unwrap();
let keys = dir.fetch_pubkeys().unwrap();
assert_eq!(keys.len(), 1);
assert_eq!(keys[0].public_hex, "bbbb");
}
#[test]
fn fs_persists_across_handles() {
let tmp = tempfile::tempdir().unwrap();
let w = make_wrap("acme", 2, "alice");
{
let mut dir = FsOrgKeyDirectory::open(tmp.path()).unwrap();
dir.publish_wrapped(&w).unwrap();
dir.publish_pubkey("alice", "cafe").unwrap();
}
let dir = FsOrgKeyDirectory::open(tmp.path()).unwrap();
assert_eq!(dir.fetch_wrapped(2, "alice").unwrap().as_ref(), Some(&w));
assert_eq!(dir.fetch_pubkeys().unwrap()[0].public_hex, "cafe");
}
#[test]
fn state_loads_from_files_missing_fields() {
let empty: OrgKeyDirectoryState = serde_json::from_str("{}").unwrap();
assert!(empty.fetch_pubkeys().is_empty());
assert!(empty.fetch_wrapped_for("alice").is_empty());
let partial: OrgKeyDirectoryState =
serde_json::from_str(r#"{"pubkeys":{"alice":"ff"}}"#).unwrap();
assert_eq!(partial.fetch_pubkeys().len(), 1);
assert!(partial.fetch_wrapped_for("alice").is_empty());
}
#[test]
fn corrupt_state_file_surfaces_as_error() {
let tmp = tempfile::tempdir().unwrap();
let mut dir = FsOrgKeyDirectory::open(tmp.path()).unwrap();
dir.publish_pubkey("alice", "cafe").unwrap(); std::fs::write(tmp.path().join("org-keys.json"), b"{ not json").unwrap();
assert!(matches!(
dir.fetch_pubkeys(),
Err(OrgKeyDirectoryError::Io(_))
));
assert!(matches!(
dir.publish_pubkey("bob", "beef"),
Err(OrgKeyDirectoryError::Io(_))
));
}
}