bitfold_peer/peer/
encoder.rs

1use bitfold_protocol::{
2    command::CommandPacket,
3    command_codec::{self, CommandEncoder},
4};
5
6use super::Peer;
7
8impl Peer {
9    /// Encodes all queued commands into a CommandPacket and returns the bytes.
10    /// Drains the command queue in the process.
11    /// Applies compression if enabled, then appends CRC32 checksum if enabled.
12    pub fn encode_queued_commands(&mut self) -> std::io::Result<Vec<u8>> {
13        let mut packet = CommandPacket::new();
14        for command in self.drain_commands() {
15            packet.add_command(command);
16        }
17
18        // Encode into a pooled scratch buffer to avoid intermediate Vec allocations per step
19        let mut scratch = self.tx_pool.allocate();
20        scratch.clear();
21        CommandEncoder::encode_packet_into(&mut scratch, &packet)?;
22
23        // Apply compression using pooled buffer to reduce allocations
24        let compression_buffer = self.compression_pool.acquire();
25        let mut final_data = command_codec::compress_with_buffer(
26            &scratch,
27            self.config.compression,
28            self.config.compression_threshold,
29            compression_buffer,
30        )?;
31
32        // Return scratch to pool
33        self.tx_pool.deallocate(scratch);
34
35        // Record packet being sent
36        self.record_packet_sent();
37
38        // Append checksum if enabled (after compression) in-place
39        if self.config.use_checksums {
40            command_codec::append_checksum_in_place(&mut final_data);
41        }
42
43        // Track bytes sent
44        self.record_data_sent(final_data.len());
45
46        Ok(final_data)
47    }
48
49    /// Encodes up to `max_size` bytes worth of queued commands into a single datagram.
50    ///
51    /// - Returns `Ok(None)` if there are no queued commands.
52    /// - Returns `Ok(Some(bytes))` where `bytes.len() <= max_size` when data was produced.
53    ///
54    /// This prevents producing UDP payloads larger than the configured receive buffer (and typical MTUs),
55    /// avoiding OS-level EMSGSIZE errors and IP fragmentation.
56    pub fn encode_queued_commands_bounded(
57        &mut self,
58        max_size: usize,
59    ) -> std::io::Result<Option<Vec<u8>>> {
60        if !self.has_queued_commands() {
61            return Ok(None);
62        }
63
64        // Allow PMTU probes to bypass the current fragment-size cap so we can
65        // actually test larger datagram sizes. Still cap to receive buffer size.
66        let first_is_pmtu_probe = matches!(
67            self.command_queue.iter().next(),
68            Some(bitfold_protocol::command::ProtocolCommand::PMTUProbe { .. })
69        );
70        let max_size =
71            if first_is_pmtu_probe { self.config.receive_buffer_max_size } else { max_size };
72
73        // Worst-case overhead outside of command bytes:
74        // - 1 byte command count header
75        // - per-command 2-byte length prefix
76        // - compression marker/header (1 byte; LZ4 adds extra 4 bytes to store original size)
77        // - optional checksum (4 bytes)
78        let compression_overhead = match self.config.compression {
79            bitfold_core::config::CompressionAlgorithm::Lz4 => 5, // 1 marker + 4 original size
80            _ => 1,                                               // 1 marker for None/Zlib
81        };
82        let checksum_overhead = if self.config.use_checksums { 4 } else { 0 };
83        let static_overhead = 1 /* command count */ + compression_overhead + checksum_overhead;
84
85        // Select as many commands as will fit within max_size when encoded
86        let mut selected_count = 0usize;
87        let mut aggregated_len = 0; // track only command bytes (static_overhead already includes command count)
88
89        // Pre-encode commands individually to know precise sizes
90        let mut per_command_sizes: Vec<usize> = Vec::new();
91        for cmd in self.command_queue.iter() {
92            let encoded = bitfold_protocol::command_codec::CommandEncoder::encode_command(cmd)?;
93            let cmd_total = 2 /* length prefix */ + encoded.len();
94
95            // Check if adding this command would exceed limit (including trailing overhead)
96            if static_overhead + aggregated_len + cmd_total > max_size {
97                break;
98            }
99
100            aggregated_len += cmd_total;
101            per_command_sizes.push(cmd_total);
102            selected_count += 1;
103        }
104
105        if selected_count == 0 {
106            // First command is too large to fit within max_size
107            // This can happen with very large data payloads or very small MTU
108            if let Some(first_cmd) = self.command_queue.iter().next() {
109                // Don't warn for PMTU probes - they're expected to exceed normal MTU
110                if !first_is_pmtu_probe {
111                    let encoded =
112                        bitfold_protocol::command_codec::CommandEncoder::encode_command(first_cmd)?;
113                    let cmd_size = 2 + encoded.len();
114                    let total_with_overhead = static_overhead + cmd_size;
115
116                    tracing::warn!(
117                        "Command too large for MTU: command type {:?}, encoded size {} bytes, total with overhead {} bytes, max allowed {} bytes. Command will remain queued.",
118                        first_cmd.command_type(),
119                        cmd_size,
120                        total_with_overhead,
121                        max_size
122                    );
123                }
124            }
125            // Nothing selected within the budget; avoid emitting oversize datagrams.
126            return Ok(None);
127        }
128
129        // Drain exactly selected_count commands, requeue the rest to preserve order
130        let drained: Vec<_> = self.drain_commands().collect();
131        let mut packet = CommandPacket::new();
132        for cmd in drained.iter().take(selected_count) {
133            packet.add_command(cmd.clone());
134        }
135        for cmd in drained.into_iter().skip(selected_count) {
136            self.enqueue_command(cmd);
137        }
138
139        // Encode into pooled scratch buffer
140        let mut scratch = self.tx_pool.allocate();
141        scratch.clear();
142        bitfold_protocol::command_codec::CommandEncoder::encode_packet_into(&mut scratch, &packet)?;
143
144        // Apply compression using pooled buffer
145        let compression_buffer = self.compression_pool.acquire();
146        let mut final_data = command_codec::compress_with_buffer(
147            &scratch,
148            self.config.compression,
149            self.config.compression_threshold,
150            compression_buffer,
151        )?;
152        self.tx_pool.deallocate(scratch);
153
154        // Record packet and bytes sent
155        self.record_packet_sent();
156
157        if self.config.use_checksums {
158            command_codec::append_checksum_in_place(&mut final_data);
159        }
160
161        // Track bytes sent (full encoded size after compression/checksum)
162        self.record_data_sent(final_data.len());
163
164        // Sanity guard: ensure we did not exceed max_size
165        if final_data.len() > max_size {
166            tracing::error!(
167                "Encoded packet exceeded max_size after compression: {} bytes > {} bytes max. Selected {} commands, pre-compression size was {} bytes (command bytes) + {} bytes (overhead). Commands will remain queued.",
168                final_data.len(),
169                max_size,
170                selected_count,
171                aggregated_len,
172                static_overhead
173            );
174            // Avoid producing oversize datagrams; keep commands queued for next attempt
175            return Ok(None);
176        }
177
178        Ok(Some(final_data))
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use std::time::Instant;
185
186    use bitfold_core::config::{CompressionAlgorithm, Config};
187    use bitfold_protocol::command::ProtocolCommand;
188
189    use super::*;
190
191    fn create_virtual_connection() -> Peer {
192        Peer::new(get_fake_addr(), &Config::default(), Instant::now())
193    }
194
195    fn get_fake_addr() -> std::net::SocketAddr {
196        "127.0.0.1:0".parse().unwrap()
197    }
198
199    #[test]
200    fn test_encode_decode_command_packet() {
201        let mut peer = create_virtual_connection();
202
203        // Enqueue several commands
204        peer.enqueue_command(ProtocolCommand::Ping { timestamp: 100 });
205        peer.enqueue_command(ProtocolCommand::SendUnreliable {
206            channel_id: 0,
207            data: vec![1, 2, 3].into(),
208        });
209        peer.enqueue_ack_command(Some(5000));
210
211        assert_eq!(peer.queued_commands_count(), 3);
212
213        // Encode to bytes
214        let bytes = peer.encode_queued_commands().unwrap();
215        assert!(!bytes.is_empty());
216        assert!(!peer.has_queued_commands()); // Should be drained
217
218        // Decode and process on another peer
219        let mut peer2 = create_virtual_connection();
220        let result = peer2.process_command_packet(&bytes, Instant::now()).unwrap();
221
222        // Should have received the unreliable data packet
223        let packets: Vec<_> = result.into_iter().collect();
224        assert_eq!(packets.len(), 1);
225
226        // peer2 should have queued responses (Pong for Ping)
227        assert!(peer2.has_queued_commands());
228    }
229
230    #[test]
231    fn test_checksum_enabled_end_to_end() {
232        let mut config = Config::default();
233        config.use_checksums = true;
234
235        let mut peer = Peer::new(get_fake_addr(), &config, Instant::now());
236        let time = Instant::now();
237
238        // Enqueue a command
239        peer.enqueue_command(ProtocolCommand::Ping { timestamp: 1234 });
240
241        // Encode with checksum
242        let encoded = peer.encode_queued_commands().unwrap();
243
244        // Should be longer due to checksum (4 bytes)
245        assert!(encoded.len() > 10); // At least command data + checksum
246
247        // Process the packet with checksum validation
248        let result = peer.process_command_packet(&encoded, time);
249        assert!(result.is_ok()); // Should succeed with valid checksum
250    }
251
252    #[test]
253    fn test_checksum_detects_corruption_in_peer() {
254        let mut config = Config::default();
255        config.use_checksums = true;
256
257        let mut peer = Peer::new(get_fake_addr(), &config, Instant::now());
258        let time = Instant::now();
259
260        // Enqueue a command
261        peer.enqueue_command(ProtocolCommand::SendUnreliable {
262            channel_id: 0,
263            data: vec![1, 2, 3, 4, 5].into(),
264        });
265
266        // Encode with checksum
267        let mut encoded = peer.encode_queued_commands().unwrap();
268
269        // Corrupt the data (but not the checksum at the end)
270        if encoded.len() > 5 {
271            encoded[0] = 99;
272        }
273
274        // Process the corrupted packet - should fail checksum validation
275        let result = peer.process_command_packet(&encoded, time);
276        assert!(result.is_err());
277    }
278
279    #[test]
280    fn test_checksum_disabled_backward_compatibility() {
281        let mut config = Config::default();
282        config.use_checksums = false; // Disabled (default)
283
284        let mut peer1 = Peer::new(get_fake_addr(), &config, Instant::now());
285        let mut peer2 = Peer::new(get_fake_addr(), &config, Instant::now());
286        let time = Instant::now();
287
288        // Enqueue a command on peer1
289        peer1.enqueue_command(ProtocolCommand::Ping { timestamp: 5678 });
290
291        // Encode without checksum
292        let encoded = peer1.encode_queued_commands().unwrap();
293
294        // Process on peer2 (also without checksum validation)
295        let result = peer2.process_command_packet(&encoded, time);
296        assert!(result.is_ok()); // Should work without checksums
297    }
298
299    #[test]
300    fn test_compression_zlib_end_to_end() {
301        let mut config = Config::default();
302        config.compression = CompressionAlgorithm::Zlib;
303        config.compression_threshold = 10;
304
305        let mut peer1 = Peer::new(get_fake_addr(), &config, Instant::now());
306        let mut peer2 = Peer::new(get_fake_addr(), &config, Instant::now());
307        let time = Instant::now();
308
309        // Enqueue a large, compressible command
310        peer1.enqueue_command(ProtocolCommand::SendUnreliable {
311            channel_id: 0,
312            data: vec![42; 300].into(), // Highly compressible
313        });
314
315        // Encode with compression
316        let encoded = peer1.encode_queued_commands().unwrap();
317
318        // Should be compressed (smaller than uncompressed would be)
319        // Uncompressed would be ~300 + overhead, compressed should be much smaller
320        assert!(encoded.len() < 100); // Compressed size
321
322        // Process on peer2
323        let result = peer2.process_command_packet(&encoded, time);
324        assert!(result.is_ok());
325    }
326
327    #[test]
328    fn test_compression_lz4_end_to_end() {
329        let mut config = Config::default();
330        config.compression = CompressionAlgorithm::Lz4;
331        config.compression_threshold = 10;
332
333        let mut peer1 = Peer::new(get_fake_addr(), &config, Instant::now());
334        let mut peer2 = Peer::new(get_fake_addr(), &config, Instant::now());
335        let time = Instant::now();
336
337        // Enqueue a large, compressible command
338        peer1.enqueue_command(ProtocolCommand::SendUnreliable {
339            channel_id: 0,
340            data: vec![99; 300].into(), // Highly compressible
341        });
342
343        // Encode with compression
344        let encoded = peer1.encode_queued_commands().unwrap();
345
346        // Should be compressed
347        assert!(encoded.len() < 100);
348
349        // Process on peer2
350        let result = peer2.process_command_packet(&encoded, time);
351        assert!(result.is_ok());
352    }
353
354    #[test]
355    fn test_compression_disabled_by_default() {
356        let config = Config::default();
357
358        let mut peer1 = Peer::new(get_fake_addr(), &config, Instant::now());
359        let mut peer2 = Peer::new(get_fake_addr(), &config, Instant::now());
360        let time = Instant::now();
361
362        // Enqueue a command
363        peer1.enqueue_command(ProtocolCommand::Ping { timestamp: 1234 });
364
365        // Encode without compression (default)
366        let encoded = peer1.encode_queued_commands().unwrap();
367
368        // First byte (after any header) should be 0 (uncompressed marker)
369        assert_eq!(encoded[0], 0);
370
371        // Process on peer2
372        let result = peer2.process_command_packet(&encoded, time);
373        assert!(result.is_ok());
374    }
375
376    #[test]
377    fn test_statistics_with_compression() {
378        let mut config = Config::default();
379        config.compression = CompressionAlgorithm::Lz4;
380        config.compression_threshold = 10;
381        let mut peer = Peer::new(get_fake_addr(), &config, Instant::now());
382
383        // Enqueue data that will be compressed
384        peer.enqueue_command(ProtocolCommand::SendReliable {
385            channel_id: 0,
386            sequence: 0,
387            ordered: true,
388            data: vec![1; 100].into(), // Repeating data compresses well
389        });
390
391        let encoded = peer.encode_queued_commands().unwrap();
392
393        // bytes_sent should track the compressed size
394        assert_eq!(peer.statistics().bytes_sent, encoded.len() as u64);
395        assert!(peer.statistics().bytes_sent < 100); // Should be compressed
396    }
397
398    #[test]
399    fn test_statistics_with_checksums() {
400        let mut config = Config::default();
401        config.use_checksums = true;
402        let mut peer = Peer::new(get_fake_addr(), &config, Instant::now());
403
404        peer.enqueue_command(ProtocolCommand::SendReliable {
405            channel_id: 0,
406            sequence: 0,
407            ordered: true,
408            data: vec![1, 2, 3].into(),
409        });
410
411        let encoded = peer.encode_queued_commands().unwrap();
412
413        // bytes_sent should include checksum overhead (4 bytes)
414        assert_eq!(peer.statistics().bytes_sent, encoded.len() as u64);
415        assert!(encoded.len() >= 4); // At least checksum size
416    }
417
418    #[test]
419    fn test_mtu_boundary_calculation() {
420        let mut config = Config::default();
421        config.use_checksums = false;
422        config.compression = bitfold_core::config::CompressionAlgorithm::None;
423        let mut peer = Peer::new(get_fake_addr(), &config, Instant::now());
424
425        // Calculate exact payload size that should fit in 1440 bytes MTU
426        // Packet structure:
427        // - 1 byte command count
428        // - 1 byte compression marker (even for None)
429        // - 2 bytes length prefix per command
430        // - SendReliable command: 1 (type) + 1 (channel) + 2 (seq) + 1 (ordered) + 2 (data len) + data
431        // Total: 1 + 1 + 2 + (1 + 1 + 2 + 1 + 2) + data = 11 + data
432        // So max data = 1440 - 11 = 1429 bytes
433        let max_payload_for_mtu = 1429;
434
435        // Enqueue a command with exactly max_payload_for_mtu bytes
436        peer.enqueue_command(ProtocolCommand::SendReliable {
437            channel_id: 0,
438            sequence: 0,
439            ordered: true,
440            data: vec![42u8; max_payload_for_mtu].into(),
441        });
442
443        // Should fit within MTU (1440 bytes)
444        let result = peer.encode_queued_commands_bounded(1440);
445        assert!(result.is_ok());
446        let encoded = result.unwrap();
447        assert!(encoded.is_some(), "Command should fit within MTU");
448
449        let bytes = encoded.unwrap();
450        assert_eq!(bytes.len(), 1440, "Encoded size should be exactly 1440 bytes");
451
452        // Verify no commands remain queued
453        assert!(!peer.has_queued_commands());
454    }
455
456    #[test]
457    fn test_mtu_boundary_too_large() {
458        let mut config = Config::default();
459        config.use_checksums = false;
460        config.compression = bitfold_core::config::CompressionAlgorithm::None;
461        let mut peer = Peer::new(get_fake_addr(), &config, Instant::now());
462
463        // Try to enqueue a command that's 1 byte too large
464        let too_large_payload = 1430; // 11 + 1430 = 1441 bytes total
465
466        peer.enqueue_command(ProtocolCommand::SendReliable {
467            channel_id: 0,
468            sequence: 0,
469            ordered: true,
470            data: vec![42u8; too_large_payload].into(),
471        });
472
473        // Should NOT fit within MTU (would be 1441 bytes)
474        let result = peer.encode_queued_commands_bounded(1440);
475        assert!(result.is_ok());
476        let encoded = result.unwrap();
477        assert!(encoded.is_none(), "Command should NOT fit within MTU");
478
479        // Command should remain queued
480        assert!(peer.has_queued_commands());
481    }
482}
bitfold_peer/peer/encoder.rs

bitfold_peer/peer/
encoder.rs
