sozu_lib/protocol/mux/hpack/
decoder.rs1use std::borrow::Cow;
11
12use super::{
13 DecoderError, decode_integer, huffman,
14 table::{DEFAULT_MAX_SIZE, DYNAMIC_TABLE_START, DynamicTable, STATIC_TABLE},
15};
16
17const MAX_LEADING_SIZE_UPDATES: usize = 2;
20
21#[derive(Clone, Copy)]
23enum Source {
24 Input(usize, usize),
26 Scratch(usize, usize),
28 Static(&'static [u8]),
30}
31
32#[derive(Clone, Copy)]
34enum Name {
35 Source(Source),
36 Dynamic(usize),
39}
40
41#[derive(Debug)]
43pub struct Decoder {
44 table: DynamicTable,
45 max_allowed_size: usize,
48 scratch: Vec<u8>,
51}
52
53impl Default for Decoder {
54 fn default() -> Self {
55 Self::new()
56 }
57}
58
59impl Decoder {
60 pub const fn new() -> Self {
63 Decoder {
64 table: DynamicTable::new(DEFAULT_MAX_SIZE),
65 max_allowed_size: usize::MAX,
66 scratch: Vec::new(),
67 }
68 }
69
70 pub fn set_max_allowed_table_size(&mut self, max_allowed_size: usize) {
73 self.max_allowed_size = max_allowed_size;
74 }
75
76 pub fn set_max_table_size(&mut self, max_size: usize) {
79 self.table.set_max_size(max_size);
80 }
81
82 #[cfg(test)]
83 pub(super) fn table(&self) -> &DynamicTable {
84 &self.table
85 }
86
87 pub fn shrink_scratch(&mut self, retain: usize) {
89 if self.scratch.capacity() > retain {
90 self.scratch = Vec::new();
91 }
92 }
93
94 pub fn decode_with_cb<F>(&mut self, block: &[u8], mut callback: F) -> Result<(), DecoderError>
102 where
103 F: FnMut(Cow<'_, [u8]>, Cow<'_, [u8]>),
104 {
105 let mut position = 0;
106 let mut fields = 0usize;
107 let mut leading_size_updates = 0;
108 let mut last_was_size_update = false;
109 while position < block.len() {
110 let first = block[position];
111 last_was_size_update = false;
112 if first & 0x80 != 0 {
113 let (index, consumed) = decode_integer(&block[position..], 7)?;
115 let (name, value) = self.lookup(index)?;
116 callback(Cow::Borrowed(name), Cow::Borrowed(value));
117 position += consumed;
118 fields += 1;
119 } else if first & 0x40 != 0 {
120 position = self.literal(block, position, 6, true, &mut callback)?;
122 fields += 1;
123 } else if first & 0x20 != 0 {
124 if fields > 0 {
126 return Err(DecoderError::SizeUpdateNotAtStart);
127 }
128 leading_size_updates += 1;
129 if leading_size_updates > MAX_LEADING_SIZE_UPDATES {
130 return Err(DecoderError::TooManySizeUpdates);
131 }
132 let (size, consumed) = decode_integer(&block[position..], 5)?;
133 if size > self.max_allowed_size {
134 return Err(DecoderError::TableSizeTooLarge);
135 }
136 self.table.set_max_size(size);
137 position += consumed;
138 last_was_size_update = true;
139 } else {
140 position = self.literal(block, position, 4, false, &mut callback)?;
143 fields += 1;
144 }
145 }
146 if last_was_size_update {
147 return Err(DecoderError::SizeUpdateAtEnd);
148 }
149 Ok(())
150 }
151
152 fn lookup(&self, index: usize) -> Result<(&[u8], &[u8]), DecoderError> {
154 match index {
155 0 => Err(DecoderError::InvalidIndex),
156 1..DYNAMIC_TABLE_START => Ok(STATIC_TABLE[index - 1]),
157 _ => self
158 .table
159 .get(index - DYNAMIC_TABLE_START)
160 .ok_or(DecoderError::InvalidIndex),
161 }
162 }
163
164 fn literal<F>(
167 &mut self,
168 block: &[u8],
169 position: usize,
170 prefix: u8,
171 indexing: bool,
172 callback: &mut F,
173 ) -> Result<usize, DecoderError>
174 where
175 F: FnMut(Cow<'_, [u8]>, Cow<'_, [u8]>),
176 {
177 self.scratch.clear();
178 let (index, consumed) = decode_integer(&block[position..], prefix)?;
179 let mut position = position + consumed;
180 let name = match index {
181 0 => {
182 let (name, next) = decode_string(block, position, &mut self.scratch)?;
183 position = next;
184 Name::Source(name)
185 }
186 1..DYNAMIC_TABLE_START => Name::Source(Source::Static(STATIC_TABLE[index - 1].0)),
187 _ => {
188 let dynamic = index - DYNAMIC_TABLE_START;
189 let (name, _) = self.table.get(dynamic).ok_or(DecoderError::InvalidIndex)?;
190 if indexing {
191 self.scratch.extend_from_slice(name);
193 Name::Source(Source::Scratch(0, self.scratch.len()))
194 } else {
195 Name::Dynamic(dynamic)
196 }
197 }
198 };
199 let (value, position) = decode_string(block, position, &mut self.scratch)?;
200 let value = resolve(value, block, &self.scratch);
201 match name {
202 Name::Source(name) => {
203 let name = resolve(name, block, &self.scratch);
204 callback(Cow::Borrowed(name), Cow::Borrowed(value));
205 if indexing {
206 self.table.insert(name, value);
207 }
208 }
209 Name::Dynamic(index) => {
210 let name = self
211 .table
212 .get(index)
213 .map(|(name, _)| name)
214 .ok_or(DecoderError::InvalidIndex)?;
215 callback(Cow::Borrowed(name), Cow::Borrowed(value));
216 }
217 }
218 Ok(position)
219 }
220}
221
222fn decode_string(
225 block: &[u8],
226 position: usize,
227 scratch: &mut Vec<u8>,
228) -> Result<(Source, usize), DecoderError> {
229 let rest = &block[position..];
230 let Some(&first) = rest.first() else {
231 return Err(DecoderError::Truncated);
232 };
233 let (length, consumed) = decode_integer(rest, 7)?;
234 if length > rest.len() - consumed {
235 return Err(DecoderError::Truncated);
236 }
237 let start = position + consumed;
238 let end = start + length;
239 if first & 0x80 == 0 {
240 return Ok((Source::Input(start, end), end));
241 }
242 let scratch_start = scratch.len();
243 huffman::decode(&block[start..end], scratch).map_err(DecoderError::Huffman)?;
244 Ok((Source::Scratch(scratch_start, scratch.len()), end))
245}
246
247fn resolve<'a>(source: Source, block: &'a [u8], scratch: &'a [u8]) -> &'a [u8] {
248 match source {
249 Source::Input(start, end) => &block[start..end],
250 Source::Scratch(start, end) => &scratch[start..end],
251 Source::Static(name) => name,
252 }
253}