fix: handle manual masks in fragment collector

Co-Authored-By: Divy Srivastava <me@littledivy.com>

Author: divybot

src/fragment.rs

@@ -118,7 +118,8 @@ impl<'f, S> FragmentCollector<S> {
       if is_closed && frame.opcode != OpCode::Close {
         return Err(WebSocketError::ConnectionClosed);
       }
-      if let Some(frame) = self.fragments.accumulate(frame)? {
+      let manual_unmask = !self.read_half.auto_apply_mask;
+      if let Some(frame) = self.fragments.accumulate(frame, manual_unmask)? {
         return Ok(frame);
       }
     }
@@ -191,7 +192,8 @@ impl<'f, S> FragmentCollectorRead<S> {
       let Some(frame) = res? else {
         continue;
       };
-      if let Some(frame) = self.fragments.accumulate(frame)? {
+      let manual_unmask = !self.read_half.auto_apply_mask;
+      if let Some(frame) = self.fragments.accumulate(frame, manual_unmask)? {
         return Ok(frame);
       }
     }
@@ -214,16 +216,20 @@ impl Fragments {
 
   pub fn accumulate<'f>(
     &mut self,
-    frame: Frame<'f>,
+    mut frame: Frame<'f>,
+    manual_unmask: bool,
   ) -> Result<Option<Frame<'f>>, WebSocketError> {
     match frame.opcode {
       OpCode::Text | OpCode::Binary => {
         if frame.fin {
           if self.fragments.is_some() {
             return Err(WebSocketError::InvalidFragment);
           }
-          return Ok(Some(Frame::new(true, frame.opcode, None, frame.payload)));
+          return Ok(Some(frame));
         } else {
+          if manual_unmask {
+            frame.unmask();
+          }
           self.fragments = match frame.opcode {
             OpCode::Text => match utf8::decode(&frame.payload) {
               Ok(text) => Some(Fragment::Text(None, text.as_bytes().to_vec())),
@@ -249,6 +255,9 @@ impl Fragments {
           return Err(WebSocketError::InvalidContinuationFrame);
         }
         Some(Fragment::Text(data, input)) => {
+          if manual_unmask {
+            frame.unmask();
+          }
           let mut tail = &frame.payload[..];
           if let Some(mut incomplete) = data.take() {
             if let Some((result, rest)) =
@@ -296,6 +305,9 @@ impl Fragments {
           }
         }
         Some(Fragment::Binary(data)) => {
+          if manual_unmask {
+            frame.unmask();
+          }
           data.extend_from_slice(&frame.payload);
           if frame.fin {
             return Ok(Some(Frame::new(

src/frame.rs

@@ -261,11 +261,15 @@ impl<'f> Frame<'f> {
   ///
   /// Note: By default, the frame payload is unmasked by `WebSocket::read_frame`.
   pub fn unmask(&mut self) {
-    if let Some(mask) = self.mask {
+    if let Some(mask) = self.mask.take() {
       crate::mask::unmask(self.payload.to_mut(), mask);
     }
   }
 
+  pub(crate) fn is_masked(&self) -> bool {
+    self.mask.is_some()
+  }
+
   /// Formats the frame header into the head buffer. Returns the size of the length field.
   ///
   /// # Panics

src/lib.rs

@@ -646,7 +646,7 @@ impl ReadHalf {
         (Ok(None), Some(Frame::pong(frame.payload)))
       }
       OpCode::Text => {
-        if frame.fin && !frame.is_utf8() {
+        if frame.fin && !frame.is_masked() && !frame.is_utf8() {
           (Err(WebSocketError::InvalidUTF8), None)
         } else {
           (Ok(Some(frame)), None)

tests/fragment_collector_manual_mask.rs

@@ -0,0 +1,128 @@
+// Copyright 2023-2026 Divy Srivastava <dj.srivastava23@gmail.com>
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use anyhow::Result;
+use fastwebsockets::FragmentCollector;
+use fastwebsockets::Frame;
+use fastwebsockets::OpCode;
+use fastwebsockets::Role;
+use fastwebsockets::WebSocket;
+use tokio::io::AsyncWriteExt;
+use tokio::io::DuplexStream;
+
+async fn write_masked_frame(
+  stream: &mut DuplexStream,
+  fin: bool,
+  opcode: OpCode,
+  mask: [u8; 4],
+  payload: &[u8],
+) -> Result<()> {
+  let mut masked = payload.to_vec();
+  fastwebsockets::unmask(&mut masked, mask);
+
+  let mut frame = Frame::new(fin, opcode, Some(mask), masked.into());
+  let mut buf = Vec::new();
+  stream.write_all(frame.write(&mut buf)).await?;
+  Ok(())
+}
+
+#[tokio::test]
+async fn unfragmented_masked_text_can_be_manually_unmasked() -> Result<()> {
+  let (mut client, server) = tokio::io::duplex(1024);
+  let mut ws = WebSocket::after_handshake(server, Role::Server);
+  ws.set_auto_apply_mask(false);
+  let mut ws = FragmentCollector::new(ws);
+
+  write_masked_frame(
+    &mut client,
+    true,
+    OpCode::Text,
+    [0xff, 0xff, 0xff, 0xff],
+    b"hello",
+  )
+  .await?;
+
+  let mut frame = ws.read_frame().await?;
+  assert_eq!(frame.opcode, OpCode::Text);
+  assert_ne!(&frame.payload[..], b"hello");
+
+  frame.unmask();
+  assert_eq!(&frame.payload[..], b"hello");
+
+  Ok(())
+}
+
+#[tokio::test]
+async fn fragmented_masked_text_is_unmasked_before_validation() -> Result<()> {
+  let (mut client, server) = tokio::io::duplex(1024);
+  let mut ws = WebSocket::after_handshake(server, Role::Server);
+  ws.set_auto_apply_mask(false);
+  let mut ws = FragmentCollector::new(ws);
+
+  write_masked_frame(
+    &mut client,
+    false,
+    OpCode::Text,
+    [0xff, 0xff, 0xff, 0xff],
+    b"hello ",
+  )
+  .await?;
+  write_masked_frame(
+    &mut client,
+    true,
+    OpCode::Continuation,
+    [0x80, 0x80, 0x80, 0x80],
+    b"world",
+  )
+  .await?;
+
+  let mut frame = ws.read_frame().await?;
+  assert_eq!(frame.opcode, OpCode::Text);
+  assert_eq!(&frame.payload[..], b"hello world");
+
+  frame.unmask();
+  assert_eq!(&frame.payload[..], b"hello world");
+
+  Ok(())
+}
+
+#[tokio::test]
+async fn fragmented_masked_binary_is_unmasked() -> Result<()> {
+  let (mut client, server) = tokio::io::duplex(1024);
+  let mut ws = WebSocket::after_handshake(server, Role::Server);
+  ws.set_auto_apply_mask(false);
+  let mut ws = FragmentCollector::new(ws);
+
+  write_masked_frame(
+    &mut client,
+    false,
+    OpCode::Binary,
+    [0x7f, 0x7f, 0x7f, 0x7f],
+    &[0, 1, 2, 3],
+  )
+  .await?;
+  write_masked_frame(
+    &mut client,
+    true,
+    OpCode::Continuation,
+    [0x55, 0x55, 0x55, 0x55],
+    &[4, 5, 6, 7],
+  )
+  .await?;
+
+  let frame = ws.read_frame().await?;
+  assert_eq!(frame.opcode, OpCode::Binary);
+  assert_eq!(&frame.payload[..], &[0, 1, 2, 3, 4, 5, 6, 7]);
+
+  Ok(())
+}