nodeset/collections/
config.rs

1use super::nodeset::ConfigurationError;
2use super::parsers::Parser;
3use super::NodeSet;
4use crate::idrange::IdRange;
5use crate::NodeSetParseError;
6use ini::Properties;
7use log::debug;
8use serde::Deserialize;
9use shellexpand::env_with_context_no_errors;
10use std::collections::HashMap;
11use std::fmt::Debug;
12use std::fmt::Display;
13use std::fs;
14use std::io::BufReader;
15use std::path::Path;
16use std::path::PathBuf;
17use std::process::Command;
18use std::sync::OnceLock;
19
20/// The default resolver used to parse NodeSet using the FromStr trait
21static GLOBAL_RESOLVER: OnceLock<Resolver> = OnceLock::new();
22
23/// Default group configuration paths
24static CONFIG_PATHS: &[&str] = &[
25    "$HOME/.local/etc/clustershell",
26    "/etc/clustershell",
27    "$XDG_CONFIG_HOME/clustershell",
28];
29
30/// An inventory of group sources used to resolve group names to node sets
31///
32/// The FromStr implementation of NodeSet uses the global resolver which can be
33/// setup to read group sources from the default configuration file as follows:
34///
35/// ```rust,no_run
36/// use nodeset::{NodeSet, Resolver};
37///
38/// fn main() -> Result<(), Box<dyn std::error::Error>> {
39///     Resolver::set_global(Resolver::from_config()?).unwrap();
40///
41///     let ns: NodeSet = "@group".parse()?;
42///
43///     Ok(())
44/// }
45/// ```
46#[derive(Debug)]
47pub struct Resolver {
48    sources: HashMap<String, Box<dyn GroupSource>>,
49    default_source: String,
50}
51
52impl Default for Resolver {
53    fn default() -> Self {
54        Self {
55            sources: HashMap::default(),
56            default_source: "local".to_string(),
57        }
58    }
59}
60
61impl Resolver {
62    /// Create a new resolver from the default configuration files
63    pub fn from_config() -> Result<Self, ConfigurationError> {
64        let mut group_config = MainGroupConfig::default();
65
66        let mut cfg_dir = None;
67        for &path in CONFIG_PATHS {
68            if let Some(file) = open_config_path(&Path::new(&path).join("groups.conf")) {
69                group_config.merge(MainGroupConfig::from_reader(BufReader::new(file))?);
70                cfg_dir = resolve_config_path(Path::new(&path));
71            }
72        }
73
74        if let Some(cfg_dir) = cfg_dir {
75            if let Some(cfg_dir) = cfg_dir.to_str() {
76                group_config.set_cfgdir(cfg_dir)?;
77            }
78        }
79
80        Resolver::from_dynamic_config(group_config)
81    }
82
83    /// Create a new resolver from a dynamic group configuration
84    ///
85    /// `set_cfgdir` must already have been called on the dynamic group configuration
86    fn from_dynamic_config(groups: MainGroupConfig) -> Result<Self, ConfigurationError> {
87        let mut resolver = Resolver {
88            sources: Default::default(),
89            default_source: groups
90                .config
91                .as_ref()
92                .and_then(|c| c.default.clone())
93                .unwrap_or_else(|| "default".to_string()),
94        };
95
96        for autodir in groups.autodirs() {
97            for path in find_files_with_ext(Path::new(&autodir), "yaml") {
98                if let Some(file) = open_config_path(&path) {
99                    let static_groups = StaticGroupConfig::from_reader(BufReader::new(file))?;
100                    resolver.add_sources(static_groups);
101                }
102            }
103        }
104        for confdir in groups.confdirs() {
105            for path in find_files_with_ext(Path::new(&confdir), "conf") {
106                if let Some(file) = open_config_path(&path) {
107                    let dynamic_groups = MainGroupConfig::from_reader(BufReader::new(file))?;
108                    resolver.add_sources(dynamic_groups);
109                }
110            }
111        }
112
113        resolver.add_sources(groups);
114
115        Ok(resolver)
116    }
117
118    /// Set the global resolver to use for parsing NodeSet using the FromStr trait
119    ///
120    /// Returns an error if the global resolver is already set
121    pub fn set_global(resolver: Resolver) -> Result<(), Resolver> {
122        GLOBAL_RESOLVER.set(resolver)?;
123
124        Ok(())
125    }
126
127    /// Get the global resolver
128    pub fn get_global() -> &'static Resolver {
129        static DEFAULT_RESOLVER: OnceLock<Resolver> = OnceLock::new();
130
131        GLOBAL_RESOLVER
132            .get()
133            .unwrap_or(DEFAULT_RESOLVER.get_or_init(Resolver::default))
134    }
135
136    /// Resolve a group name to a NodeSet
137    ///
138    /// If `source` is None, the default group source of the resolver is used.
139    pub fn resolve<T: IdRange + PartialEq + Clone + Display + Debug>(
140        &self,
141        source: Option<&str>,
142        group: &str,
143    ) -> Result<NodeSet<T>, NodeSetParseError> {
144        let source = source.unwrap_or(self.default_source.as_str());
145
146        Parser::with_resolver(self, Some(source)).parse(
147            &self
148                .sources
149                .get(source)
150                .ok_or_else(|| NodeSetParseError::Source(source.to_owned()))?
151                .map(group)?
152                .unwrap_or_default(),
153        )
154    }
155
156    /// List groups from a source
157    ///
158    /// If `source` is None, the default group source of the resolver is used.
159    pub fn list_groups<T: IdRange + PartialEq + Clone + Display + Debug>(
160        &self,
161        source: Option<&str>,
162    ) -> NodeSet<T> {
163        let source = source.unwrap_or(self.default_source.as_str());
164
165        Parser::default()
166            .parse(
167                &self
168                    .sources
169                    .get(source)
170                    .map(|s| s.list())
171                    .unwrap_or_default(),
172            )
173            .unwrap_or_default()
174    }
175
176    /// List groups from all sources
177    ///
178    /// Returns a list of tuples with the source name and the group name
179    pub fn list_all_groups<T: IdRange + PartialEq + Clone + Display + Debug>(
180        &self,
181    ) -> impl Iterator<Item = (&str, NodeSet<T>)> {
182        self.sources.iter().map(|(source, groups)| {
183            (
184                source.as_str(),
185                Parser::default().parse(&groups.list()).unwrap_or_default(),
186            )
187        })
188    }
189
190    /// List all sources
191    pub fn sources(&self) -> impl Iterator<Item = &String> {
192        self.sources.keys()
193    }
194
195    /// Returns the default group source for this resolver
196    pub fn default_source(&self) -> &str {
197        &self.default_source
198    }
199
200    pub(crate) fn add_sources(
201        &mut self,
202        sources: impl IntoIterator<Item = (String, impl GroupSource + 'static)>,
203    ) {
204        sources.into_iter().for_each(|(name, source)| {
205            self.sources.insert(name, Box::new(source));
206        });
207    }
208}
209
210/// Open a config file from a path, expanding environment variables.
211///
212/// Returns None if there was any failure
213fn open_config_path(path: &Path) -> Option<std::fs::File> {
214    resolve_config_path(path).and_then(|p| std::fs::File::open(p).ok())
215}
216
217/// Expands environment variables in a path
218///
219/// Returns None in case of non-utf8 path
220fn resolve_config_path(path: &Path) -> Option<PathBuf> {
221    let context = |s: &str| match s {
222        "HOME" => std::env::var("HOME").ok(),
223        "XDG_CONFIG_HOME" => std::env::var("XDG_CONFIG_HOME").ok().or_else(|| {
224            std::env::var("HOME")
225                .ok()
226                .map(|h| Path::new(&h).join(".config").to_str().unwrap().to_string())
227        }),
228        _ => None,
229    };
230
231    Some(PathBuf::from(
232        env_with_context_no_errors(path.to_str()?, context).as_ref(),
233    ))
234}
235
236/// Returns a list of files with a given extension in a directory.
237///
238/// Returns an empty list in case of any failure
239fn find_files_with_ext(dir: &Path, ext: &str) -> Vec<PathBuf> {
240    let mut files = vec![];
241
242    let Ok(it) = fs::read_dir(dir) else {
243        return files;
244    };
245
246    for entry in it {
247        let entry = entry.unwrap();
248        let path = entry.path();
249
250        if path.is_file() && path.extension().map(|ext| ext.to_str()) == Some(Some(ext)) {
251            files.push(path);
252        }
253    }
254
255    files
256}
257
258/// Trait for group resolution features of a group source
259pub(crate) trait GroupSource: Debug + Send + Sync {
260    fn map(&self, group: &str) -> Result<Option<String>, NodeSetParseError>;
261    fn list(&self) -> String;
262}
263
264/// Settings from the main group configuration file (groups.conf)
265#[derive(Debug, Default)]
266struct MainGroupConfig {
267    config: Option<ResolverOptions>,
268    sources: HashMap<String, DynamicGroupSource>,
269}
270
271impl MainGroupConfig {
272    fn from_reader(mut reader: impl std::io::Read) -> Result<Self, ConfigurationError> {
273        use ini::Ini;
274
275        let parser = Ini::read_from_noescape(&mut reader)?;
276        let mut config = MainGroupConfig::default();
277        for (sec, prop) in parser.iter() {
278            match sec {
279                Some("Main") => {
280                    config.config = Some(prop.try_into()?);
281                }
282                Some(sources) => {
283                    for source in sources.split(',') {
284                        config.sources.insert(
285                            source.to_string(),
286                            DynamicGroupSource::from_props(prop, source.to_string())?,
287                        );
288                    }
289                }
290                None => {
291                    if let Some(key) = prop.iter().next().map(|(k, _)| k) {
292                        return Err(ConfigurationError::UnexpectedProperty(key.to_string()));
293                    }
294                }
295            }
296        }
297
298        Ok(config)
299    }
300
301    fn autodirs(&self) -> Vec<String> {
302        self.config
303            .as_ref()
304            .map(|c| c.autodirs())
305            .unwrap_or_default()
306    }
307
308    fn confdirs(&self) -> Vec<String> {
309        self.config
310            .as_ref()
311            .map(|c| c.confdirs())
312            .unwrap_or_default()
313    }
314
315    fn set_cfgdir(&mut self, cfgdir: &str) -> Result<(), ConfigurationError> {
316        let context = |s: &str| match s {
317            "CFGDIR" => Some(cfgdir),
318            _ => None,
319        };
320
321        if let Some(config) = &mut self.config {
322            config.confdir = config
323                .confdir
324                .as_ref()
325                .map(|c| env_with_context_no_errors(&c, context).to_string());
326            config.autodir = config
327                .autodir
328                .as_ref()
329                .map(|c| env_with_context_no_errors(&c, context).to_string());
330        }
331
332        for (_, group) in self.sources.iter_mut() {
333            group.set_cfgdir(cfgdir)?;
334        }
335
336        Ok(())
337    }
338
339    /// Merge settings for another group configuration file into this one
340    fn merge(&mut self, other: Self) {
341        match (&mut self.config, other.config) {
342            (Some(ref mut main), Some(other_main)) => main.merge(other_main),
343            (None, Some(other_main)) => self.config = Some(other_main),
344            _ => (),
345        }
346        self.sources.extend(other.sources);
347    }
348}
349
350impl IntoIterator for MainGroupConfig {
351    type Item = (String, DynamicGroupSource);
352    type IntoIter = std::collections::hash_map::IntoIter<String, DynamicGroupSource>;
353
354    fn into_iter(self) -> Self::IntoIter {
355        self.sources.into_iter()
356    }
357}
358
359#[derive(Debug, Default)]
360struct ResolverOptions {
361    default: Option<String>,
362    confdir: Option<String>,
363    autodir: Option<String>,
364}
365
366impl ResolverOptions {
367    /// Merge resolver options from `other` into `self`
368    fn merge(&mut self, other: Self) {
369        if let Some(default) = other.default {
370            self.default = Some(default);
371        }
372        if let Some(confdir) = other.confdir {
373            self.confdir = Some(confdir);
374        }
375        if let Some(autodir) = other.autodir {
376            self.autodir = Some(autodir);
377        }
378    }
379
380    fn autodirs(&self) -> Vec<String> {
381        self.autodir
382            .as_ref()
383            .and_then(|autodir| shlex::split(autodir))
384            .unwrap_or_default()
385    }
386
387    fn confdirs(&self) -> Vec<String> {
388        self.confdir
389            .as_ref()
390            .and_then(|confdir| shlex::split(confdir))
391            .unwrap_or_default()
392    }
393}
394
395impl TryFrom<&Properties> for ResolverOptions {
396    type Error = ConfigurationError;
397
398    fn try_from(props: &Properties) -> Result<Self, Self::Error> {
399        let mut res = Self::default();
400
401        for (k, v) in props.iter() {
402            match k {
403                "default" => {
404                    res.default = Some(v.to_string());
405                }
406                "confdir" => {
407                    res.confdir = Some(v.to_string());
408                }
409                "autodir" => {
410                    res.autodir = Some(v.to_string());
411                }
412                _ => {
413                    return Err(ConfigurationError::UnexpectedProperty(k.to_string()));
414                }
415            }
416        }
417
418        Ok(res)
419    }
420}
421
422/// Settings from a dynamic group source (groups.conf.d/<source>.conf)
423#[derive(Debug)]
424struct DynamicGroupSource {
425    name: String,
426    map: String,
427    all: Option<String>,
428    list: Option<String>,
429}
430
431impl DynamicGroupSource {
432    fn from_props(props: &Properties, name: String) -> Result<Self, ConfigurationError> {
433        let map = props
434            .get("map")
435            .ok_or_else(|| ConfigurationError::MissingProperty("map".to_string()))?
436            .to_string();
437        let all = props.get("all").map(|s| s.to_string());
438        let list = props.get("list").map(|s| s.to_string());
439
440        Ok(Self {
441            name,
442            map,
443            all,
444            list,
445        })
446    }
447
448    fn set_cfgdir(&mut self, cfgdir: &str) -> Result<(), ConfigurationError> {
449        let context = |s: &str| match s {
450            "CFGDIR" => Some(cfgdir),
451            "SOURCE" => Some(self.name.as_str()),
452            _ => None,
453        };
454
455        self.map = env_with_context_no_errors(&self.map, context).to_string();
456        self.all = self
457            .all
458            .as_ref()
459            .map(|s| env_with_context_no_errors(s, context).to_string());
460        self.list = self
461            .list
462            .as_ref()
463            .map(|s| env_with_context_no_errors(s, context).to_string());
464
465        Ok(())
466    }
467}
468
469impl GroupSource for DynamicGroupSource {
470    fn map(&self, group: &str) -> Result<Option<String>, NodeSetParseError> {
471        let context = |s: &str| match s {
472            "GROUP" => Some(group),
473            "SOURCE" => Some(self.name.as_str()),
474            _ => None,
475        };
476        let map = env_with_context_no_errors(&self.map, context).to_string();
477
478        let output = Command::new("/bin/sh").arg("-c").arg(&map).output()?;
479
480        if !output.status.success() {
481            return Err(NodeSetParseError::Command(std::io::Error::other(format!(
482                "Command '{}' returned non-zero exit code",
483                map
484            ))));
485        }
486
487        let res = String::from_utf8_lossy(&output.stdout);
488
489        debug!(
490            "Map command '{}' for @'{}':'{}' returned: {}",
491            map, self.name, group, res
492        );
493
494        Ok(Some(res.trim().to_string()))
495    }
496
497    fn list(&self) -> String {
498        let Some(ref list_cmd) = self.list else {
499            return Default::default();
500        };
501
502        let context = |s: &str| match s {
503            "SOURCE" => Some(self.name.as_str()),
504            _ => None,
505        };
506        let list = env_with_context_no_errors(&list_cmd, context).to_string();
507
508        let output = Command::new("/bin/sh")
509            .arg("-c")
510            .arg(&list)
511            .output()
512            .unwrap();
513
514        if !output.status.success() {
515            panic!("Command '{}' returned non-zero exit code", list);
516        }
517
518        let res = String::from_utf8_lossy(&output.stdout);
519
520        debug!(
521            "List command '{}' for @'{}':* returned: {}",
522            list, self.name, res
523        );
524
525        res.trim().to_string()
526    }
527}
528
529/// Settings from a static group source configuration file (groups.d/*.yaml)
530#[derive(Deserialize, Debug)]
531struct StaticGroupConfig {
532    #[serde(flatten)]
533    sources: HashMap<String, StaticGroupSource>,
534}
535
536impl StaticGroupConfig {
537    fn from_reader(reader: impl std::io::Read) -> Result<Self, ConfigurationError> {
538        let config: Self = serde_yaml::from_reader(reader)?;
539        Ok(config)
540    }
541}
542
543impl IntoIterator for StaticGroupConfig {
544    type Item = (String, StaticGroupSource);
545    type IntoIter = std::collections::hash_map::IntoIter<String, StaticGroupSource>;
546
547    fn into_iter(self) -> Self::IntoIter {
548        self.sources.into_iter()
549    }
550}
551
552#[derive(Deserialize, Debug)]
553struct StaticGroupSource {
554    #[serde(flatten)]
555    groups: HashMap<String, SingleOrVec>,
556}
557
558#[derive(Deserialize, Debug)]
559#[serde(untagged)]
560enum SingleOrVec {
561    Single(String),
562    Vec(Vec<String>),
563}
564
565impl From<&SingleOrVec> for String {
566    fn from(s: &SingleOrVec) -> Self {
567        match s {
568            SingleOrVec::Single(s) => s.clone(),
569            SingleOrVec::Vec(v) => v.join(","),
570        }
571    }
572}
573
574impl GroupSource for StaticGroupSource {
575    fn map(&self, group: &str) -> Result<Option<String>, NodeSetParseError> {
576        Ok(self.groups.get(group).map(|v| v.into()))
577    }
578
579    fn list(&self) -> String {
580        use itertools::Itertools;
581        self.groups.keys().join(" ")
582    }
583}
584
585#[cfg(test)]
586#[derive(Debug)]
587pub(crate) struct DummySource {
588    map: HashMap<String, String>,
589}
590#[cfg(test)]
591impl DummySource {
592    pub(crate) fn new() -> Self {
593        Self {
594            map: HashMap::new(),
595        }
596    }
597
598    pub(crate) fn add(&mut self, group: &str, nodes: &str) {
599        self.map.insert(group.to_string(), nodes.to_string());
600    }
601}
602
603#[cfg(test)]
604impl GroupSource for DummySource {
605    fn map(&self, group: &str) -> Result<Option<String>, NodeSetParseError> {
606        Ok(self.map.get(group).cloned())
607    }
608
609    fn list(&self) -> String {
610        use itertools::Itertools;
611
612        self.map.keys().join(" ")
613    }
614}
615
616#[cfg(test)]
617mod tests {
618    use crate::{collections::parsers::Parser, IdRangeList};
619
620    use super::*;
621
622    #[test]
623    fn test_static_config() {
624        let config = include_str!("tests/cluster.yaml");
625        let mut resolver = Resolver::default();
626        resolver.add_sources(StaticGroupConfig::from_reader(config.as_bytes()).unwrap());
627        let parser = Parser::with_resolver(&resolver, Some("roles"));
628        assert_eq!(
629            parser.parse::<IdRangeList>("@login").unwrap().to_string(),
630            "login[1-2]"
631        );
632
633        assert_eq!(
634            parser.parse::<IdRangeList>("@*").unwrap(),
635            parser
636                .parse::<IdRangeList>(
637                    "node[0001-0288],mds[1-4],oss[0-15],server0001,login[1-2],mgmt[1-2]"
638                )
639                .unwrap()
640        );
641
642        match parser.parse::<IdRangeList>("@login:aa") {
643            Err(NodeSetParseError::Source(_)) => (),
644            e => panic!("Expected Source error, got {e:?}",),
645        }
646        assert_eq!(
647            parser
648                .parse::<IdRangeList>("@roles:cpu_only")
649                .unwrap()
650                .to_string(),
651            "node[0009-0288]"
652        );
653
654        assert_eq!(
655            parser
656                .parse::<IdRangeList>("@roles:non_existent")
657                .unwrap()
658                .to_string(),
659            ""
660        );
661
662        match parser.parse::<IdRangeList>("@non_existent:non_existent") {
663            Err(NodeSetParseError::Source(_)) => (),
664            _ => panic!("Expected Source error"),
665        }
666
667        let ns1 = parser.parse::<IdRangeList>("@rack[1-2]:hsw").unwrap();
668        let ns2 = "mgmt[1-2],oss[0-15],mds[1-4]".parse().unwrap();
669        assert_eq!(ns1, ns2);
670
671        let ns1 = parser.parse::<IdRangeList>("@rack[1-2]:*").unwrap();
672        let ns2 = "mgmt[1-2],oss[0-15],mds[1-4],node[0001-0288]"
673            .parse()
674            .unwrap();
675        assert_eq!(ns1, ns2);
676
677        let ns1 = parser.parse::<IdRangeList>("@network:net[1,3]").unwrap();
678        let ns2 = "node[10-19,30-39]".parse().unwrap();
679        assert_eq!(ns1, ns2);
680
681        assert_eq!(
682            resolver.list_groups::<IdRangeList>(Some("numerical")),
683            "1-2,03".parse::<NodeSet>().unwrap()
684        );
685
686        assert_eq!(
687            resolver
688                .resolve::<IdRangeList>(Some("numerical"), "1")
689                .unwrap(),
690            "node[10-19]".parse::<NodeSet>().unwrap()
691        );
692    }
693
694    #[test]
695    fn test_parse_dynamic_config() {
696        use tempfile::TempDir;
697
698        let config = include_str!("tests/groups.conf");
699        let mut dynamic = MainGroupConfig::from_reader(config.as_bytes()).unwrap();
700
701        let tmp_dir = TempDir::new().unwrap();
702
703        std::fs::create_dir(tmp_dir.path().join("groups.d")).unwrap();
704        std::fs::write(
705            tmp_dir.path().join("groups.d").join("local.cfg"),
706            include_str!("tests/local.cfg"),
707        )
708        .unwrap();
709
710        std::fs::create_dir(tmp_dir.path().join("groups.conf.d")).unwrap();
711        std::fs::write(
712            tmp_dir.path().join("groups.conf.d").join("multi.conf"),
713            include_str!("tests/multi.conf"),
714        )
715        .unwrap();
716
717        dynamic
718            .set_cfgdir(tmp_dir.path().to_str().unwrap())
719            .unwrap();
720
721        assert_eq!(
722            dynamic.autodirs(),
723            vec![
724                "/etc/clustershell/groups.d",
725                &format!("{}/groups.d", tmp_dir.path().to_str().unwrap())
726            ]
727        );
728        assert_eq!(
729            dynamic.confdirs(),
730            vec![
731                "/etc/clustershell/groups.conf.d",
732                &format!("{}/groups.conf.d", tmp_dir.path().to_str().unwrap())
733            ]
734        );
735
736        let resolver = Resolver::from_dynamic_config(dynamic).unwrap();
737
738        assert_eq!(
739            resolver
740                .resolve::<IdRangeList>(None, "oss")
741                .unwrap()
742                .to_string(),
743            "example[4-5]"
744        );
745
746        assert_eq!(
747            resolver
748                .resolve::<IdRangeList>(Some("local"), "mds")
749                .unwrap()
750                .to_string(),
751            "example6"
752        );
753
754        assert_eq!(
755            resolver.list_groups::<IdRangeList>(Some("local")),
756            "compute,gpu,all,adm,io,mds,oss,[1-2],03"
757                .parse::<NodeSet>()
758                .unwrap()
759        );
760
761        assert_eq!(
762            resolver.resolve::<IdRangeList>(Some("local"), "1").unwrap(),
763            "example[32-33]".parse::<NodeSet>().unwrap()
764        );
765
766        assert_eq!(
767            resolver.list_groups::<IdRangeList>(Some("rack1")),
768            "rack1_switches[1-4],rack1_nodes[1-4]"
769                .parse::<NodeSet>()
770                .unwrap()
771        );
772
773        assert_eq!(
774            resolver.list_groups::<IdRangeList>(Some("rack2")),
775            "rack2_switches[1-4],rack2_nodes[1-4]"
776                .parse::<NodeSet>()
777                .unwrap()
778        );
779
780        assert_eq!(
781            resolver
782                .resolve::<IdRangeList>(Some("rack1"), "nodes")
783                .unwrap(),
784            "rack1_nodes[1-4]".parse::<NodeSet>().unwrap()
785        );
786    }
787}
nodeset/collections/config.rs

nodeset/collections/
config.rs
