diff --git a/crates/stun-types/src/attributes/fingerprint.rs b/crates/stun-types/src/attributes/fingerprint.rs index b8ab11e..661c9cd 100644 --- a/crates/stun-types/src/attributes/fingerprint.rs +++ b/crates/stun-types/src/attributes/fingerprint.rs @@ -1,5 +1,6 @@ -use super::Attribute; +use super::{Attribute, ATTRIBUTE_HEADER_LEN}; use crate::builder::MessageBuilder; +use crate::header::STUN_HEADER_LENGTH; use crate::parse::{AttrSpan, Message}; use crate::{Error, NE}; use byteorder::ReadBytesExt; @@ -62,29 +63,39 @@ impl Attribute<'_> for Fingerprint { const TYPE: u16 = 0x8028; fn decode(_: Self::Context, msg: &mut Message, attr: AttrSpan) -> Result { - let mut value = attr.get_value(msg.buffer()); + msg.with_msg_len( + u16::try_from(attr.padding_end - STUN_HEADER_LENGTH)?, + |msg| { + let mut value = attr.get_value(msg.buffer()); - if value.len() != 4 { - return Err(Error::InvalidData("fingerprint value must be 4 bytes")); - } - - let attr_value = value.read_u32::()?; - - let data = &msg.buffer()[..attr.begin]; + if value.len() != 4 { + return Err(Error::InvalidData("fingerprint value must be 4 bytes")); + } - let crc = Self::crc32(data) ^ 0x5354554e; + let attr_value = value.read_u32::()?; + let data = &msg.buffer()[..attr.begin - ATTRIBUTE_HEADER_LEN]; + let crc = Self::crc32(data) ^ 0x5354554e; - if crc != attr_value { - return Err(Error::InvalidData("failed to verify message fingerprint")); - } + if crc != attr_value { + return Err(Error::InvalidData("failed to verify message fingerprint")); + } - Ok(Self) + Ok(Self) + }, + ) } fn encode(&self, _: Self::Context, builder: &mut MessageBuilder) -> Result<(), Error> { - let data = builder.buffer(); - let data = &data[..data.len() - 4]; + // First set the length of the message to the end of the fingerprint attribute + // 4 bytes containing type and length is already written into the buffer + let message_length_with_fingerprint_attribute = + (builder.buffer().len() + 4) - STUN_HEADER_LENGTH; + builder.set_len(message_length_with_fingerprint_attribute.try_into()?); + + // Calculate the checksum + let data = builder.buffer(); + let data = &data[..data.len() - ATTRIBUTE_HEADER_LEN]; let crc = Self::crc32(data) ^ 0x5354554e; builder.buffer().put_u32(crc); diff --git a/crates/stun-types/src/attributes/integrity.rs b/crates/stun-types/src/attributes/integrity.rs index 82a05e9..238d2db 100644 --- a/crates/stun-types/src/attributes/integrity.rs +++ b/crates/stun-types/src/attributes/integrity.rs @@ -1,5 +1,6 @@ use super::{Attribute, ATTRIBUTE_HEADER_LEN}; use crate::builder::MessageBuilder; +use crate::header::STUN_HEADER_LENGTH; use crate::parse::{AttrSpan, Message}; use crate::Error; use hmac::digest::core_api::BlockSizeUser; @@ -59,9 +60,7 @@ impl<'k> Attribute<'_> for MessageIntegrity<'k> { let hmac: SimpleHmac = SimpleHmac::new_from_slice(&ctx.0) .map_err(|_| Error::InvalidData("invalid key length"))?; - message_integrity_encode(hmac, builder); - - Ok(()) + message_integrity_encode(hmac, builder) } fn encode_len(&self) -> Result { @@ -90,9 +89,7 @@ impl<'k> Attribute<'_> for MessageIntegritySha256<'k> { let hmac: SimpleHmac = SimpleHmac::new_from_slice(&ctx.0) .map_err(|_| Error::InvalidData("invalid key length"))?; - message_integrity_encode(hmac, builder); - - Ok(()) + message_integrity_encode(hmac, builder) } fn encode_len(&self) -> Result { @@ -114,35 +111,41 @@ where // the end of the MESSAGE-INTEGRITY attribute. // The length of the message is temprorarily set to the end of the previous attribute - msg.with_msg_len(u16::try_from(attr.padding_end)?, |msg| { - // Get the digest from the received attribute - let received_digest = attr.get_value(msg.buffer()); - - // Get all bytes before the integrity attribute to calculate the hmac over - let message = &msg.buffer()[..attr.begin - ATTRIBUTE_HEADER_LEN]; - - // Calculate the expected digest, - Update::update(&mut hmac, message); - let calculated_digest = hmac.finalize().into_bytes(); - - // Compare the received and calculated digest - if calculated_digest.as_slice() != received_digest { - return Err(Error::InvalidData("failed to verify message integrity")); - } - - Ok(()) - }) + msg.with_msg_len( + u16::try_from(attr.padding_end - STUN_HEADER_LENGTH)?, + |msg| { + // Get the digest from the received attribute + let received_digest = attr.get_value(msg.buffer()); + + // Get all bytes before the integrity attribute to calculate the hmac over + let message = &msg.buffer()[..attr.begin - ATTRIBUTE_HEADER_LEN]; + + // Calculate the expected digest, + Update::update(&mut hmac, message); + let calculated_digest = hmac.finalize().into_bytes(); + + // Compare the received and calculated digest + if calculated_digest.as_slice() != received_digest { + return Err(Error::InvalidData("failed to verify message integrity")); + } + + Ok(()) + }, + ) } -fn message_integrity_encode(mut hmac: SimpleHmac, builder: &mut MessageBuilder) +fn message_integrity_encode( + mut hmac: SimpleHmac, + builder: &mut MessageBuilder, +) -> Result<(), Error> where D: Digest + BlockSizeUser, { // 4 bytes containing type and length is already written into the buffer let message_length_with_integrity_attribute = - builder.buffer().len() + ::output_size(); + (builder.buffer().len() + ::output_size()) - STUN_HEADER_LENGTH; - builder.set_len(message_length_with_integrity_attribute.try_into().unwrap()); + builder.set_len(message_length_with_integrity_attribute.try_into()?); // Calculate the digest of the message up until the previous attribute let data = builder.buffer(); @@ -151,6 +154,8 @@ where let digest = hmac.finalize().into_bytes(); builder.buffer().extend_from_slice(&digest); + + Ok(()) } #[cfg(test)] diff --git a/crates/stun-types/src/header.rs b/crates/stun-types/src/header.rs index 16d0c3c..bc10fe9 100644 --- a/crates/stun-types/src/header.rs +++ b/crates/stun-types/src/header.rs @@ -2,6 +2,8 @@ use crate::{Error, COOKIE}; use bitfield::bitfield; use std::convert::TryFrom; +pub(crate) const STUN_HEADER_LENGTH: usize = 20; + bitfield! { /// Internal bitfield representing the STUN message head pub struct MessageHead(u32);