1
//! One owner drives networking, protocol timers, media, and channel writes.
2
 
3
use serde::Serialize;
4
use std::{collections::HashMap, error::Error, net::UdpSocket, time::Instant};
5
use str0m::{
6
    Candidate, Event, Input, Output, Rtc,
7
    change::{SdpAnswer, SdpPendingOffer},
8
    channel::ChannelId,
9
    crypto::dtls::DtlsCert,
10
    media::{Direction, Frequency, MediaKind, MediaTime},
11
    net::{Protocol, Receive},
12
};
13
 
14
pub type Result<T> = std::result::Result<T, Box<dyn Error>>;
15
 
16
#[derive(Default, Debug, Serialize)]
17
pub struct Stats {
18
    pub audio_sent: u32,
19
    pub robot_sent: u32,
20
    pub spectrum_sent: u32,
21
    pub commands_received: u32,
22
    pub application_datagrams: u32,
23
    pub deliberately_dropped: u32,
24
    pub rtp_datagrams: u32,
25
}
26
 
27
pub struct Driver {
28
    rtc: Rtc,
29
    socket: UdpSocket,
30
    pending: Option<SdpPendingOffer>,
31
    mid: str0m::media::Mid,
32
    channels: HashMap<String, (u16, ChannelId)>,
33
    pub events: Vec<serde_json::Value>,
34
    pub stats: Stats,
35
    pub drop_every: u32,
36
}
37
 
38
impl Driver {
39
    pub fn new(ip: &str, certificate: DtlsCert) -> Result<(Self, String)> {
40
        let socket = UdpSocket::bind((ip, 0))?;
41
        socket.set_nonblocking(true)?;
42
        let mut rtc = crate::rtc(certificate);
43
        rtc.add_local_candidate(Candidate::host(socket.local_addr()?, "udp")?);
44
        while !matches!(rtc.poll_output()?, Output::Timeout(_)) {}
45
        let mut change = rtc.sdp_api();
46
        let mid = change.add_media(MediaKind::Audio, Direction::SendOnly, None, None, None);
47
        let bootstrap = change.add_channel_with_config(crate::channel("bootstrap", 0, true));
48
        let (offer, pending) = change.apply().ok_or("no initial offer")?;
49
        let mut driver = Self {
50
            rtc,
51
            socket,
52
            pending: Some(pending),
53
            mid,
54
            channels: HashMap::from([("bootstrap".to_owned(), (0, bootstrap))]),
55
            events: Vec::new(),
56
            stats: Stats::default(),
57
            drop_every: 0,
58
        };
59
        driver.drain()?;
60
        Ok((driver, offer.to_sdp_string()))
61
    }
62
 
63
    pub fn mid(&self) -> String {
64
        self.mid.to_string()
65
    }
66
 
67
    pub fn answer(&mut self, sdp: &str) -> Result<()> {
68
        let pending = self.pending.take().ok_or("answer already applied")?;
69
        self.rtc
70
            .sdp_api()
71
            .accept_answer(pending, SdpAnswer::from_sdp_string(sdp)?)?;
72
        self.drain()
73
    }
74
 
75
    pub fn add_channel(&mut self, label: &str, id: u16, reliable: bool) -> Result<()> {
76
        let cid = self
77
            .rtc
78
            .direct_api()
79
            .create_data_channel(crate::channel(label, id, reliable));
80
        self.channels.insert(label.to_owned(), (id, cid));
81
        self.drain()
82
    }
83
 
84
    pub fn pump(&mut self) -> Result<()> {
85
        let mut buffer = [0u8; 2048];
86
        for _ in 0..32 {
87
            let (length, source) = match self.socket.recv_from(&mut buffer) {
88
                Ok(value) => value,
89
                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => break,
90
                Err(error) => return Err(error.into()),
91
            };
92
            let receive = Receive {
93
                proto: Protocol::Udp,
94
                source,
95
                destination: self.socket.local_addr()?,
96
                contents: buffer[..length].try_into()?,
97
            };
98
            self.rtc
99
                .handle_input(Input::Receive(Instant::now(), receive))?;
100
            self.drain()?;
101
        }
102
        self.rtc.handle_input(Input::Timeout(Instant::now()))?;
103
        self.drain()
104
    }
105
 
106
    pub fn audio(&mut self, _frame: u32) -> Result<()> {
107
        let writer = self
108
            .rtc
109
            .writer(self.mid)
110
            .ok_or("audio writer unavailable")?;
111
        let pt = writer
112
            .payload_params()
113
            .next()
114
            .ok_or("audio codec missing")?
115
            .pt();
116
        writer.write(
117
            pt,
118
            Instant::now(),
119
            MediaTime::new(
120
                u64::from(self.stats.audio_sent) * 960,
121
                Frequency::FORTY_EIGHT_KHZ,
122
            ),
123
            [0xf8, 0xff, 0xfe],
124
        )?;
125
        self.stats.audio_sent += 1;
126
        self.drain()
127
    }
128
 
129
    pub fn data(&mut self, label: &str, sequence: u32) -> Result<()> {
130
        let (_, cid) = *self.channels.get(label).ok_or("unknown channel")?;
131
        let mut channel = self.rtc.channel(cid).ok_or("channel is not open")?;
132
        let accepted = if label == "robot" {
133
            let payload =
134
                serde_json::json!({"kind": "telemetry", "sequence": sequence}).to_string();
135
            channel.write(false, payload.as_bytes())?
136
        } else {
137
            channel.write(true, &sequence.to_le_bytes())?
138
        };
139
        if !accepted {
140
            return Err("channel backpressure".into());
141
        }
142
        if label == "robot" {
143
            self.stats.robot_sent += 1;
144
        } else {
145
            self.stats.spectrum_sent += 1;
146
        }
147
        self.drain()
148
    }
149
 
150
    fn drain(&mut self) -> Result<()> {
151
        loop {
152
            match self.rtc.poll_output()? {
153
                Output::Timeout(_) => break,
154
                Output::Transmit(packet) => {
155
                    if packet.contents.len() >= 2
156
                        && packet.contents[0] & 0xc0 == 0x80
157
                        && !(192..=223).contains(&packet.contents[1])
158
                    {
159
                        self.stats.rtp_datagrams += 1;
160
                    }
161
                    // DTLS 1.2 application records; enabled only after setup.
162
                    let application = packet.contents.first() == Some(&23);
163
                    if application {
164
                        self.stats.application_datagrams += 1;
165
                    }
166
                    if application
167
                        && self.drop_every != 0
168
                        && self.stats.application_datagrams % self.drop_every == 0
169
                    {
170
                        self.stats.deliberately_dropped += 1;
171
                    } else {
172
                        self.socket.send_to(&packet.contents, packet.destination)?;
173
                    }
174
                }
175
                Output::Event(Event::Connected) => {
176
                    self.events.push(serde_json::json!({"event": "connected"}))
177
                }
178
                Output::Event(Event::ChannelOpen(cid, label)) => {
179
                    let id = self
180
                        .channels
181
                        .values()
182
                        .find(|(_, value)| *value == cid)
183
                        .map(|(id, _)| *id);
184
                    self.events.push(
185
                        serde_json::json!({"event": "channel_open", "label": label, "id": id}),
186
                    );
187
                }
188
                Output::Event(Event::ChannelData(data)) => {
189
                    let is_robot = self
190
                        .channels
191
                        .get("robot")
192
                        .is_some_and(|(_, cid)| *cid == data.id);
193
                    if is_robot {
194
                        self.stats.commands_received += 1;
195
                        let mut channel = self
196
                            .rtc
197
                            .channel(data.id)
198
                            .ok_or("reply channel unavailable")?;
199
                        channel.write(data.binary, &data.data)?;
200
                    }
201
                }
202
                Output::Event(_) => {}
203
            }
204
        }
205
        Ok(())
206
    }
207
}