1 // Copyright (c) 2023 Huawei Device Co., Ltd.
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 //     http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13 
14 use std::cmp::min;
15 
16 use crate::h2::frame::{FrameFlags, FrameType, Payload, Priority, Setting};
17 use crate::h2::{Frame, Goaway, HpackEncoder, Settings};
18 
19 // TODO: Classify encoder errors per RFC specifications into categories like
20 // stream or connection errors. Identify specific error types such as
21 // Frame_size_error/Protocol Error.
22 const DEFAULT_MAX_FRAME_SIZE: usize = 16384;
23 
24 const DEFAULT_HEADER_TABLE_SIZE: usize = 4096;
25 
26 #[derive(Debug)]
27 pub enum FrameEncoderErr {
28     EncodingData,
29     UnexpectedPayloadType,
30     NoCurrentFrame,
31     InternalError,
32     HeadersNotEnd,
33 }
34 
35 #[derive(PartialEq, Debug)]
36 enum FrameEncoderState {
37     // The initial state for the frame encoder.
38     Idle,
39     // The initial state for encoding the HEADERS frame, including the frame header and the Field
40     // Block Fragment.
41     EncodingHeadersFrame,
42     // The state for encoding the payload of the HEADERS frame, including the header block
43     // fragment.
44     EncodingHeadersPayload,
45     // The state for encoding the padding octets for the HEADERS frame, if the PADDED flag is set.
46     EncodingHeadersPadding,
47     // TODO compare to max_header_list_size
48     // The state for encoding CONTINUATION frames if the header block exceeds the max_frame_size.
49     EncodingContinuationFrames,
50     // The final state, indicating that the HEADERS frame and any necessary CONTINUATION frames
51     // have been fully encoded.
52     HeadersComplete,
53     // The initial state for encoding the DATA frame, including the frame header and the Pad
54     // Length field (if PADDED flag is set).
55     EncodingDataHeader,
56     // The state for encoding the actual data payload of the DATA frame.
57     EncodingDataPayload,
58     // The state for encoding the padding octets for the DATA frame, if the PADDED flag is set.
59     EncodingDataPadding,
60     // The initial state for encoding the SETTINGS frame, including the frame header.
61     EncodingSettingsFrame,
62     // The state for encoding the SETTINGS frame payload.
63     EncodingSettingsPayload,
64     // The initial state for encoding the GOAWAY frame, including the frame header.
65     EncodingGoawayFrame,
66     // The state for encoding the GOAWAY frame payload.
67     EncodingGoawayPayload,
68     // The initial state for encoding the WINDOW_UPDATE frame, including the frame header.
69     EncodingWindowUpdateFrame,
70     // The state for encoding the WINDOW_UPDATE frame payload.
71     EncodingWindowUpdatePayload,
72     // The initial state for encoding the PRIORITY frame, including the frame header.
73     EncodingPriorityFrame,
74     // The state for encoding Priority frame payload.
75     EncodingPriorityPayload,
76     // The initial state for encoding the RST_STREAM frame, including the frame header.
77     EncodingRstStreamFrame,
78     // The state for encoding the RST_STREAM frame payload.
79     EncodingRstStreamPayload,
80     // The initial state for encoding the PING frame, including the frame header.
81     EncodingPingFrame,
82     // The state for encoding the PING frame payload.
83     EncodingPingPayload,
84     // The final state, indicating that the DATA frame has been fully encoded.
85     DataComplete,
86 }
87 
88 /// A structure for encoding frames into bytes, supporting the serialization of
89 /// HTTP/2 Frames. It maintains the state of the current frame being encoded and
90 /// also handles the fragmentation of frames.
91 pub struct FrameEncoder {
92     current_frame: Option<Frame>,
93     // `max_frame_size` is actually useless in the Encoder for headers frame and continuation
94     // frame, because the frame length does not exceed the length of the
95     // `header_payload_buffer`
96     max_frame_size: usize,
97     max_header_list_size: usize,
98     hpack_encoder: HpackEncoder,
99     state: FrameEncoderState,
100     encoded_bytes: usize,
101     data_offset: usize,
102     // `remaining_header_payload` will always be smaller than the minimum max_frame_size,
103     // because the `header_payload_buffer` length is the minimum max_frame_size.
104     remaining_header_payload: usize,
105     remaining_payload_bytes: usize,
106     is_end_stream: bool,
107     is_end_headers: bool,
108     header_payload_buffer: Vec<u8>,
109     header_payload_index: usize,
110 }
111 
112 impl FrameEncoder {
113     /// Constructs a new `FrameEncoder` with specified maximum frame size and
114     /// maximum header list size.
new(max_frame_size: usize, use_huffman: bool) -> Self115     pub fn new(max_frame_size: usize, use_huffman: bool) -> Self {
116         FrameEncoder {
117             current_frame: None,
118             max_frame_size,
119             max_header_list_size: usize::MAX,
120             hpack_encoder: HpackEncoder::new(DEFAULT_HEADER_TABLE_SIZE, use_huffman),
121             state: FrameEncoderState::Idle,
122             encoded_bytes: 0,
123             data_offset: 0,
124             remaining_header_payload: 0,
125             remaining_payload_bytes: 0,
126             is_end_stream: false,
127             is_end_headers: false,
128             // Initialized to default max frame size(16384).
129             header_payload_buffer: vec![0; DEFAULT_MAX_FRAME_SIZE],
130             header_payload_index: 0,
131         }
132     }
133 
134     /// Sets the current frame to be encoded by the `FrameEncoder`. The state of
135     /// the encoder is updated based on the payload type of the frame.
set_frame(&mut self, frame: Frame) -> Result<(), FrameEncoderErr>136     pub fn set_frame(&mut self, frame: Frame) -> Result<(), FrameEncoderErr> {
137         self.is_end_stream = frame.flags().is_end_stream();
138         self.is_end_headers = frame.flags().is_end_headers();
139         self.current_frame = Some(frame);
140         // Reset the encoded bytes counter
141         self.encoded_bytes = 0;
142 
143         // Set the initial state based on the frame payload type
144         match &self.current_frame {
145             Some(frame) => match frame.payload() {
146                 Payload::Headers(headers) => {
147                     if !self.is_end_headers {
148                         return Err(FrameEncoderErr::HeadersNotEnd);
149                     }
150                     self.hpack_encoder.set_parts(headers.get_parts());
151                     self.header_payload_index = 0;
152                     // TODO: Handle potential scenario where HPACK encoding may not be able to
153                     // complete output in one go.
154                     let payload_size = self.hpack_encoder.encode(&mut self.header_payload_buffer);
155                     self.remaining_header_payload = payload_size;
156                     self.state = FrameEncoderState::EncodingHeadersFrame;
157                 }
158                 Payload::Priority(_) => self.state = FrameEncoderState::EncodingPriorityFrame,
159                 Payload::RstStream(_) => self.state = FrameEncoderState::EncodingRstStreamFrame,
160                 Payload::Ping(_) => self.state = FrameEncoderState::EncodingPingFrame,
161                 Payload::Data(data) => {
162                     self.state = {
163                         self.data_offset = 0;
164                         self.remaining_payload_bytes = data.size();
165                         FrameEncoderState::EncodingDataHeader
166                     }
167                 }
168                 Payload::Settings(_) => self.state = FrameEncoderState::EncodingSettingsFrame,
169                 Payload::Goaway(_) => self.state = FrameEncoderState::EncodingGoawayFrame,
170                 Payload::WindowUpdate(_) => {
171                     self.state = FrameEncoderState::EncodingWindowUpdateFrame
172                 }
173                 _ => {}
174             },
175             None => self.state = FrameEncoderState::Idle,
176         }
177         Ok(())
178     }
179 
180     /// Encodes the current frame into the given buffer. The state of the
181     /// encoder determines which part of the frame is currently being encoded.
182     /// This function returns the number of bytes written to the buffer or an
183     /// error if the encoding process fails.
encode(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>184     pub fn encode(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
185         let mut written_bytes = 0;
186 
187         while written_bytes < buf.len() {
188             match self.state {
189                 FrameEncoderState::Idle
190                 | FrameEncoderState::HeadersComplete
191                 | FrameEncoderState::DataComplete => {
192                     break;
193                 }
194                 FrameEncoderState::EncodingHeadersFrame => {
195                     let bytes = self.encode_headers_frame(&mut buf[written_bytes..])?;
196                     written_bytes += bytes;
197                     if self.state == FrameEncoderState::EncodingHeadersFrame {
198                         break;
199                     }
200                 }
201                 FrameEncoderState::EncodingHeadersPayload => {
202                     let bytes = self.encode_headers_payload(&mut buf[written_bytes..])?;
203                     written_bytes += bytes;
204                     if self.state == FrameEncoderState::EncodingHeadersPayload {
205                         break;
206                     }
207                 }
208                 FrameEncoderState::EncodingHeadersPadding => {
209                     let bytes = self.encode_padding(&mut buf[written_bytes..])?;
210                     written_bytes += bytes;
211                     if self.state == FrameEncoderState::EncodingHeadersPadding {
212                         break;
213                     }
214                 }
215                 FrameEncoderState::EncodingContinuationFrames => {
216                     let bytes = self.encode_continuation_frames(&mut buf[written_bytes..])?;
217                     written_bytes += bytes;
218                     if self.state == FrameEncoderState::EncodingContinuationFrames {
219                         break;
220                     }
221                 }
222                 FrameEncoderState::EncodingDataHeader => {
223                     let bytes = self.encode_data_header(&mut buf[written_bytes..])?;
224                     written_bytes += bytes;
225                     if self.state == FrameEncoderState::EncodingDataHeader {
226                         break;
227                     }
228                 }
229                 FrameEncoderState::EncodingDataPayload => {
230                     let bytes = self.encode_data_payload(&mut buf[written_bytes..])?;
231                     written_bytes += bytes;
232                     if self.state == FrameEncoderState::EncodingDataPayload {
233                         break;
234                     }
235                 }
236                 FrameEncoderState::EncodingDataPadding => {
237                     let bytes = self.encode_padding(&mut buf[written_bytes..])?;
238                     written_bytes += bytes;
239                     if self.state == FrameEncoderState::EncodingDataPadding {
240                         break;
241                     }
242                 }
243                 FrameEncoderState::EncodingSettingsFrame => {
244                     let bytes = self.encode_settings_frame(&mut buf[written_bytes..])?;
245                     written_bytes += bytes;
246                     if self.state == FrameEncoderState::EncodingSettingsFrame {
247                         break;
248                     }
249                 }
250                 FrameEncoderState::EncodingGoawayFrame => {
251                     let bytes = self.encode_goaway_frame(&mut buf[written_bytes..])?;
252                     written_bytes += bytes;
253                     if self.state == FrameEncoderState::EncodingGoawayFrame {
254                         break;
255                     }
256                 }
257                 FrameEncoderState::EncodingWindowUpdateFrame => {
258                     let bytes = self.encode_window_update_frame(&mut buf[written_bytes..])?;
259                     written_bytes += bytes;
260                     if self.state == FrameEncoderState::EncodingWindowUpdateFrame {
261                         break;
262                     }
263                 }
264                 FrameEncoderState::EncodingSettingsPayload => {
265                     let bytes = self.encode_settings_payload(&mut buf[written_bytes..])?;
266                     written_bytes += bytes;
267                     if self.state == FrameEncoderState::EncodingSettingsPayload {
268                         break;
269                     }
270                 }
271                 FrameEncoderState::EncodingGoawayPayload => {
272                     let bytes = self.encode_goaway_payload(&mut buf[written_bytes..])?;
273                     written_bytes += bytes;
274                     if self.state == FrameEncoderState::EncodingGoawayPayload {
275                         break;
276                     }
277                 }
278                 FrameEncoderState::EncodingWindowUpdatePayload => {
279                     let bytes = self.encode_window_update_payload(&mut buf[written_bytes..])?;
280                     written_bytes += bytes;
281                     if self.state == FrameEncoderState::EncodingWindowUpdatePayload {
282                         break;
283                     }
284                 }
285                 FrameEncoderState::EncodingPriorityFrame => {
286                     let bytes = self.encode_priority_frame(&mut buf[written_bytes..])?;
287                     written_bytes += bytes;
288                     if self.state == FrameEncoderState::EncodingPriorityFrame {
289                         break;
290                     }
291                 }
292                 FrameEncoderState::EncodingPriorityPayload => {
293                     let bytes = self.encode_priority_payload(&mut buf[written_bytes..])?;
294                     written_bytes += bytes;
295                     if self.state == FrameEncoderState::EncodingPriorityPayload {
296                         break;
297                     }
298                 }
299                 FrameEncoderState::EncodingRstStreamFrame => {
300                     let bytes = self.encode_rst_stream_frame(&mut buf[written_bytes..])?;
301                     written_bytes += bytes;
302                     if self.state == FrameEncoderState::EncodingRstStreamFrame {
303                         break;
304                     }
305                 }
306                 FrameEncoderState::EncodingRstStreamPayload => {
307                     let bytes = self.encode_rst_stream_payload(&mut buf[written_bytes..])?;
308                     written_bytes += bytes;
309                     if self.state == FrameEncoderState::EncodingRstStreamPayload {
310                         break;
311                     }
312                 }
313                 FrameEncoderState::EncodingPingFrame => {
314                     let bytes = self.encode_ping_frame(&mut buf[written_bytes..])?;
315                     written_bytes += bytes;
316                     if self.state == FrameEncoderState::EncodingPingFrame {
317                         break;
318                     }
319                 }
320                 FrameEncoderState::EncodingPingPayload => {
321                     let bytes = self.encode_ping_payload(&mut buf[written_bytes..])?;
322                     written_bytes += bytes;
323                     if self.state == FrameEncoderState::EncodingPingPayload {
324                         break;
325                     }
326                 }
327             }
328         }
329 
330         Ok(written_bytes)
331     }
332 
333     /// Updates the provided setting for the current frame if it is a `Settings`
334     /// frame.
update_setting(&mut self, setting: Setting)335     pub fn update_setting(&mut self, setting: Setting) {
336         if let Some(frame) = &mut self.current_frame {
337             if let Payload::Settings(settings) = frame.payload_mut() {
338                 settings.update_setting(setting);
339             }
340         }
341     }
342 
343     /// Sets the maximum frame size for the current encoder instance.
update_max_frame_size(&mut self, size: usize)344     pub fn update_max_frame_size(&mut self, size: usize) {
345         self.max_frame_size = size;
346     }
347 
348     /// Sets the maximum header table size for the current encoder instance.
349     // TODO enable update header table size.
update_header_table_size(&mut self, size: usize)350     pub fn update_header_table_size(&mut self, size: usize) {
351         self.hpack_encoder.update_max_dynamic_table_size(size)
352     }
353 
354     // TODO enable update max header list size.
update_max_header_list_size(&mut self, size: usize)355     pub(crate) fn update_max_header_list_size(&mut self, size: usize) {
356         self.max_header_list_size = size;
357     }
358 
finish_current_frame(&mut self)359     fn finish_current_frame(&mut self) {
360         self.header_payload_index = 0;
361         self.is_end_stream = false;
362         self.current_frame = None;
363         self.is_end_headers = false;
364         self.remaining_header_payload = 0;
365         self.encoded_bytes = 0;
366     }
367 
read_rest_payload(&mut self)368     fn read_rest_payload(&mut self) {
369         self.header_payload_index = 0;
370         let payload_size = self.hpack_encoder.encode(&mut self.header_payload_buffer);
371         self.remaining_header_payload = payload_size;
372     }
373 
encode_headers_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>374     fn encode_headers_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
375         if let Some(frame) = &self.current_frame {
376             if let Payload::Headers(_) = frame.payload() {
377                 // HTTP/2 frame header size is 9 bytes.
378                 let frame_header_size = 9;
379                 let remaining_header_bytes = if self.encoded_bytes >= frame_header_size {
380                     0
381                 } else {
382                     frame_header_size - self.encoded_bytes
383                 };
384                 let bytes_to_write = remaining_header_bytes.min(buf.len());
385 
386                 self.iterate_headers_header(frame, buf, bytes_to_write)?;
387                 self.encoded_bytes += bytes_to_write;
388                 let bytes_written = bytes_to_write;
389                 let mut payload_bytes_written = 0;
390 
391                 if self.encoded_bytes >= frame_header_size {
392                     payload_bytes_written = self
393                         .write_payload(&mut buf[bytes_written..], self.remaining_header_payload);
394                     self.encoded_bytes += payload_bytes_written;
395                     self.headers_header_status();
396                 }
397 
398                 Ok(bytes_written + payload_bytes_written)
399             } else {
400                 Err(FrameEncoderErr::UnexpectedPayloadType)
401             }
402         } else {
403             Err(FrameEncoderErr::NoCurrentFrame)
404         }
405     }
406 
headers_header_status(&mut self)407     fn headers_header_status(&mut self) {
408         // Headers encoding does not need to consider max_frame_size
409         // because header_payload_buffer must be smaller than max_frame_size.
410         self.state = if self.header_payload_index < self.remaining_header_payload {
411             FrameEncoderState::EncodingHeadersPayload
412         } else if self.hpack_encoder.is_finished() {
413             // set_frame ensures that headers must be is_end_headers
414             self.finish_current_frame();
415             FrameEncoderState::HeadersComplete
416         } else {
417             self.read_rest_payload();
418             FrameEncoderState::EncodingContinuationFrames
419         }
420     }
421 
iterate_headers_header( &self, frame: &Frame, buf: &mut [u8], len: usize, ) -> Result<(), FrameEncoderErr>422     fn iterate_headers_header(
423         &self,
424         frame: &Frame,
425         buf: &mut [u8],
426         len: usize,
427     ) -> Result<(), FrameEncoderErr> {
428         for (buf_index, item) in buf.iter_mut().enumerate().take(len) {
429             let header_byte_index = self.encoded_bytes + buf_index;
430             match header_byte_index {
431                 // The first 3 bytes represent the payload length in the frame header.
432                 0..=2 => {
433                     let payload_len = self.remaining_header_payload;
434                     *item = ((payload_len >> (16 - (8 * header_byte_index))) & 0xFF) as u8;
435                 }
436                 // The 4th byte represents the frame type in the frame header.
437                 3 => {
438                     *item = FrameType::Headers as u8;
439                 }
440                 // The 5th byte represents the frame flags in the frame header.
441                 4 => {
442                     if self.hpack_encoder.is_finished() {
443                         *item = frame.flags().bits();
444                     } else {
445                         // If not finished, it is followed by a Continuation frame.
446                         *item = frame.flags().bits() & 0xFB
447                     }
448                 }
449                 // The last 4 bytes (6th to 9th) represent the stream identifier in the
450                 // frame header.
451                 5..=8 => {
452                     let stream_id_byte_index = header_byte_index - 5;
453                     *item = (frame.stream_id() >> (24 - (8 * stream_id_byte_index))) as u8;
454                 }
455                 _ => {
456                     return Err(FrameEncoderErr::InternalError);
457                 }
458             }
459         }
460         Ok(())
461     }
462 
encode_headers_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>463     fn encode_headers_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
464         if let Some(frame) = &self.current_frame {
465             if let Payload::Headers(_) = frame.payload() {
466                 if buf.is_empty() {
467                     return Ok(0);
468                 }
469 
470                 let payload_bytes_written = self.write_payload(buf, self.remaining_header_payload);
471                 self.encoded_bytes += payload_bytes_written;
472 
473                 // Updates the state based on the encoding progress
474                 self.headers_header_status();
475 
476                 Ok(payload_bytes_written)
477             } else {
478                 Err(FrameEncoderErr::UnexpectedPayloadType)
479             }
480         } else {
481             Err(FrameEncoderErr::NoCurrentFrame)
482         }
483     }
484 
encode_continuation_frames(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>485     fn encode_continuation_frames(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
486         if let Some(frame) = &self.current_frame {
487             if let Payload::Headers(_) = frame.payload() {
488                 let available_space = buf.len();
489                 let frame_header_size = 9;
490                 // TODO allow available_space < 9
491                 if available_space < frame_header_size {
492                     return Ok(0);
493                 }
494                 // Encodes CONTINUATION frame header.
495                 // And this value is always the remaining_header_payload.
496                 let continuation_frame_len = self.remaining_header_payload.min(self.max_frame_size);
497                 for (buf_index, item) in buf.iter_mut().enumerate().take(3) {
498                     *item = ((continuation_frame_len >> (16 - (8 * buf_index))) & 0xFF) as u8;
499                 }
500                 buf[3] = FrameType::Continuation as u8;
501                 let mut new_flags = FrameFlags::empty();
502                 if self.remaining_header_payload <= self.max_frame_size
503                     && self.hpack_encoder.is_finished()
504                     && self.is_end_headers
505                 {
506                     // Sets the END_HEADER flag on the last CONTINUATION frame.
507                     new_flags.set_end_headers(true);
508                 }
509                 buf[4] = new_flags.bits();
510 
511                 for buf_index in 0..4 {
512                     let stream_id_byte_index = buf_index;
513                     buf[5 + buf_index] =
514                         (frame.stream_id() >> (24 - (8 * stream_id_byte_index))) as u8;
515                 }
516 
517                 // Encodes CONTINUATION frame payload.
518                 let payload_bytes_written =
519                     self.write_payload(&mut buf[frame_header_size..], continuation_frame_len);
520                 self.encoded_bytes += payload_bytes_written;
521 
522                 // Updates the state based on the encoding progress.
523                 self.headers_header_status();
524 
525                 Ok(frame_header_size + payload_bytes_written)
526             } else {
527                 Err(FrameEncoderErr::UnexpectedPayloadType)
528             }
529         } else {
530             Err(FrameEncoderErr::NoCurrentFrame)
531         }
532     }
533 
encode_data_header(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>534     fn encode_data_header(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
535         if let Some(frame) = &self.current_frame {
536             if let Payload::Data(_) = frame.payload() {
537                 // HTTP/2 frame header size is 9 bytes.
538                 let frame_header_size = 9;
539                 let remaining_header_bytes = if self.encoded_bytes >= frame_header_size {
540                     0
541                 } else {
542                     frame_header_size - self.encoded_bytes
543                 };
544                 let bytes_to_write = remaining_header_bytes.min(buf.len());
545 
546                 self.iterate_data_header(
547                     frame,
548                     buf,
549                     min(self.remaining_payload_bytes, self.max_frame_size),
550                     bytes_to_write,
551                 )?;
552 
553                 self.encoded_bytes += bytes_to_write;
554                 if self.encoded_bytes == frame_header_size {
555                     // data frame header is finished, reset encoded_bytes.
556                     self.encoded_bytes = 0;
557                     self.state = FrameEncoderState::EncodingDataPayload;
558                 }
559                 Ok(bytes_to_write)
560             } else {
561                 Err(FrameEncoderErr::UnexpectedPayloadType)
562             }
563         } else {
564             Err(FrameEncoderErr::NoCurrentFrame)
565         }
566     }
567 
iterate_data_header( &self, frame: &Frame, buf: &mut [u8], payload_len: usize, len: usize, ) -> Result<(), FrameEncoderErr>568     fn iterate_data_header(
569         &self,
570         frame: &Frame,
571         buf: &mut [u8],
572         payload_len: usize,
573         len: usize,
574     ) -> Result<(), FrameEncoderErr> {
575         for (buf_index, item) in buf.iter_mut().enumerate().take(len) {
576             let header_byte_index = self.encoded_bytes + buf_index;
577             match header_byte_index {
578                 // The first 3 bytes represent the payload length in the frame header.
579                 0..=2 => {
580                     *item = ((payload_len >> (16 - (8 * header_byte_index))) & 0xFF) as u8;
581                 }
582                 // The 4th byte represents the frame type in the frame header.
583                 3 => {
584                     *item = frame.payload().frame_type() as u8;
585                 }
586                 // The 5th byte represents the frame flags in the frame header.
587                 4 => {
588                     if self.remaining_payload_bytes <= self.max_frame_size {
589                         *item = frame.flags().bits();
590                     } else {
591                         // When max_frame_size is exceeded, a data frame cannot send all data.
592                         *item = frame.flags().bits() & 0xFE;
593                     }
594                 }
595                 // The last 4 bytes (6th to 9th) represent the stream identifier in the
596                 // frame header.
597                 5..=8 => {
598                     let stream_id_byte_index = header_byte_index - 5;
599                     *item = (frame.stream_id() >> (24 - (8 * stream_id_byte_index))) as u8;
600                 }
601                 _ => {
602                     return Err(FrameEncoderErr::InternalError);
603                 }
604             }
605         }
606         Ok(())
607     }
608 
encode_data_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>609     fn encode_data_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
610         if let Some(frame) = self.current_frame.as_ref() {
611             if let Payload::Data(data_frame) = frame.payload() {
612                 let payload = data_frame.data();
613                 let writen_bytes = self.encode_slice(buf, payload);
614                 self.data_offset += writen_bytes;
615                 self.remaining_payload_bytes -= writen_bytes;
616 
617                 self.state = if self.remaining_payload_bytes == 0 {
618                     self.data_offset = 0;
619                     FrameEncoderState::DataComplete
620                 } else if self.data_offset == self.max_frame_size {
621                     self.data_offset = 0;
622                     FrameEncoderState::EncodingDataHeader
623                 } else {
624                     FrameEncoderState::EncodingDataPayload
625                 };
626 
627                 Ok(writen_bytes)
628             } else {
629                 Err(FrameEncoderErr::UnexpectedPayloadType)
630             }
631         } else {
632             Err(FrameEncoderErr::NoCurrentFrame)
633         }
634     }
635 
encode_padding(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>636     fn encode_padding(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
637         if let Some(frame) = &self.current_frame {
638             if frame.flags().is_padded() {
639                 let padding_len = if let Payload::Data(data_frame) = frame.payload() {
640                     data_frame.data().len()
641                 } else {
642                     return Err(FrameEncoderErr::UnexpectedPayloadType);
643                 };
644 
645                 let remaining_padding_bytes = padding_len - self.encoded_bytes;
646                 let bytes_to_write = remaining_padding_bytes.min(buf.len());
647 
648                 for item in buf.iter_mut().take(bytes_to_write) {
649                     // Padding bytes are filled with 0.
650                     *item = 0;
651                 }
652 
653                 self.encoded_bytes += bytes_to_write;
654 
655                 if self.encoded_bytes == padding_len {
656                     self.state = FrameEncoderState::DataComplete;
657                 }
658 
659                 Ok(bytes_to_write)
660             } else {
661                 Ok(0) // No padding to encode, so return 0 bytes written.
662             }
663         } else {
664             Err(FrameEncoderErr::NoCurrentFrame)
665         }
666     }
667 
encode_goaway_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>668     fn encode_goaway_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
669         if let Some(frame) = &self.current_frame {
670             if let Payload::Goaway(_) = frame.payload() {
671                 let frame_header_size = 9;
672                 let remaining_header_bytes = if self.encoded_bytes >= frame_header_size {
673                     0
674                 } else {
675                     frame_header_size - self.encoded_bytes
676                 };
677                 let bytes_to_write = remaining_header_bytes.min(buf.len());
678                 self.iterate_goaway_header(frame, buf, bytes_to_write)?;
679                 self.encoded_bytes += bytes_to_write;
680                 if self.encoded_bytes == frame_header_size {
681                     self.state = FrameEncoderState::EncodingGoawayPayload;
682                 }
683                 Ok(bytes_to_write)
684             } else {
685                 Err(FrameEncoderErr::UnexpectedPayloadType)
686             }
687         } else {
688             Err(FrameEncoderErr::NoCurrentFrame)
689         }
690     }
691 
iterate_goaway_header( &self, frame: &Frame, buf: &mut [u8], len: usize, ) -> Result<(), FrameEncoderErr>692     fn iterate_goaway_header(
693         &self,
694         frame: &Frame,
695         buf: &mut [u8],
696         len: usize,
697     ) -> Result<(), FrameEncoderErr> {
698         for (buf_index, item) in buf.iter_mut().enumerate().take(len) {
699             let header_byte_index = self.encoded_bytes + buf_index;
700             match header_byte_index {
701                 0..=2 => {
702                     if let Payload::Goaway(goaway_payload) = frame.payload() {
703                         let payload_size = goaway_payload.encoded_len();
704                         *item = ((payload_size >> (8 * (2 - header_byte_index))) & 0xFF) as u8;
705                     } else {
706                         return Err(FrameEncoderErr::UnexpectedPayloadType);
707                     }
708                 }
709                 3 => {
710                     *item = FrameType::Goaway as u8;
711                 }
712                 4 => {
713                     *item = frame.flags().bits();
714                 }
715                 5..=8 => {
716                     let stream_id_byte_index = header_byte_index - 5;
717                     *item = (frame.stream_id() >> (24 - (8 * stream_id_byte_index))) as u8;
718                 }
719                 _ => {
720                     return Err(FrameEncoderErr::InternalError);
721                 }
722             }
723         }
724         Ok(())
725     }
726 
encode_goaway_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>727     fn encode_goaway_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
728         if let Some(frame) = &self.current_frame {
729             if let Payload::Goaway(goaway) = frame.payload() {
730                 let payload_size = goaway.encoded_len();
731                 let remaining_payload_bytes =
732                     payload_size.saturating_sub(self.encoded_bytes.saturating_sub(9));
733                 let bytes_to_write = remaining_payload_bytes.min(buf.len());
734 
735                 self.iterate_goaway_payload(goaway, buf, bytes_to_write)?;
736                 self.encoded_bytes += bytes_to_write;
737                 if self.encoded_bytes == 9 + payload_size {
738                     self.state = FrameEncoderState::DataComplete;
739                 }
740 
741                 Ok(bytes_to_write)
742             } else {
743                 Err(FrameEncoderErr::UnexpectedPayloadType)
744             }
745         } else {
746             Err(FrameEncoderErr::NoCurrentFrame)
747         }
748     }
749 
iterate_goaway_payload( &self, goaway: &Goaway, buf: &mut [u8], len: usize, ) -> Result<(), FrameEncoderErr>750     fn iterate_goaway_payload(
751         &self,
752         goaway: &Goaway,
753         buf: &mut [u8],
754         len: usize,
755     ) -> Result<(), FrameEncoderErr> {
756         for (buf_index, buf_item) in buf.iter_mut().enumerate().take(len) {
757             let payload_byte_index = self.encoded_bytes - 9 + buf_index;
758             match payload_byte_index {
759                 0..=3 => {
760                     let last_stream_id_byte_index = payload_byte_index;
761                     *buf_item = (goaway.get_last_stream_id()
762                         >> (24 - (8 * last_stream_id_byte_index)))
763                         as u8;
764                 }
765                 4..=7 => {
766                     let error_code_byte_index = payload_byte_index - 4;
767                     *buf_item =
768                         (goaway.get_error_code() >> (24 - (8 * error_code_byte_index))) as u8;
769                 }
770                 _ => {
771                     let debug_data_index = payload_byte_index - 8;
772                     if debug_data_index < goaway.get_debug_data().len() {
773                         *buf_item = goaway.get_debug_data()[debug_data_index];
774                     } else {
775                         return Err(FrameEncoderErr::InternalError);
776                     }
777                 }
778             }
779         }
780         Ok(())
781     }
782 
encode_window_update_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>783     fn encode_window_update_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
784         if let Some(frame) = &self.current_frame {
785             if let Payload::WindowUpdate(_) = frame.payload() {
786                 // HTTP/2 frame header size is 9 bytes.
787                 let frame_header_size = 9;
788                 let remaining_header_bytes = if self.encoded_bytes >= frame_header_size {
789                     0
790                 } else {
791                     frame_header_size - self.encoded_bytes
792                 };
793                 let bytes_to_write = remaining_header_bytes.min(buf.len());
794                 self.iterate_window_update_header(frame, buf, bytes_to_write)?;
795                 self.encoded_bytes += bytes_to_write;
796                 if self.encoded_bytes == frame_header_size {
797                     self.state = FrameEncoderState::EncodingWindowUpdatePayload;
798                     // Resets the encoded_bytes counter here.
799                     self.encoded_bytes = 0;
800                 }
801                 Ok(bytes_to_write)
802             } else {
803                 Err(FrameEncoderErr::UnexpectedPayloadType)
804             }
805         } else {
806             Err(FrameEncoderErr::NoCurrentFrame)
807         }
808     }
809 
iterate_window_update_header( &self, frame: &Frame, buf: &mut [u8], len: usize, ) -> Result<(), FrameEncoderErr>810     fn iterate_window_update_header(
811         &self,
812         frame: &Frame,
813         buf: &mut [u8],
814         len: usize,
815     ) -> Result<(), FrameEncoderErr> {
816         for (buf_index, item) in buf.iter_mut().enumerate().take(len) {
817             let header_byte_index = self.encoded_bytes + buf_index;
818             match header_byte_index {
819                 // The first 3 bytes represent the payload length in the frame header. For
820                 // WindowUpdate frame, this is always 4 bytes.
821                 0..=1 => {
822                     *item = 0;
823                 }
824                 2 => {
825                     // Window Update frame payload size is always 4 bytes.
826                     *item = 4;
827                 }
828                 // The 4th byte represents the frame type in the frame header.
829                 3 => {
830                     *item = FrameType::WindowUpdate as u8;
831                 }
832                 // The 5th byte represents the frame flags in the frame header.
833                 4 => {
834                     *item = frame.flags().bits();
835                 }
836                 // The last 4 bytes (6th to 9th) represent the stream identifier in the
837                 // frame header.
838                 5..=8 => {
839                     let stream_id_byte_index = header_byte_index - 5;
840                     *item = (frame.stream_id() >> (24 - (8 * stream_id_byte_index))) as u8;
841                 }
842                 _ => {
843                     return Err(FrameEncoderErr::InternalError);
844                 }
845             }
846         }
847         Ok(())
848     }
849 
encode_window_update_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>850     fn encode_window_update_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
851         if let Some(frame) = &self.current_frame {
852             if let Payload::WindowUpdate(window_update) = frame.payload() {
853                 let payload_size = 4usize;
854                 let remaining_payload_bytes =
855                     payload_size.saturating_sub(self.encoded_bytes.saturating_sub(9usize));
856                 let bytes_to_write = remaining_payload_bytes.min(buf.len());
857                 for (buf_index, buf_item) in buf.iter_mut().enumerate().take(bytes_to_write) {
858                     let payload_byte_index = self
859                         .encoded_bytes
860                         .saturating_sub(9)
861                         .saturating_add(buf_index);
862                     let increment_byte_index = payload_byte_index;
863                     *buf_item =
864                         (window_update.get_increment() >> (24 - (8 * increment_byte_index))) as u8;
865                 }
866                 self.encoded_bytes += bytes_to_write;
867                 if self.encoded_bytes == payload_size {
868                     self.state = FrameEncoderState::DataComplete;
869                 }
870 
871                 Ok(bytes_to_write)
872             } else {
873                 Err(FrameEncoderErr::UnexpectedPayloadType)
874             }
875         } else {
876             Err(FrameEncoderErr::NoCurrentFrame)
877         }
878     }
879 
encode_settings_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>880     fn encode_settings_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
881         if let Some(frame) = &self.current_frame {
882             if let Payload::Settings(settings) = frame.payload() {
883                 let frame_header_size = 9;
884                 let remaining_header_bytes = if self.encoded_bytes >= frame_header_size {
885                     0
886                 } else {
887                     frame_header_size - self.encoded_bytes
888                 };
889                 let bytes_to_write = remaining_header_bytes.min(buf.len());
890                 self.iterate_settings_header(
891                     frame,
892                     buf,
893                     settings.get_settings().len() * 6,
894                     bytes_to_write,
895                 )?;
896                 self.encoded_bytes += bytes_to_write;
897                 if self.encoded_bytes == frame_header_size {
898                     self.state = FrameEncoderState::EncodingSettingsPayload;
899                 }
900                 Ok(bytes_to_write)
901             } else {
902                 Err(FrameEncoderErr::UnexpectedPayloadType)
903             }
904         } else {
905             Err(FrameEncoderErr::NoCurrentFrame)
906         }
907     }
908 
iterate_settings_header( &self, frame: &Frame, buf: &mut [u8], payload_len: usize, len: usize, ) -> Result<(), FrameEncoderErr>909     fn iterate_settings_header(
910         &self,
911         frame: &Frame,
912         buf: &mut [u8],
913         payload_len: usize,
914         len: usize,
915     ) -> Result<(), FrameEncoderErr> {
916         for buf_index in 0..len {
917             let header_byte_index = self.encoded_bytes + buf_index;
918             match header_byte_index {
919                 // The first 3 bytes represent the payload length in the frame header.
920                 0..=2 => {
921                     buf[buf_index] = ((payload_len >> (16 - (8 * buf_index))) & 0xFF) as u8;
922                 }
923                 // The 4th byte represents the frame type in the frame header.
924                 3 => {
925                     buf[3] = FrameType::Settings as u8;
926                 }
927                 // The 5th byte represents the frame flags in the frame header.
928                 4 => {
929                     buf[4] = frame.flags().bits();
930                 }
931                 // The last 4 bytes (6th to 9th) represent the stream identifier in the
932                 // frame header. For SETTINGS frames, this should
933                 // always be 0.
934                 5..=8 => {
935                     // Stream ID should be 0 for SETTINGS frames.
936                     buf[buf_index] = 0;
937                 }
938                 _ => {
939                     return Err(FrameEncoderErr::InternalError);
940                 }
941             }
942         }
943         Ok(())
944     }
945 
encode_settings_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>946     fn encode_settings_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
947         if let Some(frame) = &self.current_frame {
948             if let Payload::Settings(settings) = frame.payload() {
949                 let settings_len = settings.get_settings().len() * 6;
950                 let remaining_payload_bytes =
951                     settings_len.saturating_sub(self.encoded_bytes.saturating_sub(9));
952                 let bytes_to_write = remaining_payload_bytes.min(buf.len());
953                 self.iterate_settings_payload(settings, buf, bytes_to_write)?;
954                 self.encoded_bytes += bytes_to_write;
955                 if self.encoded_bytes == 9 + settings_len {
956                     self.state = FrameEncoderState::DataComplete;
957                 }
958 
959                 Ok(bytes_to_write)
960             } else {
961                 Err(FrameEncoderErr::UnexpectedPayloadType)
962             }
963         } else {
964             Err(FrameEncoderErr::NoCurrentFrame)
965         }
966     }
967 
iterate_settings_payload( &self, settings: &Settings, buf: &mut [u8], len: usize, ) -> Result<(), FrameEncoderErr>968     fn iterate_settings_payload(
969         &self,
970         settings: &Settings,
971         buf: &mut [u8],
972         len: usize,
973     ) -> Result<(), FrameEncoderErr> {
974         for (buf_index, buf_item) in buf.iter_mut().enumerate().take(len) {
975             let payload_byte_index = self.encoded_bytes - 9 + buf_index;
976             let setting_index = payload_byte_index / 6;
977             let setting_byte_index = payload_byte_index % 6;
978 
979             if let Some(setting) = settings.get_settings().get(setting_index) {
980                 let (id, value) = match setting {
981                     Setting::HeaderTableSize(v) => (0x1, *v),
982                     Setting::EnablePush(v) => (0x2, *v as u32),
983                     Setting::MaxConcurrentStreams(v) => (0x3, *v),
984                     Setting::InitialWindowSize(v) => (0x4, *v),
985                     Setting::MaxFrameSize(v) => (0x5, *v),
986                     Setting::MaxHeaderListSize(v) => (0x6, *v),
987                 };
988                 match setting_byte_index {
989                     0..=1 => {
990                         *buf_item = ((id >> (8 * (1 - setting_byte_index))) & 0xFF) as u8;
991                     }
992                     2..=5 => {
993                         let shift_amount = 8 * (3 - (setting_byte_index - 2));
994                         *buf_item = ((value >> shift_amount) & 0xFF) as u8;
995                     }
996                     _ => {
997                         return Err(FrameEncoderErr::InternalError);
998                     }
999                 }
1000             } else {
1001                 return Err(FrameEncoderErr::InternalError);
1002             }
1003         }
1004         Ok(())
1005     }
1006 
encode_priority_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>1007     fn encode_priority_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
1008         if let Some(frame) = &self.current_frame {
1009             if let Payload::Priority(_) = frame.payload() {
1010                 // HTTP/2 frame header size is 9 bytes.
1011                 let frame_header_size = 9;
1012                 let remaining_header_bytes = if self.encoded_bytes >= frame_header_size {
1013                     0
1014                 } else {
1015                     frame_header_size - self.encoded_bytes
1016                 };
1017                 let bytes_to_write = remaining_header_bytes.min(buf.len());
1018 
1019                 self.iterate_priority_header(frame, buf, bytes_to_write)?;
1020                 self.encoded_bytes += bytes_to_write;
1021                 if self.encoded_bytes == frame_header_size {
1022                     self.state = FrameEncoderState::EncodingPriorityPayload;
1023                 }
1024                 Ok(bytes_to_write)
1025             } else {
1026                 Err(FrameEncoderErr::UnexpectedPayloadType)
1027             }
1028         } else {
1029             Err(FrameEncoderErr::NoCurrentFrame)
1030         }
1031     }
1032 
iterate_priority_header( &self, frame: &Frame, buf: &mut [u8], len: usize, ) -> Result<(), FrameEncoderErr>1033     fn iterate_priority_header(
1034         &self,
1035         frame: &Frame,
1036         buf: &mut [u8],
1037         len: usize,
1038     ) -> Result<(), FrameEncoderErr> {
1039         for (buf_index, item) in buf.iter_mut().enumerate().take(len) {
1040             let header_byte_index = self.encoded_bytes + buf_index;
1041             match header_byte_index {
1042                 // The first 3 bytes represent the payload length in the frame header.
1043                 0..=2 => {
1044                     let payload_len = 5;
1045                     *item = ((payload_len >> (16 - (8 * header_byte_index))) & 0xFF) as u8;
1046                 }
1047                 // The 4th byte represents the frame type in the frame header.
1048                 3 => {
1049                     *item = frame.payload().frame_type() as u8;
1050                 }
1051                 // The 5th byte represents the frame flags in the frame header.
1052                 4 => {
1053                     *item = frame.flags().bits();
1054                 }
1055                 // The last 4 bytes (6th to 9th) represent the stream identifier in the
1056                 // frame header.
1057                 5..=8 => {
1058                     let stream_id_byte_index = header_byte_index - 5;
1059                     *item = (frame.stream_id() >> (24 - (8 * stream_id_byte_index))) as u8;
1060                 }
1061                 _ => {
1062                     return Err(FrameEncoderErr::InternalError);
1063                 }
1064             }
1065         }
1066         Ok(())
1067     }
1068 
encode_priority_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>1069     fn encode_priority_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
1070         if let Some(frame) = &self.current_frame {
1071             if let Payload::Priority(priority) = frame.payload() {
1072                 // HTTP/2 frame header size is 9 bytes.
1073                 let frame_header_size = 9;
1074                 let remaining_payload_bytes = 5 - (self.encoded_bytes - frame_header_size);
1075                 let bytes_to_write = remaining_payload_bytes.min(buf.len());
1076 
1077                 self.iterate_priority_payload(priority, buf, frame_header_size, bytes_to_write)?;
1078                 self.encoded_bytes += bytes_to_write;
1079                 if self.encoded_bytes == frame_header_size + 5 {
1080                     self.state = FrameEncoderState::DataComplete
1081                 }
1082 
1083                 Ok(bytes_to_write)
1084             } else {
1085                 Err(FrameEncoderErr::UnexpectedPayloadType)
1086             }
1087         } else {
1088             Err(FrameEncoderErr::NoCurrentFrame)
1089         }
1090     }
1091 
iterate_priority_payload( &self, priority: &Priority, buf: &mut [u8], frame_header_size: usize, len: usize, ) -> Result<(), FrameEncoderErr>1092     fn iterate_priority_payload(
1093         &self,
1094         priority: &Priority,
1095         buf: &mut [u8],
1096         frame_header_size: usize,
1097         len: usize,
1098     ) -> Result<(), FrameEncoderErr> {
1099         for (buf_index, buf_item) in buf.iter_mut().enumerate().take(len) {
1100             let payload_byte_index = self
1101                 .encoded_bytes
1102                 .saturating_sub(frame_header_size)
1103                 .saturating_add(buf_index);
1104             match payload_byte_index {
1105                 0 => {
1106                     *buf_item = (priority.get_exclusive() as u8) << 7
1107                         | ((priority.get_stream_dependency() >> 24) & 0x7F) as u8;
1108                 }
1109                 1..=3 => {
1110                     let stream_dependency_byte_index = payload_byte_index - 1;
1111                     *buf_item = (priority.get_stream_dependency()
1112                         >> (16 - (8 * stream_dependency_byte_index)))
1113                         as u8;
1114                 }
1115                 4 => {
1116                     // The last byte is the weight.
1117                     *buf_item = priority.get_weight();
1118                 }
1119                 _ => {
1120                     return Err(FrameEncoderErr::InternalError);
1121                 }
1122             }
1123         }
1124         Ok(())
1125     }
1126 
encode_rst_stream_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>1127     fn encode_rst_stream_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
1128         if let Some(frame) = &self.current_frame {
1129             let frame_header_size = 9;
1130             if self.encoded_bytes >= frame_header_size {
1131                 return Ok(0);
1132             }
1133 
1134             let bytes_to_write = (frame_header_size - self.encoded_bytes).min(buf.len());
1135 
1136             for (buf_index, item) in buf.iter_mut().enumerate().take(bytes_to_write) {
1137                 let header_byte_index = self.encoded_bytes + buf_index;
1138                 match header_byte_index {
1139                     0..=2 => {
1140                         let payload_len = 4;
1141                         *item = ((payload_len >> (16 - (8 * buf_index))) & 0xFF) as u8;
1142                     }
1143                     3 => {
1144                         *item = FrameType::RstStream as u8;
1145                     }
1146                     4 => {
1147                         *item = frame.flags().bits();
1148                     }
1149                     5..=8 => {
1150                         let stream_id = frame.stream_id();
1151                         *item = ((stream_id >> (24 - (8 * (buf_index - 5)))) & 0xFF) as u8;
1152                     }
1153                     _ => {
1154                         return Err(FrameEncoderErr::InternalError);
1155                     }
1156                 }
1157             }
1158             self.encoded_bytes += bytes_to_write;
1159 
1160             if self.encoded_bytes == frame_header_size {
1161                 self.state = FrameEncoderState::EncodingRstStreamPayload;
1162             }
1163 
1164             Ok(bytes_to_write)
1165         } else {
1166             Err(FrameEncoderErr::NoCurrentFrame)
1167         }
1168     }
1169 
encode_rst_stream_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>1170     fn encode_rst_stream_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
1171         if let Some(frame) = &self.current_frame {
1172             if let Payload::RstStream(rst_stream) = frame.payload() {
1173                 let frame_header_size = 9;
1174                 if self.encoded_bytes < frame_header_size {
1175                     return Ok(0);
1176                 }
1177 
1178                 let payload_size = 4;
1179                 let encoded_payload_bytes = self.encoded_bytes - frame_header_size;
1180 
1181                 if encoded_payload_bytes >= payload_size {
1182                     return Ok(0);
1183                 }
1184 
1185                 let bytes_to_write = (payload_size - encoded_payload_bytes).min(buf.len());
1186 
1187                 for (buf_index, item) in buf.iter_mut().enumerate().take(bytes_to_write) {
1188                     let payload_byte_index = encoded_payload_bytes + buf_index;
1189                     *item =
1190                         ((rst_stream.error_code() >> (24 - (8 * payload_byte_index))) & 0xFF) as u8;
1191                 }
1192 
1193                 self.encoded_bytes += bytes_to_write;
1194 
1195                 if self.encoded_bytes == frame_header_size + payload_size {
1196                     self.state = FrameEncoderState::DataComplete;
1197                 }
1198 
1199                 Ok(bytes_to_write)
1200             } else {
1201                 Err(FrameEncoderErr::UnexpectedPayloadType)
1202             }
1203         } else {
1204             Err(FrameEncoderErr::NoCurrentFrame)
1205         }
1206     }
1207 
encode_ping_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>1208     fn encode_ping_frame(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
1209         if let Some(frame) = &self.current_frame {
1210             if let Payload::Ping(_) = frame.payload() {
1211                 let frame_header_size = 9;
1212                 let remaining_header_bytes = if self.encoded_bytes >= frame_header_size {
1213                     0
1214                 } else {
1215                     frame_header_size - self.encoded_bytes
1216                 };
1217                 let bytes_to_write = remaining_header_bytes.min(buf.len());
1218                 self.iterate_ping_header(frame, buf, bytes_to_write)?;
1219                 self.encoded_bytes += bytes_to_write;
1220                 if self.encoded_bytes == frame_header_size {
1221                     self.state = FrameEncoderState::EncodingPingPayload;
1222                 }
1223                 Ok(bytes_to_write)
1224             } else {
1225                 Err(FrameEncoderErr::UnexpectedPayloadType)
1226             }
1227         } else {
1228             Err(FrameEncoderErr::NoCurrentFrame)
1229         }
1230     }
1231 
iterate_ping_header( &self, frame: &Frame, buf: &mut [u8], len: usize, ) -> Result<(), FrameEncoderErr>1232     fn iterate_ping_header(
1233         &self,
1234         frame: &Frame,
1235         buf: &mut [u8],
1236         len: usize,
1237     ) -> Result<(), FrameEncoderErr> {
1238         for buf_index in 0..len {
1239             let header_byte_index = self.encoded_bytes + buf_index;
1240             match header_byte_index {
1241                 // The first 3 bytes represent the payload length in the frame header.
1242                 0..=2 => {
1243                     // PING payload is always 8 bytes.
1244                     let payload_len = 8;
1245                     buf[buf_index] = ((payload_len >> (16 - (8 * buf_index))) & 0xFF) as u8;
1246                 }
1247                 // The 4th byte represents the frame type in the frame header.
1248                 3 => {
1249                     buf[3] = FrameType::Ping as u8;
1250                 }
1251                 // The 5th byte represents the frame flags in the frame header.
1252                 4 => {
1253                     buf[4] = frame.flags().bits();
1254                 }
1255                 // The last 4 bytes (6th to 9th) represent the stream identifier in the
1256                 // frame header. For PING frames, this should always
1257                 // be 0.
1258                 5..=8 => {
1259                     // Stream ID should be 0 for PING frames.
1260                     buf[buf_index] = 0;
1261                 }
1262                 _ => {
1263                     return Err(FrameEncoderErr::InternalError);
1264                 }
1265             }
1266         }
1267         Ok(())
1268     }
1269 
encode_ping_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr>1270     fn encode_ping_payload(&mut self, buf: &mut [u8]) -> Result<usize, FrameEncoderErr> {
1271         if let Some(frame) = &self.current_frame {
1272             if let Payload::Ping(ping) = frame.payload() {
1273                 // PING payload is always 8 bytes.
1274                 let payload_size = 8usize;
1275                 let remaining_payload_bytes =
1276                     payload_size.saturating_sub(self.encoded_bytes.saturating_sub(9usize));
1277                 let bytes_to_write = remaining_payload_bytes.min(buf.len());
1278                 for (buf_index, buf_item) in buf.iter_mut().enumerate().take(bytes_to_write) {
1279                     let payload_byte_index = self
1280                         .encoded_bytes
1281                         .saturating_sub(9)
1282                         .saturating_add(buf_index);
1283                     *buf_item = ping.data[payload_byte_index];
1284                 }
1285                 self.encoded_bytes += bytes_to_write;
1286                 if self.encoded_bytes == 9 + 8 {
1287                     self.state = FrameEncoderState::DataComplete;
1288                 }
1289 
1290                 Ok(bytes_to_write)
1291             } else {
1292                 Err(FrameEncoderErr::UnexpectedPayloadType)
1293             }
1294         } else {
1295             Err(FrameEncoderErr::NoCurrentFrame)
1296         }
1297     }
1298 
encode_slice(&self, buf: &mut [u8], data: &[u8]) -> usize1299     fn encode_slice(&self, buf: &mut [u8], data: &[u8]) -> usize {
1300         let start = self.data_offset;
1301         let allow_to_write = (self.max_frame_size - start).min(self.remaining_payload_bytes);
1302         let bytes_to_write = allow_to_write.min(buf.len());
1303 
1304         buf[..bytes_to_write].copy_from_slice(&data[start..start + bytes_to_write]);
1305         bytes_to_write
1306     }
1307 
1308     // Helper method for writing the payload from the buffer to the output buffer.
write_payload(&mut self, buf: &mut [u8], payload_len: usize) -> usize1309     fn write_payload(&mut self, buf: &mut [u8], payload_len: usize) -> usize {
1310         let bytes_to_write = buf.len().min(payload_len - self.header_payload_index);
1311         buf[..bytes_to_write].copy_from_slice(
1312             &self.header_payload_buffer
1313                 [self.header_payload_index..self.header_payload_index + bytes_to_write],
1314         );
1315         self.header_payload_index += bytes_to_write;
1316         bytes_to_write
1317     }
1318 }
1319 
1320 #[cfg(test)]
1321 mod ut_frame_encoder {
1322     use super::*;
1323     use crate::error::HttpError;
1324     use crate::h2::frame::{
1325         Data, FrameFlags, Goaway, Headers, Ping, Priority, RstStream, Settings, WindowUpdate,
1326     };
1327     use crate::h2::{Parts, PseudoHeaders};
1328 
1329     /// UT test cases for `FrameEncoder` encoding `DATA` frame.
1330     ///
1331     /// # Brief
1332     /// 1. Creates a `FrameEncoder`.
1333     /// 2. Creates a `Frame` with `Payload::Data`.
1334     /// 3. Sets the frame for the encoder.
1335     /// 4. Encodes the frame using buffer segments.
1336     /// 5. Checks whether the result is correct.
1337     #[test]
ut_data_frame_encoding()1338     fn ut_data_frame_encoding() {
1339         let mut encoder = FrameEncoder::new(4096, false);
1340         let data_payload = b"hhhhhhhhhhhhhhhhhhhhhhhhhhhhhhh".to_vec();
1341 
1342         let data_frame = Frame::new(
1343             131,
1344             FrameFlags::new(0),
1345             Payload::Data(Data::new(data_payload.clone())),
1346         );
1347 
1348         encoder.set_frame(data_frame).unwrap();
1349 
1350         let mut first_buf = [0u8; 2];
1351         let mut second_buf = [0u8; 38];
1352 
1353         let first_encoded = encoder.encode(&mut first_buf).unwrap();
1354         let second_encoded = encoder.encode(&mut second_buf).unwrap();
1355 
1356         assert_eq!(first_encoded, 2);
1357         assert_eq!(second_encoded, 38);
1358         assert_eq!(first_buf[0], 0);
1359         assert_eq!(first_buf[1], 0);
1360         assert_eq!(second_buf[0], data_payload.len() as u8);
1361         assert_eq!(second_buf[6], 131);
1362 
1363         for &item in second_buf.iter().skip(7).take(30) {
1364             assert_eq!(item, 104);
1365         }
1366     }
1367 
1368     /// UT test cases for `FrameEncoder` encoding `HEADERS` frame.
1369     ///
1370     /// # Brief
1371     /// 1. Creates a `FrameEncoder`.
1372     /// 2. Creates a `Frame` with `Payload::Headers`.
1373     /// 3. Sets the frame for the encoder.
1374     /// 4. Encodes the frame using buffer segments.
1375     /// 5. Checks whether the result is correct.
1376     #[test]
ut_headers_frame_encoding()1377     fn ut_headers_frame_encoding() {
1378         let mut frame_encoder = FrameEncoder::new(4096, false);
1379 
1380         let mut new_parts = Parts::new();
1381         new_parts.pseudo.set_method(Some("GET".to_string()));
1382         new_parts.pseudo.set_scheme(Some("https".to_string()));
1383         new_parts.pseudo.set_path(Some("/code".to_string()));
1384         new_parts
1385             .pseudo
1386             .set_authority(Some("example.com".to_string()));
1387         let mut frame_flag = FrameFlags::empty();
1388         frame_flag.set_end_headers(true);
1389         frame_flag.set_end_stream(true);
1390         let frame = Frame::new(1, frame_flag, Payload::Headers(Headers::new(new_parts)));
1391 
1392         // Set the current frame for the encoder
1393         frame_encoder.set_frame(frame).unwrap();
1394 
1395         let mut buf = vec![0; 50];
1396         let first_encoded = frame_encoder.encode(&mut buf).unwrap();
1397         assert_eq!(first_encoded, 22 + 9);
1398         assert_eq!(buf[0], 0);
1399         assert_eq!(buf[2], 22);
1400         assert_eq!(buf[3], 0x1);
1401         assert_eq!(buf[4], 5);
1402         assert_eq!(buf[8], 1);
1403 
1404         assert_eq!(frame_encoder.state, FrameEncoderState::HeadersComplete);
1405     }
1406 
1407     /// UT test cases for `FrameEncoder` encoding `SETTINGS` frame.
1408     ///
1409     /// # Brief
1410     /// 1. Creates a `FrameEncoder`.
1411     /// 2. Creates a `Frame` with `Payload::Settings`.
1412     /// 3. Sets the frame for the encoder.
1413     /// 4. Encodes the frame using buffer segments.
1414     /// 5. Checks whether the result is correct.
1415     #[test]
ut_settings_frame_encoding()1416     fn ut_settings_frame_encoding() {
1417         let mut encoder = FrameEncoder::new(4096, false);
1418         let settings_payload = vec![
1419             Setting::HeaderTableSize(4096),
1420             Setting::EnablePush(true),
1421             Setting::MaxConcurrentStreams(100),
1422             Setting::InitialWindowSize(65535),
1423             Setting::MaxFrameSize(16384),
1424             Setting::MaxHeaderListSize(8192),
1425         ];
1426 
1427         let settings = Settings::new(settings_payload.clone());
1428 
1429         let settings_frame = Frame::new(0, FrameFlags::new(0), Payload::Settings(settings));
1430 
1431         let mut first_buf = [0u8; 9];
1432         let mut second_buf = [0u8; 30];
1433         let mut third_buf = [0u8; 6];
1434         encoder.set_frame(settings_frame).unwrap();
1435 
1436         let first_encoded = encoder.encode(&mut first_buf).unwrap();
1437         assert_eq!(encoder.state, FrameEncoderState::EncodingSettingsPayload);
1438         let second_encoded = encoder.encode(&mut second_buf).unwrap();
1439         let third_encoded = encoder.encode(&mut third_buf).unwrap();
1440 
1441         assert_eq!(encoder.state, FrameEncoderState::DataComplete);
1442         // Updated expected value for first_encoded
1443         assert_eq!(first_encoded, 9);
1444         assert_eq!(second_encoded, 30);
1445         // Updated expected value for third_encoded
1446         assert_eq!(third_encoded, 6);
1447 
1448         // Validate the encoded settings
1449         let mut expected_encoded_settings = [0u8; 60];
1450         for (i, setting) in settings_payload.iter().enumerate() {
1451             let offset = i * 6;
1452             let (id, value) = match setting {
1453                 Setting::HeaderTableSize(v) => (0x1, *v),
1454                 Setting::EnablePush(v) => (0x2, *v as u32),
1455                 Setting::MaxConcurrentStreams(v) => (0x3, *v),
1456                 Setting::InitialWindowSize(v) => (0x4, *v),
1457                 Setting::MaxFrameSize(v) => (0x5, *v),
1458                 Setting::MaxHeaderListSize(v) => (0x6, *v),
1459             };
1460             expected_encoded_settings[offset] = (id >> 8) as u8;
1461             expected_encoded_settings[offset + 1] = (id & 0xFF) as u8;
1462             expected_encoded_settings[offset + 2] = (value >> 24) as u8;
1463             expected_encoded_settings[offset + 3] = ((value >> 16) & 0xFF) as u8;
1464             expected_encoded_settings[offset + 4] = ((value >> 8) & 0xFF) as u8;
1465             expected_encoded_settings[offset + 5] = (value & 0xFF) as u8;
1466         }
1467 
1468         let actual_encoded_settings = [&second_buf[..], &third_buf[..]].concat();
1469         for i in 0..35 {
1470             assert_eq!(expected_encoded_settings[i], actual_encoded_settings[i]);
1471         }
1472     }
1473 
1474     /// UT test cases for `FrameEncoder` encoding `PING` frame.
1475     ///
1476     /// # Brief
1477     /// 1. Creates a `FrameEncoder`.
1478     /// 2. Creates a `Frame` with `Payload::Ping`.
1479     /// 3. Sets the frame for the encoder.
1480     /// 4. Encodes the frame using buffer segments.
1481     /// 5. Checks whether the result is correct.
1482     #[test]
ut_ping_frame_encoding()1483     fn ut_ping_frame_encoding() {
1484         let mut encoder = FrameEncoder::new(4096, false);
1485         let ping_payload = [1, 2, 3, 4, 5, 6, 7, 8];
1486 
1487         let ping_frame = Frame::new(
1488             0,
1489             FrameFlags::new(0),
1490             Payload::Ping(Ping { data: ping_payload }),
1491         );
1492 
1493         encoder.set_frame(ping_frame).unwrap();
1494 
1495         let mut first_buf = [0u8; 9];
1496         let mut second_buf = [0u8; 8];
1497 
1498         let first_encoded = encoder.encode(&mut first_buf).unwrap();
1499         let second_encoded = encoder.encode(&mut second_buf).unwrap();
1500 
1501         assert_eq!(first_encoded, 9);
1502         assert_eq!(second_encoded, 8);
1503 
1504         assert_eq!(first_buf[0], 0);
1505         assert_eq!(first_buf[1], 0);
1506         // payload length
1507         assert_eq!(first_buf[2], 8);
1508         assert_eq!(first_buf[3], FrameType::Ping as u8);
1509         // flags
1510         assert_eq!(first_buf[4], 0);
1511         // stream id
1512         assert_eq!(first_buf[5], 0);
1513         // stream id
1514         assert_eq!(first_buf[6], 0);
1515         // stream id
1516         assert_eq!(first_buf[7], 0);
1517         // stream id
1518         assert_eq!(first_buf[8], 0);
1519 
1520         for i in 0..8 {
1521             assert_eq!(second_buf[i], ping_payload[i]);
1522         }
1523     }
1524 
1525     /// UT test case for FrameEncoder encoding a sequence of frames: Headers,
1526     /// Data, Headers.
1527     ///
1528     /// # Brief
1529     /// 1. Creates a FrameEncoder.
1530     /// 2. Creates multiple frames including Headers and Data frames.
1531     /// 3. Sets each frame for the encoder and encodes them using buffer
1532     ///    segments.
1533     /// 4. Checks whether the encoding results are correct.
1534     #[test]
ut_continue_frame_encoding()1535     fn ut_continue_frame_encoding() {
1536         let mut encoder = FrameEncoder::new(4096, false);
1537 
1538         let mut new_parts = Parts::new();
1539         new_parts.pseudo.set_method(Some("GET".to_string()));
1540         new_parts.pseudo.set_scheme(Some("https".to_string()));
1541         new_parts.pseudo.set_path(Some("/code".to_string()));
1542         new_parts
1543             .pseudo
1544             .set_authority(Some("example.com".to_string()));
1545         let mut frame_flag = FrameFlags::empty();
1546         frame_flag.set_end_headers(true);
1547         frame_flag.set_end_stream(false);
1548         let frame_1 = Frame::new(
1549             1,
1550             frame_flag.clone(),
1551             Payload::Headers(Headers::new(new_parts.clone())),
1552         );
1553 
1554         let data_payload = b"hhhhhhhhhhhhhhhhhhhhhhhhhhhhhhh".to_vec();
1555         let data_frame = Frame::new(
1556             1,
1557             FrameFlags::new(1),
1558             Payload::Data(Data::new(data_payload)),
1559         );
1560 
1561         let frame_2 = Frame::new(
1562             1,
1563             frame_flag.clone(),
1564             Payload::Headers(Headers::new(new_parts.clone())),
1565         );
1566 
1567         encoder.set_frame(frame_1).unwrap();
1568         let mut first_buf = [0u8; 50];
1569         let first_encoding = encoder.encode(&mut first_buf).unwrap();
1570 
1571         encoder.set_frame(data_frame).unwrap();
1572         let mut second_buf = [0u8; 50];
1573         let second_encoding = encoder.encode(&mut second_buf).unwrap();
1574 
1575         encoder.set_frame(frame_2).unwrap();
1576         let mut third_buf = [0u8; 50];
1577         let third_encoding = encoder.encode(&mut third_buf).unwrap();
1578 
1579         assert_eq!(first_encoding, 31);
1580         assert_eq!(second_encoding, 40);
1581         assert_eq!(third_encoding, 13);
1582 
1583         assert_eq!(first_buf[2], 22);
1584         assert_eq!(second_buf[2], 31);
1585         assert_eq!(third_buf[2], 4);
1586     }
1587 
1588     /// UT test cases for `FrameEncoder` encoding `RST_STREAM` frame.
1589     ///
1590     /// # Brief
1591     /// 1. Creates a `FrameEncoder`.
1592     /// 2. Creates a `Frame` with `Payload::RstStream`.
1593     /// 3. Sets the frame for the encoder.
1594     /// 4. Encodes the frame using buffer segments.
1595     /// 5. Checks whether the result is correct.
1596     #[test]
ut_rst_stream_frame_encoding()1597     fn ut_rst_stream_frame_encoding() {
1598         let mut frame_encoder = FrameEncoder::new(4096, false);
1599 
1600         let error_code = 12345678;
1601         let rst_stream_payload = Payload::RstStream(RstStream::new(error_code));
1602 
1603         let frame_flags = FrameFlags::empty();
1604         let frame = Frame::new(
1605             // Stream ID can be non-zero for RST_STREAM frames
1606             1,
1607             frame_flags,
1608             rst_stream_payload,
1609         );
1610 
1611         // Set the current frame for the encoder
1612         frame_encoder.set_frame(frame).unwrap();
1613 
1614         let mut buf = vec![0; 50];
1615         let first_encoded = frame_encoder.encode(&mut buf).unwrap();
1616         // 9 bytes for header, 4 bytes for payload
1617         assert_eq!(first_encoded, 9 + 4);
1618         assert_eq!(buf[0], 0);
1619         // payload length should be 4 for RST_STREAM frames
1620         assert_eq!(buf[2], 4);
1621         assert_eq!(buf[3], FrameType::RstStream as u8);
1622         // frame flags should be 0 for RST_STREAM frames
1623         assert_eq!(buf[4], 0);
1624         // stream ID should be 1 for this test case
1625         assert_eq!(buf[8], 1);
1626 
1627         // Check if encoded error code is correct
1628         assert_eq!(&buf[9..13], &error_code.to_be_bytes());
1629 
1630         assert_eq!(frame_encoder.state, FrameEncoderState::DataComplete);
1631     }
1632 
1633     /// UT test cases for `FrameEncoder` encoding `WINDOW_UPDATE` frame.
1634     ///
1635     /// # Brief
1636     /// 1. Creates a `FrameEncoder`.
1637     /// 2. Creates a `Frame` with `Payload::WindowUpdate`.
1638     /// 3. Sets the frame for the encoder.
1639     /// 4. Encodes the frame using buffer segments.
1640     /// 5. Checks whether the result is correct.
1641     #[test]
ut_window_update_frame_encoding()1642     fn ut_window_update_frame_encoding() {
1643         let mut frame_encoder = FrameEncoder::new(4096, false);
1644 
1645         let window_size_increment = 12345678;
1646         let window_update_payload = Payload::WindowUpdate(WindowUpdate::new(window_size_increment));
1647 
1648         let frame_flags = FrameFlags::empty();
1649         let frame = Frame::new(
1650             // Stream ID can be zero for WINDOW_UPDATE frames.
1651             0,
1652             frame_flags,
1653             window_update_payload,
1654         );
1655 
1656         // Sets the current frame for the encoder.
1657         frame_encoder.set_frame(frame).unwrap();
1658 
1659         let mut buf = vec![0; 50];
1660         let first_encoded = frame_encoder.encode(&mut buf).unwrap();
1661         // 9 bytes for header, 4 bytes for payload.
1662         assert_eq!(first_encoded, 9 + 4);
1663         assert_eq!(buf[0], 0);
1664         // Payload length should be 4 for WINDOW_UPDATE frames.
1665         assert_eq!(buf[2], 4);
1666         assert_eq!(buf[3], FrameType::WindowUpdate as u8);
1667         // Frame flags should be 0 for WINDOW_UPDATE frames.
1668         assert_eq!(buf[4], 0);
1669         // Stream ID should be 0 for this test case.
1670         assert_eq!(buf[8], 0);
1671 
1672         // Checks if encoded window size increment is correct.
1673         assert_eq!(&buf[9..13], &window_size_increment.to_be_bytes());
1674 
1675         assert_eq!(frame_encoder.state, FrameEncoderState::DataComplete);
1676     }
1677 
1678     /// UT test case for FrameEncoder encoding `PRIORITY` frame.
1679     ///
1680     /// # Brief
1681     /// 1. Creates a FrameEncoder.
1682     /// 2. Creates a Frame with Payload::Priority.
1683     /// 3. Sets the frame for the encoder.
1684     /// 4. Encodes the frame using buffer segments.
1685     /// 5. Checks whether the result is correct.
1686     #[test]
ut_priority_frame_encoding()1687     fn ut_priority_frame_encoding() {
1688         let mut encoder = FrameEncoder::new(4096, false);
1689         // Maximum value for a 31-bit integer
1690         let stream_dependency = 0x7FFFFFFF;
1691         let priority_payload = Priority::new(true, stream_dependency, 15);
1692 
1693         let priority_frame =
1694             Frame::new(131, FrameFlags::new(0), Payload::Priority(priority_payload));
1695 
1696         encoder.set_frame(priority_frame).unwrap();
1697 
1698         let mut buf = [0u8; 14];
1699 
1700         let encoded = encoder.encode(&mut buf).unwrap();
1701 
1702         assert_eq!(encoded, 14);
1703         // Payload length (most significant byte)
1704         assert_eq!(buf[0], 0);
1705         // Payload length (middle byte)
1706         assert_eq!(buf[1], 0);
1707         // Payload length (least significant byte)
1708         assert_eq!(buf[2], 5);
1709         // Frame flags
1710         assert_eq!(buf[3], FrameType::Priority as u8);
1711         assert_eq!(buf[4], 0);
1712         // Stream ID (most significant byte)
1713         assert_eq!(buf[5], 0);
1714         // Stream ID (middle bytes)
1715         assert_eq!(buf[6], 0);
1716         // Stream ID (middle bytes)
1717         assert_eq!(buf[7], 0);
1718         // Stream ID (least significant byte)
1719         assert_eq!(buf[8], 131);
1720         // Exclusive flag and most significant byte of stream dependency
1721         assert_eq!(buf[9], (0x80 | ((stream_dependency >> 24) & 0x7F)) as u8);
1722         // Stream dependency (middle bytes)
1723         assert_eq!(buf[10], ((stream_dependency >> 16) & 0xFF) as u8);
1724         // Stream dependency (middle bytes)
1725         assert_eq!(buf[11], ((stream_dependency >> 8) & 0xFF) as u8);
1726         // Stream dependency (least significant byte)
1727         assert_eq!(buf[12], (stream_dependency & 0xFF) as u8);
1728         // Weight
1729         assert_eq!(buf[13], 15);
1730     }
1731 
1732     /// UT test cases for `FrameEncoder` encoding `GOAWAY` frame.
1733     ///
1734     /// # Brief
1735     /// 1. Creates a `FrameEncoder`.
1736     /// 2. Creates a `Frame` with `Payload::Goaway`.
1737     /// 3. Sets the frame for the encoder.
1738     /// 4. Encodes the frame and its payload using buffer segments.
1739     /// 5. Checks whether the result is correct.
1740     #[test]
ut_goaway_frame_encoding()1741     fn ut_goaway_frame_encoding() {
1742         // 1. Creates a `FrameEncoder`.
1743         let mut encoder = FrameEncoder::new(4096, false);
1744 
1745         // 2. Creates a `Frame` with `Payload::Goaway`.
1746         let last_stream_id = 1;
1747         let error_code = 1;
1748         let debug_data = vec![1, 2, 3, 4, 5];
1749         let goaway_frame = Frame::new(
1750             131,
1751             FrameFlags::new(0),
1752             Payload::Goaway(Goaway::new(error_code, last_stream_id, debug_data.clone())),
1753         );
1754 
1755         // 3. Sets the frame for the encoder.
1756         encoder.set_frame(goaway_frame).unwrap();
1757 
1758         // 4. Encodes the frame and its payload using buffer segments.
1759         let mut first_buf = [0u8; 9];
1760         let mut second_buf = [0u8; 13];
1761         let first_encoded = encoder.encode(&mut first_buf).unwrap();
1762         let second_encoded = encoder.encode(&mut second_buf).unwrap();
1763 
1764         // 5. Checks whether the result is correct.
1765         assert_eq!(first_encoded, 9);
1766         assert_eq!(second_encoded, 13);
1767 
1768         // Validate the encoded GOAWAY frame.
1769         let mut expected_encoded_goaway = [0u8; 13];
1770         expected_encoded_goaway[0..4].copy_from_slice(&(last_stream_id as u32).to_be_bytes());
1771         expected_encoded_goaway[4..8].copy_from_slice(&(error_code).to_be_bytes());
1772 
1773         expected_encoded_goaway[8..13].copy_from_slice(&debug_data[..]);
1774 
1775         // payload length should be 13 bytes
1776         assert_eq!(first_buf[0..3], [0u8, 0, 13]);
1777         assert_eq!(first_buf[3], FrameType::Goaway as u8);
1778         // flags
1779         assert_eq!(first_buf[4], 0);
1780 
1781         // Validate the encoded Last-Stream-ID, Error Code, and debug data
1782         assert_eq!(second_buf[..], expected_encoded_goaway[..]);
1783     }
1784 
1785     /// UT test cases for `FrameEncoder::update_max_frame_size`.
1786     ///
1787     /// # Brief
1788     /// 1. Creates a `FrameEncoder`.
1789     /// 2. Updates the maximum frame size.
1790     /// 3. Checks whether the maximum frame size was updated correctly.
1791     #[test]
ut_update_max_frame_size()1792     fn ut_update_max_frame_size() {
1793         let mut encoder = FrameEncoder::new(4096, false);
1794         encoder.update_max_frame_size(8192);
1795         assert_eq!(encoder.max_frame_size, 8192);
1796     }
1797 
1798     /// UT test cases for `FrameEncoder::update_setting`.
1799     ///
1800     /// # Brief
1801     /// 1. Creates a `FrameEncoder`.
1802     /// 2. Creates a `Setting` variant.
1803     /// 3. Creates a `Frame` with `Payload::Settings`.
1804     /// 4. Sets the frame for the encoder.
1805     /// 5. Updates the setting.
1806     /// 6. Checks whether the setting was updated correctly.
1807     #[test]
ut_update_setting()1808     fn ut_update_setting() {
1809         let mut encoder = FrameEncoder::new(4096, false);
1810         let settings_payload = vec![Setting::MaxFrameSize(4096)];
1811         let settings = Settings::new(settings_payload);
1812         let settings_frame = Frame::new(0, FrameFlags::new(0), Payload::Settings(settings));
1813 
1814         encoder.set_frame(settings_frame).unwrap();
1815         let new_setting = Setting::MaxFrameSize(8192);
1816         encoder.update_setting(new_setting.clone());
1817 
1818         if let Some(frame) = &mut encoder.current_frame {
1819             if let Payload::Settings(settings) = frame.payload_mut() {
1820                 let updated_settings = settings.get_settings();
1821                 assert!(updated_settings.iter().any(|s| *s == new_setting));
1822             }
1823         }
1824     }
1825 
1826     /// UT test cases for `FrameEncoder` encoding continuation frames.
1827     ///
1828     /// # Brief
1829     /// 1. Creates a `FrameEncoder`.
1830     /// 2. Creates a `Frame` with `Payload::Headers` and sets the flags.
1831     /// 3. Sets the frame for the encoder.
1832     /// 4. Encodes the continuation frames using a buffer.
1833     /// 5. Checks whether the result is correct.
1834     #[test]
ut_encode_continuation_frames()1835     fn ut_encode_continuation_frames() {
1836         let mut frame_encoder = FrameEncoder::new(4096, false);
1837         let mut new_parts = Parts::new();
1838         assert!(new_parts.is_empty());
1839         new_parts.pseudo.set_method(Some("GET".to_string()));
1840         new_parts.pseudo.set_scheme(Some("https".to_string()));
1841         new_parts.pseudo.set_path(Some("/code".to_string()));
1842         new_parts
1843             .pseudo
1844             .set_authority(Some("example.com".to_string()));
1845 
1846         let mut frame_flag = FrameFlags::empty();
1847         frame_flag.set_end_headers(true);
1848         frame_flag.set_end_stream(false);
1849         let frame = Frame::new(
1850             1,
1851             frame_flag.clone(),
1852             Payload::Headers(Headers::new(new_parts.clone())),
1853         );
1854 
1855         frame_encoder.set_frame(frame).unwrap();
1856         frame_encoder.state = FrameEncoderState::EncodingContinuationFrames;
1857         let mut buf = [0u8; 5000];
1858 
1859         assert!(frame_encoder.encode_continuation_frames(&mut buf).is_ok());
1860 
1861         let mut frame_flag = FrameFlags::empty();
1862         frame_flag.set_end_headers(true);
1863         let frame = Frame::new(
1864             1,
1865             frame_flag,
1866             Payload::Headers(Headers::new(new_parts.clone())),
1867         );
1868 
1869         frame_encoder.set_frame(frame).unwrap();
1870         frame_encoder.state = FrameEncoderState::EncodingContinuationFrames;
1871         assert!(frame_encoder.encode_continuation_frames(&mut buf).is_ok());
1872 
1873         let mut frame_flag = FrameFlags::empty();
1874         frame_flag.set_end_headers(true);
1875         let frame = Frame::new(1, frame_flag, Payload::Ping(Ping::new([0; 8])));
1876 
1877         frame_encoder.set_frame(frame).unwrap();
1878         frame_encoder.state = FrameEncoderState::EncodingContinuationFrames;
1879         assert!(frame_encoder.encode_continuation_frames(&mut buf).is_err());
1880     }
1881 
1882     /// UT test cases for `FrameEncoder` encoding padded data.
1883     ///
1884     /// # Brief
1885     /// 1. Creates a `FrameEncoder`.
1886     /// 2. Creates a `Frame` with `Payload::Data` and sets the flags.
1887     /// 3. Sets the frame for the encoder.
1888     /// 4. Encodes the padding using a buffer.
1889     /// 5. Checks whether the result is correct.
1890     #[test]
ut_encode_padding()1891     fn ut_encode_padding() {
1892         let mut frame_encoder = FrameEncoder::new(4096, false);
1893 
1894         // Creates a padded data frame.
1895         let mut frame_flags = FrameFlags::empty();
1896         frame_flags.set_end_headers(true);
1897         frame_flags.set_padded(true);
1898         let data_payload = vec![0u8; 500];
1899         let data_frame = Frame::new(
1900             1,
1901             frame_flags.clone(),
1902             Payload::Data(Data::new(data_payload)),
1903         );
1904 
1905         // Sets the frame to the frame_encoder and test padding encoding.
1906         frame_encoder.set_frame(data_frame).unwrap();
1907         frame_encoder.state = FrameEncoderState::EncodingDataPadding;
1908         let mut buf = [0u8; 600];
1909         assert!(frame_encoder.encode_padding(&mut buf).is_ok());
1910 
1911         let headers_payload = Payload::Headers(Headers::new(Parts::new()));
1912         let headers_frame = Frame::new(1, frame_flags.clone(), headers_payload);
1913         frame_encoder.set_frame(headers_frame).unwrap();
1914         frame_encoder.state = FrameEncoderState::EncodingDataPadding;
1915         assert!(frame_encoder.encode_padding(&mut buf).is_err());
1916 
1917         frame_encoder.current_frame = None;
1918         assert!(frame_encoder.encode_padding(&mut buf).is_err());
1919     }
1920 
1921     /// UT test cases for `FrameEncoder` encoding data frame.
1922     ///
1923     /// # Brief
1924     /// 1. Creates a `FrameEncoder`.
1925     /// 2. Creates a `Frame` smaller than max_frame_size with `Payload::Data`
1926     ///    and sets the flags.
1927     /// 3. Sets the frame for the encoder.
1928     /// 4. Encode the data frame with a buf.
1929     /// 5. Checks whether the result is correct.
1930     #[test]
ut_encode_small_data_frame()1931     fn ut_encode_small_data_frame() {
1932         let mut encoder = FrameEncoder::new(100, false);
1933         let data_payload = vec![b'a'; 10];
1934         let mut buf = [0u8; 10];
1935         encode_small_frame(&mut encoder, &mut buf, data_payload.clone());
1936 
1937         // buf is larger than frame but smaller than max_frame_size;
1938         let mut buf = [0u8; 50];
1939         encode_small_frame(&mut encoder, &mut buf, data_payload.clone());
1940 
1941         // buf is larger than max_frame_size;
1942         let mut buf = [0u8; 200];
1943         encode_small_frame(&mut encoder, &mut buf, data_payload);
1944     }
1945 
encode_small_frame(encoder: &mut FrameEncoder, buf: &mut [u8], data_payload: Vec<u8>)1946     fn encode_small_frame(encoder: &mut FrameEncoder, buf: &mut [u8], data_payload: Vec<u8>) {
1947         let data_frame = Frame::new(
1948             1,
1949             FrameFlags::new(1),
1950             Payload::Data(Data::new(data_payload)),
1951         );
1952         encoder.set_frame(data_frame).unwrap();
1953 
1954         let mut result = [b'0'; 10 + 9];
1955         let total = assert_encoded_data(encoder, &mut result, buf);
1956         assert_eq!(total, 10 + 9);
1957         assert_eq!(result[4], 0x1);
1958         assert_eq!(encoder.state, FrameEncoderState::DataComplete);
1959     }
1960 
1961     /// UT test cases for `FrameEncoder` encoding data frame.
1962     ///
1963     /// # Brief
1964     /// 1. Creates a `FrameEncoder`.
1965     /// 2. Creates a `Frame` larger than max_frame_size with `Payload::Data` and
1966     ///    sets the flags.
1967     /// 3. Sets the frame for the encoder.
1968     /// 4. Encode the data frame with a buf.
1969     /// 5. Checks whether the result is correct.
1970     #[test]
ut_encode_large_data_frame()1971     fn ut_encode_large_data_frame() {
1972         let mut encoder = FrameEncoder::new(100, false);
1973         let data_payload = vec![b'a'; 1024];
1974         let mut buf = [0u8; 10];
1975 
1976         let data_frame = Frame::new(
1977             1,
1978             FrameFlags::new(0),
1979             Payload::Data(Data::new(data_payload.clone())),
1980         );
1981         let mut result = [b'0'; 1024 + 9 * 11];
1982         encoder.set_frame(data_frame).unwrap();
1983 
1984         let total = assert_encoded_data(&mut encoder, &mut result, &mut buf);
1985         // This is because the data frame flag is 0.
1986         assert_eq!(result[4 + 10 * (9 + 100)], 0x0);
1987         assert_eq!(total, 1024 + 9 * 11);
1988         for index in 0..=9 {
1989             assert_eq!(result[4 + index * (9 + 100)], 0x0);
1990         }
1991         assert_eq!(encoder.state, FrameEncoderState::DataComplete);
1992 
1993         // finished
1994         let data_frame = Frame::new(
1995             1,
1996             FrameFlags::new(1),
1997             Payload::Data(Data::new(data_payload)),
1998         );
1999         let mut result = [b'0'; 1024 + 9 * 11];
2000         encoder.set_frame(data_frame).unwrap();
2001 
2002         let total = assert_encoded_data(&mut encoder, &mut result, &mut buf);
2003         // This is because the data frame flag is 0.
2004         assert_eq!(result[4 + 10 * (9 + 100)], 0x1);
2005         assert_eq!(total, 1024 + 9 * 11);
2006         for index in 0..=9 {
2007             assert_eq!(result[4 + index * (9 + 100)], 0x0);
2008         }
2009         assert_eq!(encoder.state, FrameEncoderState::DataComplete);
2010     }
2011 
assert_encoded_data(encoder: &mut FrameEncoder, result: &mut [u8], buf: &mut [u8]) -> usize2012     fn assert_encoded_data(encoder: &mut FrameEncoder, result: &mut [u8], buf: &mut [u8]) -> usize {
2013         let mut total = 0;
2014         loop {
2015             let size = encoder.encode(buf).unwrap();
2016             result[total..total + size].copy_from_slice(&buf[..size]);
2017             total += size;
2018             if size == 0 {
2019                 break;
2020             }
2021         }
2022         total
2023     }
2024 }
2025