git.delta.rocks / jrsonnet / refs/commits / faec7071817b

difftreelog

source

crates/fleet-base/src/fleetdata.rs10.5 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, skip_serializing_if = "Vec::is_empty")]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, Debug)]125pub struct FleetSecretPart {126	pub raw: SecretData,127}128129#[derive(Serialize, Deserialize, Clone, Debug)]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, Debug)]148#[serde(rename_all = "camelCase")]149#[must_use]150pub struct FleetSecretDistribution {151	#[serde(default)]152	pub owners: BTreeSet<String>,153	#[serde(flatten)]154	pub secret: FleetSecretData,155156	#[serde(default, skip_serializing, alias = "managed")]157	pub _deprecated_managed: bool,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			owners: BTreeSet::from_iter([owner.clone()]),248			secret,249250			_deprecated_managed: true,251		});252		OccupiedDistEntry {253			distributions,254			owner,255			idx,256		}257	}258}259260enum DistEntry<'d> {261	Vacant(VacantDistEntry<'d>),262	Occupied(OccupiedDistEntry<'d>),263}264impl DistEntry<'_> {265	fn remove(self) -> Self {266		match self {267			DistEntry::Vacant(_) => self,268			DistEntry::Occupied(o) => Self::Vacant(o.remove()),269		}270	}271	fn set(self, secret: FleetSecretData) -> Self {272		Self::Occupied(match self {273			DistEntry::Vacant(e) => e.set(secret),274			DistEntry::Occupied(e) => e.set(secret),275		})276	}277}278279impl Serialize for FleetSecretDistributions {280	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>281	where282		S: serde::Serializer,283	{284		let mut found_hosts = BTreeSet::new();285		for ele in self.0.iter() {286			if ele.owners.is_empty() {287				panic!("consistency: secret distribution has no defined owners");288			}289			for ele in ele.owners.iter() {290				if !found_hosts.insert(ele) {291					panic!(292						"consistency: secret distribution contains duplicate entry for the same host",293					);294				}295			}296		}297		match self.0.len() {298			0 => panic!("consistency: empty distributions"),299			1 => self.0[0].serialize(serializer),300			_ => self.0.serialize(serializer),301		}302	}303}304impl<'de> Deserialize<'de> for FleetSecretDistributions {305	fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>306	where307		D: serde::Deserializer<'de>,308	{309		#[derive(Deserialize)]310		#[serde(untagged)]311		enum Distributions {312			One(FleetSecretDistribution),313			Many(Vec<FleetSecretDistribution>),314		}315		let d = Distributions::deserialize(deserializer)?;316		let ds = match d {317			Distributions::One(d) => vec![d],318			Distributions::Many(ds) => ds,319		};320		if ds.is_empty() {321			return Err(de::Error::custom("consistency: empty distributions"));322		}323		let mut found_hosts = BTreeSet::new();324		for ele in ds.iter() {325			if ele.owners.is_empty() {326				return Err(de::Error::custom(327					"consistency: secret distribution has no defined owners",328				));329			}330			for ele in ele.owners.iter() {331				if !found_hosts.insert(ele) {332					return Err(de::Error::custom(333						"consistency: secret distribution contains duplicate entry for the same host",334					));335				}336			}337		}338		Ok(Self(ds))339	}340}341342#[derive(Serialize, Deserialize, Default)]343pub struct FleetSecrets(BTreeMap<String, FleetSecretDistributions>);344345impl FleetSecrets {346	pub fn keys(&self) -> btree_map::Keys<String, FleetSecretDistributions> {347		self.0.keys()348	}349350	pub fn keys_for_owner(&self, owner: &str) -> impl Iterator<Item = &String> {351		self.0352			.iter()353			.filter(|(_, d)| d.contains(owner))354			.map(|(n, _)| n)355	}356357	pub fn drop_owner_no_reencrypt(&mut self, secret: &str, owner: &str) -> bool {358		let Entry::Occupied(mut dists) = self.0.entry(secret.to_owned()) else {359			return false;360		};361		let DistEntry::Occupied(dist) = dists.get_mut().entry(owner.to_owned()) else {362			return false;363		};364365		dist.remove();366367		if dists.get().0.is_empty() {368			dists.remove();369		};370371		true372	}373	pub fn set_single_data(&mut self, secret: String, owner: String, data: FleetSecretData) {374		let e = self375			.0376			.entry(secret.to_owned())377			.or_insert_with(|| FleetSecretDistributions(Default::default()));378		e.entry(owner.to_owned()).set(data);379	}380	pub fn set_data(&mut self, secret: String, data: FleetSecretDistribution) {381		match self.0.entry(secret) {382			Entry::Vacant(e) => {383				e.insert(FleetSecretDistributions(vec![data]));384			}385			Entry::Occupied(mut e) => {386				let dists = e.get_mut();387				dists.extend(data)388			}389		}390	}391	pub fn get_single(&self, secret: &str, owner: &str) -> Option<&FleetSecretDistribution> {392		let secret = self.0.get(secret)?;393		secret.get(owner)394	}395	pub fn get(&self, secret: &str) -> Option<&FleetSecretDistributions> {396		self.0.get(secret)397	}398399	pub fn contains_for_owner(&self, secret: &str, owner: &str) -> bool {400		let Some(secret) = self.0.get(secret) else {401			return false;402		};403		secret.contains(owner)404	}405	pub fn contains(&self, secret: &str) -> bool {406		self.0.contains_key(secret)407	}408	pub fn remove(&mut self, secret: &str) {409		self.0.remove(secret);410	}411412	fn merge_from_hosts(413		&mut self,414		host_secrets: BTreeMap<String, BTreeMap<String, FleetSecretDistribution>>,415	) {416		for (host, host_secrets) in host_secrets {417			for (secret_name, mut secret_data) in host_secrets {418				secret_data.owners.insert(host.clone());419				self.set_data(secret_name, secret_data);420			}421		}422	}423}424425#[derive(Debug)]426pub struct Expectations {427	pub owners: BTreeSet<String>,428	pub generation_data: serde_json::Value,429	pub parts: BTreeMap<String, GeneratorPart>,430}431#[derive(Deserialize, Debug, Clone)]432pub struct GeneratorPart {433	pub encrypted: bool,434}