1use std::{2 collections::{BTreeMap, BTreeSet},3 env::current_dir,4 ffi::OsString,5 str::FromStr,6 sync::{Arc, Mutex},7};89use anyhow::{Context, Result, bail};10use chrono::Utc;11use nix_eval::{12 FetchSettings, FlakeLockFlags, FlakeReference, FlakeReferenceParseFlags, FlakeSettings, Value,13 gc_now, nix_go, util::assert_warn,14};15use nom::{16 Parser,17 bytes::complete::take_while1,18 character::complete::char,19 combinator::{map, opt},20 multi::separated_list1,21 sequence::{preceded, separated_pair},22};2324use crate::{25 fleetdata::FleetData,26 host::{Config, ConfigHost, FleetConfigInternals},27 primops::{PRIMOPS_DATA, init_primops},28};2930#[derive(Clone)]31pub enum HostItem {32 Host {33 name: String,34 attrs: BTreeMap<String, String>,35 },36 Tag {37 name: String,38 attrs: BTreeMap<String, String>,39 },40}41fn host_item_parser(input: &str) -> Result<HostItem, String> {42 fn err_to_string(err: nom::Err<nom::error::Error<&str>>) -> String {43 err.to_string()44 }4546 let (input, is_tag) = map(opt(char('@')), |c| c.is_some())47 .parse_complete(input)48 .map_err(err_to_string)?;49 let (input, name) = map(50 take_while1(|v| v != ',' && v != '?' && v != '@'),51 str::to_owned,52 )53 .parse_complete(input)54 .map_err(err_to_string)?;5556 let kw_item = separated_pair(57 map(take_while1(|v| v != '&' && v != '='), str::to_owned),58 char('='),59 map(take_while1(|v| v != '&'), str::to_owned),60 );61 let kw = map(separated_list1(char('&'), kw_item), |vec| {62 vec.into_iter().collect::<BTreeMap<_, _>>()63 });64 let mut opt_kw = map(opt(preceded(char('?'), kw)), Option::unwrap_or_default);6566 let (input, attrs) = opt_kw.parse_complete(input).map_err(err_to_string)?;6768 if !input.is_empty() {69 return Err(format!("unexpected trailing input: {input:?}"));70 }71 Ok(if is_tag {72 HostItem::Tag { name, attrs }73 } else {74 HostItem::Host { name, attrs }75 })76}777879#[derive(clap::Parser, Clone)]80pub struct FleetOpts {81 82 #[clap(long, number_of_values = 1, value_parser = host_item_parser)]83 pub only: Vec<HostItem>,8485 86 #[clap(long, number_of_values = 1)]87 pub skip: Vec<String>,8889 90 91 #[clap(long, default_value_t = hostname::get().expect("unknown hostname").to_str().expect("hostname is not utf-8").to_owned())]92 pub localhost: String,9394 95 96 #[clap(long, default_value = env!("NIX_SYSTEM"))]97 pub local_system: String,9899 100 101 102 103 #[clap(long)]104 pub fail_fast: bool,105}106107impl FleetOpts {108 pub fn filter_skipped(109 &self,110 hosts: impl IntoIterator<Item = ConfigHost>,111 ) -> Result<Vec<ConfigHost>> {112 let mut out = Vec::new();113 for host in hosts {114 if self.should_skip(&host)? {115 continue;116 }117 out.push(host);118 }119 Ok(out)120 }121 pub fn should_skip(&self, host: &ConfigHost) -> Result<bool> {122 if self.skip.iter().any(|h| h as &str == host.name) {123 return Ok(true);124 }125 if self.only.is_empty() {126 return Ok(false);127 }128 let mut have_group_matches = false;129 for item in self.only.iter() {130 match item {131 HostItem::Host { name, .. } if *name == host.name => {132 return Ok(false);133 }134 HostItem::Tag { .. } => {135 have_group_matches = true;136 }137 _ => {}138 }139 }140 if have_group_matches {141 let host_tags = host.tags()?;142 for item in self.only.iter() {143 match item {144 HostItem::Tag { name, .. } if host_tags.contains(name) => {145 return Ok(false);146 }147 _ => {}148 }149 }150 }151 Ok(true)152 }153 pub fn action_attr<T: FromStr>(&self, host: &ConfigHost, attr: &str) -> Result<Option<T>>154 where155 T::Err: Sync,156 anyhow::Error: From<T::Err>,157 {158 let str = self.action_attr_str(host, attr)?;159 Ok(str.map(|v| T::from_str(&v)).transpose()?)160 }161 pub fn action_attr_str(&self, host: &ConfigHost, attr: &str) -> Result<Option<String>> {162 if self.only.is_empty() {163 return Ok(None);164 }165 let mut have_group_matches = false;166 for item in self.only.iter() {167 match item {168 HostItem::Host { name, attrs }169 if *name == host.name && attrs.contains_key(attr) =>170 {171 return Ok(attrs.get(attr).cloned());172 }173 HostItem::Tag { attrs, .. } if attrs.contains_key(attr) => {174 have_group_matches = true;175 }176 _ => {}177 }178 }179 if have_group_matches {180 let host_tags = host.tags()?;181 for item in self.only.iter() {182 match item {183 HostItem::Tag { name, attrs }184 if host_tags.contains(name) && attrs.contains_key(attr) =>185 {186 return Ok(attrs.get(attr).cloned());187 }188 _ => {}189 }190 }191 }192 Ok(None)193 }194 pub fn is_local(&self, host: &str) -> bool {195 self.localhost == host196 }197198 199 pub fn build(&self, nix_args: Vec<OsString>, assert: bool) -> Result<Config> {200 let cwd = current_dir()?;201 let mut directory = cwd.clone();202 let mut fleet_data_path = directory.join("fleet.nix");203 while !fleet_data_path.is_file() {204 205 fleet_data_path.pop();206 if !directory.pop() || !fleet_data_path.pop() {207 bail!(208 "fleet.nix not found at {} or any of the parent directories",209 cwd.display()210 );211 }212 fleet_data_path.push("fleet.nix");213 }214 let bytes =215 std::fs::read_to_string(&fleet_data_path).context("reading fleet state (fleet.nix)")?;216 let data = Arc::new(FleetData::from_str(&bytes)?);217218 init_primops();219220 let mut fetch_settings = FetchSettings::new();221 fetch_settings.set(c"warn-dirty", c"false");222223 let mut flake_settings = FlakeSettings::new()?;224 let mut parse = FlakeReferenceParseFlags::new(&flake_settings)?;225 226 parse.set_base_dir("/")?;227228 let (mut flake, _) = FlakeReference::new(229 directory230 .to_str()231 .ok_or_else(|| anyhow::anyhow!("fleet dir should have utf-8 path"))?,232 &flake_settings,233 &parse,234 &fetch_settings,235 )?;236237 let lock = FlakeLockFlags::new(&flake_settings)?;238239 let flake = flake.lock(&fetch_settings, &flake_settings, &lock)?;240241 let flake = flake.get_attrs(&mut flake_settings)?;242243 let builtins_field = Value::eval("builtins")?;244245 let fleet_root = flake.get_field("fleetConfigurations")?;246 let fleet_field = nix_go!(fleet_root.default(Obj {}));247248 let config_field = nix_go!(fleet_field.config);249250 if assert {251 assert_warn("fleet config evaluation", &config_field)252 .context("failed to verify assertions")?;253 }254255 let import = nix_go!(builtins_field.import);256 let overlays = nix_go!(config_field.nixpkgs.overlays);257 let nixpkgs = nix_go!(config_field.nixpkgs.buildUsing);258 let nixpkgs_imported = nix_go!(import(nixpkgs));259260 let default_pkgs = nix_go!(nixpkgs_imported(Obj {261 overlays,262 system: self.local_system.clone(),263 }));264265 if cfg!(debug_assertions) {266 gc_now();267 }268 let config = Config(Arc::new(FleetConfigInternals {269 270 prefer_identities: BTreeSet::new(),271 now: Utc::now(),272273 directory,274 data,275 flake_outputs: flake,276 local_system: self.local_system.clone(),277 nix_args,278 config_field,279 default_pkgs,280 nixpkgs,281 localhost: self.localhost.to_owned(),282 }));283284 PRIMOPS_DATA285 .set(config.clone())286 .map_err(|_| ())287 .expect("only one fleet config may exist per process");288 Ok(config)289 }290}