1use alloc::vec::Vec;
14use byteorder::{BigEndian, ByteOrder};
15use core::net::SocketAddr;
16use stun_types::message::Message;
17use stun_types::message::MessageHeader;
18use stun_types::message::MessageType;
19use stun_types::message::MessageWriteMutSlice;
20use stun_types::message::MessageWriteVec;
21use stun_types::message::TransactionId;
22use stun_types::prelude::AttributeExt;
23use stun_types::prelude::MessageWrite;
24use stun_types::prelude::MessageWriteExt;
25
26use stun_proto::agent::Transmit;
27use stun_types::TransportType;
28
29use crate::attribute::Data as AData;
30use crate::attribute::XorPeerAddress;
31use crate::message::{DATA, SEND};
32
33#[derive(Debug)]
35pub struct TransmitBuild<T: DelayedTransmitBuild + core::fmt::Debug> {
36 pub data: T,
38 pub transport: TransportType,
40 pub from: SocketAddr,
42 pub to: SocketAddr,
44}
45
46impl<T: DelayedTransmitBuild + core::fmt::Debug> TransmitBuild<T> {
47 pub fn new(data: T, transport: TransportType, from: SocketAddr, to: SocketAddr) -> Self {
49 Self {
50 data,
51 transport,
52 from,
53 to,
54 }
55 }
56
57 pub fn build(self) -> Transmit<Vec<u8>> {
59 Transmit {
60 data: self.data.build(),
61 transport: self.transport,
62 from: self.from,
63 to: self.to,
64 }
65 }
66
67 pub fn write_into(self, dest: &mut [u8]) -> Transmit<&mut [u8]> {
69 let len = self.data.write_into(dest);
70 Transmit {
71 data: &mut dest[..len],
72 transport: self.transport,
73 from: self.from,
74 to: self.to,
75 }
76 }
77}
78
79pub trait DelayedTransmitBuild {
81 fn build(self) -> Vec<u8>;
83 fn len(&self) -> usize;
85 fn is_empty(&self) -> bool {
87 self.len() == 0
88 }
89 fn write_into(self, data: &mut [u8]) -> usize;
93}
94
95impl DelayedTransmitBuild for Vec<u8> {
96 fn len(&self) -> usize {
97 self.len()
98 }
99 fn build(self) -> Vec<u8> {
100 self
101 }
102 fn is_empty(&self) -> bool {
103 self.is_empty()
104 }
105 fn write_into(self, data: &mut [u8]) -> usize {
106 data[..self.len()].copy_from_slice(&self);
107 self.len()
108 }
109}
110
111impl DelayedTransmitBuild for &[u8] {
112 fn len(&self) -> usize {
113 (**self).len()
114 }
115 fn build(self) -> Vec<u8> {
116 self.to_vec()
117 }
118 fn is_empty(&self) -> bool {
119 (**self).is_empty()
120 }
121 fn write_into(self, data: &mut [u8]) -> usize {
122 data[..self.len()].copy_from_slice(self);
123 self.len()
124 }
125}
126
127#[derive(Debug)]
129pub struct DelayedChannel<T: AsRef<[u8]> + core::fmt::Debug> {
130 data: T,
131 channel_id: u16,
132}
133
134impl<T: AsRef<[u8]> + core::fmt::Debug> DelayedChannel<T> {
135 pub fn new(channel_id: u16, data: T) -> Self {
137 Self { channel_id, data }
138 }
139}
140
141impl<T: AsRef<[u8]> + core::fmt::Debug> DelayedChannel<T> {
142 fn write_header_into(&self, len: u16, dest: &mut [u8]) {
143 BigEndian::write_u16(&mut dest[..2], self.channel_id);
144 BigEndian::write_u16(&mut dest[2..4], len);
145 }
146}
147
148impl<T: AsRef<[u8]> + core::fmt::Debug> DelayedTransmitBuild for DelayedChannel<T> {
149 fn len(&self) -> usize {
150 self.data.as_ref().len() + 4
151 }
152
153 fn build(self) -> Vec<u8> {
154 let data = self.data.as_ref();
155 let data_len = data.len();
156 let mut header = [0; 4];
157 self.write_header_into(data_len as u16, &mut header);
158 let mut out = Vec::with_capacity(4 + data_len);
159 out.extend(header.as_slice());
160 out.extend_from_slice(data);
161 out
162 }
163
164 fn write_into(self, dest: &mut [u8]) -> usize {
165 let data = self.data.as_ref();
166 let data_len = data.len();
167 self.write_header_into(data_len as u16, dest);
168 dest[4..4 + data_len].copy_from_slice(data);
169 data_len + 4
170 }
171}
172
173#[derive(Debug)]
175pub struct DelayedMessage<T: AsRef<[u8]> + core::fmt::Debug> {
176 data: T,
177 peer_addr: SocketAddr,
178 for_client: bool,
179}
180
181impl<T: AsRef<[u8]> + core::fmt::Debug> DelayedMessage<T> {
182 pub fn for_client(peer_addr: SocketAddr, data: T) -> Self {
184 Self {
185 peer_addr,
186 data,
187 for_client: true,
188 }
189 }
190 pub fn for_server(peer_addr: SocketAddr, data: T) -> Self {
192 Self {
193 peer_addr,
194 data,
195 for_client: false,
196 }
197 }
198}
199
200impl<T: AsRef<[u8]> + core::fmt::Debug> DelayedTransmitBuild for DelayedMessage<T> {
201 fn len(&self) -> usize {
202 let xor_peer_addr = XorPeerAddress::new(self.peer_addr, 0.into());
203 let data = AData::new(self.data.as_ref());
204 MessageHeader::LENGTH + xor_peer_addr.padded_len() + data.padded_len()
205 }
206
207 fn build(self) -> Vec<u8> {
208 let transaction_id = TransactionId::generate();
209 let method = if self.for_client { DATA } else { SEND };
210 let mut msg = Message::builder(
211 MessageType::from_class_method(
212 stun_proto::types::message::MessageClass::Indication,
213 method,
214 ),
215 transaction_id,
216 MessageWriteVec::with_capacity(self.len()),
217 );
218 let xor_peer_address = XorPeerAddress::new(self.peer_addr, transaction_id);
219 msg.add_attribute(&xor_peer_address).unwrap();
220 let data = AData::new(self.data.as_ref());
221 msg.add_attribute(&data).unwrap();
222 msg.finish()
223 }
224
225 fn write_into(self, dest: &mut [u8]) -> usize {
226 let transaction_id = TransactionId::generate();
227 let mut msg = Message::builder(
228 MessageType::from_class_method(
229 stun_proto::types::message::MessageClass::Indication,
230 SEND,
231 ),
232 transaction_id,
233 MessageWriteMutSlice::new(dest),
234 );
235 let xor_peer_address = XorPeerAddress::new(self.peer_addr, transaction_id);
236 msg.add_attribute(&xor_peer_address).unwrap();
237 let data = AData::new(self.data.as_ref());
238 msg.add_attribute(&data).unwrap();
239 msg.finish()
240 }
241}
242
243#[cfg(test)]
244mod tests {
245 use super::*;
246
247 use alloc::vec;
248
249 #[test]
250 fn test_delayed_vecu8() {
251 let data = vec![7; 7];
252 assert_eq!(DelayedTransmitBuild::len(&data), data.len());
253 assert_eq!(DelayedTransmitBuild::build(data.clone()), data);
254 assert!(!DelayedTransmitBuild::is_empty(&data));
255 assert!(DelayedTransmitBuild::is_empty(&Vec::new()));
256 let mut out = vec![0; 8];
257 assert_eq!(DelayedTransmitBuild::write_into(data.clone(), &mut out), 7);
258 }
259
260 #[test]
261 fn test_delayed_u8slice() {
262 let data = [7; 7];
263 assert_eq!(DelayedTransmitBuild::len(&data.as_slice()), data.len());
264 assert_eq!(DelayedTransmitBuild::build(data.as_slice()), data);
265 assert!(!DelayedTransmitBuild::is_empty(&data.as_slice()));
266 assert!(DelayedTransmitBuild::is_empty(&[].as_slice()));
267 let mut out = [0; 8];
268 assert_eq!(
269 DelayedTransmitBuild::write_into(data.as_slice(), &mut out),
270 7
271 );
272 }
273}