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//! Transmit structures and helpers.
10
11use 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/// A piece of data that needs to be built before it can be transmitted.
32#[derive(Debug)]
33pub struct TransmitBuild<T: DelayedTransmitBuild + core::fmt::Debug> {
34    /// The data blob
35    pub data: T,
36    /// The transport for the transmission
37    pub transport: TransportType,
38    /// The source address of the transmission
39    pub from: SocketAddr,
40    /// The destination address of the transmission
41    pub to: SocketAddr,
42}
43
44impl<T: DelayedTransmitBuild + core::fmt::Debug> TransmitBuild<T> {
45    /// Construct a new [`Transmit`] with the specifid data and 5-tuple.
46    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    /// Write the [`TransmitBuild`] to a new `Vec<u8>`.
56    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    /// Write the [`TransmitBuild`] into the provided destination buffer.
66    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
77/// A trait for delaying building a byte sequence for transmission
78pub trait DelayedTransmitBuild {
79    /// Write the packet in to a new Vec.
80    fn build(self) -> Vec<u8>;
81    /// The length (in bytes) of the produced data.
82    fn len(&self) -> usize;
83    /// Whether the resulting data would be empty.
84    fn is_empty(&self) -> bool {
85        self.len() == 0
86    }
87    /// Write the data into a provided output buffer.
88    ///
89    /// Returns the number of bytes written.
90    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/// A `Transmit` that will construct a channel message towards a TURN client or server.
126#[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    /// Construct a new [`DelayedChannel`] with the specified channel ID and data.
134    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/// A `Transmit` that will construct a STUN message towards a client with the relevant data.
172#[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    /// Construct a [`Message`] aimed at being delievered to a TURN client.
181    pub fn for_client(peer_addr: SocketAddr, data: T) -> Self {
182        Self {
183            peer_addr,
184            data,
185            for_client: true,
186        }
187    }
188    /// Construct a [`Message`] aimed at being delievered to a TURN server.
189    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}