pb_mapper_protocol/
forward.rs

1use bytes::Bytes;
2use snafu::ResultExt;
3use std::future::Future;
4use std::pin::Pin;
5use std::time::Duration;
6use tokio::io::{AsyncReadExt, AsyncWriteExt};
7use tokio::time::Instant;
8
9use super::{
10    CodecMessageReader, CodecMessageWriter, MessageReader, MessageWriter, NormalMessageReader,
11    NormalMessageWriter,
12};
13use crate::buffer::{BufferReader, BufferedReader};
14use crate::data::DataCodec;
15use pb_mapper_core::checksum::AesKeyType;
16use pb_mapper_core::codec::{Decryptor, Encryptor};
17use pb_mapper_core::config::duration_from_env;
18use pb_mapper_core::error::{FwdNetworkWriteWithNormalSnafu, Result};
19use uni_stream::stream::{StreamSplit, TcpStreamImpl, UdpStreamImpl};
20use uni_stream::udp::{UdpStreamReadHalf, UdpStreamWriteHalf};
21
22pub trait ForwardReader {
23    async fn read(&mut self) -> Result<&'_ [u8]>;
24}
25
26pub trait ForwardWriter {
27    async fn write(&mut self, src: &[u8]) -> Result<()>;
28
29    // Gracefully close the write side to support half-closed TCP streams.
30    async fn shutdown(&mut self);
31}
32
33pub trait DatagramReader {
34    async fn recv(&mut self) -> Result<Bytes>;
35}
36
37pub trait DatagramWriter {
38    async fn send(&mut self, src: &[u8]) -> Result<()>;
39}
40
41const DEFAULT_TUNNEL_IDLE_TIMEOUT: Duration = Duration::from_secs(60 * 60);
42const DEFAULT_HALF_CLOSE_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
43const PB_MAPPER_TUNNEL_IDLE_TIMEOUT: &str = "PB_MAPPER_TUNNEL_IDLE_TIMEOUT";
44const PB_MAPPER_HALF_CLOSE_IDLE_TIMEOUT: &str = "PB_MAPPER_HALF_CLOSE_IDLE_TIMEOUT";
45
46tokio::task_local! {
47    static TUNNEL_IDLE_OVERRIDE: Duration;
48}
49
50/// Apply an idle deadline to forwarding polled by this future only.
51///
52/// Activity in either direction resets the deadline. This does not change
53/// process-wide defaults, half-close draining, or independently spawned tasks.
54/// A zero duration disables the idle deadline.
55pub async fn with_tunnel_idle_timeout<F: Future>(timeout: Duration, forward: F) -> F::Output {
56    TUNNEL_IDLE_OVERRIDE.scope(timeout, forward).await
57}
58
59#[derive(Debug, Clone, Copy)]
60struct ForwardTimeoutConfig {
61    tunnel_idle_timeout: Duration,
62    half_close_idle_timeout: Duration,
63}
64
65impl ForwardTimeoutConfig {
66    fn from_env() -> Self {
67        Self {
68            tunnel_idle_timeout: TUNNEL_IDLE_OVERRIDE
69                .try_with(|timeout| *timeout)
70                .unwrap_or_else(|_| {
71                    duration_from_env(PB_MAPPER_TUNNEL_IDLE_TIMEOUT, DEFAULT_TUNNEL_IDLE_TIMEOUT)
72                }),
73            half_close_idle_timeout: duration_from_env(
74                PB_MAPPER_HALF_CLOSE_IDLE_TIMEOUT,
75                DEFAULT_HALF_CLOSE_IDLE_TIMEOUT,
76            ),
77        }
78    }
79}
80
81pub struct NormalForwardReader<'a, T> {
82    buffered_reader: BufferReader<'a, T>,
83}
84
85impl<'a, T: AsyncReadExt + Unpin + Send> NormalForwardReader<'a, T> {
86    pub fn new(reader: &'a mut T) -> Self {
87        Self {
88            buffered_reader: BufferReader::new(reader),
89        }
90    }
91}
92
93impl<'a, T: AsyncReadExt + Unpin + Send> ForwardReader for NormalForwardReader<'a, T> {
94    async fn read(&mut self) -> Result<&'_ [u8]> {
95        self.buffered_reader.read().await
96    }
97}
98
99pub struct NormalDatagramReader<'a, T: AsyncReadExt + Unpin> {
100    reader: NormalMessageReader<'a, T>,
101}
102
103impl<'a, T: AsyncReadExt + Unpin + Send> NormalDatagramReader<'a, T> {
104    pub fn new(reader: &'a mut T) -> Self {
105        Self {
106            reader: NormalMessageReader::new(reader),
107        }
108    }
109
110    pub fn with_checksum_key(self, key: AesKeyType) -> Self {
111        Self {
112            reader: self.reader.with_checksum_key(key),
113        }
114    }
115}
116
117impl<'a, T: AsyncReadExt + Unpin + Send> DatagramReader for NormalDatagramReader<'a, T> {
118    async fn recv(&mut self) -> Result<Bytes> {
119        let msg = self.reader.read_msg().await?;
120        Ok(Bytes::copy_from_slice(msg))
121    }
122}
123
124pub struct NormalForwardWriter<'a, T> {
125    writer: &'a mut T,
126}
127
128impl<'a, T: AsyncWriteExt + Unpin + Send> NormalForwardWriter<'a, T> {
129    pub fn new(writer: &'a mut T) -> Self {
130        Self { writer }
131    }
132
133    async fn write_inner(&mut self, src: &[u8]) -> Result<()> {
134        self.writer
135            .write_all(src)
136            .await
137            .context(FwdNetworkWriteWithNormalSnafu)
138    }
139}
140
141impl<'a, T: AsyncWriteExt + Unpin + Send> ForwardWriter for NormalForwardWriter<'a, T> {
142    async fn write(&mut self, src: &[u8]) -> Result<()> {
143        self.write_inner(src).await
144    }
145
146    async fn shutdown(&mut self) {
147        let _ = self.writer.shutdown().await;
148    }
149}
150
151pub struct NormalDatagramWriter<'a, T: AsyncWriteExt + Unpin> {
152    writer: NormalMessageWriter<'a, T>,
153}
154
155impl<'a, T: AsyncWriteExt + Unpin + Send> NormalDatagramWriter<'a, T> {
156    pub fn new(writer: &'a mut T) -> Self {
157        Self {
158            writer: NormalMessageWriter::new(writer),
159        }
160    }
161
162    pub fn with_checksum_key(self, key: AesKeyType) -> Self {
163        Self {
164            writer: self.writer.with_checksum_key(key),
165        }
166    }
167}
168
169impl<'a, T: AsyncWriteExt + Unpin + Send> DatagramWriter for NormalDatagramWriter<'a, T> {
170    async fn send(&mut self, src: &[u8]) -> Result<()> {
171        self.writer.write_msg(src).await
172    }
173}
174
175pub struct CodecForwardReader<'a, T: AsyncReadExt + Unpin + Send, D: Decryptor>(
176    CodecMessageReader<'a, T, D>,
177);
178
179impl<'a, T: AsyncReadExt + Send + Unpin, D: Decryptor> CodecForwardReader<'a, T, D> {
180    pub fn new(reader: &'a mut T, decryptor: D) -> Self {
181        Self(CodecMessageReader::new(reader, decryptor))
182    }
183
184    pub fn with_checksum_key(self, key: AesKeyType) -> Self {
185        Self(self.0.with_checksum_key(key))
186    }
187}
188
189impl<'a, T: AsyncReadExt + Send + Unpin, D: Decryptor> ForwardReader
190    for CodecForwardReader<'a, T, D>
191{
192    async fn read(&mut self) -> Result<&'_ [u8]> {
193        self.0.read_stream().await
194    }
195}
196
197pub struct CodecDatagramReader<'a, T: AsyncReadExt + Unpin + Send, D: Decryptor>(
198    CodecMessageReader<'a, T, D>,
199);
200
201impl<'a, T: AsyncReadExt + Send + Unpin, D: Decryptor> CodecDatagramReader<'a, T, D> {
202    pub fn new(reader: &'a mut T, decryptor: D) -> Self {
203        Self(CodecMessageReader::new(reader, decryptor))
204    }
205
206    pub fn with_checksum_key(self, key: AesKeyType) -> Self {
207        Self(self.0.with_checksum_key(key))
208    }
209}
210
211impl<'a, T: AsyncReadExt + Send + Unpin, D: Decryptor> DatagramReader
212    for CodecDatagramReader<'a, T, D>
213{
214    async fn recv(&mut self) -> Result<Bytes> {
215        let msg = self.0.read_msg().await?;
216        Ok(Bytes::copy_from_slice(msg))
217    }
218}
219
220pub struct CodecForwardWriter<'a, T: AsyncWriteExt + Send + Unpin, E: Encryptor>(
221    CodecMessageWriter<'a, T, E>,
222);
223
224impl<'a, T: AsyncWriteExt + Send + Unpin, E: Encryptor> CodecForwardWriter<'a, T, E> {
225    pub fn new(writer: &'a mut T, encryptor: E) -> Self {
226        Self(CodecMessageWriter::new(writer, encryptor))
227    }
228
229    pub fn with_checksum_key(self, key: AesKeyType) -> Self {
230        Self(self.0.with_checksum_key(key))
231    }
232}
233
234impl<'a, T: AsyncWriteExt + Send + Unpin, E: Encryptor> ForwardWriter
235    for CodecForwardWriter<'a, T, E>
236{
237    /// SAFETY: Same as [`CodecMessageWriter`]
238    async fn write(&mut self, src: &[u8]) -> Result<()> {
239        self.0.write_msg(src).await
240    }
241
242    async fn shutdown(&mut self) {
243        let _ = self.0.shutdown().await;
244    }
245}
246
247pub struct CodecDatagramWriter<'a, T: AsyncWriteExt + Send + Unpin, E: Encryptor>(
248    CodecMessageWriter<'a, T, E>,
249);
250
251impl<'a, T: AsyncWriteExt + Send + Unpin, E: Encryptor> CodecDatagramWriter<'a, T, E> {
252    pub fn new(writer: &'a mut T, encryptor: E) -> Self {
253        Self(CodecMessageWriter::new(writer, encryptor))
254    }
255
256    pub fn with_checksum_key(self, key: AesKeyType) -> Self {
257        Self(self.0.with_checksum_key(key))
258    }
259}
260
261impl<'a, T: AsyncWriteExt + Send + Unpin, E: Encryptor> DatagramWriter
262    for CodecDatagramWriter<'a, T, E>
263{
264    async fn send(&mut self, src: &[u8]) -> Result<()> {
265        let buf = src.to_vec();
266        self.0.write_msg(&buf).await
267    }
268}
269
270pub async fn copy<R: ForwardReader, W: ForwardWriter>(
271    mut reader: R,
272    mut writer: W,
273) -> Result<usize> {
274    let mut length: usize = 0;
275    loop {
276        let src = reader.read().await?;
277        let n = src.len();
278        if n == 0 {
279            break;
280        }
281        writer.write(src).await?;
282        length += n;
283    }
284    writer.shutdown().await;
285    Ok(length)
286}
287
288pub async fn transfer_datagrams<R: DatagramReader, W: DatagramWriter>(
289    label: &'static str,
290    mut reader: R,
291    mut writer: W,
292) -> Result<usize> {
293    let mut _length: usize = 0;
294    loop {
295        let src = reader.recv().await?;
296        let n = src.len();
297        tracing::debug!("datagram forward {label} {n} bytes");
298        writer.send(&src).await?;
299        _length += n;
300    }
301}
302
303pub async fn start_forward<
304    ClientReader: ForwardReader,
305    ClientWriter: ForwardWriter,
306    ServerReader: ForwardReader,
307    ServerWriter: ForwardWriter,
308>(
309    client_reader: ClientReader,
310    client_writer: ClientWriter,
311    server_reader: ServerReader,
312    server_writer: ServerWriter,
313) {
314    start_forward_with_config(
315        client_reader,
316        client_writer,
317        server_reader,
318        server_writer,
319        ForwardTimeoutConfig::from_env(),
320    )
321    .await
322}
323
324/// Both directions are polled by the same forwarding task. A short, uncontended
325/// stack-owned lock makes the futures Send without allocating or queueing progress.
326struct ForwardActivity {
327    state: parking_lot::Mutex<(usize, Instant)>,
328}
329impl ForwardActivity {
330    fn new() -> Self {
331        Self {
332            state: parking_lot::Mutex::new((0, Instant::now())),
333        }
334    }
335    fn record(&self, n: usize) {
336        let mut state = self.state.lock();
337        state.0 += n;
338        state.1 = Instant::now();
339    }
340    fn snapshot(&self) -> (usize, Instant) {
341        *self.state.lock()
342    }
343}
344
345async fn copy_with_activity<R: ForwardReader, W: ForwardWriter>(
346    mut reader: R,
347    mut writer: W,
348    activity: &ForwardActivity,
349) -> Result<usize> {
350    let mut length = 0;
351    loop {
352        let src = reader.read().await?;
353        if src.is_empty() {
354            break;
355        }
356        writer.write(src).await?;
357        length += src.len();
358        activity.record(src.len());
359    }
360    writer.shutdown().await;
361    Ok(length)
362}
363
364async fn start_forward_with_config<
365    CR: ForwardReader,
366    CW: ForwardWriter,
367    SR: ForwardReader,
368    SW: ForwardWriter,
369>(
370    client_reader: CR,
371    client_writer: CW,
372    server_reader: SR,
373    server_writer: SW,
374    config: ForwardTimeoutConfig,
375) {
376    let client_activity = ForwardActivity::new();
377    let server_activity = ForwardActivity::new();
378    let client_to_server = copy_with_activity(client_reader, server_writer, &client_activity);
379    let server_to_client = copy_with_activity(server_reader, client_writer, &server_activity);
380    tokio::pin!(client_to_server, server_to_client);
381    let idle = tokio::time::sleep(config.tunnel_idle_timeout);
382    let half_close = tokio::time::sleep(config.half_close_idle_timeout);
383    tokio::pin!(idle, half_close);
384    let mut client_result = None;
385    let mut server_result = None;
386    let mut half_closed_at = None;
387    loop {
388        let client_done = client_result.is_some();
389        let server_done = server_result.is_some();
390        if client_done && server_done {
391            break;
392        }
393        tokio::select! {
394            biased;
395            result = &mut client_to_server, if !client_done => {
396                let failed = result.is_err();
397                client_result = Some(result);
398                half_closed_at = Some(Instant::now());
399                reset_sleep(&mut half_close, config.half_close_idle_timeout);
400                if failed { break; }
401            }
402            result = &mut server_to_client, if !server_done => {
403                let failed = result.is_err();
404                server_result = Some(result);
405                half_closed_at = Some(Instant::now());
406                reset_sleep(&mut half_close, config.half_close_idle_timeout);
407                if failed { break; }
408            }
409            _ = &mut idle, if !config.tunnel_idle_timeout.is_zero() && half_closed_at.is_none() => {
410                let last = client_activity.snapshot().1.max(server_activity.snapshot().1);
411                if last.elapsed() >= config.tunnel_idle_timeout { break; }
412                reset_sleep_at(&mut idle, config.tunnel_idle_timeout, last);
413            }
414            _ = &mut half_close, if !config.half_close_idle_timeout.is_zero() && half_closed_at.is_some() => {
415                let last = client_activity.snapshot().1.max(server_activity.snapshot().1).max(half_closed_at.unwrap_or_else(Instant::now));
416                if last.elapsed() >= config.half_close_idle_timeout { break; }
417                reset_sleep_at(&mut half_close, config.half_close_idle_timeout, last);
418            }
419        }
420    }
421    handle_forward_final_result(
422        client_result,
423        client_activity.snapshot().0,
424        "client->server",
425    );
426    handle_forward_final_result(
427        server_result,
428        server_activity.snapshot().0,
429        "server->client",
430    );
431}
432
433fn reset_sleep(sleep: &mut Pin<&mut tokio::time::Sleep>, timeout: Duration) {
434    if !timeout.is_zero() {
435        sleep.as_mut().reset(Instant::now() + timeout);
436    }
437}
438
439fn reset_sleep_at(
440    sleep: &mut Pin<&mut tokio::time::Sleep>,
441    timeout: Duration,
442    activity_at: Instant,
443) {
444    if !timeout.is_zero() {
445        sleep.as_mut().reset(activity_at + timeout);
446    }
447}
448
449fn handle_forward_final_result(result: Option<Result<usize>>, len: usize, detail: &'static str) {
450    if let Some(result) = result {
451        handle_forward_result(result, detail);
452    } else {
453        tracing::debug!("forward stopped before peer closed; we send {len} bytes,detail:{detail}");
454    }
455}
456
457pub async fn start_datagram_forward<
458    ClientReader: DatagramReader,
459    ClientWriter: DatagramWriter,
460    ServerReader: DatagramReader,
461    ServerWriter: DatagramWriter,
462>(
463    client_reader: ClientReader,
464    client_writer: ClientWriter,
465    server_reader: ServerReader,
466    server_writer: ServerWriter,
467) {
468    let client_to_server = transfer_datagrams("udp->tcp", client_reader, server_writer);
469    let server_to_client = transfer_datagrams("tcp->udp", server_reader, client_writer);
470    tokio::select! {
471        result = client_to_server =>{
472            handle_forward_result( result,"udp->tcp");
473        },
474        result = server_to_client =>{
475            handle_forward_result( result,"tcp->udp");
476        }
477    }
478}
479
480fn handle_forward_result(result: Result<usize>, detail: &'static str) {
481    match result {
482        Ok(len) => tracing::info!("forward finish! we send {len} bytes,detail:{detail}"),
483        Err(e) => {
484            // Treat peer-initiated shutdowns as expected to avoid noisy error logs.
485            if e.is_expected_disconnect() {
486                tracing::debug!("forward closed by peer:{e},detail:{detail}");
487            } else {
488                tracing::error!("got forward error:{e},detail:{detail}");
489            }
490        }
491    }
492}
493
494impl DatagramReader for UdpStreamReadHalf {
495    async fn recv(&mut self) -> Result<Bytes> {
496        self.recv_datagram()
497            .await
498            .map_err(|e| pb_mapper_core::error::Error::MsgForward {
499                action: "read",
500                source: e,
501            })
502    }
503}
504
505impl DatagramWriter for UdpStreamWriteHalf<'_> {
506    async fn send(&mut self, src: &[u8]) -> Result<()> {
507        self.send_datagram(src)
508            .await
509            .map_err(|e| pb_mapper_core::error::Error::MsgForward {
510                action: "write",
511                source: e,
512            })
513    }
514}
515
516pub trait StreamForward: StreamSplit + Sized {
517    fn forward_local_to_remote<'a, R, W>(
518        codec_key: Option<AesKeyType>,
519        framing_key: AesKeyType,
520        local_reader: Self::ReaderRef<'a>,
521        local_writer: Self::WriterRef<'a>,
522        remote_reader: R,
523        remote_writer: W,
524    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>
525    where
526        R: AsyncReadExt + Unpin + Send + 'a,
527        W: AsyncWriteExt + Unpin + Send + 'a;
528
529    /// Custom transports keep the legacy data format unless they implement the
530    /// negotiated forwarding method and explicitly advertise support.
531    fn supports_data_v2() -> bool {
532        false
533    }
534
535    /// Forward using the format explicitly selected in the authenticated setup.
536    fn forward_local_to_remote_with_codec<'a, R, W>(
537        codec: Option<DataCodec>,
538        framing_key: AesKeyType,
539        local_reader: Self::ReaderRef<'a>,
540        local_writer: Self::WriterRef<'a>,
541        remote_reader: R,
542        remote_writer: W,
543    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>
544    where
545        R: AsyncReadExt + Unpin + Send + 'a,
546        W: AsyncWriteExt + Unpin + Send + 'a,
547    {
548        if codec.is_some_and(|codec| codec.protocol().is_some()) {
549            return Box::pin(async {
550                Err(crate::secure::protocol_error(
551                    "transport does not implement data-v2",
552                ))
553            });
554        }
555        Self::forward_local_to_remote(
556            codec.map(DataCodec::key),
557            framing_key,
558            local_reader,
559            local_writer,
560            remote_reader,
561            remote_writer,
562        )
563    }
564}
565
566impl StreamForward for TcpStreamImpl {
567    fn supports_data_v2() -> bool {
568        true
569    }
570
571    fn forward_local_to_remote<'a, R, W>(
572        codec_key: Option<AesKeyType>,
573        framing_key: AesKeyType,
574        local_reader: Self::ReaderRef<'a>,
575        local_writer: Self::WriterRef<'a>,
576        remote_reader: R,
577        remote_writer: W,
578    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>
579    where
580        R: AsyncReadExt + Unpin + Send + 'a,
581        W: AsyncWriteExt + Unpin + Send + 'a,
582    {
583        Self::forward_local_to_remote_with_codec(
584            codec_key.map(DataCodec::legacy),
585            framing_key,
586            local_reader,
587            local_writer,
588            remote_reader,
589            remote_writer,
590        )
591    }
592
593    fn forward_local_to_remote_with_codec<'a, R, W>(
594        codec_key: Option<DataCodec>,
595        framing_key: AesKeyType,
596        local_reader: Self::ReaderRef<'a>,
597        local_writer: Self::WriterRef<'a>,
598        remote_reader: R,
599        remote_writer: W,
600    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>
601    where
602        R: AsyncReadExt + Unpin + Send + 'a,
603        W: AsyncWriteExt + Unpin + Send + 'a,
604    {
605        Box::pin(async move {
606            // Both the accepted subscriber socket and the publisher's target
607            // socket pass here. Small RPC frames must not wait for Nagle's
608            // coalescing after the relay legs have already disabled it.
609            if let Err(error) = local_reader.as_ref().set_nodelay(true) {
610                tracing::warn!(%error, "failed to disable Nagle on local TCP stream");
611            }
612            let mut local_reader = local_reader;
613            let mut local_writer = local_writer;
614            let mut remote_reader = remote_reader;
615            let mut remote_writer = remote_writer;
616            match codec_key {
617                Some(key) => {
618                    let (decoder, encoder) = key.endpoint_codecs()?;
619                    start_forward(
620                        NormalForwardReader::new(&mut local_reader),
621                        NormalForwardWriter::new(&mut local_writer),
622                        CodecForwardReader::new(&mut remote_reader, decoder)
623                            .with_checksum_key(framing_key),
624                        CodecForwardWriter::new(&mut remote_writer, encoder)
625                            .with_checksum_key(framing_key),
626                    )
627                    .await;
628                }
629                None => {
630                    start_forward(
631                        NormalForwardReader::new(&mut local_reader),
632                        NormalForwardWriter::new(&mut local_writer),
633                        NormalForwardReader::new(&mut remote_reader),
634                        NormalForwardWriter::new(&mut remote_writer),
635                    )
636                    .await;
637                }
638            }
639            Ok(())
640        })
641    }
642}
643
644impl StreamForward for UdpStreamImpl {
645    fn supports_data_v2() -> bool {
646        true
647    }
648
649    fn forward_local_to_remote<'a, R, W>(
650        codec_key: Option<AesKeyType>,
651        framing_key: AesKeyType,
652        local_reader: Self::ReaderRef<'a>,
653        local_writer: Self::WriterRef<'a>,
654        remote_reader: R,
655        remote_writer: W,
656    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>
657    where
658        R: AsyncReadExt + Unpin + Send + 'a,
659        W: AsyncWriteExt + Unpin + Send + 'a,
660    {
661        Self::forward_local_to_remote_with_codec(
662            codec_key.map(DataCodec::legacy),
663            framing_key,
664            local_reader,
665            local_writer,
666            remote_reader,
667            remote_writer,
668        )
669    }
670
671    fn forward_local_to_remote_with_codec<'a, R, W>(
672        codec_key: Option<DataCodec>,
673        framing_key: AesKeyType,
674        local_reader: Self::ReaderRef<'a>,
675        local_writer: Self::WriterRef<'a>,
676        remote_reader: R,
677        remote_writer: W,
678    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>
679    where
680        R: AsyncReadExt + Unpin + Send + 'a,
681        W: AsyncWriteExt + Unpin + Send + 'a,
682    {
683        Box::pin(async move {
684            let mut remote_reader = remote_reader;
685            let mut remote_writer = remote_writer;
686            match codec_key {
687                Some(key) => {
688                    let (decoder, encoder) = key.endpoint_codecs()?;
689                    start_datagram_forward(
690                        local_reader,
691                        local_writer,
692                        CodecDatagramReader::new(&mut remote_reader, decoder)
693                            .with_checksum_key(framing_key),
694                        CodecDatagramWriter::new(&mut remote_writer, encoder)
695                            .with_checksum_key(framing_key),
696                    )
697                    .await;
698                }
699                None => {
700                    start_datagram_forward(
701                        local_reader,
702                        local_writer,
703                        NormalDatagramReader::new(&mut remote_reader)
704                            .with_checksum_key(framing_key),
705                        NormalDatagramWriter::new(&mut remote_writer)
706                            .with_checksum_key(framing_key),
707                    )
708                    .await;
709                }
710            }
711            Ok(())
712        })
713    }
714}
715
716#[cfg(test)]
717mod tests {
718    #[tokio::test]
719    async fn encrypted_stream_fin_requires_complete_authenticated_frames() {
720        use crate::{data::DATA_PROTOCOL_V2, secure::HeaderProtocol};
721        for protocol in [None, Some(DATA_PROTOCOL_V2)] {
722            let codec = DataCodec::negotiate([7; 32], protocol, HeaderProtocol::V2).unwrap();
723            let (_, encoder) = codec.endpoint_codecs().unwrap();
724            let mut wire = Vec::new();
725            CodecMessageWriter::for_session_key(&mut wire, encoder, [8; 32])
726                .write_msg(b"response tail")
727                .await
728                .unwrap();
729            let (decoder, _) = codec.relay_codecs().unwrap();
730            let mut input = wire.as_slice();
731            let mut reader =
732                CodecForwardReader::new(&mut input, decoder).with_checksum_key([8; 32]);
733            assert_eq!(reader.read().await.unwrap(), b"response tail");
734            assert!(reader.read().await.unwrap().is_empty());
735
736            // EOF within either the header or encrypted body is never FIN.
737            for end in [3, wire.len() - 1] {
738                let (decoder, _) = codec.relay_codecs().unwrap();
739                let mut input = &wire[..end];
740                assert!(
741                    CodecForwardReader::new(&mut input, decoder)
742                        .with_checksum_key([8; 32])
743                        .read()
744                        .await
745                        .is_err()
746                );
747            }
748            *wire.last_mut().unwrap() ^= 1;
749            let (decoder, _) = codec.relay_codecs().unwrap();
750            let mut input = wire.as_slice();
751            assert!(
752                CodecForwardReader::new(&mut input, decoder)
753                    .with_checksum_key([8; 32])
754                    .read()
755                    .await
756                    .is_err()
757            );
758
759            // A complete authenticated empty record is valid; missing control
760            // and datagram frames still use the strict MessageReader API.
761            let (_, encoder) = codec.endpoint_codecs().unwrap();
762            let mut empty = Vec::new();
763            CodecMessageWriter::for_session_key(&mut empty, encoder, [8; 32])
764                .write_msg(&[])
765                .await
766                .unwrap();
767            let (decoder, _) = codec.relay_codecs().unwrap();
768            let mut input = empty.as_slice();
769            let mut reader = CodecMessageReader::for_session_key(&mut input, decoder, [8; 32]);
770            assert!(reader.read_msg().await.unwrap().is_empty());
771            assert!(reader.read_msg().await.is_err());
772        }
773    }
774
775    #[tokio::test]
776    async fn tcp_forward_disables_nagle_on_the_local_leg() {
777        use tokio::net::{TcpListener, TcpStream};
778        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
779        let (local, accepted) = tokio::join!(
780            TcpStream::connect(listener.local_addr().unwrap()),
781            listener.accept()
782        );
783        let socket = local.unwrap().into_std().unwrap();
784        let observer = socket.try_clone().unwrap();
785        assert!(!observer.nodelay().unwrap());
786        let mut local = TcpStreamImpl::new(TcpStream::from_std(socket).unwrap());
787        let (mut peer, _) = accepted.unwrap();
788        let (remote, mut echo) = tokio::io::duplex(64);
789        let forwarding = tokio::spawn(async move {
790            let (reader, writer) = local.split();
791            let (remote_reader, remote_writer) = tokio::io::split(remote);
792            TcpStreamImpl::forward_local_to_remote(
793                None,
794                [0; 32],
795                reader,
796                writer,
797                remote_reader,
798                remote_writer,
799            )
800            .await
801            .unwrap();
802        });
803        peer.write_all(b"x").await.unwrap();
804        assert_eq!(
805            tokio::time::timeout(Duration::from_secs(1), echo.read_u8())
806                .await
807                .unwrap()
808                .unwrap(),
809            b'x'
810        );
811        assert!(observer.nodelay().unwrap());
812        forwarding.abort();
813    }
814
815    use std::collections::VecDeque;
816    use std::io;
817    use std::sync::Arc;
818
819    use parking_lot::Mutex;
820    use std::time::Duration;
821
822    use super::*;
823    use pb_mapper_core::config::parse_duration;
824    use pb_mapper_core::error::Error;
825    use tokio::sync::Notify;
826
827    enum ReadAction {
828        Data(Vec<u8>),
829        Eof,
830        Pending,
831        Error(io::ErrorKind),
832    }
833
834    struct ScriptedReader {
835        actions: VecDeque<ReadAction>,
836        current: Vec<u8>,
837    }
838
839    impl ScriptedReader {
840        fn new(actions: impl IntoIterator<Item = ReadAction>) -> Self {
841            Self {
842                actions: actions.into_iter().collect(),
843                current: Vec::new(),
844            }
845        }
846    }
847
848    impl ForwardReader for ScriptedReader {
849        async fn read(&mut self) -> Result<&'_ [u8]> {
850            match self.actions.pop_front().unwrap_or(ReadAction::Pending) {
851                ReadAction::Data(data) => {
852                    self.current = data;
853                    Ok(&self.current)
854                }
855                ReadAction::Eof => {
856                    self.current.clear();
857                    Ok(&self.current)
858                }
859                ReadAction::Pending => std::future::pending().await,
860                ReadAction::Error(kind) => Err(Error::MsgForward {
861                    action: "read",
862                    source: io::Error::new(kind, "scripted read error"),
863                }),
864            }
865        }
866    }
867
868    #[derive(Default)]
869    struct WriterState {
870        chunks: Vec<Vec<u8>>,
871        shutdowns: usize,
872    }
873
874    #[derive(Clone, Default)]
875    struct ScriptedWriter {
876        state: Arc<Mutex<WriterState>>,
877    }
878
879    impl ScriptedWriter {
880        fn chunks(&self) -> Vec<Vec<u8>> {
881            self.state.lock().chunks.clone()
882        }
883
884        fn shutdowns(&self) -> usize {
885            self.state.lock().shutdowns
886        }
887    }
888
889    impl ForwardWriter for ScriptedWriter {
890        async fn write(&mut self, src: &[u8]) -> Result<()> {
891            self.state.lock().chunks.push(src.to_vec());
892            Ok(())
893        }
894
895        async fn shutdown(&mut self) {
896            self.state.lock().shutdowns += 1;
897        }
898    }
899
900    struct EofAfterWriteStartsReader {
901        write_started: Arc<Notify>,
902        returned_eof: bool,
903        empty: Vec<u8>,
904    }
905
906    impl EofAfterWriteStartsReader {
907        fn new(write_started: Arc<Notify>) -> Self {
908            Self {
909                write_started,
910                returned_eof: false,
911                empty: Vec::new(),
912            }
913        }
914    }
915
916    impl ForwardReader for EofAfterWriteStartsReader {
917        async fn read(&mut self) -> Result<&'_ [u8]> {
918            if self.returned_eof {
919                return std::future::pending().await;
920            }
921            self.write_started.notified().await;
922            self.returned_eof = true;
923            Ok(&self.empty)
924        }
925    }
926
927    #[derive(Clone)]
928    struct DelayedWriter {
929        state: Arc<Mutex<WriterState>>,
930        write_started: Arc<Notify>,
931        delay: Duration,
932    }
933
934    impl DelayedWriter {
935        fn new(write_started: Arc<Notify>, delay: Duration) -> Self {
936            Self {
937                state: Arc::new(Mutex::new(WriterState::default())),
938                write_started,
939                delay,
940            }
941        }
942
943        fn chunks(&self) -> Vec<Vec<u8>> {
944            self.state.lock().chunks.clone()
945        }
946    }
947
948    impl ForwardWriter for DelayedWriter {
949        async fn write(&mut self, src: &[u8]) -> Result<()> {
950            self.write_started.notify_one();
951            tokio::time::sleep(self.delay).await;
952            self.state.lock().chunks.push(src.to_vec());
953            Ok(())
954        }
955
956        async fn shutdown(&mut self) {
957            self.state.lock().shutdowns += 1;
958        }
959    }
960
961    #[test]
962    fn parse_duration_accepts_suffixes_and_plain_seconds() {
963        assert_eq!(parse_duration("42"), Some(Duration::from_secs(42)));
964        assert_eq!(parse_duration("500ms"), Some(Duration::from_millis(500)));
965        assert_eq!(parse_duration("2s"), Some(Duration::from_secs(2)));
966        assert_eq!(parse_duration("3m"), Some(Duration::from_secs(180)));
967        assert_eq!(parse_duration("1h"), Some(Duration::from_secs(3600)));
968        assert_eq!(parse_duration(""), None);
969        assert_eq!(parse_duration("bad"), None);
970        assert_eq!(parse_duration("18446744073709551615h"), None);
971    }
972
973    #[tokio::test]
974    async fn half_close_idle_timeout_closes_stalled_peer() {
975        let client_reader = ScriptedReader::new([ReadAction::Eof]);
976        let client_writer = ScriptedWriter::default();
977        let server_reader = ScriptedReader::new([ReadAction::Pending]);
978        let server_writer = ScriptedWriter::default();
979        let server_writer_state = server_writer.clone();
980
981        tokio::time::timeout(
982            Duration::from_millis(200),
983            start_forward_with_config(
984                client_reader,
985                client_writer,
986                server_reader,
987                server_writer,
988                ForwardTimeoutConfig {
989                    tunnel_idle_timeout: Duration::from_secs(60 * 60),
990                    half_close_idle_timeout: Duration::from_millis(20),
991                },
992            ),
993        )
994        .await
995        .expect("half-closed tunnel did not stop after half-close idle timeout");
996
997        assert_eq!(server_writer_state.shutdowns(), 1);
998    }
999
1000    #[tokio::test]
1001    async fn expected_disconnect_stops_waiting_for_pending_peer() {
1002        let client_reader =
1003            ScriptedReader::new([ReadAction::Error(io::ErrorKind::ConnectionReset)]);
1004        let client_writer = ScriptedWriter::default();
1005        let server_reader = ScriptedReader::new([ReadAction::Pending]);
1006        let server_writer = ScriptedWriter::default();
1007
1008        tokio::time::timeout(
1009            Duration::from_millis(200),
1010            start_forward_with_config(
1011                client_reader,
1012                client_writer,
1013                server_reader,
1014                server_writer,
1015                ForwardTimeoutConfig {
1016                    tunnel_idle_timeout: Duration::from_secs(60 * 60),
1017                    half_close_idle_timeout: Duration::from_secs(60),
1018                },
1019            ),
1020        )
1021        .await
1022        .expect("expected disconnect did not stop the tunnel");
1023    }
1024
1025    #[tokio::test]
1026    async fn half_closed_tunnel_drains_peer_before_timeout() {
1027        let client_reader = ScriptedReader::new([ReadAction::Eof]);
1028        let client_writer = ScriptedWriter::default();
1029        let client_writer_state = client_writer.clone();
1030        let server_reader =
1031            ScriptedReader::new([ReadAction::Data(b"response".to_vec()), ReadAction::Eof]);
1032        let server_writer = ScriptedWriter::default();
1033
1034        tokio::time::timeout(
1035            Duration::from_millis(200),
1036            start_forward_with_config(
1037                client_reader,
1038                client_writer,
1039                server_reader,
1040                server_writer,
1041                ForwardTimeoutConfig {
1042                    tunnel_idle_timeout: Duration::from_secs(60 * 60),
1043                    half_close_idle_timeout: Duration::from_millis(200),
1044                },
1045            ),
1046        )
1047        .await
1048        .expect("half-closed tunnel failed to drain the peer");
1049
1050        assert_eq!(client_writer_state.chunks(), vec![b"response".to_vec()]);
1051        assert_eq!(client_writer_state.shutdowns(), 1);
1052    }
1053
1054    #[tokio::test]
1055    async fn delayed_tail_write_survives_peer_half_close() {
1056        let write_started = Arc::new(Notify::new());
1057        let client_reader = EofAfterWriteStartsReader::new(write_started.clone());
1058        let client_writer = DelayedWriter::new(write_started, Duration::from_millis(20));
1059        let client_writer_state = client_writer.clone();
1060        let tail = vec![0x5a; 499];
1061        let server_reader = ScriptedReader::new([ReadAction::Data(tail.clone()), ReadAction::Eof]);
1062        let server_writer = ScriptedWriter::default();
1063
1064        tokio::time::timeout(
1065            Duration::from_millis(300),
1066            start_forward_with_config(
1067                client_reader,
1068                client_writer,
1069                server_reader,
1070                server_writer,
1071                ForwardTimeoutConfig {
1072                    tunnel_idle_timeout: Duration::from_secs(60 * 60),
1073                    half_close_idle_timeout: Duration::from_millis(200),
1074                },
1075            ),
1076        )
1077        .await
1078        .expect("delayed response tail was lost after the peer half-closed");
1079
1080        assert_eq!(client_writer_state.chunks(), vec![tail]);
1081    }
1082
1083    #[tokio::test]
1084    async fn open_tunnel_idle_timeout_closes_inactive_tunnel() {
1085        let client_reader = ScriptedReader::new([ReadAction::Pending]);
1086        let client_writer = ScriptedWriter::default();
1087        let server_reader = ScriptedReader::new([ReadAction::Pending]);
1088        let server_writer = ScriptedWriter::default();
1089
1090        tokio::time::timeout(
1091            Duration::from_millis(200),
1092            start_forward_with_config(
1093                client_reader,
1094                client_writer,
1095                server_reader,
1096                server_writer,
1097                ForwardTimeoutConfig {
1098                    tunnel_idle_timeout: Duration::from_millis(20),
1099                    half_close_idle_timeout: Duration::from_secs(60),
1100                },
1101            ),
1102        )
1103        .await
1104        .expect("inactive open tunnel did not stop after tunnel idle timeout");
1105    }
1106}
pb_mapper_protocol/forward.rs

pb_mapper_protocol/
forward.rs
