Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

message: add has_any_attribute() #25

Merged
merged 1 commit into from
Oct 10, 2024
Merged

message: add has_any_attribute() #25

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 63 additions & 42 deletions stun-types/src/message.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1468,15 +1468,7 @@
.map(|attr| attr.padded_len())
.sum::<usize>();
let mut ret = vec![0; MessageHeader::LENGTH + attr_size];
self.msg_type.write_into(&mut ret[..2]);
let transaction: u128 = self.transaction_id.into();
let tid = (MAGIC_COOKIE as u128) << 96 | transaction & 0xffff_ffff_ffff_ffff_ffff_ffff;
BigEndian::write_u128(&mut ret[4..20], tid);
BigEndian::write_u16(&mut ret[2..4], attr_size as u16);
let mut offset = MessageHeader::LENGTH;
for attr in &self.attributes {
offset += attr.write_into(&mut ret[offset..]).unwrap();
}
let _ = self.write_into(&mut ret);
ret
}

Expand Down Expand Up @@ -1574,16 +1566,32 @@
credentials: &MessageIntegrityCredentials,
algorithm: IntegrityAlgorithm,
) -> Result<(), StunWriteError> {
if self.has_attribute(MessageIntegrity::TYPE) && algorithm == IntegrityAlgorithm::Sha1 {
return Err(StunWriteError::AttributeExists(MessageIntegrity::TYPE));
}
if self.has_attribute(MessageIntegritySha256::TYPE) {
return Err(StunWriteError::AttributeExists(
MessageIntegritySha256::TYPE,
));
let mut atypes = [AttributeType::new(0); 3];
let mut i = 0;
atypes[i] = match algorithm {
IntegrityAlgorithm::Sha1 => MessageIntegrity::TYPE,
IntegrityAlgorithm::Sha256 => MessageIntegritySha256::TYPE,
};
i += 1;
if algorithm == IntegrityAlgorithm::Sha1 {
atypes[i] = MessageIntegritySha256::TYPE;
i += 1;
}
if self.has_attribute(Fingerprint::TYPE) {
return Err(StunWriteError::FingerprintExists);
atypes[i] = Fingerprint::TYPE;
i += 1;

match self.has_any_attribute(&atypes[..i]) {
// can't validly add generic attributes after message integrity or fingerprint
Some(MessageIntegrity::TYPE) => {

Check warning on line 1585 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1585 

Added line #L1585 was not covered by tests
return Err(StunWriteError::AttributeExists(MessageIntegrity::TYPE))
}
Some(MessageIntegritySha256::TYPE) => {

Check warning on line 1588 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1588 

Added line #L1588 was not covered by tests
return Err(StunWriteError::AttributeExists(
MessageIntegritySha256::TYPE,

Check warning on line 1590 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1590 

Added line #L1590 was not covered by tests
));
}
Some(Fingerprint::TYPE) => return Err(StunWriteError::FingerprintExists),
_ => (),

Check warning on line 1594 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1594 

Added line #L1594 was not covered by tests
}

self.add_message_integrity_unchecked(credentials, algorithm);
Expand Down Expand Up @@ -1720,18 +1728,20 @@
}
_ => (),
}
if self.has_attribute(ty) {
return Err(StunWriteError::AttributeExists(ty));
}
// can't validly add generic attributes after message integrity or fingerprint
if self.has_attribute(MessageIntegrity::TYPE) {
return Err(StunWriteError::MessageIntegrityExists);
}
if self.has_attribute(MessageIntegritySha256::TYPE) {
return Err(StunWriteError::MessageIntegrityExists);
}
if self.has_attribute(Fingerprint::TYPE) {
return Err(StunWriteError::FingerprintExists);
match self.has_any_attribute(&[
ty,
MessageIntegrity::TYPE,
MessageIntegritySha256::TYPE,
Fingerprint::TYPE,

Check warning on line 1735 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1732-L1735 

Added lines #L1732 - L1735 were not covered by tests
]) {
// can't validly add generic attributes after message integrity or fingerprint
Some(MessageIntegrity::TYPE) => return Err(StunWriteError::MessageIntegrityExists),
Some(MessageIntegritySha256::TYPE) => {

Check warning on line 1739 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1739 

Added line #L1739 was not covered by tests
return Err(StunWriteError::MessageIntegrityExists)
}
Some(Fingerprint::TYPE) => return Err(StunWriteError::FingerprintExists),
Some(typ) if typ == ty => return Err(StunWriteError::AttributeExists(ty)),
_ => (),

Check warning on line 1744 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1743-L1744 

Added lines #L1743 - L1744 were not covered by tests
}
self.attributes.push(AttrOrRaw::Attr(attr));
self.attribute_types.push(ty);
Expand Down Expand Up @@ -1777,18 +1787,20 @@
}
_ => (),
}
if self.has_attribute(ty) {
return Err(StunWriteError::AttributeExists(ty));
}
// can't validly add generic attributes after message integrity or fingerprint
if self.has_attribute(MessageIntegrity::TYPE) {
return Err(StunWriteError::MessageIntegrityExists);
}
if self.has_attribute(MessageIntegritySha256::TYPE) {
return Err(StunWriteError::MessageIntegrityExists);
}
if self.has_attribute(Fingerprint::TYPE) {
return Err(StunWriteError::FingerprintExists);
match self.has_any_attribute(&[
ty,
MessageIntegrity::TYPE,
MessageIntegritySha256::TYPE,
Fingerprint::TYPE,

Check warning on line 1794 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1791-L1794 

Added lines #L1791 - L1794 were not covered by tests
]) {
// can't validly add generic attributes after message integrity or fingerprint
Some(MessageIntegrity::TYPE) => return Err(StunWriteError::MessageIntegrityExists),
Some(MessageIntegritySha256::TYPE) => {
return Err(StunWriteError::MessageIntegrityExists)

Check warning on line 1799 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1797-L1799 

Added lines #L1797 - L1799 were not covered by tests
}
Some(Fingerprint::TYPE) => return Err(StunWriteError::FingerprintExists),

Check warning on line 1801 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1801 

Added line #L1801 was not covered by tests
Some(typ) if typ == ty => return Err(StunWriteError::AttributeExists(ty)),
_ => (),

Check warning on line 1803 in stun-types/src/message.rs

View check run for this annotation

Codecov / codecov/patch

stun-types/src/message.rs#L1803 

Added line #L1803 was not covered by tests
}
self.attributes.push(AttrOrRaw::Raw(attr));
self.attribute_types.push(ty);
Expand All @@ -1799,6 +1811,15 @@
pub fn has_attribute(&self, atype: AttributeType) -> bool {
self.attribute_types.iter().any(|&ty| ty == atype)
}

/// Return whether this [`MessageBuilder`] contains any of the provided attributes and
/// returns the attribute found.
pub fn has_any_attribute(&self, atypes: &[AttributeType]) -> Option<AttributeType> {
self.attribute_types
.iter()
.find(|&ty| atypes.contains(ty))
.cloned()
}
}

#[cfg(test)]
Expand Down
Loading