1use std::{2 collections::BTreeMap,3 env::current_dir,4 ffi::OsString,5 str::FromStr,6 sync::{Arc, Mutex},7};89use anyhow::{Context, Result, bail};10use nix_eval::{11 FetchSettings, FlakeLockFlags, FlakeReference, FlakeReferenceParseFlags, FlakeSettings, Value,12 gc_now, nix_go, util::assert_warn,13};14use nom::{15 Parser,16 bytes::complete::take_while1,17 character::complete::char,18 combinator::{map, opt},19 multi::separated_list1,20 sequence::{preceded, separated_pair},21};2223use crate::{24 fleetdata::FleetData,25 host::{Config, ConfigHost, FleetConfigInternals},26};2728#[derive(Clone)]29pub enum HostItem {30 Host {31 name: String,32 attrs: BTreeMap<String, String>,33 },34 Tag {35 name: String,36 attrs: BTreeMap<String, String>,37 },38}39fn host_item_parser(input: &str) -> Result<HostItem, String> {40 fn err_to_string(err: nom::Err<nom::error::Error<&str>>) -> String {41 err.to_string()42 }4344 let (input, is_tag) = map(opt(char('@')), |c| c.is_some())45 .parse_complete(input)46 .map_err(err_to_string)?;47 let (input, name) = map(48 take_while1(|v| v != ',' && v != '?' && v != '@'),49 str::to_owned,50 )51 .parse_complete(input)52 .map_err(err_to_string)?;5354 let kw_item = separated_pair(55 map(take_while1(|v| v != '&' && v != '='), str::to_owned),56 char('='),57 map(take_while1(|v| v != '&'), str::to_owned),58 );59 let kw = map(separated_list1(char('&'), kw_item), |vec| {60 vec.into_iter().collect::<BTreeMap<_, _>>()61 });62 let mut opt_kw = map(opt(preceded(char('?'), kw)), Option::unwrap_or_default);6364 let (input, attrs) = opt_kw.parse_complete(input).map_err(err_to_string)?;6566 if !input.is_empty() {67 return Err(format!("unexpected trailing input: {input:?}"));68 }69 Ok(if is_tag {70 HostItem::Tag { name, attrs }71 } else {72 HostItem::Host { name, attrs }73 })74}757677#[derive(clap::Parser, Clone)]78pub struct FleetOpts {79 80 #[clap(long, number_of_values = 1, value_parser = host_item_parser)]81 pub only: Vec<HostItem>,8283 84 #[clap(long, number_of_values = 1)]85 pub skip: Vec<String>,8687 88 89 #[clap(long, default_value_t = hostname::get().expect("unknown hostname").to_str().expect("hostname is not utf-8").to_owned())]90 pub localhost: String,9192 93 94 #[clap(long, default_value = env!("NIX_SYSTEM"))]95 pub local_system: String,9697 98 99 100 101 #[clap(long)]102 pub fail_fast: bool,103}104105impl FleetOpts {106 pub async fn filter_skipped(107 &self,108 hosts: impl IntoIterator<Item = ConfigHost>,109 ) -> Result<Vec<ConfigHost>> {110 let mut out = Vec::new();111 for host in hosts {112 if self.should_skip(&host).await? {113 continue;114 }115 out.push(host);116 }117 Ok(out)118 }119 pub async fn should_skip(&self, host: &ConfigHost) -> Result<bool> {120 if self.skip.iter().any(|h| h as &str == host.name) {121 return Ok(true);122 }123 if self.only.is_empty() {124 return Ok(false);125 }126 let mut have_group_matches = false;127 for item in self.only.iter() {128 match item {129 HostItem::Host { name, .. } if *name == host.name => {130 return Ok(false);131 }132 HostItem::Tag { .. } => {133 have_group_matches = true;134 }135 _ => {}136 }137 }138 if have_group_matches {139 let host_tags = host.tags().await?;140 for item in self.only.iter() {141 match item {142 HostItem::Tag { name, .. } if host_tags.contains(name) => {143 return Ok(false);144 }145 _ => {}146 }147 }148 }149 Ok(true)150 }151 pub async fn action_attr<T: FromStr>(&self, host: &ConfigHost, attr: &str) -> Result<Option<T>>152 where153 T::Err: Sync,154 anyhow::Error: From<T::Err>,155 {156 let str = self.action_attr_str(host, attr).await?;157 Ok(str.map(|v| T::from_str(&v)).transpose()?)158 }159 pub async fn action_attr_str(&self, host: &ConfigHost, attr: &str) -> Result<Option<String>> {160 if self.only.is_empty() {161 return Ok(None);162 }163 let mut have_group_matches = false;164 for item in self.only.iter() {165 match item {166 HostItem::Host { name, attrs }167 if *name == host.name && attrs.contains_key(attr) =>168 {169 return Ok(attrs.get(attr).cloned());170 }171 HostItem::Tag { attrs, .. } if attrs.contains_key(attr) => {172 have_group_matches = true;173 }174 _ => {}175 }176 }177 if have_group_matches {178 let host_tags = host.tags().await?;179 for item in self.only.iter() {180 match item {181 HostItem::Tag { name, attrs }182 if host_tags.contains(name) && attrs.contains_key(attr) =>183 {184 return Ok(attrs.get(attr).cloned());185 }186 _ => {}187 }188 }189 }190 Ok(None)191 }192 pub fn is_local(&self, host: &str) -> bool {193 self.localhost == host194 }195196 197 pub async fn build(&self, nix_args: Vec<OsString>, assert: bool) -> Result<Config> {198 let cwd = current_dir()?;199 let mut directory = cwd.clone();200 let mut fleet_data_path = directory.join("fleet.nix");201 while !fleet_data_path.is_file() {202 203 fleet_data_path.pop();204 if !directory.pop() || !fleet_data_path.pop() {205 bail!(206 "fleet.nix not found at {} or any of the parent directories",207 cwd.display()208 );209 }210 fleet_data_path.push("fleet.nix");211 }212 let bytes =213 std::fs::read_to_string(&fleet_data_path).context("reading fleet state (fleet.nix)")?;214 let data: Mutex<FleetData> = nixlike::parse_str(&bytes)?;215216 let mut fetch_settings = FetchSettings::new();217 fetch_settings.set(c"warn-dirty", c"false");218219 let mut flake_settings = FlakeSettings::new()?;220 let mut parse = FlakeReferenceParseFlags::new(&flake_settings)?;221 222 parse.set_base_dir("/")?;223224 let (mut flake, _) = FlakeReference::new(225 directory226 .to_str()227 .ok_or_else(|| anyhow::anyhow!("fleet dir should have utf-8 path"))?,228 &flake_settings,229 &parse,230 &fetch_settings,231 )?;232233 let lock = FlakeLockFlags::new(&flake_settings)?;234235 let flake = flake.lock(&fetch_settings, &flake_settings, &lock)?;236237 let flake = flake.get_attrs(&mut flake_settings)?;238239 let builtins_field = Value::eval("builtins")?;240241 let fleet_root = flake.get_field("fleetConfigurations")?;242 let data_val = Value::serialized(&data)?;243 let fleet_field = nix_go!(fleet_root.default(data_val));244245 let config_field = nix_go!(fleet_field.config);246247 if assert {248 assert_warn("fleet config evaluation", &config_field)249 .await250 .context("failed to verify assertions")?;251 }252253 let import = nix_go!(builtins_field.import);254 let overlays = nix_go!(config_field.nixpkgs.overlays);255 let nixpkgs = nix_go!(config_field.nixpkgs.buildUsing);256 let nixpkgs_imported = nix_go!(import(nixpkgs));257258 let default_pkgs = nix_go!(nixpkgs_imported(Obj {259 overlays,260 system: self.local_system.clone(),261 }));262263 if cfg!(debug_assertions) {264 gc_now();265 }266267 Ok(Config(Arc::new(FleetConfigInternals {268 directory,269 data,270 flake_outputs: flake,271 local_system: self.local_system.clone(),272 nix_args,273 config_field,274 default_pkgs,275 nixpkgs,276 localhost: self.localhost.to_owned(),277 })))278 }279}