1use 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#[derive(Debug)]
73pub enum TidMode {
74 Fixed(u16),
76 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#[derive(Debug, Clone, Copy, PartialEq)]
92pub enum LengthMode {
93 Standard,
96 PduOnly,
99}
100
101#[derive(Debug, Clone)]
106pub struct TcpConfig {
107 pub tid: TidMode,
108 pub unit_id_in_body: bool,
111 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 pub fn standard() -> Self {
128 Self::default()
129 }
130
131 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
141fn 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
150pub 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); 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
172fn is_tcp_header_corrupt(buf: &[u8]) -> bool {
187 if buf.len() < MBAP_HEADER_SIZE {
188 return false; }
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
195const MAX_TCP_BUF_BYTES: usize = MAX_TCP_ADU_SIZE * 4; pub 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 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
232fn try_extract_tcp_pdu(
245 tid: u16,
246 unit_id: u8,
247 body: &[u8],
248 consumed: usize,
249) -> Option<(u16, u8, Bytes, usize)> {
250 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 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 return Some((tid, unit_id, pdu, consumed));
271 }
272
273 None
274}
275
276pub 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 tap: Option<Arc<dyn WireTap>>,
306 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 pub async fn connect(addr: SocketAddr) -> std::io::Result<Self> {
332 Self::connect_with_timeout(addr, Duration::from_secs(5)).await
333 }
334
335 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 pub fn with_config(mut self, config: TcpConfig) -> Self {
371 self.tcp_config = config;
372 self
373 }
374
375 pub fn with_gateway_mode(self) -> Self {
377 self.with_config(TcpConfig::gateway())
378 }
379
380 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 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 let mut scratch = [0u8; MAX_ADU_SIZE];
404 send_recv::drain_stale_data(&mut inner.stream, &mut scratch).await?;
405
406 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 inner.read_buf.clear();
421 let deadline = Instant::now() + self.timeout;
422
423 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 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 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 if tcp_cfg.unit_id_in_body && pdu.len() > 1 && pdu[0] == slave_id {
491 pdu = pdu.slice(1..);
492 }
493
494 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
527pub 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
556pub struct TcpServer {
577 listener: tokio::net::TcpListener,
578 tcp_config: TcpConfig,
579}
580
581impl TcpServer {
582 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 pub fn with_config(mut self, config: TcpConfig) -> Self {
597 self.tcp_config = config;
598 self
599 }
600
601 pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
603 self.listener.local_addr()
604 }
605
606 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 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#[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 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 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 let mut buf = [0u8; MBAP_HEADER_SIZE + 1];
710 buf[0] = 0x00;
711 buf[1] = 0x01; buf[2] = 0x00;
713 buf[3] = 0x00; buf[4] = 0x01;
715 buf[5] = 0x05; buf[6] = 0x01; buf[7] = 0x03; assert!(try_parse_tcp_frame(&buf).is_none());
719 }
720
721 #[test]
722 fn try_parse_tcp_frame_payload_len_zero() {
723 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 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 let buf = [
741 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x01, 0x46, 0x00, ];
747 assert!(try_parse_tcp_frame(&buf).is_none());
748 }
749
750 #[test]
751 fn try_parse_tcp_frame_valid_minimal() {
752 let buf = [
754 0x00, 0x01, 0x00, 0x00, 0x00, 0x05, 0x01, 0x01, 0x00, 0x00, 0x00, 0x01, ];
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 let buf = [
772 0x00, 0x01, 0x00, 0x00, 0x00, 0x06, 0x01, 0x01, 0x03, 0x00, 0x00, 0x00, 0x01, ];
778 let (tid, slave, pdu, _) = try_parse_tcp_frame(&buf).unwrap();
779 assert_eq!(tid, 1);
780 assert_eq!(slave, 1);
782 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 let buf = [
793 0x00, 0x01, 0x00, 0x00, 0x00, 0x05, 0x01, 0x03, 0x00, 0x00, 0x00, 0x01, ];
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 assert_eq!(consumed, 12);
805 }
806
807 #[test]
808 fn try_parse_tcp_frame_not_enough_data() {
809 let buf = [
812 0x00, 0x01, 0x00, 0x00, 0x00, 0x05, 0x01, 0x03, 0x00, 0x00, ];
815 assert!(try_parse_tcp_frame(&buf).is_none());
816 }
817
818 #[test]
819 fn try_parse_tcp_frame_max_valid_length() {
820 let mut buf = vec![0u8; MBAP_HEADER_SIZE + 260];
822 buf[0] = 0x00;
823 buf[1] = 0x01; buf[2] = 0x00;
825 buf[3] = 0x00; buf[4] = 0x01;
827 buf[5] = 0x04; buf[6] = 0x01; buf[7] = 0x03; assert!(try_parse_tcp_frame(&buf).is_some());
831 }
832}