Respect peer MAX_FRAME_SIZE

Signed-off-by: Eloi DEMOLIS <[email protected]>
sozu-proxy · Jan 23, 2025 · ad8fb54 · ad8fb54
1 parent e1cc9c0
commit ad8fb54
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 26 deletions.
diff --git a/lib/src/protocol/mux/converter.rs b/lib/src/protocol/mux/converter.rs
@@ -1,3 +1,5 @@
+use std::cmp::min;
+
 use kawa::{
     AsBuffer, Block, BlockConverter, Chunk, Flags, Kawa, Pair, ParsingErrorKind, ParsingPhase,
     StatusLine, Store,
@@ -13,6 +15,7 @@ use crate::protocol::{
 };
 
 pub struct H2BlockConverter<'a> {
+    pub max_frame_size: usize,
     pub window: i32,
     pub stream_id: StreamId,
     pub encoder: &'a mut hpack::Encoder<'static>,
@@ -121,19 +124,25 @@ impl<'a, T: AsBuffer> BlockConverter<T> for H2BlockConverter<'a> {
             Block::Chunk(Chunk { data }) => {
                 let mut header = [0; 9];
                 let payload_len = data.len();
-                let (data, payload_len, can_continue) = if self.window >= payload_len as i32 {
-                    // the window is wide enought to send the entire chunk
-                    (data, payload_len as u32, true)
-                } else if self.window > 0 {
-                    // we split the chunk to fit in the window
-                    let (before, after) = data.split(self.window as usize);
-                    kawa.blocks.push_front(Block::Chunk(Chunk { data: after }));
-                    (before, self.window as u32, false)
-                } else {
-                    // the window can't take any more bytes, return the chunk to the blocks
-                    kawa.blocks.push_front(Block::Chunk(Chunk { data }));
-                    return false;
-                };
+                let (data, payload_len, can_continue) =
+                    if self.window >= payload_len as i32 && self.max_frame_size >= payload_len {
+                        // the window is wide enought to send the entire chunk
+                        (data, payload_len as u32, true)
+                    } else if self.window > 0 {
+                        // we split the chunk to fit in the window
+                        let payload_len = min(self.max_frame_size, self.window as usize);
+                        let (before, after) = data.split(payload_len);
+                        kawa.blocks.push_front(Block::Chunk(Chunk { data: after }));
+                        (
+                            before,
+                            payload_len as u32,
+                            self.max_frame_size < self.window as usize,
+                        )
+                    } else {
+                        // the window can't take any more bytes, return the chunk to the blocks
+                        kawa.blocks.push_front(Block::Chunk(Chunk { data }));
+                        return false;
+                    };
                 self.window -= payload_len as i32;
                 gen_frame_header(
                     &mut header,
@@ -158,19 +167,37 @@ impl<'a, T: AsBuffer> BlockConverter<T> for H2BlockConverter<'a> {
                 if end_header {
                     let payload = std::mem::take(&mut self.out);
                     let mut header = [0; 9];
-                    let flags = if end_stream { 1 } else { 0 } | if end_header { 4 } else { 0 };
-                    gen_frame_header(
-                        &mut header,
-                        &FrameHeader {
-                            payload_len: payload.len() as u32,
-                            frame_type: FrameType::Headers,
-                            flags,
-                            stream_id: self.stream_id,
-                        },
-                    )
-                    .unwrap();
-                    kawa.push_out(Store::from_slice(&header));
-                    kawa.push_out(Store::Alloc(payload.into_boxed_slice(), 0));
+                    let chunks = payload.chunks(self.max_frame_size);
+                    let n_chunks = chunks.len();
+                    for (i, chunk) in chunks.enumerate() {
+                        let flags = if i == 0 && end_stream { 1 } else { 0 }
+                            | if i + 1 == n_chunks { 4 } else { 0 };
+                        if i == 0 {
+                            gen_frame_header(
+                                &mut header,
+                                &FrameHeader {
+                                    payload_len: chunk.len() as u32,
+                                    frame_type: FrameType::Headers,
+                                    flags,
+                                    stream_id: self.stream_id,
+                                },
+                            )
+                            .unwrap();
+                        } else {
+                            gen_frame_header(
+                                &mut header,
+                                &FrameHeader {
+                                    payload_len: chunk.len() as u32,
+                                    frame_type: FrameType::Continuation,
+                                    flags,
+                                    stream_id: self.stream_id,
+                                },
+                            )
+                            .unwrap();
+                        }
+                        kawa.push_out(Store::from_slice(&header));
+                        kawa.push_out(Store::from_slice(chunk));
+                    }
                 } else if end_stream {
                     let mut header = [0; 9];
                     gen_frame_header(

diff --git a/lib/src/protocol/mux/h2.rs b/lib/src/protocol/mux/h2.rs
@@ -595,6 +595,7 @@ impl<Front: SocketHandler> ConnectionH2<Front> {
                 }
 
                 let mut converter = converter::H2BlockConverter {
+                    max_frame_size: self.peer_settings.settings_max_frame_size as usize,
                     window: 0,
                     stream_id: 0,
                     encoder: &mut self.encoder,