utils_box/connections/
tcp_client.rs

1//! # TCP Client utility
2//! A small TCP utility to connect and manipulate TCP connections to a server.
3//! Useful for sending and receiving data via a TCP socket.
4
5use anyhow::{Result, bail};
6use std::{
7    io::{Read, Write},
8    net::{Shutdown, TcpStream},
9};
10
11use crate::{log_error, log_info, log_trace};
12
13pub static BUFFER_SIZE: usize = 500;
14
15pub struct TcpClient {
16    server_ip: String,
17    server_port: u16,
18    tcp_stream: TcpStream,
19}
20
21impl TcpClient {
22    /// Create a new TCP connection to the specified server and connect to it
23    pub fn new(server_ip: String, server_port: u16) -> Result<Self> {
24        let tcp_stream = TcpStream::connect(format!("{server_ip}:{server_port}"))?;
25
26        Ok(Self {
27            server_ip,
28            server_port,
29            tcp_stream,
30        })
31    }
32
33    /// Get the server IP and port information
34    pub fn server_info(&self) -> (String, u16) {
35        (self.server_ip.clone(), self.server_port)
36    }
37
38    pub fn send(&mut self, data: &[u8]) -> Result<()> {
39        match self.tcp_stream.write(data) {
40            Ok(size) => {
41                log_trace!(
42                    "[TcpClient][{}:{}] Send [{} Bytes] SUCCESSFULLY!",
43                    self.server_ip,
44                    self.server_port,
45                    size
46                );
47                Ok(())
48            }
49            Err(e) => {
50                log_error!(
51                    "[TcpClient][{}:{}] Send FAILED with [{}]",
52                    self.server_ip,
53                    self.server_port,
54                    e
55                );
56                bail!(e)
57            }
58        }
59    }
60
61    pub fn receive(&mut self) -> Result<Vec<u8>> {
62        let mut data: Vec<u8> = vec![0; BUFFER_SIZE];
63
64        match self.tcp_stream.read(&mut data) {
65            Ok(data_len) => {
66                log_trace!(
67                    "[TcpClient][{}:{}] Received [{} Bytes] SUCCESSFULLY!",
68                    self.server_ip,
69                    self.server_port,
70                    data_len
71                );
72                Ok(data[..data_len].to_vec())
73            }
74            Err(e) => {
75                log_error!(
76                    "[TcpClient][{}:{}] Receive FAILED with [{}]",
77                    self.server_ip,
78                    self.server_port,
79                    e
80                );
81                bail!(e)
82            }
83        }
84    }
85}
86
87impl Drop for TcpClient {
88    fn drop(&mut self) {
89        match self.tcp_stream.shutdown(Shutdown::Both) {
90            Ok(_) => log_info!(
91                "[TcpClient] Disconnected from [{}:{}]",
92                self.server_ip,
93                self.server_port
94            ),
95            Err(_) => log_error!(
96                "[TcpClient] FAILED to Grecefully Disconnect from [{}:{}]. Dropping...",
97                self.server_ip,
98                self.server_port
99            ),
100        }
101    }
102}
utils_box/connections/tcp_client.rs

utils_box/connections/
tcp_client.rs
