git.delta.rocks / jrsonnet / refs/commits / 5a5b360a3403

difftreelog

source

crates/fleet-base/src/fleetdata.rs10.2 KiBsourcehistory
1use std::{2	collections::{3		BTreeMap, BTreeSet,4		btree_map::{self, Entry},5	},6	io::{self, Cursor},7	ops::Deref,8};910use age::Recipient;11use chrono::{DateTime, Utc};12use fleet_shared::SecretData;13use rand::{14	distr::{Alphanumeric, SampleString as _},15	rng,16};17use serde::{18	Deserialize, Serialize,19	de::{self, Error},20};21use serde_json::Value;22use tracing::info;2324#[derive(Serialize, Deserialize, Default)]25#[serde(rename_all = "camelCase")]26pub struct HostData {27	#[serde(default)]28	#[serde(skip_serializing_if = "String::is_empty")]29	pub encryption_key: String,30}3132const VERSION: &str = "0.1.0";33pub struct FleetDataVersion;34impl Serialize for FleetDataVersion {35	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>36	where37		S: serde::Serializer,38	{39		VERSION.serialize(serializer)40	}41}42impl<'de> Deserialize<'de> for FleetDataVersion {43	fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>44	where45		D: serde::Deserializer<'de>,46	{47		let version = String::deserialize(deserializer)?;48		if version != VERSION {49			return Err(D::Error::custom(format!(50				"fleet.nix data version mismatch, expected {VERSION}, got {version}.\nFollow the docs for migration instruction"51			)));52		}53		Ok(Self)54	}55}5657fn generate_gc_prefix() -> String {58	let id = Alphanumeric.sample_string(&mut rng(), 8);59	format!("fleet-gc-{id}")60}6162#[derive(Serialize, Deserialize)]63#[serde(rename_all = "camelCase")]64pub struct ManagerKey {65	pub name: String,66	pub key: String,67}6869#[derive(Serialize, Deserialize)]70#[serde(rename_all = "camelCase")]71pub struct FleetData {72	pub version: FleetDataVersion,73	#[serde(default = "generate_gc_prefix")]74	pub gc_root_prefix: String,7576	#[serde(default)]77	pub manager_keys: Vec<ManagerKey>,7879	#[serde(default)]80	pub hosts: BTreeMap<String, HostData>,8182	#[serde(default, alias = "shared_secrets")]83	pub secrets: FleetSecrets,8485	// extra_name => anything86	#[serde(default)]87	#[serde(skip_serializing_if = "BTreeMap::is_empty")]88	pub extra: BTreeMap<String, Value>,8990	#[serde(default)]91	#[serde(skip_serializing_if = "BTreeMap::is_empty")]92	host_secrets: BTreeMap<String, BTreeMap<String, FleetSecretDistribution>>,93}94impl FleetData {95	pub fn from_str(s: &str) -> anyhow::Result<Self> {96		let mut data: Self = nixlike::parse_str(s)?;97		if !data.host_secrets.is_empty() {98			info!("migrating host secrets into shared secrets structure");99			data.secrets100				.merge_from_hosts(std::mem::take(&mut data.host_secrets));101		}102		Ok(data)103	}104}105106/// Returns None if recipients.is_empty()107pub fn encrypt_secret_data<'r>(108	recipients: impl IntoIterator<Item = &'r Box<dyn Recipient>>,109	data: Vec<u8>,110) -> Option<SecretData> {111	let mut encrypted = vec![];112	let mut encryptor = age::Encryptor::with_recipients(recipients.into_iter().map(|v| &**v))113		.ok()?114		.wrap_output(&mut encrypted)115		.expect("in memory write");116	io::copy(&mut Cursor::new(data), &mut encryptor).expect("in memory copy");117	encryptor.finish().expect("in memory flush");118	Some(SecretData {119		data: encrypted,120		encrypted: true,121	})122}123124#[derive(Serialize, Deserialize, Clone)]125pub struct FleetSecretPart {126	pub raw: SecretData,127}128129#[derive(Serialize, Deserialize, Clone)]130#[serde(rename_all = "camelCase")]131#[must_use]132pub struct FleetSecretData {133	#[serde(default = "Utc::now")]134	pub created_at: DateTime<Utc>,135	#[serde(default)]136	#[serde(skip_serializing_if = "Option::is_none", alias = "expire_at")]137	pub expires_at: Option<DateTime<Utc>>,138139	#[serde(flatten)]140	pub parts: BTreeMap<String, FleetSecretPart>,141142	#[serde(default)]143	#[serde(skip_serializing_if = "Value::is_null")]144	pub generation_data: Value,145}146147#[derive(Serialize, Deserialize, Clone)]148#[serde(rename_all = "camelCase")]149#[must_use]150pub struct FleetSecretDistribution {151	#[serde(default)]152	#[serde(skip_serializing_if = "Option::is_none")]153	pub managed: Option<bool>,154	#[serde(default)]155	pub owners: BTreeSet<String>,156	#[serde(flatten)]157	pub secret: FleetSecretData,158}159160#[derive(Clone)]161#[must_use]162pub struct FleetSecretDistributions(Vec<FleetSecretDistribution>);163164impl Deref for FleetSecretDistributions {165	type Target = [FleetSecretDistribution];166167	fn deref(&self) -> &Self::Target {168		self.0.as_slice()169	}170}171172impl FleetSecretDistributions {173	pub fn owners(&self) -> impl Iterator<Item = &String> {174		self.0.iter().flat_map(|v| v.owners.iter())175	}176	#[allow(177		clippy::len_without_is_empty,178		reason = "should not be empty for a long time"179	)]180	pub fn len(&self) -> usize {181		self.0.len()182	}183184	pub fn get(&self, owner: &str) -> Option<&FleetSecretDistribution> {185		self.0.iter().find(|d| d.owners.contains(owner))186	}187	fn entry(&mut self, owner: String) -> DistEntry<'_> {188		let Some(idx) = self.0.iter().position(|d| d.owners.contains(&owner)) else {189			return DistEntry::Vacant(VacantDistEntry {190				distributions: self,191				owner,192			});193		};194		DistEntry::Occupied(OccupiedDistEntry {195			distributions: self,196			idx,197			owner,198		})199	}200	fn extend(&mut self, dist: FleetSecretDistribution) {201		for owner in &dist.owners {202			self.entry(owner.to_owned()).remove();203		}204		self.0.push(dist);205	}206	pub fn contains(&self, owner: &str) -> bool {207		self.0.iter().any(|d| d.owners.contains(owner))208	}209}210211struct OccupiedDistEntry<'d> {212	distributions: &'d mut FleetSecretDistributions,213	idx: usize,214	owner: String,215}216impl<'d> OccupiedDistEntry<'d> {217	fn remove(self) -> VacantDistEntry<'d> {218		let dist = &mut self.distributions.0[self.idx];219		assert!(220			dist.owners.remove(&self.owner),221			"entry exists, as we have its reference"222		);223		if dist.owners.is_empty() {224			self.distributions.0.remove(self.idx);225		}226		VacantDistEntry {227			distributions: self.distributions,228			owner: self.owner,229		}230	}231	fn set(self, secret: FleetSecretData) -> Self {232		self.remove().set(secret)233	}234}235struct VacantDistEntry<'d> {236	distributions: &'d mut FleetSecretDistributions,237	owner: String,238}239impl<'d> VacantDistEntry<'d> {240	fn set(self, secret: FleetSecretData) -> OccupiedDistEntry<'d> {241		let Self {242			distributions,243			owner,244		} = self;245		let idx = distributions.0.len();246		distributions.0.push(FleetSecretDistribution {247			managed: None,248			owners: BTreeSet::from_iter([owner.clone()]),249			secret,250		});251		OccupiedDistEntry {252			distributions,253			owner,254			idx,255		}256	}257}258259enum DistEntry<'d> {260	Vacant(VacantDistEntry<'d>),261	Occupied(OccupiedDistEntry<'d>),262}263impl DistEntry<'_> {264	fn remove(self) -> Self {265		match self {266			DistEntry::Vacant(_) => self,267			DistEntry::Occupied(o) => Self::Vacant(o.remove()),268		}269	}270	fn set(self, secret: FleetSecretData) -> Self {271		Self::Occupied(match self {272			DistEntry::Vacant(e) => e.set(secret),273			DistEntry::Occupied(e) => e.set(secret),274		})275	}276}277278impl Serialize for FleetSecretDistributions {279	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>280	where281		S: serde::Serializer,282	{283		let mut found_hosts = BTreeSet::new();284		for ele in self.0.iter() {285			if ele.owners.is_empty() {286				panic!("consistency: secret distribution has no defined owners");287			}288			for ele in ele.owners.iter() {289				if !found_hosts.insert(ele) {290					panic!(291						"consistency: secret distribution contains duplicate entry for the same host",292					);293				}294			}295		}296		match self.0.len() {297			0 => panic!("consistency: empty distributions"),298			1 => self.0[0].serialize(serializer),299			_ => self.0.serialize(serializer),300		}301	}302}303impl<'de> Deserialize<'de> for FleetSecretDistributions {304	fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>305	where306		D: serde::Deserializer<'de>,307	{308		#[derive(Deserialize)]309		#[serde(untagged)]310		enum Distributions {311			One(FleetSecretDistribution),312			Many(Vec<FleetSecretDistribution>),313		}314		let d = Distributions::deserialize(deserializer)?;315		let ds = match d {316			Distributions::One(d) => vec![d],317			Distributions::Many(ds) => ds,318		};319		if ds.is_empty() {320			return Err(de::Error::custom("consistency: empty distributions"));321		}322		let mut found_hosts = BTreeSet::new();323		for ele in ds.iter() {324			if ele.owners.is_empty() {325				return Err(de::Error::custom(326					"consistency: secret distribution has no defined owners",327				));328			}329			for ele in ele.owners.iter() {330				if !found_hosts.insert(ele) {331					return Err(de::Error::custom(332						"consistency: secret distribution contains duplicate entry for the same host",333					));334				}335			}336		}337		Ok(Self(ds))338	}339}340341#[derive(Serialize, Deserialize, Default)]342pub struct FleetSecrets(BTreeMap<String, FleetSecretDistributions>);343344impl FleetSecrets {345	pub fn keys(&self) -> btree_map::Keys<String, FleetSecretDistributions> {346		self.0.keys()347	}348349	pub fn keys_for_owner(&self, owner: &str) -> impl Iterator<Item = &String> {350		self.0351			.iter()352			.filter(|(_, d)| d.contains(owner))353			.map(|(n, _)| n)354	}355356	pub fn drop_owner_no_reencrypt(&mut self, secret: &str, owner: &str) -> bool {357		let Entry::Occupied(mut dists) = self.0.entry(secret.to_owned()) else {358			return false;359		};360		let DistEntry::Occupied(dist) = dists.get_mut().entry(owner.to_owned()) else {361			return false;362		};363364		dist.remove();365366		if dists.get().0.is_empty() {367			dists.remove();368		};369370		true371	}372	pub fn set_single_data(&mut self, secret: String, owner: String, data: FleetSecretData) {373		let e = self374			.0375			.entry(secret.to_owned())376			.or_insert_with(|| FleetSecretDistributions(Default::default()));377		e.entry(owner.to_owned()).set(data);378	}379	pub fn set_data(&mut self, secret: String, data: FleetSecretDistribution) {380		match self.0.entry(secret) {381			Entry::Vacant(e) => {382				e.insert(FleetSecretDistributions(vec![data]));383			}384			Entry::Occupied(mut e) => {385				let dists = e.get_mut();386				dists.extend(data)387			}388		}389	}390	pub fn get_single(&self, secret: &str, owner: &str) -> Option<&FleetSecretDistribution> {391		let secret = self.0.get(secret)?;392		secret.get(owner)393	}394	pub fn get(&self, secret: &str) -> Option<&FleetSecretDistributions> {395		self.0.get(secret)396	}397398	pub fn contains_for_owner(&self, secret: &str, owner: &str) -> bool {399		let Some(secret) = self.0.get(secret) else {400			return false;401		};402		secret.contains(owner)403	}404	pub fn contains(&self, secret: &str) -> bool {405		self.0.contains_key(secret)406	}407	pub fn remove(&mut self, secret: &str) {408		self.0.remove(secret);409	}410411	fn merge_from_hosts(412		&mut self,413		host_secrets: BTreeMap<String, BTreeMap<String, FleetSecretDistribution>>,414	) {415		for (host, host_secrets) in host_secrets {416			for (secret_name, mut secret_data) in host_secrets {417				secret_data.owners.insert(host.clone());418				self.set_data(secret_name, secret_data);419			}420		}421	}422}