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