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