Skip to main content

oms_modbus/transport/
tcp.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Modbus TCP client and server with MBAP framing.
3//!
4//! # Quick start
5//!
6//! ```no_run
7//! use std::net::{IpAddr, Ipv4Addr, SocketAddr};
8//! use std::sync::Arc;
9//! use std::time::Duration;
10//! use oms_modbus::*;
11//!
12//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
13//! // ── Server ──────────────────────────────────────────────────
14//! let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 502);
15//! let store = Arc::new(SlaveStore::with_holding_registers(&[(0, 1234)]));
16//! let server = tcp::TcpServer::bind(addr).await?;
17//! tokio::spawn(async move { server.serve_forever(store).await.ok(); });
18//!
19//! // ── Client ──────────────────────────────────────────────────
20//! let client = tcp::TcpClient::connect_with_timeout(addr, Duration::from_secs(3)).await?;
21//! let regs = client.read_holding_registers(1, 0, 1).await?;
22//! # Ok(())
23//! # }
24//! ```
25//!
26//! # Advanced MBAP configuration
27//!
28//! Use [`TcpConfig`] for non-standard devices and RTU-over-TCP gateways:
29//!
30//! ```no_run
31//! use oms_modbus::tcp::{TcpConfig, TidMode, LengthMode};
32//! # use std::net::SocketAddr;
33//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
34//! # let addr = "127.0.0.1:502".parse()?;
35//! let client = tcp::TcpClient::connect(addr).await?.with_config(TcpConfig {
36//!     tid: TidMode::Auto,                // auto-increment transaction ID
37//!     unit_id_in_body: true,             // UID also in PDU body
38//!     length_mode: LengthMode::Standard,
39//! });
40//! # Ok(())
41//! # }
42//! ```
43
44use std::net::SocketAddr;
45use std::sync::atomic::{AtomicU16, Ordering};
46use std::sync::Arc;
47use std::time::{Duration, Instant};
48
49use std::panic::AssertUnwindSafe;
50
51use async_trait::async_trait;
52use bytes::{BufMut, Bytes, BytesMut};
53use futures_util::FutureExt;
54use tokio::io::{AsyncReadExt, AsyncWriteExt};
55use tokio::net::TcpStream;
56use tokio::sync::Mutex;
57
58use crate::bus_timing::BusTiming;
59use crate::client::ModbusClient;
60use crate::error::ModbusError;
61use crate::error::*;
62use crate::frame::{Request, Response};
63use crate::options::ClientOptions;
64use crate::transport::send_recv;
65use crate::transport::sniff_io::SniffIo;
66use crate::transport::{MAX_ADU_SIZE, MAX_TCP_ADU_SIZE, MBAP_HEADER_SIZE, MBAP_PREFIX_SIZE};
67use crate::wire_tap::WireTap;
68
69// ── TCP configuration ───────────────────────────────────────────────────
70
71/// Transaction ID generation mode.
72#[derive(Debug)]
73pub enum TidMode {
74    /// Always use the same value (default: 0).
75    Fixed(u16),
76    /// Auto-increment on each request, wrapping at u16::MAX.
77    /// Starts from 1.
78    Auto,
79}
80
81impl Clone for TidMode {
82    fn clone(&self) -> Self {
83        match self {
84            TidMode::Fixed(v) => TidMode::Fixed(*v),
85            TidMode::Auto => TidMode::Auto,
86        }
87    }
88}
89
90/// MBAP Length field calculation.
91#[derive(Debug, Clone, Copy, PartialEq)]
92pub enum LengthMode {
93    /// Standard per spec: Length = 1 (UID at byte 6) + PDU bytes.
94    /// Compatible with tokio-modbus and most devices.
95    Standard,
96    /// Length = PDU bytes only (excludes UID at byte 6).
97    /// Some non-standard devices use this mode.
98    PduOnly,
99}
100
101/// TCP-specific MBAP framing configuration.
102///
103/// Controls how the MBAP header is built. Defaults match the
104/// Modbus TCP specification (TID=0, UID in header only, Standard length).
105#[derive(Debug, Clone)]
106pub struct TcpConfig {
107    pub tid: TidMode,
108    /// If true, the unit_id is also written as the first byte of the PDU body
109    /// (in addition to MBAP header byte 6). Required by some TCP-to-RTU gateways.
110    pub unit_id_in_body: bool,
111    /// How the MBAP Length field is calculated.
112    pub length_mode: LengthMode,
113}
114
115impl Default for TcpConfig {
116    fn default() -> Self {
117        Self {
118            tid: TidMode::Fixed(0),
119            unit_id_in_body: false,
120            length_mode: LengthMode::Standard,
121        }
122    }
123}
124
125impl TcpConfig {
126    /// Standard Modbus TCP: TID=0, UID in MBAP header only, Standard length.
127    pub fn standard() -> Self {
128        Self::default()
129    }
130
131    /// Gateway-friendly: auto-increment TID from 1, UID in PDU body, Standard length.
132    pub fn gateway() -> Self {
133        Self {
134            tid: TidMode::Auto,
135            unit_id_in_body: true,
136            length_mode: LengthMode::Standard,
137        }
138    }
139}
140
141// ── TCP frame encoding / decoding ────────────────────────────────────────
142
143fn next_tid(mode: &TidMode, counter: &AtomicU16) -> u16 {
144    match mode {
145        TidMode::Fixed(v) => *v,
146        TidMode::Auto => counter.fetch_add(1, Ordering::Relaxed),
147    }
148}
149
150/// Encode [slave_id, pdu...] into a TCP MBAP frame.
151pub fn encode_tcp_frame(data: &[u8], buf: &mut BytesMut, config: &TcpConfig, tid: u16) {
152    let (unit_id, pdu) = if data.is_empty() {
153        (1u8, &[] as &[u8])
154    } else {
155        (data[0], &data[1..])
156    };
157    let extra = if config.unit_id_in_body { 1u16 } else { 0u16 };
158    let len = match config.length_mode {
159        LengthMode::Standard => pdu.len() as u16 + 1 + extra,
160        LengthMode::PduOnly => pdu.len() as u16 + extra,
161    };
162    buf.put_u16(tid);
163    buf.put_u16(0); // Protocol ID
164    buf.put_u16(len);
165    buf.put_u8(unit_id);
166    if config.unit_id_in_body {
167        buf.put_u8(unit_id);
168    }
169    buf.extend_from_slice(pdu);
170}
171
172/// Try to parse one TCP frame from a byte buffer.
173///
174/// Returns `Some((tid, slave_id, pdu, bytes_consumed))` if MBAP header is complete
175/// and frame data is sufficient. The `tid` (Transaction Identifier) is extracted
176/// from the MBAP header for the server to echo back in its response.
177/// Validates Protocol ID == 0 and max frame size.
178///
179/// Handles both normal frames (PDU byte 0 = function code) and `unit_id_in_body`
180/// Check whether the first 7 bytes of `buf` contain a corrupt MBAP header.
181///
182/// Returns `false` when there aren't enough bytes yet or the header looks
183/// valid (proto_id == 0, payload within range).  Returns `true` only when a
184/// full header is present but structurally invalid — this is the signal to
185/// discard the buffer rather than wait for more data.
186fn is_tcp_header_corrupt(buf: &[u8]) -> bool {
187    if buf.len() < MBAP_HEADER_SIZE {
188        return false; // need more data
189    }
190    let proto_id = u16::from_be_bytes([buf[2], buf[3]]);
191    let payload_len = u16::from_be_bytes([buf[4], buf[5]]) as usize;
192    proto_id != 0 || !(1..=MAX_TCP_ADU_SIZE).contains(&payload_len)
193}
194
195/// Safety cap: clear the buffer if it grows beyond 4× the maximum legal ADU.
196/// A well-formed stream should never exceed this; oversized buffers indicate
197/// a corrupted length field or a peer that isn't speaking Modbus TCP.
198const MAX_TCP_BUF_BYTES: usize = MAX_TCP_ADU_SIZE * 4; // 1040 bytes
199
200/// frames (PDU byte 0 = duplicate slave ID, byte 1 = function code). When the
201/// first body byte matches the MBAP Unit ID, both interpretations are tested
202/// and the one that parses as a valid [`Request`] is returned.
203pub fn try_parse_tcp_frame(buf: &[u8]) -> Option<(u16, u8, Bytes, usize)> {
204    if buf.len() < MBAP_HEADER_SIZE {
205        return None;
206    }
207    let tid = u16::from_be_bytes([buf[0], buf[1]]);
208    let payload_len = u16::from_be_bytes([buf[4], buf[5]]) as usize;
209    let proto_id = u16::from_be_bytes([buf[2], buf[3]]);
210    if proto_id != 0 || payload_len > MAX_TCP_ADU_SIZE {
211        return None;
212    }
213
214    let unit_id = buf[6];
215    for &fl in &[
216        MBAP_HEADER_SIZE + payload_len,
217        MBAP_PREFIX_SIZE + payload_len,
218    ] {
219        // Guard: when payload_len < 1, fl can be < MBAP_HEADER_SIZE,
220        // which would panic on buf[MBAP_HEADER_SIZE..fl].
221        if fl > MBAP_HEADER_SIZE && buf.len() >= fl {
222            let body = &buf[MBAP_HEADER_SIZE..fl];
223            if body.is_empty() {
224                continue;
225            }
226            return try_extract_tcp_pdu(tid, unit_id, body, fl);
227        }
228    }
229    None
230}
231
232/// Given a TCP frame body (bytes after the MBAP header), try to
233/// extract the PDU. Handles both normal and `unit_id_in_body` framing.
234///
235/// Priority: when byte 0 matches the MBAP Unit ID AND byte 1 is also a valid
236/// function code, prefer the `unit_id_in_body` interpretation (skip byte 0).
237/// This is correct because:
238/// - `unit_id_in_body` clients put the slave ID in byte 0, which is always in
239///   the valid FC range (1–247). Function codes are also 1–127. When both
240///   interpretations parse as a valid Request, we must prefer unit_id_in_body.
241/// - Normal clients have byte 0 = FC. The Unit ID in the MBAP header (byte 6)
242///   is independent. The only time byte 0 == unit_id AND byte 1 is a valid FC
243///   is when a unit_id_in_body client sent the frame.
244fn try_extract_tcp_pdu(
245    tid: u16,
246    unit_id: u8,
247    body: &[u8],
248    consumed: usize,
249) -> Option<(u16, u8, Bytes, usize)> {
250    // Case 1: unit_id_in_body — byte 0 matches MBAP UID and byte 1 is a valid FC.
251    if body.len() >= 2 && body[0] == unit_id && crate::frame::is_known_function_code(body[1]) {
252        let pdu = Bytes::copy_from_slice(&body[1..]);
253        if crate::frame::Request::try_from(pdu.clone()).is_ok() {
254            return Some((tid, body[0], pdu, consumed));
255        }
256    }
257
258    // Case 2: normal framing — byte 0 is the function code.
259    // When body[0] == unit_id but Case 1 fell through, we're in an
260    // ambiguous situation: body[0] could be a UID byte (unit_id_in_body
261    // with unknown FC) or a normal FC that coincidentally equals the UID.
262    // Try parsing as a Request; if it fails, reject rather than guessing.
263    if crate::frame::is_known_function_code(body[0]) {
264        let pdu = Bytes::copy_from_slice(body);
265        if body[0] == unit_id && Request::try_from(pdu.clone()).is_err() {
266            return None;
267        }
268        // Even if Request::try_from fails, return the PDU — the caller will
269        // log the invalid request rather than dropping bytes.
270        return Some((tid, unit_id, pdu, consumed));
271    }
272
273    None
274}
275
276// ── TCP client ──────────────────────────────────────────────────────────
277
278/// Modbus TCP client — `Send + Sync + Clone` via internal `Mutex`.
279///
280/// Call `.with_reconnect(max, backoff)` to enable auto-reconnect on transport
281/// errors. Without it, transport errors are returned directly (no retry).
282///
283/// # Example
284///
285/// ```no_run
286/// use std::time::Duration;
287/// use oms_modbus::*;
288///
289/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
290/// let addr = "192.168.1.10:502".parse()?;
291/// let client = tcp::TcpClient::connect_with_timeout(addr, Duration::from_secs(3)).await?;
292/// let regs = client.read_holding_registers(1, 0, 10).await?;
293/// client.write_single_register(1, 0, 42).await?;
294/// # Ok(())
295/// # }
296/// ```
297pub struct TcpClient {
298    inner: Mutex<TcpInner>,
299    addr: SocketAddr,
300    timeout: Duration,
301    reconnect: Option<crate::reconnect::ReconnectConfig>,
302    tcp_config: TcpConfig,
303    tcp_counter: AtomicU16,
304    /// Preserved for reconnect rebuilds — active tap lives in SniffIo.
305    tap: Option<Arc<dyn WireTap>>,
306    /// Preserved for reconnect rebuilds — active timing lives in SniffIo.
307    bus_timing: Option<Arc<BusTiming>>,
308}
309
310struct TcpInner {
311    stream: SniffIo<TcpStream>,
312    write_buf: BytesMut,
313    read_buf: BytesMut,
314}
315
316impl std::fmt::Debug for TcpClient {
317    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
318        let mut d = f.debug_struct("TcpClient");
319        d.field("addr", &self.addr);
320        d.field("timeout", &self.timeout);
321        if let Some(cfg) = &self.reconnect {
322            d.field("reconnect_max_retries", &cfg.max_retries());
323            d.field("reconnect_interval", &cfg.interval());
324        }
325        d.finish()
326    }
327}
328
329impl TcpClient {
330    /// Connect to a Modbus TCP server (default 5-second timeout).
331    pub async fn connect(addr: SocketAddr) -> std::io::Result<Self> {
332        Self::connect_with_timeout(addr, Duration::from_secs(5)).await
333    }
334
335    /// Connect with an explicit connection timeout.
336    pub async fn connect_with_timeout(
337        addr: SocketAddr,
338        timeout: Duration,
339    ) -> std::io::Result<Self> {
340        Self::connect_with_config(addr, timeout, TcpConfig::default()).await
341    }
342
343    async fn connect_with_config(
344        addr: SocketAddr,
345        timeout: Duration,
346        tcp_config: TcpConfig,
347    ) -> std::io::Result<Self> {
348        let stream = TcpStream::connect(addr).await?;
349        stream.set_nodelay(true)?;
350        Ok(Self {
351            inner: Mutex::new(TcpInner {
352                stream: SniffIo::new(stream, None, None),
353                write_buf: BytesMut::with_capacity(MAX_ADU_SIZE),
354                read_buf: BytesMut::with_capacity(MAX_ADU_SIZE),
355            }),
356            addr,
357            timeout,
358            reconnect: None,
359            tcp_config: tcp_config.clone(),
360            tcp_counter: AtomicU16::new(1),
361            tap: None,
362            bus_timing: None,
363        })
364    }
365
366    /// Set MBAP framing configuration.
367    ///
368    /// Controls transaction ID, unit ID placement, and length calculation.
369    /// Use [`TcpConfig::gateway`] for RTU-over-TCP gateways.
370    pub fn with_config(mut self, config: TcpConfig) -> Self {
371        self.tcp_config = config;
372        self
373    }
374
375    /// Shortcut for gateway mode: auto TID, UID in body, standard length.
376    pub fn with_gateway_mode(self) -> Self {
377        self.with_config(TcpConfig::gateway())
378    }
379
380    /// Enable auto-reconnect on transport errors.
381    /// `max_retries = 0` means infinite retry. Protocol errors are not retried.
382    pub fn with_reconnect(mut self, max_retries: u32, backoff: Duration) -> Self {
383        self.reconnect = Some(crate::reconnect::ReconnectConfig::new(max_retries, backoff));
384        self
385    }
386
387    /// ── TCP receive: two-phase MBAP ─────────────────────────────────────
388    ///
389    /// Phase 1: read MBAP header, parse Length field.
390    /// Phase 2: read remaining PDU bytes.
391    /// No CRC — TCP checksum handles integrity.
392    async fn send_recv(
393        &self,
394        slave_id: u8,
395        request: &Request<'_>,
396    ) -> Result<Response, ModbusError> {
397        let mut inner = self.inner.lock().await;
398        let inner = &mut *inner;
399
400        // TCP is full-duplex, but stale bytes from a previous response or
401        // unsolicited server data can accumulate in the receive buffer.
402        // Drain before sending so the next read returns the correct response.
403        let mut scratch = [0u8; MAX_ADU_SIZE];
404        send_recv::drain_stale_data(&mut inner.stream, &mut scratch).await?;
405
406        // ── Send ────────────────────────────────────────────────────────
407        let tcp_cfg = self.tcp_config.clone();
408        let tid = next_tid(&tcp_cfg.tid, &self.tcp_counter);
409        send_recv::send_frame(
410            &mut inner.stream,
411            &mut inner.write_buf,
412            slave_id,
413            self.timeout,
414            request,
415            |data, buf| encode_tcp_frame(data, buf, &tcp_cfg, tid),
416        )
417        .await?;
418
419        // ── Receive ──────────────────────────────────────────────────────
420        inner.read_buf.clear();
421        let deadline = Instant::now() + self.timeout;
422
423        // Phase 1: read MBAP header (7 bytes)
424        send_recv::read_at_least(
425            &mut inner.stream,
426            &mut inner.read_buf,
427            deadline,
428            MBAP_HEADER_SIZE,
429        )
430        .await?;
431
432        if inner.read_buf.len() < MBAP_HEADER_SIZE {
433            return Err(ModbusError::timeout(TCP_RECV_TIMEOUT));
434        }
435
436        let proto_id = u16::from_be_bytes([inner.read_buf[2], inner.read_buf[3]]);
437        if proto_id != 0 {
438            return Err(ModbusError::protocol("TCP: invalid Protocol ID"));
439        }
440
441        let payload_len = u16::from_be_bytes([inner.read_buf[4], inner.read_buf[5]]) as usize;
442        if payload_len > MAX_TCP_ADU_SIZE {
443            return Err(ModbusError::protocol("TCP: MBAP Length exceeds max ADU"));
444        }
445
446        // Phase 2: read PDU bytes.
447        // The MBAP Length field is ambiguous:
448        //   Standard (LengthMode::Standard): Length = PDU bytes + 1 (includes UID)
449        //   Non-standard (LengthMode::PduOnly): Length = PDU bytes only
450        // Frame sizes: min = MBAP_PREFIX_SIZE + payload_len (standard), max = MBAP_HEADER_SIZE + payload_len (PduOnly)
451        let min_total = MBAP_PREFIX_SIZE + payload_len;
452        send_recv::read_at_least(&mut inner.stream, &mut inner.read_buf, deadline, min_total)
453            .await?;
454
455        // If we only have min_total bytes, one more byte may be needed for PduOnly.
456        let max_total = MBAP_HEADER_SIZE + payload_len;
457        if inner.read_buf.len() < max_total {
458            let remaining = deadline.saturating_duration_since(Instant::now());
459            if !remaining.is_zero() {
460                match tokio::time::timeout(
461                    remaining.min(Duration::from_millis(200)),
462                    inner.stream.read(&mut scratch),
463                )
464                .await
465                {
466                    Ok(Ok(n)) if n > 0 => {
467                        inner.read_buf.extend_from_slice(&scratch[..n]);
468                    }
469                    _ => {}
470                }
471            }
472        }
473
474        let available = inner.read_buf.len();
475        let mut pdu = if available >= max_total {
476            Bytes::copy_from_slice(&inner.read_buf[MBAP_HEADER_SIZE..max_total])
477        } else if available >= min_total {
478            Bytes::copy_from_slice(&inner.read_buf[MBAP_HEADER_SIZE..min_total])
479        } else {
480            return Err(ModbusError::timeout(TCP_RECV_TIMEOUT));
481        };
482
483        if pdu.is_empty() {
484            return Err(ModbusError::protocol(TCP_EMPTY_RESP));
485        }
486
487        // Strip unit_id_in_body prefix if configured — the MBAP header
488        // already carries the Unit ID, but gateway devices may echo it
489        // as the first byte of the body.
490        if tcp_cfg.unit_id_in_body && pdu.len() > 1 && pdu[0] == slave_id {
491            pdu = pdu.slice(1..);
492        }
493
494        // TCP does NOT verify slave_id — Unit ID in MBAP is informational
495        Response::try_from(pdu)
496            .map_err(|e| ModbusError::protocol(format!("{PDU_DECODE_ERROR} {e}")))
497    }
498}
499
500#[async_trait]
501impl ModbusClient for TcpClient {
502    async fn call(&self, slave: u8, request: Request<'_>) -> Result<Response, ModbusError> {
503        let request = request.into_owned();
504        let slave_id = slave;
505
506        send_recv::run_with_reconnect(
507            self.reconnect.as_ref(),
508            || self.send_recv(slave_id, &request),
509            || async {
510                let mut inner = self.inner.lock().await;
511                if let Ok(stream) = TcpStream::connect(self.addr).await {
512                    stream.set_nodelay(true).ok();
513                    inner.stream = SniffIo::new(stream, self.tap.clone(), self.bus_timing.clone());
514                    inner.write_buf.clear();
515                    inner.read_buf.clear();
516                    true
517                } else {
518                    false
519                }
520            },
521            ModbusError::connection,
522        )
523        .await
524    }
525}
526
527/// Create a TCP client with optional capture and reconnect.
528pub async fn with_options(addr: SocketAddr, opts: ClientOptions) -> std::io::Result<TcpClient> {
529    let tap = opts.tap().cloned();
530    let timing = opts.bus_timing.clone();
531    let stream = TcpStream::connect(addr).await?;
532    stream.set_nodelay(true)?;
533    let mut sniff = SniffIo::new(stream, tap, timing);
534    if let Some(cap) = opts.data_channel_capacity {
535        sniff = sniff.with_channel_capacity(cap);
536    }
537    if let Some(cap) = opts.tap_channel_capacity {
538        sniff = sniff.with_tap_channel_capacity(cap);
539    }
540    Ok(TcpClient {
541        inner: Mutex::new(TcpInner {
542            stream: sniff,
543            write_buf: BytesMut::with_capacity(MAX_ADU_SIZE),
544            read_buf: BytesMut::with_capacity(MAX_ADU_SIZE),
545        }),
546        addr,
547        timeout: opts.timeout,
548        reconnect: opts.reconnect,
549        tcp_config: TcpConfig::default(),
550        tcp_counter: AtomicU16::new(1),
551        tap: opts.tap().cloned(),
552        bus_timing: opts.bus_timing.clone(),
553    })
554}
555
556// ── TCP server ──────────────────────────────────────────────────────────
557
558/// Modbus TCP server. Spawns one tokio task per connection.
559///
560/// # Example
561///
562/// ```no_run
563/// use std::net::{IpAddr, Ipv4Addr, SocketAddr};
564/// use std::sync::Arc;
565/// use oms_modbus::*;
566///
567/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
568/// let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 502);
569/// let store = Arc::new(SlaveStore::with_holding_registers(&[(0, 1234)]));
570/// let server = tcp::TcpServer::bind(addr).await?;
571/// println!("Listening on {}", server.local_addr()?);
572/// // server.serve_forever(store).await?;  // blocks forever
573/// # Ok(())
574/// # }
575/// ```
576pub struct TcpServer {
577    listener: tokio::net::TcpListener,
578    tcp_config: TcpConfig,
579}
580
581impl TcpServer {
582    /// Bind to a TCP address and prepare to accept connections.
583    pub async fn bind(addr: SocketAddr) -> std::io::Result<Self> {
584        let listener = tokio::net::TcpListener::bind(addr).await?;
585        Ok(Self {
586            listener,
587            tcp_config: TcpConfig::default(),
588        })
589    }
590
591    /// Set the MBAP framing configuration for response encoding.
592    ///
593    /// Default is [`TcpConfig::standard`]. Use [`TcpConfig::gateway`] for
594    /// RTU-over-TCP gateways, or construct a custom config with
595    /// [`LengthMode::PduOnly`] for non-standard devices.
596    pub fn with_config(mut self, config: TcpConfig) -> Self {
597        self.tcp_config = config;
598        self
599    }
600
601    /// The socket address this server is bound to.
602    pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
603        self.listener.local_addr()
604    }
605
606    /// Accept connections indefinitely, processing each with the given service.
607    ///
608    /// Handler panics are caught via [`futures_util::FutureExt::catch_unwind`]
609    /// and logged rather than silently swallowed.
610    pub async fn serve_forever<S>(self, service: S) -> std::io::Result<()>
611    where
612        S: crate::server::Service + Send + Sync + Clone + 'static,
613    {
614        loop {
615            let (stream, addr) = self.listener.accept().await?;
616            log::info!("TCP server: new connection from {}", addr);
617            let svc = service.clone();
618            let tcp_cfg = self.tcp_config.clone();
619            tokio::spawn(async move {
620                let result = AssertUnwindSafe(handle_tcp_connection(stream, svc, tcp_cfg))
621                    .catch_unwind()
622                    .await;
623                match result {
624                    Ok(Ok(())) => {}
625                    Ok(Err(e)) => {
626                        log::warn!("TCP server: connection {} error: {}", addr, e)
627                    }
628                    Err(_) => {
629                        log::error!("TCP server: handler task panicked");
630                    }
631                }
632            });
633        }
634    }
635}
636
637async fn handle_tcp_connection<S>(
638    stream: TcpStream,
639    service: S,
640    tcp_cfg: TcpConfig,
641) -> std::io::Result<()>
642where
643    S: crate::server::Service + Send + Sync + 'static,
644{
645    let mut sniff = SniffIo::new(stream, None, None);
646    let mut buf = BytesMut::with_capacity(MAX_ADU_SIZE);
647    let mut rsp_buf = BytesMut::with_capacity(MAX_ADU_SIZE);
648    let mut frame_buf = BytesMut::with_capacity(MAX_ADU_SIZE);
649    loop {
650        let mut tmp = [0u8; MAX_ADU_SIZE];
651        match sniff.read(&mut tmp).await {
652            Ok(0) => break,
653            Ok(n) => {
654                buf.extend_from_slice(&tmp[..n]);
655                while let Some((tid, slave_id, pdu, consumed)) = try_parse_tcp_frame(&buf) {
656                    if let Some(rsp_data) =
657                        send_recv::process_server_request(&pdu, slave_id, &service, &mut rsp_buf)
658                            .await
659                    {
660                        frame_buf.clear();
661                        encode_tcp_frame(&rsp_data, &mut frame_buf, &tcp_cfg, tid);
662                        if sniff.write_all(&frame_buf).await.is_err() {
663                            return Ok(());
664                        }
665                    }
666                    let _ = buf.split_to(consumed);
667                }
668                // If unparseable bytes remain and the header is corrupt
669                // (or the buffer is absurdly large), discard them so the
670                // next valid frame isn't contaminated.
671                if is_tcp_header_corrupt(&buf) || buf.len() > MAX_TCP_BUF_BYTES {
672                    buf.clear();
673                }
674            }
675            Err(_) => break,
676        }
677    }
678    Ok(())
679}
680
681// ── Tests ──────────────────────────────────────────────────────────────────
682
683#[cfg(test)]
684mod tests {
685    use super::*;
686
687    #[test]
688    fn try_parse_tcp_frame_empty_buffer() {
689        assert!(try_parse_tcp_frame(&[]).is_none());
690    }
691
692    #[test]
693    fn try_parse_tcp_frame_incomplete_header() {
694        // Less than 7 bytes (MBAP_HEADER_SIZE)
695        assert!(try_parse_tcp_frame(&[0x00, 0x01, 0x00]).is_none());
696        assert!(try_parse_tcp_frame(&[0x00, 0x01, 0x00, 0x00, 0x00, 0x01]).is_none());
697    }
698
699    #[test]
700    fn try_parse_tcp_frame_nonzero_protocol_id() {
701        // Proto ID = 0x0001 (must be 0)
702        let buf = [0x00, 0x01, 0x00, 0x01, 0x00, 0x02, 0x01, 0x03, 0x00];
703        assert!(try_parse_tcp_frame(&buf).is_none());
704    }
705
706    #[test]
707    fn try_parse_tcp_frame_payload_too_large() {
708        // payload_len > MAX_TCP_ADU_SIZE (260)
709        let mut buf = [0u8; MBAP_HEADER_SIZE + 1];
710        buf[0] = 0x00;
711        buf[1] = 0x01; // TID
712        buf[2] = 0x00;
713        buf[3] = 0x00; // Proto = 0
714        buf[4] = 0x01;
715        buf[5] = 0x05; // Length = 261 (> 260)
716        buf[6] = 0x01; // UID
717        buf[7] = 0x03; // FC (body)
718        assert!(try_parse_tcp_frame(&buf).is_none());
719    }
720
721    #[test]
722    fn try_parse_tcp_frame_payload_len_zero() {
723        // payload_len = 0: fl values are 7 and 6, both ≤ MBAP_HEADER_SIZE.
724        // Guard prevents panic, returns None (invalid MBAP frame).
725        let buf = [0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x01];
726        assert!(try_parse_tcp_frame(&buf).is_none());
727    }
728
729    #[test]
730    fn try_parse_tcp_frame_payload_len_one_no_pdu() {
731        // payload_len = 1 → body is just [UID], not a valid FC.
732        // UID=0x46 is NOT a known function code → try_extract_tcp_pdu returns None.
733        let buf = [0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x46, 0x46];
734        assert!(try_parse_tcp_frame(&buf).is_none());
735    }
736
737    #[test]
738    fn try_parse_tcp_frame_unknown_function_code() {
739        // MBAP header complete, body exists, but byte 0 is NOT a known FC (0x46)
740        let buf = [
741            0x00, 0x01, // TID
742            0x00, 0x00, // Proto = 0
743            0x00, 0x02, // Length = 2
744            0x01, // UID
745            0x46, 0x00, // Unknown FC 0x46
746        ];
747        assert!(try_parse_tcp_frame(&buf).is_none());
748    }
749
750    #[test]
751    fn try_parse_tcp_frame_valid_minimal() {
752        // Minimal valid frame: ReadCoils request
753        let buf = [
754            0x00, 0x01, // TID
755            0x00, 0x00, // Proto = 0
756            0x00, 0x05, // Length = 5
757            0x01, // UID
758            0x01, 0x00, 0x00, 0x00, 0x01, // FC=1, addr=0, count=1
759        ];
760        let (tid, slave, pdu, consumed) = try_parse_tcp_frame(&buf).unwrap();
761        assert_eq!(tid, 1);
762        assert_eq!(slave, 1);
763        assert_eq!(pdu.len(), 5);
764        assert_eq!(consumed, 12);
765    }
766
767    #[test]
768    fn try_parse_tcp_frame_with_unit_id_in_body() {
769        // Gateway-style: UID in body byte 0 matches MBAP UID
770        // Body: [UID=1, FC=3, addr_hi, addr_lo, count_hi, count_lo]
771        let buf = [
772            0x00, 0x01, // TID
773            0x00, 0x00, // Proto = 0
774            0x00, 0x06, // Length = 6
775            0x01, // UID in MBAP
776            0x01, 0x03, 0x00, 0x00, 0x00, 0x01, // body: UID=1, FC=3, addr=0, count=1
777        ];
778        let (tid, slave, pdu, _) = try_parse_tcp_frame(&buf).unwrap();
779        assert_eq!(tid, 1);
780        // Slave should come from body[0], PDU should skip the duplicate UID
781        assert_eq!(slave, 1);
782        // PDU is body[1..] = [FC, addr, count]
783        assert_eq!(pdu[0], 0x03);
784        assert_eq!(pdu.len(), 5);
785    }
786
787    #[test]
788    fn try_parse_tcp_frame_pdu_only_length() {
789        // PduOnly: Length = 5 (PDU only, no +1 for UID)
790        // min_total = MBAP_PREFIX_SIZE + 5 = 11, max_total = MBAP_HEADER_SIZE + 5 = 12
791        // So second iteration of the `for` loop finds the match.
792        let buf = [
793            0x00, 0x01, // TID
794            0x00, 0x00, // Proto = 0
795            0x00, 0x05, // Length = 5 (PduOnly)
796            0x01, // UID
797            0x03, 0x00, 0x00, 0x00, 0x01, // PDU: FC=3, addr=0, count=1
798        ];
799        let (tid, slave, pdu, consumed) = try_parse_tcp_frame(&buf).unwrap();
800        assert_eq!(tid, 1);
801        assert_eq!(slave, 1);
802        assert_eq!(pdu.len(), 5);
803        // consumed = MBAP_HEADER_SIZE + payload_len = 12 (full frame)
804        assert_eq!(consumed, 12);
805    }
806
807    #[test]
808    fn try_parse_tcp_frame_not_enough_data() {
809        // Correct MBAP header with payload_len=5, but buffer only has 10 bytes
810        // (MBAP_HEADER_SIZE + 3 < MBAP_HEADER_SIZE + 5)
811        let buf = [
812            0x00, 0x01, 0x00, 0x00, 0x00, 0x05, 0x01, // MBAP (7 bytes)
813            0x03, 0x00, 0x00, // Only 3 PDU bytes (need 5)
814        ];
815        assert!(try_parse_tcp_frame(&buf).is_none());
816    }
817
818    #[test]
819    fn try_parse_tcp_frame_max_valid_length() {
820        // payload_len = MAX_TCP_ADU_SIZE (260) → valid size
821        let mut buf = vec![0u8; MBAP_HEADER_SIZE + 260];
822        buf[0] = 0x00;
823        buf[1] = 0x01; // TID
824        buf[2] = 0x00;
825        buf[3] = 0x00; // Proto = 0
826        buf[4] = 0x01;
827        buf[5] = 0x04; // Length = 260
828        buf[6] = 0x01; // UID
829        buf[7] = 0x03; // FC=3 (first byte of body, ensures known FC)
830        assert!(try_parse_tcp_frame(&buf).is_some());
831    }
832}