From 50d1371e52ed2883b606ac80b32c0758665bdb5b Mon Sep 17 00:00:00 2001 From: Alexey Galakhov Date: Wed, 11 Jul 2018 00:42:58 +0200 Subject: [PATCH] frame: cleanup frame API Signed-off-by: Alexey Galakhov --- src/protocol/frame/frame.rs | 131 ++++++++++-------------------------- src/protocol/mod.rs | 17 +++-- 2 files changed, 44 insertions(+), 104 deletions(-) diff --git a/src/protocol/frame/frame.rs b/src/protocol/frame/frame.rs index 89f82f3..cee9b00 100644 --- a/src/protocol/frame/frame.rs +++ b/src/protocol/frame/frame.rs @@ -50,7 +50,7 @@ pub struct FrameHeader { /// WebSocket protocol opcode. pub opcode: OpCode, /// A frame mask, if any. - pub mask: Option<[u8; 4]>, + mask: Option<[u8; 4]>, } impl Default for FrameHeader { @@ -120,6 +120,13 @@ impl FrameHeader { Ok(()) } + + /// Generate a random frame mask and store this in the header. + /// + /// Of course this does not change frame contents. It just generates a mask. + pub(crate) fn set_random_mask(&mut self) { + self.mask = Some(generate_mask()) + } } impl FrameHeader { @@ -218,34 +225,16 @@ impl Frame { self.header.len(length as u64) + length } - /// Test whether the frame is a final frame. - #[inline] - pub fn is_final(&self) -> bool { - self.header.is_final - } - - /// Test whether the first reserved bit is set. - #[inline] - pub fn has_rsv1(&self) -> bool { - self.header.rsv1 - } - - /// Test whether the second reserved bit is set. - #[inline] - pub fn has_rsv2(&self) -> bool { - self.header.rsv2 - } - - /// Test whether the third reserved bit is set. + /// Get a reference to the frame's header. #[inline] - pub fn has_rsv3(&self) -> bool { - self.header.rsv3 + pub fn header(&self) -> &FrameHeader { + &self.header } - /// Get the OpCode of the frame. + /// Get a mutable reference to the frame's header. #[inline] - pub fn opcode(&self) -> OpCode { - self.header.opcode + pub fn header_mut(&mut self) -> &mut FrameHeader { + &mut self.header } /// Get a reference to the frame's payload. @@ -254,86 +243,34 @@ impl Frame { &self.payload } - // Test whether the frame is masked. - #[doc(hidden)] - #[inline] - pub fn is_masked(&self) -> bool { - self.header.mask.is_some() - } - - // Get an optional reference to the frame's mask. - #[doc(hidden)] - #[allow(dead_code)] - #[inline] - pub fn mask(&self) -> Option<&[u8; 4]> { - self.header.mask.as_ref() - } - - /// Make this frame a final frame. - #[allow(dead_code)] + /// Get a mutable reference to the frame's payload. #[inline] - pub fn set_final(&mut self, is_final: bool) -> &mut Frame { - self.header.is_final = is_final; - self - } - - /// Set the first reserved bit. - #[inline] - pub fn set_rsv1(&mut self, has_rsv1: bool) -> &mut Frame { - self.header.rsv1 = has_rsv1; - self - } - - /// Set the second reserved bit. - #[inline] - pub fn set_rsv2(&mut self, has_rsv2: bool) -> &mut Frame { - self.header.rsv2 = has_rsv2; - self - } - - /// Set the third reserved bit. - #[inline] - pub fn set_rsv3(&mut self, has_rsv3: bool) -> &mut Frame { - self.header.rsv3 = has_rsv3; - self - } - - /// Set the OpCode. - #[allow(dead_code)] - #[inline] - pub fn set_opcode(&mut self, opcode: OpCode) -> &mut Frame { - self.header.opcode = opcode; - self + pub fn payload_mut(&mut self) -> &mut Vec { + &mut self.payload } - /// Edit the frame's payload. - #[allow(dead_code)] + /// Test whether the frame is masked. #[inline] - pub fn payload_mut(&mut self) -> &mut Vec { - &mut self.payload + pub(crate) fn is_masked(&self) -> bool { + self.header.mask.is_some() } - // Generate a new mask for this frame. - // - // This method simply generates and stores the mask. It does not change the payload data. - // Instead, the payload data will be masked with the generated mask when the frame is sent - // to the other endpoint. - #[doc(hidden)] + /// Generate a random mask for the frame. + /// + /// This just generates a mask, payload is not changed. The actual masking is performed + /// either on `format()` or on `apply_mask()` call. #[inline] - pub fn set_mask(&mut self) -> &mut Frame { - self.header.mask = Some(generate_mask()); - self + pub(crate) fn set_random_mask(&mut self) { + self.header.set_random_mask() } - // This method unmasks the payload and should only be called on frames that are actually - // masked. In other words, those frames that have just been received from a client endpoint. - #[doc(hidden)] + /// This method unmasks the payload and should only be called on frames that are actually + /// masked. In other words, those frames that have just been received from a client endpoint. #[inline] - pub fn remove_mask(&mut self) { - self.header.mask.and_then(|mask| { - Some(apply_mask(&mut self.payload, &mask)) - }); - self.header.mask = None; + pub(crate) fn apply_mask(&mut self) { + if let Some(mask) = self.header.mask.take() { + apply_mask(&mut self.payload, &mask) + } } /// Consume the frame into its payload as binary. @@ -350,7 +287,7 @@ impl Frame { /// Consume the frame into a closing frame. #[inline] - pub fn into_close(self) -> Result>> { + pub(crate) fn into_close(self) -> Result>> { match self.payload.len() { 0 => Ok(None), 1 => Err(Error::Protocol("Invalid close sequence".into())), @@ -435,7 +372,7 @@ impl Frame { /// Write a frame out to a buffer pub fn format(mut self, output: &mut impl Write) -> Result<()> { self.header.format(self.payload.len() as u64, output)?; - self.remove_mask(); + self.apply_mask(); output.write_all(self.payload())?; Ok(()) } diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 3436e20..8124e73 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -246,8 +246,11 @@ impl WebSocket { // the negotiated extensions defines the meaning of such a nonzero // value, the receiving endpoint MUST _Fail the WebSocket // Connection_. - if frame.has_rsv1() || frame.has_rsv2() || frame.has_rsv3() { - return Err(Error::Protocol("Reserved bits are non-zero".into())) + { + let hdr = frame.header(); + if hdr.rsv1 || hdr.rsv2 || hdr.rsv3 { + return Err(Error::Protocol("Reserved bits are non-zero".into())) + } } match self.role { @@ -255,7 +258,7 @@ impl WebSocket { if frame.is_masked() { // A server MUST remove masking for data frames received from a client // as described in Section 5.3. (RFC 6455) - frame.remove_mask() + frame.apply_mask() } else { // The server MUST close the connection upon receiving a // frame that is not masked. (RFC 6455) @@ -270,13 +273,13 @@ impl WebSocket { } } - match frame.opcode() { + match frame.header().opcode { OpCode::Control(ctl) => { match ctl { // All control frames MUST have a payload length of 125 bytes or less // and MUST NOT be fragmented. (RFC 6455) - _ if !frame.is_final() => { + _ if !frame.header().is_final => { Err(Error::Protocol("Fragmented control frame".into())) } _ if frame.payload().len() > 125 => { @@ -309,7 +312,7 @@ impl WebSocket { } OpCode::Data(data) => { - let fin = frame.is_final(); + let fin = frame.header().is_final; match data { OpData::Continue => { if let Some(ref mut msg) = self.incomplete { @@ -423,7 +426,7 @@ impl WebSocket { Role::Client => { // 5. If the data is being sent by the client, the frame(s) MUST be // masked as defined in Section 5.3. (RFC 6455) - frame.set_mask(); + frame.set_random_mask(); } } let res = self.socket.write_frame(frame);