warg_server/policy/content/
wasm.rs

1use super::{ContentPolicy, ContentPolicyError, ContentPolicyResult, ContentStreamPolicy};
2use warg_crypto::hash::AnyHash;
3use wasmparser::{
4    Chunk, Encoding, FuncValidatorAllocations, Parser, ValidPayload, Validator, WasmFeatures,
5};
6
7/// A policy that ensures all uploaded content is valid WebAssembly.
8pub struct WasmContentPolicy {
9    allow_modules: bool,
10    allow_components: bool,
11}
12
13impl WasmContentPolicy {
14    /// Creates a new WebAssembly content policy.
15    pub fn new() -> Self {
16        Self::default()
17    }
18
19    /// Disallows WebAssembly modules from being acceptable content.
20    pub fn disallow_modules(mut self) -> Self {
21        self.allow_modules = false;
22        self
23    }
24
25    /// Disallows WebAssembly components from being acceptable content.
26    pub fn disallow_components(mut self) -> Self {
27        self.allow_components = false;
28        self
29    }
30}
31
32impl Default for WasmContentPolicy {
33    fn default() -> Self {
34        Self {
35            allow_modules: true,
36            allow_components: true,
37        }
38    }
39}
40
41impl ContentPolicy for WasmContentPolicy {
42    fn new_stream_policy(
43        &self,
44        _digest: &AnyHash,
45    ) -> ContentPolicyResult<Box<dyn ContentStreamPolicy>> {
46        Ok(Box::new(WasmContentStreamPolicy {
47            buffer: Vec::new(),
48            parser: Parser::new(0),
49            stack: Vec::new(),
50            validator: wasmparser::Validator::new_with_features(WasmFeatures::default()),
51            allocs: FuncValidatorAllocations::default(),
52            allow_modules: self.allow_modules,
53            allow_components: self.allow_components,
54        }))
55    }
56}
57
58struct WasmContentStreamPolicy {
59    buffer: Vec<u8>,
60    parser: Parser,
61    stack: Vec<Parser>,
62    validator: Validator,
63    allocs: FuncValidatorAllocations,
64    allow_modules: bool,
65    allow_components: bool,
66}
67
68impl WasmContentStreamPolicy {
69    fn process(&mut self, bytes: &[u8], eof: bool) -> ContentPolicyResult<()> {
70        // Extend the buffer if we need to; otherwise, parse the given slice
71        let buf = if !self.buffer.is_empty() {
72            self.buffer.extend(bytes);
73            self.buffer.as_slice()
74        } else {
75            bytes
76        };
77
78        let mut offset = 0;
79        loop {
80            let (payload, consumed) = match self.parser.parse(&buf[offset..], eof).map_err(|e| {
81                ContentPolicyError::Rejection(format!("content is not valid WebAssembly: {e}"))
82            })? {
83                Chunk::NeedMoreData(_) => {
84                    // If the buffer is empty and there's still data in the given slice,
85                    // copy the remaining data to the buffer.
86                    // If there's still data remaining in the buffer, copy it to the
87                    // beginning of the buffer and truncate it.
88                    // Otherwise, clear the buffer.
89                    if self.buffer.is_empty() && offset < bytes.len() {
90                        self.buffer.extend_from_slice(&bytes[offset..]);
91                    } else if offset < self.buffer.len() {
92                        self.buffer.copy_within(offset.., 0);
93                        self.buffer.truncate(self.buffer.len() - offset);
94                    } else {
95                        self.buffer.clear();
96                    }
97                    return Ok(());
98                }
99
100                Chunk::Parsed { consumed, payload } => (payload, consumed),
101            };
102
103            offset += consumed;
104
105            match &payload {
106                wasmparser::Payload::Version {
107                    encoding: Encoding::Module,
108                    ..
109                } if !self.allow_modules => {
110                    return Err(ContentPolicyError::Rejection(
111                        "WebAssembly modules are not allowed".to_string(),
112                    ))
113                }
114                wasmparser::Payload::Version {
115                    encoding: Encoding::Component,
116                    ..
117                } if !self.allow_components => {
118                    return Err(ContentPolicyError::Rejection(
119                        "WebAssembly components are not allowed".to_string(),
120                    ))
121                }
122                _ => {}
123            }
124
125            match self.validator.payload(&payload).map_err(|e| {
126                ContentPolicyError::Rejection(format!("content is not valid WebAssembly: {e}"))
127            })? {
128                ValidPayload::Ok => {}
129                ValidPayload::Parser(p) => {
130                    self.stack.push(self.parser.clone());
131                    self.parser = p;
132                }
133                ValidPayload::Func(func, body) => {
134                    let allocs = std::mem::take(&mut self.allocs);
135                    let mut validator = func.into_validator(allocs);
136                    validator.validate(&body).map_err(|e| {
137                        ContentPolicyError::Rejection(format!(
138                            "content is not valid WebAssembly: {e}"
139                        ))
140                    })?;
141                    self.allocs = validator.into_allocations();
142                }
143                ValidPayload::End(_) => {
144                    if let Some(parser) = self.stack.pop() {
145                        self.parser = parser;
146                    } else {
147                        return Ok(());
148                    }
149                }
150            }
151        }
152    }
153}
154
155impl ContentStreamPolicy for WasmContentStreamPolicy {
156    fn check(&mut self, bytes: &[u8]) -> ContentPolicyResult<()> {
157        self.process(bytes, false)
158    }
159
160    fn finalize(&mut self) -> ContentPolicyResult<()> {
161        self.process(&[], true)
162    }
163}
warg_server/policy/content/wasm.rs

warg_server/policy/content/
wasm.rs
