rust_http2/
frame.rs

1use hpack::Encoder;
2
3pub enum FrameType {
4    DATA = 0x00,
5    HEADERS = 0x01,
6    PRIORITY = 0x02,
7    RSTSTREAM = 0x03,
8    SETTINGS = 0x04,
9    PUSHPROMISE = 0x05,
10    PING = 0x06,
11    GOAWAY = 0x07,
12    WINDOWUPDATE = 0x08,
13    CONTINUATION = 0x09,
14}
15
16impl TryFrom<u8> for FrameType {
17    type Error = &'static str;
18
19    fn try_from(value: u8) -> Result<Self, Self::Error> {
20        match value {
21            0x00 => Ok(FrameType::DATA),
22            0x01 => Ok(FrameType::HEADERS),
23            0x02 => Ok(FrameType::PRIORITY),
24            0x03 => Ok(FrameType::RSTSTREAM),
25            0x04 => Ok(FrameType::SETTINGS),
26            0x05 => Ok(FrameType::PUSHPROMISE),
27            0x06 => Ok(FrameType::PING),
28            0x07 => Ok(FrameType::GOAWAY),
29            0x08 => Ok(FrameType::WINDOWUPDATE),
30            0x09 => Ok(FrameType::CONTINUATION),
31            _ => Err("Invalid frame type"),
32        }
33    }
34}
35
36impl TryInto<u8> for FrameType {
37    type Error = &'static str;
38
39    fn try_into(self) -> Result<u8, Self::Error> {
40        match self {
41            FrameType::DATA => Ok(0x00),
42            FrameType::HEADERS => Ok(0x01),
43            FrameType::PRIORITY => Ok(0x02),
44            FrameType::RSTSTREAM => Ok(0x03),
45            FrameType::SETTINGS => Ok(0x04),
46            FrameType::PUSHPROMISE => Ok(0x05),
47            FrameType::PING => Ok(0x06),
48            FrameType::GOAWAY => Ok(0x07),
49            FrameType::WINDOWUPDATE => Ok(0x08),
50            FrameType::CONTINUATION => Ok(0x09),
51        }
52    }
53}
54
55pub enum FrameFlags {
56    ENDSTREAM = 0x01,
57    ENDHEADERS = 0x04,
58    PADDED = 0x08,
59    PRIORITY = 0x20,
60    NONE = 0x00,
61}
62
63impl TryFrom<u8> for FrameFlags {
64    type Error = &'static str;
65
66    fn try_from(value: u8) -> Result<Self, Self::Error> {
67        match value {
68            0x01 => Ok(FrameFlags::ENDSTREAM),
69            0x04 => Ok(FrameFlags::ENDHEADERS),
70            0x08 => Ok(FrameFlags::PADDED),
71            0x20 => Ok(FrameFlags::PRIORITY),
72            _ => Ok(FrameFlags::NONE),
73        }
74    }
75}
76
77impl TryInto<u8> for FrameFlags {
78    type Error = &'static str;
79
80    fn try_into(self) -> Result<u8, Self::Error> {
81        match self {
82            FrameFlags::ENDSTREAM => Ok(0x01),
83            FrameFlags::ENDHEADERS => Ok(0x04),
84            FrameFlags::PADDED => Ok(0x08),
85            FrameFlags::PRIORITY => Ok(0x20),
86            FrameFlags::NONE => Ok(0x00),
87        }
88    }
89}
90
91pub struct FrameWriter {
92    pub frame_type: FrameType,
93    pub flags: FrameFlags,
94    pub stream_id: u32,
95    pub payload_len: u32,
96    pub payload: Vec<u8>,
97}
98
99impl FrameWriter {
100    pub fn new(frame_type: FrameType, flags: FrameFlags, stream_id: u32, payload: Vec<u8>) -> Self {
101        FrameWriter {
102            frame_type,
103            flags,
104            stream_id,
105            payload_len: payload.len() as u32,
106            payload,
107        }
108    }
109
110    pub fn new_frame_writer(
111        frame_type: FrameType,
112        flags: FrameFlags,
113        stream_id: u32
114    ) -> Result<Self, Box<dyn std::error::Error>> {
115        let mut headers = FrameHeaders::new();
116        headers.add_header(":method".to_string(), "GET".to_string());
117        headers.add_header(":path".to_string(), "/".to_string());
118        headers.add_header(":scheme".to_string(), "https".to_string());
119        headers.add_header(":authority".to_string(), "example.com".to_string());
120
121        let payload = headers.serialize();
122
123        // println!("[*] Serialized headers: {:?}", payload);
124
125        Ok(FrameWriter {
126            frame_type,
127            flags,
128            stream_id,
129            payload_len: payload.len() as u32,
130            payload: payload,
131        })
132    }
133
134    // pub fn payload_from_frame_type(
135    //     frame_type: FrameType
136    // ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
137    //     let payload = match frame_type {
138    //         FrameType::HEADERS => {
139    //             let mut headers = FrameHeaders::new();
140    //             headers.add_header(":method".to_string(), "GET".to_string());
141    //             headers.add_header(":path".to_string(), "/".to_string());
142    //             headers.add_header(":scheme".to_string(), "https".to_string());
143    //             headers.add_header(":authority".to_string(), "example.com".to_string());
144
145    //             headers.serialize()
146    //         }
147
148    //         _ => Vec::new(),
149    //     };
150
151    //     Ok(payload)
152    // }
153
154    pub fn serialize(self) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
155        let mut serialized = Vec::with_capacity(9 + (self.payload_len as usize));
156
157        serialized.extend_from_slice(&self.payload_len.to_be_bytes()[1..]);
158        serialized.push(FrameType::try_into(self.frame_type)?);
159        serialized.push(FrameFlags::try_into(self.flags)?);
160        serialized.extend_from_slice(
161            &[
162                (self.stream_id >> 24) as u8,
163                (self.stream_id >> 16) as u8,
164                (self.stream_id >> 8) as u8,
165                self.stream_id as u8,
166            ]
167        );
168        // serialized.extend_from_slice(&self.stream_id.to_be_bytes()[1..]);
169        serialized.extend_from_slice(&self.payload);
170
171        Ok(serialized)
172    }
173}
174
175pub struct FrameReader {
176    pub frame_type: FrameType,
177    pub flags: FrameFlags,
178    pub stream_id: u32,
179    pub payload_len: u32,
180    pub payload: Vec<u8>,
181}
182
183impl FrameReader {
184    // pub async fn parse_frame(buf: &[u8]) -> Result<FrameReader, Box<dyn std::error::Error>> {
185    //     FrameReader::deserialize(buf).await
186    // }
187
188    pub async fn read_frame(buf: &[u8]) -> Result<Self, Box<dyn std::error::Error>> {
189        let payload_len = ((buf[0] as u32) << 16) | ((buf[1] as u32) << 8) | (buf[2] as u32);
190
191        if buf.len() < 9 + (payload_len as usize) {
192            return Err("Buffer too small".into());
193        }
194
195        let frame_type = FrameType::try_from(buf[3])?;
196        let flags = FrameFlags::try_from(buf[4])?;
197        let stream_id = u32::from_be_bytes([buf[5], buf[6], buf[7], buf[8]]) & 0x7fff_ffff;
198        let payload = buf[9..9 + (payload_len as usize)].to_vec();
199
200        Ok(FrameReader {
201            frame_type,
202            flags,
203            stream_id,
204            payload_len,
205            payload,
206        })
207    }
208
209    pub async fn parse_settings_payload(&self) -> Vec<Setting> {
210        let payload = &self.payload;
211
212        if self.payload_len == 0 {
213            return Vec::new();
214        }
215
216        let mut i = 0;
217        let mut settings = Vec::new();
218        while i < payload.len() {
219            if i + 6 > payload.len() {
220                break;
221            }
222
223            let id = u16::from_be_bytes([payload[i], payload[i + 1]]);
224            let value = u32::from_be_bytes([
225                payload[i + 2],
226                payload[i + 3],
227                payload[i + 4],
228                payload[i + 5],
229            ]);
230
231            settings.push(Setting::new(id, value));
232
233            i += 6;
234        }
235
236        settings
237    }
238}
239
240pub struct HeadersBuilder {
241    headers: Vec<(String, String)>,
242}
243
244impl HeadersBuilder {
245    pub fn new() -> Self {
246        HeadersBuilder { headers: Vec::new() }
247    }
248
249    pub fn add_header(mut self, name: String, value: String) -> Self {
250        self.headers.push((name, value));
251        self
252    }
253
254    pub fn build(
255        self,
256        stream_id: u32
257        // flags: FrameFlags
258    ) -> Result<FrameWriter, Box<dyn std::error::Error>> {
259        let headers = FrameHeaders::from_pairs(self.headers);
260        let payload = headers.serialize();
261
262        Ok(FrameWriter::new(FrameType::HEADERS, FrameFlags::ENDHEADERS, stream_id, payload))
263    }
264}
265
266pub struct FrameHeaders {
267    headers: Vec<(String, String)>,
268}
269
270impl FrameHeaders {
271    pub fn new() -> Self {
272        FrameHeaders { headers: Vec::new() }
273    }
274
275    pub fn from_pairs(headers: Vec<(String, String)>) -> Self {
276        FrameHeaders { headers }
277    }
278
279    pub fn add_header(&mut self, name: String, value: String) {
280        self.headers.push((name, value));
281    }
282
283    pub fn serialize(&self) -> Vec<u8> {
284        let mut encoder = Encoder::new();
285        let headers: Vec<(&[u8], &[u8])> = self.headers
286            .iter()
287            .map(|(name, value)| (name.as_bytes(), value.as_bytes()))
288            .collect();
289        encoder.encode(headers)
290    }
291}
292
293pub struct Setting {
294    pub id: u16,
295    pub value: u32,
296}
297
298impl Setting {
299    pub fn new(id: u16, value: u32) -> Self {
300        Setting { id, value }
301    }
302}
303
304pub struct Settings {
305    pub max_concurrent_streams: Setting,
306    pub initial_window_size: Setting,
307    pub max_frame_size: Setting,
308    pub max_header_list_size: Setting,
309}
310
311impl Settings {
312    pub async fn new() -> Self {
313        Settings {
314            max_concurrent_streams: Setting::new(0x03, 100),
315            initial_window_size: Setting::new(0x04, 65535),
316            max_frame_size: Setting::new(0x05, 16384),
317            max_header_list_size: Setting::new(0x06, 262144),
318        }
319    }
320
321    pub async fn serialize(&self) -> Vec<u8> {
322        let mut payload = Vec::new();
323
324        for setting in [
325            &self.max_concurrent_streams,
326            &self.initial_window_size,
327            &self.max_frame_size,
328            &self.max_header_list_size,
329        ] {
330            payload.extend_from_slice(&setting.id.to_be_bytes());
331            payload.extend_from_slice(&setting.value.to_be_bytes());
332        }
333
334        let mut serialized = Vec::with_capacity(9 + payload.len());
335
336        serialized.extend_from_slice(&(payload.len() as u32).to_be_bytes()[1..]);
337        serialized.extend_from_slice(&[4, 0]);
338        serialized.extend_from_slice(&(0u32).to_be_bytes());
339        // serialized.extend_from_slice(&[(1 >> 24) as u8, (1 >> 16) as u8, (1 >> 8) as u8, 1 as u8]);
340
341        serialized.extend_from_slice(&payload);
342
343        serialized
344    }
345}
rust_http2/frame.rs

rust_http2/
frame.rs
