cc_lb_runtime_protocol/
identity.rs

1use cc_lb_plugin_wire::identity::{
2    CC_LB_PLUGIN_MAGIC, CC_LB_PLUGIN_SECTION_NAME, IdentityError, PluginIdentity,
3};
4use cc_lb_plugin_wire::limits;
5use thiserror::Error;
6use wasmparser::{Parser, Payload};
7
8pub fn read_identity(wasm_bytes: &[u8]) -> Result<PluginIdentity, IdentityReadError> {
9    let mut identity = None;
10
11    for payload in Parser::new(0).parse_all(wasm_bytes) {
12        let payload = payload?;
13        let Payload::CustomSection(section) = payload else {
14            continue;
15        };
16        if section.name() != CC_LB_PLUGIN_SECTION_NAME {
17            continue;
18        }
19        if identity.is_some() {
20            return Err(IdentityReadError::DuplicateCustomSection);
21        }
22        identity = Some(read_section_payload(section.data())?);
23    }
24
25    identity.ok_or(IdentityReadError::MissingCustomSection)
26}
27
28fn read_section_payload(data: &[u8]) -> Result<PluginIdentity, IdentityReadError> {
29    if data.len() > limits::CUSTOM_SECTION_MAX_SIZE {
30        return Err(IdentityReadError::SectionTooLarge { size: data.len() });
31    }
32
33    let identity: PluginIdentity = serde_json::from_slice(data)?;
34    identity.validate().map_err(|error| match error {
35        IdentityError::MagicMismatch => IdentityReadError::MagicMismatch {
36            expected: CC_LB_PLUGIN_MAGIC,
37            found: identity.magic,
38        },
39        error => IdentityReadError::Validation(error),
40    })?;
41    Ok(identity)
42}
43
44#[non_exhaustive]
45#[derive(Debug, Error)]
46pub enum IdentityReadError {
47    #[error("invalid wasm: {0}")]
48    WasmParseError(#[from] wasmparser::BinaryReaderError),
49    #[error("missing {CC_LB_PLUGIN_SECTION_NAME} custom section")]
50    MissingCustomSection,
51    #[error("duplicate {CC_LB_PLUGIN_SECTION_NAME} custom section")]
52    DuplicateCustomSection,
53    #[error("{CC_LB_PLUGIN_SECTION_NAME} custom section is {size} bytes, max {max}", max = limits::CUSTOM_SECTION_MAX_SIZE)]
54    SectionTooLarge { size: usize },
55    #[error("magic mismatch in {CC_LB_PLUGIN_SECTION_NAME} custom section")]
56    MagicMismatch { expected: [u8; 8], found: [u8; 8] },
57    #[error("malformed {CC_LB_PLUGIN_SECTION_NAME} payload: {0}")]
58    MalformedPayload(#[from] serde_json::Error),
59    #[error("invalid plugin identity: {0}")]
60    Validation(IdentityError),
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use cc_lb_plugin_wire::{identity::CC_LB_PLUGIN_MAGIC, limits};
67    use serde_json::{Value, json};
68
69    #[test]
70    fn reads_identity_from_custom_section() {
71        let wasm = wasm_with_custom_sections(&[(
72            CC_LB_PLUGIN_SECTION_NAME,
73            valid_identity_payload().as_bytes(),
74        )]);
75
76        let identity = read_identity(&wasm).expect("identity is read");
77
78        assert_eq!(identity.magic, CC_LB_PLUGIN_MAGIC);
79        assert_eq!(identity.abi_envelope, 1);
80        assert_eq!(identity.plugin_name, "test-plugin");
81        assert_eq!(identity.plugin_version, "1.0.0");
82    }
83
84    #[test]
85    fn rejects_missing_identity_section() {
86        let wasm = wasm_with_custom_sections(&[]);
87
88        let error = read_identity(&wasm).expect_err("identity section is required");
89
90        assert!(matches!(error, IdentityReadError::MissingCustomSection));
91    }
92
93    #[test]
94    fn rejects_duplicate_identity_sections() {
95        let payload = valid_identity_payload();
96        let wasm = wasm_with_custom_sections(&[
97            (CC_LB_PLUGIN_SECTION_NAME, payload.as_bytes()),
98            (CC_LB_PLUGIN_SECTION_NAME, payload.as_bytes()),
99        ]);
100
101        let error = read_identity(&wasm).expect_err("duplicate section is rejected");
102
103        assert!(matches!(error, IdentityReadError::DuplicateCustomSection));
104    }
105
106    #[test]
107    fn rejects_section_over_size_limit() {
108        let oversized = vec![b' '; limits::CUSTOM_SECTION_MAX_SIZE + 1];
109        let wasm = wasm_with_custom_sections(&[(CC_LB_PLUGIN_SECTION_NAME, &oversized)]);
110
111        let error = read_identity(&wasm).expect_err("oversized section is rejected");
112
113        assert!(matches!(
114            error,
115            IdentityReadError::SectionTooLarge {
116                size,
117            } if size == limits::CUSTOM_SECTION_MAX_SIZE + 1
118        ));
119    }
120
121    #[test]
122    fn rejects_bad_eight_byte_magic() {
123        let payload = identity_payload(json!([0, 0, 0, 0, 0, 0, 0, 0]));
124        let wasm = wasm_with_custom_sections(&[(CC_LB_PLUGIN_SECTION_NAME, payload.as_bytes())]);
125
126        let error = read_identity(&wasm).expect_err("bad magic is rejected");
127
128        match error {
129            IdentityReadError::MagicMismatch { expected, found } => {
130                assert_eq!(expected, CC_LB_PLUGIN_MAGIC);
131                assert_eq!(found, [0; 8]);
132            }
133            other => panic!("expected magic mismatch, got {other:?}"),
134        }
135    }
136
137    #[test]
138    fn rejects_four_byte_ascii_magic() {
139        let payload = identity_payload(json!([67, 67, 76, 66]));
140        let wasm = wasm_with_custom_sections(&[(CC_LB_PLUGIN_SECTION_NAME, payload.as_bytes())]);
141
142        let error = read_identity(&wasm).expect_err("4-byte magic is rejected");
143
144        assert!(matches!(error, IdentityReadError::MalformedPayload(_)));
145    }
146
147    #[test]
148    fn rejects_extra_identity_field() {
149        let payload = json!({
150            "magic": CC_LB_PLUGIN_MAGIC,
151            "abi_envelope": 1,
152            "plugin_name": "test-plugin",
153            "plugin_version": "1.0.0",
154            "extra": true,
155        })
156        .to_string();
157        let wasm = wasm_with_custom_sections(&[(CC_LB_PLUGIN_SECTION_NAME, payload.as_bytes())]);
158
159        let error = read_identity(&wasm).expect_err("extra field is rejected");
160
161        assert!(matches!(error, IdentityReadError::MalformedPayload(_)));
162    }
163
164    #[test]
165    fn rejects_invalid_wasm_bytes() {
166        let error = read_identity(b"not wasm").expect_err("invalid wasm is rejected");
167
168        assert!(matches!(error, IdentityReadError::WasmParseError(_)));
169    }
170
171    fn valid_identity_payload() -> String {
172        identity_payload(json!(CC_LB_PLUGIN_MAGIC))
173    }
174
175    fn identity_payload(magic: Value) -> String {
176        json!({
177            "magic": magic,
178            "abi_envelope": 1,
179            "plugin_name": "test-plugin",
180            "plugin_version": "1.0.0",
181        })
182        .to_string()
183    }
184
185    fn wasm_with_custom_sections(sections: &[(&str, &[u8])]) -> Vec<u8> {
186        let mut wasm = Vec::from([0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00]);
187        for (name, data) in sections {
188            wasm.push(0);
189            let mut payload = Vec::new();
190            encode_u32(name.len() as u32, &mut payload);
191            payload.extend_from_slice(name.as_bytes());
192            payload.extend_from_slice(data);
193            encode_u32(payload.len() as u32, &mut wasm);
194            wasm.extend_from_slice(&payload);
195        }
196        wasm
197    }
198
199    fn encode_u32(mut value: u32, output: &mut Vec<u8>) {
200        loop {
201            let mut byte = (value & 0x7f) as u8;
202            value >>= 7;
203            if value != 0 {
204                byte |= 0x80;
205            }
206            output.push(byte);
207            if value == 0 {
208                break;
209            }
210        }
211    }
212}
cc_lb_runtime_protocol/identity.rs

cc_lb_runtime_protocol/
identity.rs
