virtio_accel_coreml/
artifact.rs

1#[cfg(any(target_os = "macos", test))]
2use std::collections::BTreeSet;
3
4const MAGIC: [u8; 4] = *b"CMLP";
5const VERSION_MAJOR: u8 = 1;
6const VERSION_MINOR: u8 = 0;
7const HEADER_BYTES: usize = 16;
8const ENTRY_HEADER_BYTES: usize = 8;
9pub(crate) const MAX_ARTIFACT_BYTES: u64 = 64 * 1024;
10const MAX_MODEL_PATH_BYTES: usize = 4 * 1024;
11const MAX_FEATURE_NAME_BYTES: usize = 1024;
12const MAX_MAPPINGS: usize = 256;
13
14/// A Core ML model feature's role in one virtio-accel binding.
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16#[repr(u8)]
17pub enum FeatureRole {
18    Input = 1,
19    Output = 2,
20}
21
22impl FeatureRole {
23    #[cfg(any(target_os = "macos", test))]
24    fn from_wire(value: u8) -> Result<Self, DecodeError> {
25        match value {
26            1 => Ok(Self::Input),
27            2 => Ok(Self::Output),
28            _ => Err(DecodeError::Invalid),
29        }
30    }
31}
32
33/// Invalid Core ML path artifact construction.
34#[derive(Clone, Copy, Debug, PartialEq, Eq)]
35pub enum ArtifactBuildError {
36    EmptyModelPath,
37    ModelPathTooLong,
38    EmptyFeatureName,
39    FeatureNameTooLong,
40    DuplicateFeature,
41    EmptyMappings,
42    TooManyMappings,
43    ArtifactTooLarge,
44}
45
46impl std::fmt::Display for ArtifactBuildError {
47    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        write!(formatter, "{self:?}")
49    }
50}
51
52impl std::error::Error for ArtifactBuildError {}
53
54#[derive(Clone, Debug, PartialEq, Eq)]
55pub(crate) struct FeatureMapping {
56    pub slot: u32,
57    pub role: FeatureRole,
58    pub name: String,
59}
60
61/// Builder for the provider-owned Core ML path artifact.
62///
63/// The model path is relative to the host-selected model root. Every nonoptional model input and
64/// output must be mapped exactly once. Mapping one input and one compatible output to the same slot
65/// creates an in-place `ReadWrite` binding.
66#[derive(Clone, Debug, PartialEq, Eq)]
67pub struct CoreMlArtifact {
68    model_path: String,
69    mappings: Vec<FeatureMapping>,
70}
71
72impl CoreMlArtifact {
73    pub fn new(model_path: impl Into<String>) -> Result<Self, ArtifactBuildError> {
74        let model_path = model_path.into();
75        if model_path.is_empty() {
76            return Err(ArtifactBuildError::EmptyModelPath);
77        }
78        if model_path.len() > MAX_MODEL_PATH_BYTES || model_path.len() > u16::MAX as usize {
79            return Err(ArtifactBuildError::ModelPathTooLong);
80        }
81        Ok(Self {
82            model_path,
83            mappings: Vec::new(),
84        })
85    }
86
87    pub fn map_input(
88        self,
89        slot: u32,
90        feature_name: impl Into<String>,
91    ) -> Result<Self, ArtifactBuildError> {
92        self.map(slot, FeatureRole::Input, feature_name.into())
93    }
94
95    pub fn map_output(
96        self,
97        slot: u32,
98        feature_name: impl Into<String>,
99    ) -> Result<Self, ArtifactBuildError> {
100        self.map(slot, FeatureRole::Output, feature_name.into())
101    }
102
103    pub fn encode(&self) -> Result<Vec<u8>, ArtifactBuildError> {
104        if self.mappings.is_empty() {
105            return Err(ArtifactBuildError::EmptyMappings);
106        }
107        if self.mappings.len() > MAX_MAPPINGS || self.mappings.len() > u16::MAX as usize {
108            return Err(ArtifactBuildError::TooManyMappings);
109        }
110        let mut bytes = Vec::with_capacity(
111            HEADER_BYTES
112                + self.model_path.len()
113                + self
114                    .mappings
115                    .iter()
116                    .map(|mapping| ENTRY_HEADER_BYTES + mapping.name.len())
117                    .sum::<usize>(),
118        );
119        bytes.extend_from_slice(&MAGIC);
120        bytes.push(VERSION_MAJOR);
121        bytes.push(VERSION_MINOR);
122        bytes.extend_from_slice(&0u16.to_le_bytes());
123        bytes.extend_from_slice(&(self.model_path.len() as u16).to_le_bytes());
124        bytes.extend_from_slice(&(self.mappings.len() as u16).to_le_bytes());
125        bytes.extend_from_slice(&0u32.to_le_bytes());
126        bytes.extend_from_slice(self.model_path.as_bytes());
127        for mapping in &self.mappings {
128            bytes.extend_from_slice(&mapping.slot.to_le_bytes());
129            bytes.push(mapping.role as u8);
130            bytes.push(0);
131            bytes.extend_from_slice(&(mapping.name.len() as u16).to_le_bytes());
132            bytes.extend_from_slice(mapping.name.as_bytes());
133        }
134        if bytes.len() as u64 > MAX_ARTIFACT_BYTES {
135            return Err(ArtifactBuildError::ArtifactTooLarge);
136        }
137        Ok(bytes)
138    }
139
140    fn map(
141        mut self,
142        slot: u32,
143        role: FeatureRole,
144        name: String,
145    ) -> Result<Self, ArtifactBuildError> {
146        if name.is_empty() {
147            return Err(ArtifactBuildError::EmptyFeatureName);
148        }
149        if name.len() > MAX_FEATURE_NAME_BYTES || name.len() > u16::MAX as usize {
150            return Err(ArtifactBuildError::FeatureNameTooLong);
151        }
152        if self
153            .mappings
154            .iter()
155            .any(|mapping| mapping.role == role && mapping.name == name)
156        {
157            return Err(ArtifactBuildError::DuplicateFeature);
158        }
159        if self.mappings.len() == MAX_MAPPINGS {
160            return Err(ArtifactBuildError::TooManyMappings);
161        }
162        self.mappings.push(FeatureMapping { slot, role, name });
163        Ok(self)
164    }
165}
166
167#[derive(Clone, Copy, Debug, PartialEq, Eq)]
168#[cfg(any(target_os = "macos", test))]
169pub(crate) enum DecodeError {
170    Invalid,
171    OutOfBounds,
172    ResourceLimit,
173}
174
175#[derive(Debug)]
176#[cfg(any(target_os = "macos", test))]
177pub(crate) struct DecodedArtifact {
178    pub model_path: String,
179    pub mappings: Vec<FeatureMapping>,
180}
181
182#[cfg(any(target_os = "macos", test))]
183pub(crate) fn decode<S: virtio_accel_core::ByteSource + ?Sized>(
184    source: &S,
185) -> Result<DecodedArtifact, DecodeError> {
186    if source.len() > MAX_ARTIFACT_BYTES {
187        return Err(DecodeError::ResourceLimit);
188    }
189    if source.len() < HEADER_BYTES as u64 {
190        return Err(DecodeError::Invalid);
191    }
192
193    let mut header = [0; HEADER_BYTES];
194    source
195        .read_at(0, &mut header)
196        .map_err(|_| DecodeError::OutOfBounds)?;
197    if header[..4] != MAGIC
198        || header[4] != VERSION_MAJOR
199        || header[5] != VERSION_MINOR
200        || header[6..8] != [0, 0]
201        || header[12..16] != [0, 0, 0, 0]
202    {
203        return Err(DecodeError::Invalid);
204    }
205    let path_len = u16::from_le_bytes([header[8], header[9]]) as usize;
206    let mapping_count = u16::from_le_bytes([header[10], header[11]]) as usize;
207    if path_len == 0 || path_len > MAX_MODEL_PATH_BYTES || mapping_count > MAX_MAPPINGS {
208        return Err(DecodeError::Invalid);
209    }
210
211    let mut cursor = HEADER_BYTES as u64;
212    let mut path = vec![0; path_len];
213    source
214        .read_at(cursor, &mut path)
215        .map_err(|_| DecodeError::OutOfBounds)?;
216    cursor += path_len as u64;
217    let model_path = String::from_utf8(path).map_err(|_| DecodeError::Invalid)?;
218
219    let mut mappings = Vec::new();
220    mappings
221        .try_reserve_exact(mapping_count)
222        .map_err(|_| DecodeError::ResourceLimit)?;
223    let mut features = BTreeSet::new();
224    for _ in 0..mapping_count {
225        let mut entry = [0; ENTRY_HEADER_BYTES];
226        source
227            .read_at(cursor, &mut entry)
228            .map_err(|_| DecodeError::OutOfBounds)?;
229        cursor += ENTRY_HEADER_BYTES as u64;
230        if entry[5] != 0 {
231            return Err(DecodeError::Invalid);
232        }
233        let slot = u32::from_le_bytes(entry[..4].try_into().unwrap());
234        let role = FeatureRole::from_wire(entry[4])?;
235        let name_len = u16::from_le_bytes([entry[6], entry[7]]) as usize;
236        if name_len == 0 || name_len > MAX_FEATURE_NAME_BYTES {
237            return Err(DecodeError::Invalid);
238        }
239        let mut name = vec![0; name_len];
240        source
241            .read_at(cursor, &mut name)
242            .map_err(|_| DecodeError::OutOfBounds)?;
243        cursor += name_len as u64;
244        let name = String::from_utf8(name).map_err(|_| DecodeError::Invalid)?;
245        if !features.insert((role as u8, name.clone())) {
246            return Err(DecodeError::Invalid);
247        }
248        mappings.push(FeatureMapping { slot, role, name });
249    }
250    if cursor != source.len() || mappings.is_empty() {
251        return Err(DecodeError::Invalid);
252    }
253    Ok(DecodedArtifact {
254        model_path,
255        mappings,
256    })
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    #[test]
264    fn path_artifact_round_trips() {
265        let artifact = CoreMlArtifact::new("models/twice.mlmodel")
266            .unwrap()
267            .map_input(7, "x")
268            .unwrap()
269            .map_output(7, "y")
270            .unwrap();
271        let bytes = artifact.encode().unwrap();
272        let decoded = decode(bytes.as_slice()).unwrap();
273        assert_eq!(decoded.model_path, "models/twice.mlmodel");
274        assert_eq!(decoded.mappings, artifact.mappings);
275    }
276
277    #[test]
278    fn malformed_envelopes_are_rejected() {
279        assert_eq!(
280            CoreMlArtifact::new("model.mlmodel").unwrap().encode(),
281            Err(ArtifactBuildError::EmptyMappings)
282        );
283        let artifact = CoreMlArtifact::new("model.mlmodel")
284            .unwrap()
285            .map_input(0, "x")
286            .unwrap();
287        let mut bytes = artifact.encode().unwrap();
288        bytes[6] = 1;
289        assert_eq!(decode(bytes.as_slice()).unwrap_err(), DecodeError::Invalid);
290
291        let duplicate = CoreMlArtifact {
292            model_path: "model.mlmodel".into(),
293            mappings: vec![
294                FeatureMapping {
295                    slot: 0,
296                    role: FeatureRole::Input,
297                    name: "x".into(),
298                },
299                FeatureMapping {
300                    slot: 1,
301                    role: FeatureRole::Input,
302                    name: "x".into(),
303                },
304            ],
305        }
306        .encode()
307        .unwrap();
308        assert_eq!(
309            decode(duplicate.as_slice()).unwrap_err(),
310            DecodeError::Invalid
311        );
312    }
313}
virtio_accel_coreml/artifact.rs

virtio_accel_coreml/
artifact.rs
