Skip to content

Commit 52b0a2b

Browse files
ystreetystreet
committed
message: implement write_into() a slice
Improves write throughput with supported attributes by around 25%
1 parent 0eef9c8 commit 52b0a2b

19 files changed

+817
-244
lines changed

‎stun-proto/examples/stund.rs

+6-6
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,10 @@ fn handle_binding_request<'a>(
4545
}
4646

4747
let mut response = Message::builder_success(msg);
48-
response.add_attribute(&XorMappedAddress::new(from, msg.transaction_id()))?;
48+
let xor_addr = XorMappedAddress::new(from, msg.transaction_id());
49+
response.add_attribute(&xor_addr)?;
4950
response.add_fingerprint()?;
50-
Ok(response)
51+
Ok(response.into_owned())
5152
}
5253

5354
fn handle_incoming_data<'a>(
@@ -75,10 +76,9 @@ fn handle_incoming_data<'a>(
7576
}
7677
} else {
7778
let mut response = Message::builder_error(&msg);
78-
response
79-
.add_attribute(&ErrorCode::new(400, "Bad Request").unwrap())
80-
.unwrap();
81-
return Some((response, from));
79+
let error = ErrorCode::new(400, "Bad Request").unwrap();
80+
response.add_attribute(&error).unwrap();
81+
return Some((response.into_owned(), from));
8282
}
8383
None
8484
}

‎stun-proto/src/agent.rs

+14-39
Original file line numberDiff line numberDiff line change
@@ -791,12 +791,8 @@ pub(crate) mod tests {
791791

792792
println!("generate response");
793793
let mut response = Message::builder_success(&request);
794-
response
795-
.add_attribute(&XorMappedAddress::new(
796-
transmit.from,
797-
request.transaction_id(),
798-
))
799-
.unwrap();
794+
let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
795+
response.add_attribute(&xor_addr).unwrap();
800796
response
801797
.add_message_integrity(&remote_credentials.into(), IntegrityAlgorithm::Sha1)
802798
.unwrap();
@@ -865,12 +861,8 @@ pub(crate) mod tests {
865861
let request = Message::from_bytes(&transmit.data).unwrap();
866862

867863
let mut response = Message::builder_success(&request);
868-
response
869-
.add_attribute(&XorMappedAddress::new(
870-
transmit.from,
871-
request.transaction_id(),
872-
))
873-
.unwrap();
864+
let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
865+
response.add_attribute(&xor_addr).unwrap();
874866

875867
let data = response.build();
876868
let to = transmit.to;
@@ -905,12 +897,8 @@ pub(crate) mod tests {
905897
let request = Message::from_bytes(&transmit.data).unwrap();
906898

907899
let mut response = Message::builder_success(&request);
908-
response
909-
.add_attribute(&XorMappedAddress::new(
910-
transmit.from,
911-
request.transaction_id(),
912-
))
913-
.unwrap();
900+
let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
901+
response.add_attribute(&xor_addr).unwrap();
914902

915903
let data = response.build();
916904
let to = transmit.to;
@@ -944,12 +932,8 @@ pub(crate) mod tests {
944932
let request = Message::from_bytes(&transmit.data).unwrap();
945933

946934
let mut response = Message::builder_success(&request);
947-
response
948-
.add_attribute(&XorMappedAddress::new(
949-
transmit.from,
950-
request.transaction_id(),
951-
))
952-
.unwrap();
935+
let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
936+
response.add_attribute(&xor_addr).unwrap();
953937

954938
let data = response.build();
955939
let to = transmit.to;
@@ -984,12 +968,8 @@ pub(crate) mod tests {
984968
let request = Message::from_bytes(&transmit.data).unwrap();
985969

986970
let mut response = Message::builder_success(&request);
987-
response
988-
.add_attribute(&XorMappedAddress::new(
989-
transmit.from,
990-
request.transaction_id(),
991-
))
992-
.unwrap();
971+
let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
972+
response.add_attribute(&xor_addr).unwrap();
993973
// wrong credentials, should be `remote_credentials`
994974
response
995975
.add_message_integrity(&local_credentials.into(), IntegrityAlgorithm::Sha1)
@@ -1021,12 +1001,8 @@ pub(crate) mod tests {
10211001
let request = Message::from_bytes(&transmit.data).unwrap();
10221002

10231003
let mut response = Message::builder_success(&request);
1024-
response
1025-
.add_attribute(&XorMappedAddress::new(
1026-
transmit.from,
1027-
request.transaction_id(),
1028-
))
1029-
.unwrap();
1004+
let xor_addr = XorMappedAddress::new(transmit.from, request.transaction_id());
1005+
response.add_attribute(&xor_addr).unwrap();
10301006

10311007
let data = response.build();
10321008
let to = transmit.to;
@@ -1117,9 +1093,8 @@ pub(crate) mod tests {
11171093
let request = Message::from_bytes(transmit.data()).unwrap();
11181094

11191095
let mut response = Message::builder_success(&request);
1120-
response
1121-
.add_attribute(&XorMappedAddress::new(transmit.from, transaction_id))
1122-
.unwrap();
1096+
let xor_addr = XorMappedAddress::new(transmit.from, transaction_id);
1097+
response.add_attribute(&xor_addr).unwrap();
11231098

11241099
assert!(matches!(
11251100
agent.send(msg, remote_addr, Instant::now()),

‎stun-types/Cargo.toml

+1
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ sha2 = "0.10"
2222
tracing.workspace = true
2323
arbitrary ={ workspace = true, optional = true }
2424
thiserror.workspace = true
25+
smallvec = "1"
2526

2627
[dev-dependencies]
2728
tracing-subscriber = "0.3"

‎stun-types/benches/message_write.rs

+92-2
Original file line numberDiff line numberDiff line change
@@ -13,18 +13,23 @@ use stun_types::message::{
1313
TransactionId, BINDING,
1414
};
1515

16-
fn builder_with_attribute<'a>(attr: impl Into<RawAttribute<'a>>) -> MessageBuilder<'a> {
16+
fn builder_with_attribute(attr: &dyn AttributeWrite) -> MessageBuilder {
1717
let mut msg = Message::builder_request(BINDING);
1818
msg.add_attribute(attr).unwrap();
1919
msg
2020
}
2121

22-
fn build_with_attribute<'a>(attr: impl Into<RawAttribute<'a>>) {
22+
fn build_with_attribute(attr: &dyn AttributeWrite) {
2323
let mut msg = Message::builder_request(BINDING);
2424
msg.add_attribute(attr).unwrap();
2525
let _data = msg.build();
2626
}
2727

28+
fn write_into_with_attribute(attr: &dyn AttributeWrite, dest: &mut [u8]) {
29+
let msg = builder_with_attribute(attr);
30+
msg.write_into(dest).unwrap();
31+
}
32+
2833
fn bench_message_write(c: &mut Criterion) {
2934
let software = Software::new("stun-types").unwrap();
3035
let addr = "192.168.10.200:9876".parse().unwrap();
@@ -122,6 +127,91 @@ fn bench_message_write(c: &mut Criterion) {
122127
},
123128
);
124129
group.finish();
130+
131+
let mut group = c.benchmark_group("Message/WriteInto");
132+
let mut scratch = vec![0; 1 << 8];
133+
134+
group.throughput(criterion::Throughput::Bytes(
135+
builder_with_attribute(&software).build().len() as u64,
136+
));
137+
group.bench_with_input(
138+
BenchmarkId::from_parameter("Software"),
139+
&software,
140+
|b, software| b.iter(|| write_into_with_attribute(software, &mut scratch)),
141+
);
142+
group.bench_with_input(
143+
BenchmarkId::from_parameter("Attributes/9"),
144+
&(
145+
&software,
146+
&xor_mapped_address,
147+
&nonce,
148+
&alt_server,
149+
&alt_domain,
150+
&priority,
151+
&controlled,
152+
&controlling,
153+
&use_candidate,
154+
),
155+
|b, attrs| {
156+
b.iter(|| {
157+
let mut msg = builder_with_attribute(attrs.0);
158+
msg.add_attribute(attrs.1).unwrap();
159+
msg.add_attribute(attrs.2).unwrap();
160+
msg.add_attribute(attrs.3).unwrap();
161+
msg.add_attribute(attrs.4).unwrap();
162+
msg.add_attribute(attrs.5).unwrap();
163+
msg.add_attribute(attrs.6).unwrap();
164+
msg.add_attribute(attrs.7).unwrap();
165+
msg.add_attribute(attrs.8).unwrap();
166+
msg.write_into(&mut scratch).unwrap();
167+
})
168+
},
169+
);
170+
group.bench_with_input(
171+
BenchmarkId::from_parameter("XorMappedAddress"),
172+
&xor_mapped_address,
173+
|b, xor_mapped_address| {
174+
b.iter(|| write_into_with_attribute(xor_mapped_address, &mut scratch));
175+
},
176+
);
177+
group.bench_with_input(
178+
BenchmarkId::from_parameter("XorMappedAddress+Fingerprint"),
179+
&xor_mapped_address,
180+
|b, xor_mapped_address| {
181+
b.iter(|| {
182+
let mut msg = builder_with_attribute(xor_mapped_address);
183+
msg.add_fingerprint().unwrap();
184+
msg.write_into(&mut scratch).unwrap();
185+
})
186+
},
187+
);
188+
group.bench_with_input(
189+
BenchmarkId::from_parameter("XorMappedAddress+ShortTermIntegritySha1+Fingerprint"),
190+
&(&xor_mapped_address, &short_term_integrity),
191+
|b, &(xor_mapped_address, short_term_integrity)| {
192+
b.iter(|| {
193+
let mut msg = builder_with_attribute(xor_mapped_address);
194+
msg.add_message_integrity(short_term_integrity, IntegrityAlgorithm::Sha1)
195+
.unwrap();
196+
msg.add_fingerprint().unwrap();
197+
msg.write_into(&mut scratch).unwrap();
198+
})
199+
},
200+
);
201+
group.bench_with_input(
202+
BenchmarkId::from_parameter("XorMappedAddress+ShortTermIntegritySha256+Fingerprint"),
203+
&(&xor_mapped_address, &short_term_integrity),
204+
|b, &(xor_mapped_address, short_term_integrity)| {
205+
b.iter(|| {
206+
let mut msg = builder_with_attribute(xor_mapped_address);
207+
msg.add_message_integrity(short_term_integrity, IntegrityAlgorithm::Sha256)
208+
.unwrap();
209+
msg.add_fingerprint().unwrap();
210+
msg.write_into(&mut scratch).unwrap();
211+
})
212+
},
213+
);
214+
group.finish();
125215
}
126216

127217
criterion_group!(message_parse, bench_message_write);

‎stun-types/src/attribute/address.rs

+27-10
Original file line numberDiff line numberDiff line change
@@ -70,20 +70,14 @@ impl MappedSocketAddr {
7070
/// Convert this [`MappedSocketAddr`] into a [`RawAttribute`]
7171
pub fn to_raw<'a>(&self, atype: AttributeType) -> RawAttribute<'a> {
7272
match self.addr {
73-
SocketAddr::V4(addr) => {
73+
SocketAddr::V4(_addr) => {
7474
let mut buf = [0; 8];
75-
buf[1] = AddressFamily::IPV4.to_byte();
76-
BigEndian::write_u16(&mut buf[2..4], addr.port());
77-
let octets = u32::from(*addr.ip());
78-
BigEndian::write_u32(&mut buf[4..8], octets);
75+
self.write_into_unchecked(&mut buf);
7976
RawAttribute::new(atype, &buf).into_owned()
8077
}
81-
SocketAddr::V6(addr) => {
78+
SocketAddr::V6(_addr) => {
8279
let mut buf = [0; 20];
83-
buf[1] = AddressFamily::IPV6.to_byte();
84-
BigEndian::write_u16(&mut buf[2..4], addr.port());
85-
let octets = u128::from(*addr.ip());
86-
BigEndian::write_u128(&mut buf[4..20], octets);
80+
self.write_into_unchecked(&mut buf);
8781
RawAttribute::new(atype, &buf).into_owned()
8882
}
8983
}
@@ -122,6 +116,25 @@ impl MappedSocketAddr {
122116
pub fn addr(&self) -> SocketAddr {
123117
self.addr
124118
}
119+
120+
pub(crate) fn write_into_unchecked(&self, dest: &mut [u8]) {
121+
match self.addr {
122+
SocketAddr::V4(addr) => {
123+
dest[0] = 0x0;
124+
dest[1] = AddressFamily::IPV4.to_byte();
125+
BigEndian::write_u16(&mut dest[2..4], addr.port());
126+
let octets = u32::from(*addr.ip());
127+
BigEndian::write_u32(&mut dest[4..8], octets);
128+
}
129+
SocketAddr::V6(addr) => {
130+
dest[0] = 0x0;
131+
dest[1] = AddressFamily::IPV6.to_byte();
132+
BigEndian::write_u16(&mut dest[2..4], addr.port());
133+
let octets = u128::from(*addr.ip());
134+
BigEndian::write_u128(&mut dest[4..20], octets);
135+
}
136+
}
137+
}
125138
}
126139

127140
impl std::fmt::Display for MappedSocketAddr {
@@ -190,6 +203,10 @@ impl XorSocketAddr {
190203
pub(crate) fn addr(&self, transaction: TransactionId) -> SocketAddr {
191204
XorSocketAddr::xor_addr(self.addr.addr(), transaction)
192205
}
206+
207+
pub(crate) fn write_into_unchecked(&self, dest: &mut [u8]) {
208+
self.addr.write_into_unchecked(dest)
209+
}
193210
}
194211

195212
impl std::fmt::Display for XorSocketAddr {

‎stun-types/src/attribute/alternate.rs

+35-9
Original file line numberDiff line numberDiff line change
@@ -11,25 +11,37 @@ use std::net::SocketAddr;
1111

1212
use crate::message::StunParseError;
1313

14-
use super::{Attribute, AttributeType, MappedSocketAddr, RawAttribute};
14+
use super::{
15+
Attribute, AttributeExt, AttributeStaticType, AttributeType, AttributeWrite, AttributeWriteExt,
16+
MappedSocketAddr, RawAttribute,
17+
};
1518

1619
/// The AlternateServer [`Attribute`]
1720
#[derive(Debug, Clone, PartialEq, Eq)]
1821
pub struct AlternateServer {
1922
addr: MappedSocketAddr,
2023
}
2124

22-
impl Attribute for AlternateServer {
25+
impl AttributeStaticType for AlternateServer {
2326
const TYPE: AttributeType = AttributeType(0x8023);
27+
}
28+
impl Attribute for AlternateServer {
29+
fn get_type(&self) -> AttributeType {
30+
Self::TYPE
31+
}
2432

2533
fn length(&self) -> u16 {
2634
self.addr.length()
2735
}
2836
}
2937

30-
impl<'a> From<&AlternateServer> for RawAttribute<'a> {
31-
fn from(value: &AlternateServer) -> RawAttribute<'a> {
32-
value.addr.to_raw(AlternateServer::TYPE)
38+
impl AttributeWrite for AlternateServer {
39+
fn to_raw(&self) -> RawAttribute {
40+
self.addr.to_raw(AlternateServer::TYPE)
41+
}
42+
fn write_into_unchecked(&self, dest: &mut [u8]) {
43+
self.write_header_unchecked(dest);
44+
self.addr.write_into_unchecked(&mut dest[4..]);
3345
}
3446
}
3547

@@ -87,9 +99,14 @@ pub struct AlternateDomain {
8799
domain: String,
88100
}
89101

90-
impl Attribute for AlternateDomain {
102+
impl AttributeStaticType for AlternateDomain {
91103
const TYPE: AttributeType = AttributeType(0x8003);
104+
}
92105

106+
impl Attribute for AlternateDomain {
107+
fn get_type(&self) -> AttributeType {
108+
Self::TYPE
109+
}
93110
fn length(&self) -> u16 {
94111
self.domain.len() as u16
95112
}
@@ -107,9 +124,18 @@ impl<'a> TryFrom<&RawAttribute<'a>> for AlternateDomain {
107124
})
108125
}
109126
}
110-
impl<'a> From<&AlternateDomain> for RawAttribute<'a> {
111-
fn from(value: &AlternateDomain) -> RawAttribute<'a> {
112-
RawAttribute::new(AlternateDomain::TYPE, value.domain.as_bytes()).into_owned()
127+
impl AttributeWrite for AlternateDomain {
128+
fn to_raw(&self) -> RawAttribute {
129+
RawAttribute::new(AlternateDomain::TYPE, self.domain.as_bytes())
130+
}
131+
fn write_into_unchecked(&self, dest: &mut [u8]) {
132+
let len = self.padded_len();
133+
self.write_header_unchecked(dest);
134+
dest[4..4 + self.domain.len()].copy_from_slice(self.domain.as_bytes());
135+
let offset = 4 + self.domain.len();
136+
if len > offset {
137+
dest[offset..len].fill(0);
138+
}
113139
}
114140
}
115141

0 commit comments

Comments
 (0)