Skip to main content

turn_types/
transmit.rs

1// Copyright (C) 2025 Matthew Waters <matthew@centricular.com>
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8//
9// SPDX-License-Identifier: MIT OR Apache-2.0
10
11//! Transmit structures and helpers.
12
13use 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/// A piece of data that needs to be built before it can be transmitted.
34#[derive(Debug)]
35pub struct TransmitBuild<T: DelayedTransmitBuild + core::fmt::Debug> {
36    /// The data blob
37    pub data: T,
38    /// The transport for the transmission
39    pub transport: TransportType,
40    /// The source address of the transmission
41    pub from: SocketAddr,
42    /// The destination address of the transmission
43    pub to: SocketAddr,
44}
45
46impl<T: DelayedTransmitBuild + core::fmt::Debug> TransmitBuild<T> {
47    /// Construct a new [`Transmit`] with the specifid data and 5-tuple.
48    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    /// Write the [`TransmitBuild`] to a new `Vec<u8>`.
58    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    /// Write the [`TransmitBuild`] into the provided destination buffer.
68    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
79/// A trait for delaying building a byte sequence for transmission
80pub trait DelayedTransmitBuild {
81    /// Write the packet in to a new Vec.
82    fn build(self) -> Vec<u8>;
83    /// The length (in bytes) of the produced data.
84    fn len(&self) -> usize;
85    /// Whether the resulting data would be empty.
86    fn is_empty(&self) -> bool {
87        self.len() == 0
88    }
89    /// Write the data into a provided output buffer.
90    ///
91    /// Returns the number of bytes written.
92    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/// A `Transmit` that will construct a channel message towards a TURN client or server.
128#[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    /// Construct a new [`DelayedChannel`] with the specified channel ID and data.
136    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/// A `Transmit` that will construct a STUN message towards a client with the relevant data.
174#[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    /// Construct a [`Message`] aimed at being delievered to a TURN client.
183    pub fn for_client(peer_addr: SocketAddr, data: T) -> Self {
184        Self {
185            peer_addr,
186            data,
187            for_client: true,
188        }
189    }
190    /// Construct a [`Message`] aimed at being delievered to a TURN server.
191    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}