agentos_kernel/
command_registry.rs

1use crate::vfs::{VfsError, VfsResult, VirtualFileSystem};
2use std::collections::BTreeMap;
3
4const COMMAND_STUB: &[u8] = b"#!/bin/sh\n# kernel command stub\n";
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct CommandDriver {
8    name: String,
9    commands: Vec<String>,
10}
11
12impl CommandDriver {
13    pub fn new<N, I, C>(name: N, commands: I) -> Self
14    where
15        N: Into<String>,
16        I: IntoIterator<Item = C>,
17        C: Into<String>,
18    {
19        Self {
20            name: name.into(),
21            commands: commands.into_iter().map(Into::into).collect(),
22        }
23    }
24
25    pub fn name(&self) -> &str {
26        &self.name
27    }
28
29    pub fn commands(&self) -> &[String] {
30        &self.commands
31    }
32
33    fn validate_commands(&self) -> VfsResult<()> {
34        for command in &self.commands {
35            validate_command_name(command)?;
36        }
37
38        Ok(())
39    }
40}
41
42#[derive(Debug, Default, Clone)]
43pub struct CommandRegistry {
44    commands: BTreeMap<String, CommandDriver>,
45    warnings: Vec<String>,
46}
47
48impl CommandRegistry {
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    pub fn register(&mut self, driver: CommandDriver) -> VfsResult<()> {
54        driver.validate_commands()?;
55
56        // Registering one logical driver is replacement, not an append. This is
57        // required for transactional registries (for example host callbacks):
58        // rolling a driver back to its previous command set must make aliases
59        // introduced by the failed update unresolvable.
60        self.commands
61            .retain(|_, existing| existing.name() != driver.name());
62
63        for command in &driver.commands {
64            if let Some(existing) = self.commands.get(command) {
65                self.warnings.push(format!(
66                    "command \"{command}\" overridden: {} -> {}",
67                    existing.name(),
68                    driver.name()
69                ));
70            }
71
72            self.commands.insert(command.clone(), driver.clone());
73        }
74
75        Ok(())
76    }
77
78    pub fn warnings(&self) -> &[String] {
79        &self.warnings
80    }
81
82    pub fn resolve(&self, command: &str) -> Option<&CommandDriver> {
83        self.commands.get(command)
84    }
85
86    pub fn list(&self) -> BTreeMap<String, String> {
87        self.commands
88            .iter()
89            .map(|(command, driver)| (command.clone(), driver.name().to_owned()))
90            .collect()
91    }
92
93    pub fn populate_bin<F>(&self, vfs: &mut F) -> VfsResult<()>
94    where
95        F: VirtualFileSystem,
96    {
97        self.populate_commands(vfs, self.commands.keys())
98    }
99
100    pub fn populate_driver_bin<F>(&self, vfs: &mut F, driver: &CommandDriver) -> VfsResult<()>
101    where
102        F: VirtualFileSystem,
103    {
104        self.populate_commands(vfs, driver.commands())
105    }
106
107    fn populate_commands<F, I, S>(&self, vfs: &mut F, commands: I) -> VfsResult<()>
108    where
109        F: VirtualFileSystem,
110        I: IntoIterator<Item = S>,
111        S: AsRef<str>,
112    {
113        let commands = commands
114            .into_iter()
115            .map(|command| {
116                validate_command_name(command.as_ref())?;
117                Ok(command.as_ref().to_owned())
118            })
119            .collect::<VfsResult<Vec<_>>>()?;
120
121        if !vfs.exists("/bin") {
122            vfs.mkdir("/bin", true)?;
123        }
124
125        for command in commands {
126            let path = format!("/bin/{command}");
127            if !vfs.exists(&path) {
128                vfs.write_file(&path, COMMAND_STUB.to_vec())?;
129                let _ = vfs.chmod(&path, 0o755);
130            }
131        }
132
133        Ok(())
134    }
135}
136
137fn validate_command_name(command: &str) -> VfsResult<()> {
138    if command.is_empty()
139        || command == "."
140        || command == ".."
141        || command.contains('/')
142        || command.contains('\0')
143    {
144        return Err(VfsError::new(
145            "EINVAL",
146            format!("invalid command name {command:?}"),
147        ));
148    }
149
150    Ok(())
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[test]
158    fn registering_same_driver_replaces_its_command_set() {
159        let mut registry = CommandRegistry::new();
160        registry
161            .register(CommandDriver::new("bindings", ["old", "temporary"]))
162            .expect("register initial driver");
163        registry
164            .register(CommandDriver::new("bindings", ["old"]))
165            .expect("replace driver commands");
166
167        assert_eq!(
168            registry.resolve("old").map(CommandDriver::name),
169            Some("bindings")
170        );
171        assert!(
172            registry.resolve("temporary").is_none(),
173            "aliases removed by a driver refresh must not remain executable"
174        );
175    }
176}
agentos_kernel/command_registry.rs

agentos_kernel/
command_registry.rs
