1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

//! # TCP Client utility
//! A small TCP utility to connect and manipulate TCP connections to a server.
//! Useful for sending and receiving data via a TCP socket.

use anyhow::{bail, Result};
use std::{
    io::{Read, Write},
    net::{Shutdown, TcpStream},
};

use crate::{log_error, log_info, log_trace};

pub static BUFFER_SIZE: usize = 500;

pub struct TcpClient {
    server_ip: String,
    server_port: u16,
    tcp_stream: TcpStream,
}

impl TcpClient {
    /// Create a new TCP connection to the specified server and connect to it
    pub fn new(server_ip: String, server_port: u16) -> Result<Self> {
        let tcp_stream = TcpStream::connect(format!("{}:{}", server_ip, server_port))?;

        Ok(Self {
            server_ip,
            server_port,
            tcp_stream,
        })
    }

    /// Get the server IP and port information
    pub fn server_info(&self) -> (String, u16) {
        (self.server_ip.clone(), self.server_port)
    }

    pub fn send(&mut self, data: &[u8]) -> Result<()> {
        match self.tcp_stream.write(data) {
            Ok(_) => {
                log_trace!(
                    "[TcpClient][{}:{}] Send [{} Bytes] SUCCESSFULLY!",
                    self.server_ip,
                    self.server_port,
                    data.len()
                );
                Ok(())
            }
            Err(e) => {
                log_error!(
                    "[TcpClient][{}:{}] Send FAILED with [{}]",
                    self.server_ip,
                    self.server_port,
                    e
                );
                bail!(e)
            }
        }
    }

    pub fn receive(&mut self) -> Result<Vec<u8>> {
        let mut data: Vec<u8> = vec![0; BUFFER_SIZE];

        match self.tcp_stream.read(&mut data) {
            Ok(data_len) => {
                log_trace!(
                    "[TcpClient][{}:{}] Received [{} Bytes] SUCCESSFULLY!",
                    self.server_ip,
                    self.server_port,
                    data_len
                );
                Ok(data[..data_len].to_vec())
            }
            Err(e) => {
                log_error!(
                    "[TcpClient][{}:{}] Receive FAILED with [{}]",
                    self.server_ip,
                    self.server_port,
                    e
                );
                bail!(e)
            }
        }
    }
}

impl Drop for TcpClient {
    fn drop(&mut self) {
        match self.tcp_stream.shutdown(Shutdown::Both) {
            Ok(_) => log_info!(
                "[TcpClient] Disconnected from [{}:{}]",
                self.server_ip,
                self.server_port
            ),
            Err(_) => log_error!(
                "[TcpClient] FAILED to Grecefully Disconnect from [{}:{}]. Dropping...",
                self.server_ip,
                self.server_port
            ),
        }
    }
}