container_device_interface/
annotations.rs

1use anyhow::{anyhow, Context, Result};
2use std::collections::HashMap;
3use std::vec::Vec;
4
5use crate::parser;
6
7const ANNOTATION_PREFIX: &str = "cdi.k8s.io/";
8const MAX_NAME_LEN: usize = 63;
9
10// UpdateAnnotations updates annotations with a plugin-specific CDI device
11// injection request for the given devices. Upon any error a non-nil error
12// is returned and annotations are left intact. By convention plugin should
13// be in the format of "vendor.device-type".
14#[allow(dead_code)]
15pub(crate) fn update_annotations(
16    option_annotations: Option<HashMap<String, String>>,
17    plugin_name: &str,
18    device_id: &str,
19    devices: Vec<String>,
20) -> Result<HashMap<String, String>> {
21    let mut annotations = option_annotations.unwrap_or_default();
22
23    let key = annotation_key(plugin_name, device_id).context("CDI annotation key failed")?;
24    if annotations.contains_key(&key) {
25        return Err(anyhow!("CDI annotation key collision, key {:?} used", key));
26    }
27    let value = annotation_value(devices).context("CDI annotation value failed")?;
28
29    annotations.insert(key, value);
30
31    Ok(annotations.clone())
32}
33
34// ParseAnnotations parses annotations for CDI device injection requests.
35// The keys and devices from all such requests are collected into slices
36// which are returned as the result. All devices are expected to be fully
37// qualified CDI device names. If any device fails this check empty slices
38// are returned along with a non-nil error. The annotations are expected
39// to be formatted by, or in a compatible fashion to UpdateAnnotations().
40#[allow(dead_code)]
41pub fn parse_annotations(
42    annotations: &HashMap<String, String>,
43) -> Result<(Vec<String>, Vec<String>)> {
44    let mut keys: Vec<String> = Vec::new();
45    let mut devices: Vec<String> = Vec::new();
46
47    for (k, v) in annotations {
48        if !k.starts_with(ANNOTATION_PREFIX) {
49            continue;
50        }
51
52        for device in v.split(',') {
53            if !parser::is_qualified_name(device) {
54                return Err(anyhow!("invalid CDI device name {}", device));
55            }
56            devices.push(device.to_string());
57        }
58        keys.push(k.to_string());
59    }
60
61    Ok((keys, devices))
62}
63
64// AnnotationKey returns a unique annotation key for an device allocation
65// by a K8s device plugin. pluginName should be in the format of
66// "vendor.device-type". deviceID is the ID of the device the plugin is
67// allocating. It is used to make sure that the generated key is unique
68// even if multiple allocations by a single plugin needs to be annotated.
69#[allow(dead_code)]
70pub(crate) fn annotation_key(plugin_name: &str, device_id: &str) -> Result<String> {
71    if plugin_name.is_empty() {
72        return Err(anyhow!("invalid plugin name, empty"));
73    }
74    if device_id.is_empty() {
75        return Err(anyhow!("invalid deviceID, empty"));
76    }
77
78    let name = format!(
79        "{}_{}",
80        plugin_name.to_owned(),
81        &device_id.replace('/', "_")
82    );
83    if name.len() > MAX_NAME_LEN {
84        return Err(anyhow!("invalid plugin+deviceID {:?}, too long", name));
85    }
86
87    if !name.starts_with(|c: char| c.is_alphanumeric()) {
88        return Err(anyhow!(
89            "invalid name {:?}, first '{}' should be alphanumeric",
90            name.as_str(),
91            name.chars().next().unwrap(),
92        ));
93    }
94
95    if !name
96        .chars()
97        .skip(1)
98        .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '.')
99    {
100        return Err(anyhow!(
101            "invalid name {:?}, invalid character '{}'",
102            name.as_str(),
103            name.chars()
104                .find(|c| !c.is_alphanumeric() && *c != '_' && *c != '-' && *c != '.')
105                .unwrap(),
106        ));
107    }
108
109    if !name.ends_with(|c: char| c.is_alphanumeric()) {
110        return Err(anyhow!(
111            "invalid name {:?}, last '{}' should be alphanumeric",
112            name.as_str(),
113            name.chars().next_back().unwrap(),
114        ));
115    }
116
117    Ok(format!("{}{}", ANNOTATION_PREFIX, name))
118}
119
120// AnnotationValue returns an annotation value for the given devices.
121#[allow(dead_code)]
122pub(crate) fn annotation_value(devices: Vec<String>) -> Result<String> {
123    devices.iter().try_for_each(|device| {
124        // Assuming parser::parse_qualified_name expects a &String and returns Result<(), Error>
125        match crate::parser::parse_qualified_name(device) {
126            Ok(_) => Ok(()),
127            Err(e) => Err(e),
128        }
129    })?;
130
131    let device_strs: Vec<&str> = devices.iter().map(AsRef::as_ref).collect();
132    let value = device_strs.join(",");
133
134    Ok(value)
135}
136
137#[cfg(test)]
138mod tests {
139    use std::collections::HashMap;
140
141    use crate::annotations::{
142        annotation_key, annotation_value, parse_annotations, ANNOTATION_PREFIX,
143    };
144
145    #[test]
146    fn test_parse_annotations() {
147        let mut cdi_devices = HashMap::new();
148
149        cdi_devices.insert(
150            "cdi.k8s.io/vfio17".to_string(),
151            "nvidia.com/gpu=0".to_string(),
152        );
153        cdi_devices.insert(
154            "cdi.k8s.io/vfio18".to_string(),
155            "nvidia.com/gpu=1".to_string(),
156        );
157        cdi_devices.insert(
158            "cdi.k8s.io/vfio19".to_string(),
159            "nvidia.com/gpu=all".to_string(),
160        );
161
162        // one vendor, multiple devices
163        cdi_devices.insert(
164            "vendor.class_device".to_string(),
165            "vendor.com/class=device1,vendor.com/class=device2,vendor.com/class=device3"
166                .to_string(),
167        );
168
169        assert!(parse_annotations(&cdi_devices).is_ok());
170        let (keys, devices) = parse_annotations(&cdi_devices).unwrap();
171        assert_eq!(keys.len(), 3);
172        assert_eq!(devices.len(), 3);
173    }
174
175    #[test]
176    fn test_annotation_value() {
177        let devices = vec![
178            "nvidia.com/gpu=0".to_string(),
179            "nvidia.com/gpu=1".to_string(),
180        ];
181
182        assert!(annotation_value(devices.clone()).is_ok());
183        assert_eq!(
184            annotation_value(devices.clone()).unwrap(),
185            "nvidia.com/gpu=0,nvidia.com/gpu=1"
186        );
187    }
188
189    #[test]
190    fn test_annotation_key() {
191        struct TestCase {
192            plugin_name: String,
193            device_id: String,
194            key_result: String,
195        }
196
197        let test_cases = vec![
198            // valid, with special characters
199            TestCase {
200                plugin_name: "v-e.n_d.or.cl-as_s".to_owned(),
201                device_id: "d_e-v-i-c_e".to_owned(),
202                key_result: format!(
203                    "{}{}_{}",
204                    ANNOTATION_PREFIX, "v-e.n_d.or.cl-as_s", "d_e-v-i-c_e"
205                ),
206            },
207            // valid, with /'s replaced in devID
208            TestCase {
209                plugin_name: "v-e.n_d.or.cl-as_s".to_owned(),
210                device_id: "d-e/v/i/c-e".to_owned(),
211                key_result: format!(
212                    "{}{}_{}",
213                    ANNOTATION_PREFIX, "v-e.n_d.or.cl-as_s", "d-e_v_i_c-e"
214                ),
215            },
216            TestCase {
217                // valid, simple
218                plugin_name: "vendor.class".to_owned(),
219                device_id: "device".to_owned(),
220                key_result: format!("{}{}_{}", ANNOTATION_PREFIX, "vendor.class", "device"),
221            },
222        ];
223
224        for case in test_cases {
225            let plugin_name = &case.plugin_name;
226            let device_id = &case.device_id;
227            assert!(annotation_key(plugin_name, device_id).is_ok());
228            assert_eq!(
229                annotation_key(plugin_name, device_id).unwrap(),
230                case.key_result.clone()
231            );
232        }
233
234        let test_cases_err = vec![
235            // invalid, non-alphanumeric first character
236            TestCase {
237                plugin_name: "_vendor.class".to_owned(),
238                device_id: "device".to_owned(),
239                key_result: "".to_owned(),
240            },
241            // "invalid, non-alphanumeric last character"
242            TestCase {
243                plugin_name: "vendor.class".to_owned(),
244                device_id: "device_".to_owned(),
245                key_result: "".to_owned(),
246            },
247            // invalid, plugin contains invalid characters
248            TestCase {
249                plugin_name: "ven.dor-cl+ass".to_owned(),
250                device_id: "d_e-v-i-c_e".to_owned(),
251                key_result: "".to_owned(),
252            },
253            // "invalid, devID contains invalid characters"
254            TestCase {
255                plugin_name: "vendor.class".to_owned(),
256                device_id: "dev+ice".to_owned(),
257                key_result: "".to_owned(),
258            },
259            // invalid, too plugin long
260            TestCase {
261                plugin_name: "123456789012345678901234567890123456789012345678901234567".to_owned(),
262                device_id: "device".to_owned(),
263                key_result: "".to_owned(),
264            },
265        ];
266
267        for case in test_cases_err {
268            let plugin_name = &case.plugin_name;
269            let device_id = &case.device_id;
270            assert!(annotation_key(plugin_name, device_id).is_err());
271        }
272    }
273}
container_device_interface/annotations.rs

container_device_interface/
annotations.rs
