git.delta.rocks / fleet / refs/commits / cd4874aa6417

difftreelog

refactor db data

Yaroslav Bolyukin2021-09-18parent: #5eda055.patch.diff
in: push-kyumtlkprzyotrunk

4 files changed

deletedsrc/db/db.rsdiffbeforeafterboth
--- a/src/db/db.rs
+++ /dev/null
@@ -1,117 +0,0 @@
-//! Small .toml based readable data store
-
-use anyhow::{Context, Result};
-use serde::{de::DeserializeOwned, Serialize};
-use std::{
-	cell::Cell,
-	collections::HashSet,
-	io::Write,
-	ops::{Deref, DerefMut},
-	path::Path,
-	path::PathBuf,
-	sync::{Arc, Mutex},
-};
-
-struct DbInternal {
-	root: PathBuf,
-	locked_paths: HashSet<PathBuf>,
-}
-
-pub trait DbData: DeserializeOwned + Serialize + Default {
-	const DB_NAME: &'static str;
-
-	fn open(db: &Db) -> Result<DbFile<Self>> {
-		db.db::<Self>()
-	}
-}
-
-#[derive(Clone)]
-pub struct Db(Arc<Mutex<DbInternal>>);
-impl Db {
-	pub fn new(root: impl AsRef<Path>) -> Result<Self> {
-		let root: &Path = root.as_ref();
-		std::fs::create_dir_all(&root).context("db root")?;
-		Ok(Db(Arc::new(Mutex::new(DbInternal {
-			root: root.to_owned(),
-			locked_paths: HashSet::new(),
-		}))))
-	}
-
-	pub fn db<T: DbData>(&self) -> Result<DbFile<T>> {
-		let name = T::DB_NAME;
-		assert!(!name.contains("/") && !name.contains("\\"));
-		let mut db = self.0.lock().unwrap();
-		let mut data_path = db.root.clone();
-		data_path.push(format!("{}.toml", name));
-
-		if !db.locked_paths.insert(data_path.clone()) {
-			anyhow::bail!("file is already open");
-		}
-
-		let data = if data_path.exists() {
-			let raw_data = std::fs::read(&data_path).context("reading file")?;
-			toml::from_slice(&raw_data).context("parsing file")?
-		} else {
-			T::default()
-		};
-
-		Ok(DbFile {
-			db: self.clone(),
-			root: db.root.clone(),
-			path: data_path,
-			data,
-			dirty: Cell::new(false),
-		})
-	}
-}
-
-pub struct DbFile<T: DbData> {
-	db: Db,
-	root: PathBuf,
-	path: PathBuf,
-	data: T,
-	dirty: Cell<bool>,
-}
-
-impl<T: DbData> Deref for DbFile<T> {
-	type Target = T;
-
-	fn deref(&self) -> &Self::Target {
-		&self.data
-	}
-}
-
-impl<T: DbData> DerefMut for DbFile<T> {
-	fn deref_mut(&mut self) -> &mut Self::Target {
-		self.dirty.set(true);
-		&mut self.data
-	}
-}
-
-impl<T: DbData> DbFile<T> {
-	pub fn write(&self) -> Result<()> {
-		if !self.dirty.get() {
-			return Ok(());
-		}
-		let mut temp = tempfile::Builder::new()
-			.prefix("~")
-			.suffix(".toml")
-			.tempfile_in(&self.root)?;
-		let mut out = String::new();
-		let mut serializer = toml::Serializer::new(&mut out);
-		serializer.pretty_array(true).pretty_string(true);
-		self.data.serialize(&mut serializer)?;
-		temp.write_all(&out.as_bytes())?;
-		temp.persist(&self.path)?;
-		self.dirty.set(false);
-		Ok(())
-	}
-}
-
-impl<T: DbData> Drop for DbFile<T> {
-	fn drop(&mut self) {
-		let mut db = self.db.0.lock().unwrap();
-		self.write().unwrap();
-		db.locked_paths.remove(&self.path);
-	}
-}
addedsrc/db/dbr.rsdiffbeforeafterboth
--- /dev/null
+++ b/src/db/dbr.rs
@@ -0,0 +1,117 @@
+//! Small .toml based readable data store
+
+use anyhow::{Context, Result};
+use serde::{de::DeserializeOwned, Serialize};
+use std::{
+	cell::Cell,
+	collections::HashSet,
+	io::Write,
+	ops::{Deref, DerefMut},
+	path::Path,
+	path::PathBuf,
+	sync::{Arc, Mutex},
+};
+
+struct DbInternal {
+	root: PathBuf,
+	locked_paths: HashSet<PathBuf>,
+}
+
+pub trait DbData: DeserializeOwned + Serialize + Default {
+	const DB_NAME: &'static str;
+
+	fn open(db: &Db) -> Result<DbFile<Self>> {
+		db.db::<Self>()
+	}
+}
+
+#[derive(Clone)]
+pub struct Db(Arc<Mutex<DbInternal>>);
+impl Db {
+	pub fn new(root: impl AsRef<Path>) -> Result<Self> {
+		let root: &Path = root.as_ref();
+		std::fs::create_dir_all(&root).context("db root")?;
+		Ok(Db(Arc::new(Mutex::new(DbInternal {
+			root: root.to_owned(),
+			locked_paths: HashSet::new(),
+		}))))
+	}
+
+	pub fn db<T: DbData>(&self) -> Result<DbFile<T>> {
+		let name = T::DB_NAME;
+		assert!(!name.contains('/') && !name.contains('\\'));
+		let mut db = self.0.lock().unwrap();
+		let mut data_path = db.root.clone();
+		data_path.push(format!("{}.toml", name));
+
+		if !db.locked_paths.insert(data_path.clone()) {
+			anyhow::bail!("file is already open");
+		}
+
+		let data = if data_path.exists() {
+			let raw_data = std::fs::read(&data_path).context("reading file")?;
+			toml::from_slice(&raw_data).context("parsing file")?
+		} else {
+			T::default()
+		};
+
+		Ok(DbFile {
+			db: self.clone(),
+			root: db.root.clone(),
+			path: data_path,
+			data,
+			dirty: Cell::new(false),
+		})
+	}
+}
+
+pub struct DbFile<T: DbData> {
+	db: Db,
+	root: PathBuf,
+	path: PathBuf,
+	data: T,
+	dirty: Cell<bool>,
+}
+
+impl<T: DbData> Deref for DbFile<T> {
+	type Target = T;
+
+	fn deref(&self) -> &Self::Target {
+		&self.data
+	}
+}
+
+impl<T: DbData> DerefMut for DbFile<T> {
+	fn deref_mut(&mut self) -> &mut Self::Target {
+		self.dirty.set(true);
+		&mut self.data
+	}
+}
+
+impl<T: DbData> DbFile<T> {
+	pub fn write(&self) -> Result<()> {
+		if !self.dirty.get() {
+			return Ok(());
+		}
+		let mut temp = tempfile::Builder::new()
+			.prefix("~")
+			.suffix(".toml")
+			.tempfile_in(&self.root)?;
+		let mut out = String::new();
+		let mut serializer = toml::Serializer::new(&mut out);
+		serializer.pretty_array(true).pretty_string(true);
+		self.data.serialize(&mut serializer)?;
+		temp.write_all(out.as_bytes())?;
+		temp.persist(&self.path)?;
+		self.dirty.set(false);
+		Ok(())
+	}
+}
+
+impl<T: DbData> Drop for DbFile<T> {
+	fn drop(&mut self) {
+		let mut db = self.db.0.lock().unwrap();
+		self.write().unwrap();
+		db.locked_paths.remove(&self.path);
+	}
+}
modifiedsrc/db/mod.rsdiffbeforeafterboth
--- a/src/db/mod.rs
+++ b/src/db/mod.rs
@@ -1,4 +1,4 @@
-mod db;
+mod dbr;
 pub mod secret;
 
-pub use db::*;
+pub use dbr::*;
modifiedsrc/db/secret.rsdiffbeforeafterboth
before · src/db/secret.rs
1use crate::{command::CommandExt, host::FleetConfig, nix::SECRETS_ATTRIBUTE};2use anyhow::{bail, Result};3use log::info;4use serde::{Deserialize, Deserializer, Serialize, Serializer};5use std::{6	collections::{BTreeMap, BTreeSet, HashMap},7	process::Command,8	time::Instant,9	time::SystemTime,10};11use time::{Duration, PrimitiveDateTime};1213use super::db::DbData;1415#[derive(Serialize, Deserialize, Debug)]16pub struct SecretListData {17	pub owners: BTreeSet<String>,18	#[serde(rename = "expireIn")]19	renew_in: Option<u64>,20}21pub fn list_secrets() -> Result<HashMap<String, SecretListData>> {22	Command::new("nix")23		.inherit_stdio()24		.arg("eval")25		.arg(SECRETS_ATTRIBUTE)26		.arg("--apply")27		.arg(28			r#"29				s: (builtins.mapAttrs (n: {owners, expireIn, ...}: {30					inherit owners expireIn;31				}) s)32			"#,33		)34		.run_json()35}3637struct ReadableDate(PrimitiveDateTime);38impl Serialize for ReadableDate {39	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>40	where41		S: Serializer,42	{43		serializer.serialize_str(&self.0.to_string())44	}45}46impl<'de> Deserialize<'de> for ReadableDate {47	fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>48	where49		D: Deserializer<'de>,50	{51		Ok(Self(52			PrimitiveDateTime::parse(String::deserialize(deserializer)?, "%F %T").unwrap(),53		))54	}55}56impl From<PrimitiveDateTime> for ReadableDate {57	fn from(d: PrimitiveDateTime) -> Self {58		Self(d)59	}60}61impl From<ReadableDate> for PrimitiveDateTime {62	fn from(d: ReadableDate) -> Self {63		d.064	}65}6667#[derive(serde::Serialize, serde::Deserialize)]68struct SecretData {69	created_at: ReadableDate,70	renew_at: Option<ReadableDate>,71	owners: BTreeSet<String>,7273	public_data: BTreeMap<String, String>,74	private_files: BTreeMap<String, String>,75}76impl SecretData {77	fn should_renew(&self) -> bool {78		if let Some(renew_at) = &self.renew_at {79			let now: PrimitiveDateTime = SystemTime::now().into();80			renew_at.0 <= now81		} else {82			false83		}84	}85	fn is_valid(&self, data: &SecretListData) -> bool {86		self.owners == data.owners87	}88}8990#[derive(serde::Serialize, serde::Deserialize)]91struct NixDataValue {92	data: BTreeMap<String, String>,93}9495#[derive(serde::Serialize, serde::Deserialize)]96struct NixData {97	secrets: BTreeMap<String, NixDataValue>,98}99100#[derive(serde::Serialize, serde::Deserialize, Default)]101pub struct SecretDb {102	secrets: BTreeMap<String, SecretData>,103}104impl DbData for SecretDb {105	const DB_NAME: &'static str = "secrets";106}107108impl SecretDb {109	// Secrets are generated on machine running fleet command110	pub fn generate_secret(111		&mut self,112		fleet_config: FleetConfig,113		secret: &str,114		data: &SecretListData,115	) -> Result<()> {116		let mut rage_keys = String::new();117		for (i, owner) in data.owners.iter().enumerate() {118			if i != 0 {119				rage_keys.push(' ');120			}121			rage_keys.push_str("--recipient \"");122			// rage_keys.push_str(&keys.get_host_key(&owner)?);123			rage_keys.push('"')124		}125		let created_at: PrimitiveDateTime = SystemTime::now().into();126		let renew_at = data127			.renew_in128			.map(|hours| created_at + Duration::hours(hours as i64));129		let built = tempfile::tempdir()?;130		Command::new("nix")131			.inherit_stdio()132			.arg("build")133			.arg(format!("{}.{}.generator", SECRETS_ATTRIBUTE, secret))134			.arg("--no-link")135			.arg("--out-link")136			.arg(built.path())137			.arg("--impure")138			.env("RAGE_KEYS", rage_keys)139			.env("IMPURITY_SOURCE", format!("{:?}", Instant::now()))140			.run()?;141		let path = built.path().to_owned();142		let mut secret_data = SecretData {143			created_at: created_at.into(),144			renew_at: renew_at.map(|v| v.into()),145			owners: data.owners.clone(),146			public_data: BTreeMap::new(),147			private_files: BTreeMap::new(),148		};149		for file in std::fs::read_dir(path)? {150			let entry = file?;151			if !entry.file_type()?.is_file() {152				bail!("Secret generator should produce files, not directories");153			}154			let name = entry.file_name();155			let name = name156				.to_str()157				.ok_or(anyhow::anyhow!("file name should be utf-8"))?;158			let value = String::from_utf8(std::fs::read(entry.path())?)?;159			if let Some(name) = name.strip_prefix("pub_") {160				secret_data.public_data.insert(name.into(), value);161			} else {162				secret_data.private_files.insert(name.into(), value);163			}164		}165		self.secrets.insert(secret.into(), secret_data);166		Ok(())167	}168	pub fn need_to_generate(&self, secret: &str, data: &SecretListData) -> Result<bool> {169		let secret = self.secrets.get(secret);170		if secret.is_none() {171			return Ok(true);172		}173		let secret = secret.unwrap();174175		if secret.should_renew() {176			return Ok(true);177		}178179		if !secret.is_valid(&data) {180			return Ok(true);181		}182183		Ok(false)184	}185	pub fn ensure_generated(186		&mut self,187		// keys: &KeyDb,188		secret: &str,189		data: &SecretListData,190	) -> Result<()> {191		if self.need_to_generate(secret, data)? {192			info!("Generating secret {}", secret);193			// self.generate_secret(keys, secret, data)?;194		}195196		Ok(())197	}198	pub fn generate_nix_data(&self) -> Result<String> {199		let mut out = BTreeMap::new();200		for (host, secrets) in &self.secrets {201			out.insert(202				host.to_owned(),203				NixDataValue {204					data: secrets205						.public_data206						.clone()207						.iter()208						.map(|(k, v)| (k.to_owned(), v.trim().to_owned()))209						.collect(),210				},211			);212		}213		Ok(serde_json::to_string(&out)?)214	}215216	pub fn has_secret(&self, secret: &str) -> bool {217		self.secrets.contains_key(secret)218	}219220	pub fn remove_secret(&mut self, secret: &str) {221		self.secrets.remove(secret);222	}223}
after · src/db/secret.rs
1use crate::{command::CommandExt, host::FleetOpts, nix::SECRETS_ATTRIBUTE};2use anyhow::{bail, Context, Result};3use log::info;4use serde::{Deserialize, Deserializer, Serialize, Serializer};5use std::{6	collections::{BTreeMap, BTreeSet, HashMap},7	process::Command,8	time::Instant,9	time::SystemTime,10};11use time::{Duration, PrimitiveDateTime};1213use super::DbData;1415#[derive(Serialize, Deserialize, Debug)]16pub struct SecretListData {17	pub owners: BTreeSet<String>,18	#[serde(rename = "expireIn")]19	renew_in: Option<u64>,20}21pub fn list_secrets() -> Result<HashMap<String, SecretListData>> {22	Command::new("nix")23		.inherit_stdio()24		.arg("eval")25		.arg(SECRETS_ATTRIBUTE)26		.arg("--apply")27		.arg(28			r#"29				s: (builtins.mapAttrs (n: {owners, expireIn, ...}: {30					inherit owners expireIn;31				}) s)32			"#,33		)34		.arg("--json")35		.run_json()36		.context("while getting secret list")37}3839struct ReadableDate(PrimitiveDateTime);40impl Serialize for ReadableDate {41	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>42	where43		S: Serializer,44	{45		serializer.serialize_str(&self.0.to_string())46	}47}48impl<'de> Deserialize<'de> for ReadableDate {49	fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>50	where51		D: Deserializer<'de>,52	{53		Ok(Self(54			PrimitiveDateTime::parse(String::deserialize(deserializer)?, "%F %T").unwrap(),55		))56	}57}58impl From<PrimitiveDateTime> for ReadableDate {59	fn from(d: PrimitiveDateTime) -> Self {60		Self(d)61	}62}63impl From<ReadableDate> for PrimitiveDateTime {64	fn from(d: ReadableDate) -> Self {65		d.066	}67}6869#[derive(serde::Serialize, serde::Deserialize)]70struct SecretData {71	created_at: ReadableDate,72	renew_at: Option<ReadableDate>,73	owners: BTreeSet<String>,7475	public_data: BTreeMap<String, String>,76	private_files: BTreeMap<String, String>,77}78impl SecretData {79	fn should_renew(&self) -> bool {80		if let Some(renew_at) = &self.renew_at {81			let now: PrimitiveDateTime = SystemTime::now().into();82			renew_at.0 <= now83		} else {84			false85		}86	}87	fn is_valid(&self, data: &SecretListData) -> bool {88		self.owners == data.owners89	}90}9192#[derive(serde::Serialize, serde::Deserialize)]93struct NixDataValue {94	data: BTreeMap<String, String>,95}9697#[derive(serde::Serialize, serde::Deserialize)]98struct NixData {99	secrets: BTreeMap<String, NixDataValue>,100}101102#[derive(serde::Serialize, serde::Deserialize, Default)]103pub struct SecretDb {104	secrets: BTreeMap<String, SecretData>,105}106impl DbData for SecretDb {107	const DB_NAME: &'static str = "secrets";108}109110impl SecretDb {111	// Secrets are generated on machine running fleet command112	pub fn generate_secret(113		&mut self,114		_fleet_config: &FleetOpts,115		secret: &str,116		data: &SecretListData,117	) -> Result<()> {118		let mut rage_keys = String::new();119		for (i, _owner) in data.owners.iter().enumerate() {120			if i != 0 {121				rage_keys.push(' ');122			}123			rage_keys.push_str("--recipient \"");124			// rage_keys.push_str(&fleet_config.clone().build()?.host(owner)?.key()?);125			//rage_keys.push_str(&keys.get_host_key(&owner)?);126			rage_keys.push('"')127		}128		let created_at: PrimitiveDateTime = SystemTime::now().into();129		let renew_at = data130			.renew_in131			.map(|hours| created_at + Duration::hours(hours as i64));132		let built = tempfile::tempdir()?;133		Command::new("nix")134			.inherit_stdio()135			.arg("build")136			.arg(format!("{}.{}.generator", SECRETS_ATTRIBUTE, secret))137			.arg("--no-link")138			.arg("--out-link")139			.arg(built.path())140			.arg("--impure")141			.env("RAGE_KEYS", rage_keys)142			.env("IMPURITY_SOURCE", format!("{:?}", Instant::now()))143			.run()?;144		let path = built.path().to_owned();145		let mut secret_data = SecretData {146			created_at: created_at.into(),147			renew_at: renew_at.map(|v| v.into()),148			owners: data.owners.clone(),149			public_data: BTreeMap::new(),150			private_files: BTreeMap::new(),151		};152		for file in std::fs::read_dir(path)? {153			let entry = file?;154			if !entry.file_type()?.is_file() {155				bail!("Secret generator should produce files, not directories");156			}157			let name = entry.file_name();158			let name = name159				.to_str()160				.ok_or_else(|| anyhow::anyhow!("file name should be utf-8"))?;161			let value = String::from_utf8(std::fs::read(entry.path())?)?;162			if let Some(name) = name.strip_prefix("pub_") {163				secret_data.public_data.insert(name.into(), value);164			} else {165				secret_data.private_files.insert(name.into(), value);166			}167		}168		self.secrets.insert(secret.into(), secret_data);169		Ok(())170	}171	pub fn need_to_generate(&self, secret: &str, data: &SecretListData) -> Result<bool> {172		let secret = self.secrets.get(secret);173		if secret.is_none() {174			return Ok(true);175		}176		let secret = secret.unwrap();177178		if secret.should_renew() {179			return Ok(true);180		}181182		if !secret.is_valid(data) {183			return Ok(true);184		}185186		Ok(false)187	}188	pub fn ensure_generated(189		&mut self,190		// keys: &KeyDb,191		fleet_config: &FleetOpts,192		secret: &str,193		data: &SecretListData,194	) -> Result<()> {195		if self.need_to_generate(secret, data)? {196			info!("Generating secret {}", secret);197			self.generate_secret(fleet_config, secret, data)?;198		}199200		Ok(())201	}202	pub fn generate_nix_data(&self) -> Result<String> {203		let mut out = BTreeMap::new();204		for (host, secrets) in &self.secrets {205			out.insert(206				host.to_owned(),207				NixDataValue {208					data: secrets209						.public_data210						.clone()211						.iter()212						.map(|(k, v)| (k.to_owned(), v.trim().to_owned()))213						.collect(),214				},215			);216		}217		Ok(serde_json::to_string(&out)?)218	}219220	pub fn has_secret(&self, secret: &str) -> bool {221		self.secrets.contains_key(secret)222	}223224	pub fn remove_secret(&mut self, secret: &str) {225		self.secrets.remove(secret);226	}227}