1use std::{2 collections::BTreeMap,3 env::current_dir,4 ffi::OsString,5 str::FromStr,6 sync::{Arc, Mutex},7};89use anyhow::{bail, Context, Result};10use clap::Parser;11use nix_eval::{nix_go, util::assert_warn, NixSessionPool, Value};12use nom::{13 bytes::complete::take_while1,14 character::complete::char,15 combinator::{map, opt},16 multi::separated_list1,17 sequence::{preceded, separated_pair},18};1920use crate::{21 fleetdata::FleetData,22 host::{Config, ConfigHost, FleetConfigInternals},23};2425#[derive(Clone)]26pub enum HostItem {27 Host {28 name: String,29 attrs: BTreeMap<String, String>,30 },31 Tag {32 name: String,33 attrs: BTreeMap<String, String>,34 },35}36fn host_item_parser(input: &str) -> Result<HostItem, String> {37 fn err_to_string(err: nom::Err<nom::error::Error<&str>>) -> String {38 err.to_string()39 }4041 let (input, is_tag) = map(opt(char('@')), |c| c.is_some())(input).map_err(err_to_string)?;42 let (input, name) = map(43 take_while1(|v| v != ',' && v != '?' && v != '@'),44 str::to_owned,45 )(input)46 .map_err(err_to_string)?;4748 let kw_item = separated_pair(49 map(take_while1(|v| v != '&' && v != '='), str::to_owned),50 char('='),51 map(take_while1(|v| v != '&'), str::to_owned),52 );53 let kw = map(separated_list1(char('&'), kw_item), |vec| {54 vec.into_iter().collect::<BTreeMap<_, _>>()55 });56 let mut opt_kw = map(opt(preceded(char('?'), kw)), Option::unwrap_or_default);5758 let (input, attrs) = opt_kw(input).map_err(err_to_string)?;5960 if !input.is_empty() {61 return Err(format!("unexpected trailing input: {input:?}"));62 }63 Ok(if is_tag {64 HostItem::Tag { name, attrs }65 } else {66 HostItem::Host { name, attrs }67 })68}697071#[derive(Parser, Clone)]72pub struct FleetOpts {73 74 #[clap(long, number_of_values = 1, value_parser = host_item_parser)]75 pub only: Vec<HostItem>,7677 78 #[clap(long, number_of_values = 1)]79 pub skip: Vec<String>,8081 82 83 #[clap(long, default_value_t = hostname::get().expect("unknown hostname").to_str().expect("hostname is not utf-8").to_owned())]84 pub localhost: String,8586 87 88 #[clap(long, default_value = env!("NIX_SYSTEM"))]89 pub local_system: String,9091 92 93 94 95 #[clap(long)]96 pub fail_fast: bool,97}9899impl FleetOpts {100 pub async fn filter_skipped(101 &self,102 hosts: impl IntoIterator<Item = ConfigHost>,103 ) -> Result<Vec<ConfigHost>> {104 let mut out = Vec::new();105 for host in hosts {106 if self.should_skip(&host).await? {107 continue;108 }109 out.push(host);110 }111 Ok(out)112 }113 pub async fn should_skip(&self, host: &ConfigHost) -> Result<bool> {114 if self.skip.iter().any(|h| h as &str == host.name) {115 return Ok(true);116 }117 if self.only.is_empty() {118 return Ok(false);119 }120 let mut have_group_matches = false;121 for item in self.only.iter() {122 match item {123 HostItem::Host { name, .. } if *name == host.name => {124 return Ok(false);125 }126 HostItem::Tag { .. } => {127 have_group_matches = true;128 }129 _ => {}130 }131 }132 if have_group_matches {133 let host_tags = host.tags().await?;134 for item in self.only.iter() {135 match item {136 HostItem::Tag { name, .. } if host_tags.contains(name) => {137 return Ok(false);138 }139 _ => {}140 }141 }142 }143 Ok(true)144 }145 pub async fn action_attr<T: FromStr>(&self, host: &ConfigHost, attr: &str) -> Result<Option<T>>146 where147 T::Err: Sync,148 anyhow::Error: From<T::Err>,149 {150 let str = self.action_attr_str(host, attr).await?;151 Ok(str.map(|v| T::from_str(&v)).transpose()?)152 }153 pub async fn action_attr_str(&self, host: &ConfigHost, attr: &str) -> Result<Option<String>> {154 if self.only.is_empty() {155 return Ok(None);156 }157 let mut have_group_matches = false;158 for item in self.only.iter() {159 match item {160 HostItem::Host { name, attrs }161 if *name == host.name && attrs.contains_key(attr) =>162 {163 return Ok(attrs.get(attr).cloned());164 }165 HostItem::Tag { attrs, .. } if attrs.contains_key(attr) => {166 have_group_matches = true;167 }168 _ => {}169 }170 }171 if have_group_matches {172 let host_tags = host.tags().await?;173 for item in self.only.iter() {174 match item {175 HostItem::Tag { name, attrs }176 if host_tags.contains(name) && attrs.contains_key(attr) =>177 {178 return Ok(attrs.get(attr).cloned());179 }180 _ => {}181 }182 }183 }184 Ok(None)185 }186 pub fn is_local(&self, host: &str) -> bool {187 self.localhost == host188 }189190 191 pub async fn build(&self, nix_args: Vec<OsString>, assert: bool) -> Result<Config> {192 let cwd = current_dir()?;193 let mut directory = cwd.clone();194 let mut fleet_data_path = directory.join("fleet.nix");195 while !fleet_data_path.is_file() {196 197 fleet_data_path.pop();198 if !directory.pop() || !fleet_data_path.pop() {199 bail!(200 "fleet.nix not found at {} or any of the parent directories",201 cwd.display()202 );203 }204 fleet_data_path.push("fleet.nix");205 }206 let bytes =207 std::fs::read_to_string(&fleet_data_path).context("reading fleet state (fleet.nix)")?;208 let data: Mutex<FleetData> = nixlike::parse_str(&bytes)?;209210 let pool = NixSessionPool::new(211 directory.as_os_str().to_owned(),212 nix_args.clone(),213 self.local_system.clone(),214 self.fail_fast,215 )216 .await?;217 let nix_session = pool.get().await?;218219 let builtins_field = Value::binding(nix_session.clone(), "builtins").await?;220221 let fleet_root = Value::binding(nix_session.clone(), "fleetConfigurations").await?;222 let fleet_field = nix_go!(fleet_root.default({ data }));223224 let config_field = nix_go!(fleet_field.config);225226 if assert {227 assert_warn("fleet config evaluation", &config_field).await?;228 }229230 let import = nix_go!(builtins_field.import);231 let overlays = nix_go!(config_field.nixpkgs.overlays);232 let nixpkgs = nix_go!(config_field.nixpkgs.buildUsing);233 let nixpkgs_imported = nix_go!(nixpkgs | import);234235 let default_pkgs = nix_go!(nixpkgs_imported(Obj {236 overlays,237 system: self.local_system.clone(),238 }));239240 Ok(Config(Arc::new(FleetConfigInternals {241 nix_session,242 directory,243 data,244 local_system: self.local_system.clone(),245 nix_args,246 config_field,247 default_pkgs,248 nixpkgs,249 localhost: self.localhost.to_owned(),250 })))251 }252}