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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
//! Implementation of the [RCON](https://wiki.vg/RCON) protocol.

mod packet;

use self::packet::{RconPacket, RconPacketType};
use crate::errors::RconProtocolError;
use bytes::{BufMut, BytesMut};
use tokio::{
    io::{self, AsyncReadExt, AsyncWriteExt, Error},
    net::TcpStream,
};

/// Struct that stores the connection and other state of the RCON protocol with the server.
///
/// # Examples
///
/// ```no_run
/// use mc_query::rcon::RconClient;
/// use tokio::io::Result;
///
/// #[tokio::main]
/// async fn main() -> Result<()> {
///     let mut client = RconClient::new("localhost", 25575).await?;
///     client.authenticate("password").await?;
///
///     let output = client.run_command("time set day").await?;
///     println!("{output}");
///
///     Ok(())
/// }
/// ```
#[derive(Debug)]
pub struct RconClient {
    socket: TcpStream,
}

impl RconClient {
    /// Construct an [RconClient] that connects to the given host and port.
    /// Note: to authenticate use the `authenticate` method, this method does not take a password.
    ///
    /// # Arguments
    /// * `host` - A string slice that holds the hostname of the server to connect to.
    /// * `port` - The port to connect to.
    pub async fn new(host: &str, port: u16) -> io::Result<Self> {
        let connection = TcpStream::connect(format!("{host}:{port}")).await?;

        Ok(Self { socket: connection })
    }

    /// Disconnect from the server and close the RCON connection.
    pub async fn disconnect(mut self) -> io::Result<()> {
        self.socket.shutdown().await
    }

    /// Authenticate with the server, with the given password.
    ///
    /// If authentication fails, this method will return [RconProtocolError::AuthFailed].
    ///
    /// # Arguments
    /// * `password` - A string slice that holds the RCON password.
    pub async fn authenticate(&mut self, password: &str) -> io::Result<()> {
        let packet =
            RconPacket::new(1, RconPacketType::Login, password.to_string()).map_err(Error::from)?;

        self.write_packet(packet).await?;

        let packet = self.read_packet().await?;

        if !matches!(packet.packet_type, RconPacketType::RunCommand) {
            return Err(RconProtocolError::InvalidPacketType.into());
        }

        if packet.request_id == -1 {
            return Err(RconProtocolError::AuthFailed.into());
        } else if packet.request_id != 1 {
            return Err(RconProtocolError::RequestIdMismatch.into());
        }

        Ok(())
    }

    /// Run the given command on the server and return the result.
    ///
    /// # Arguments
    /// * `command` - A string slice that holds the command to run. Must be ASCII and under 1446 bytes in length.
    pub async fn run_command(&mut self, command: &str) -> io::Result<String> {
        let packet = RconPacket::new(1, RconPacketType::RunCommand, command.to_string())
            .map_err(Error::from)?;

        self.write_packet(packet).await?;

        let mut full_payload = "".to_string();

        loop {
            let recieved = self.read_packet().await?;

            if recieved.request_id == -1 {
                return Err(RconProtocolError::AuthFailed.into());
            } else if recieved.request_id != 1 {
                return Err(RconProtocolError::RequestIdMismatch.into());
            }

            full_payload.push_str(&recieved.payload);

            // wiki says this method of determining if this is the end of the
            // response is not 100% reliable, but this is the best solution imo
            // if this ends up being a problem, this can be changed later
            if recieved.payload.len() < 4096 {
                break;
            }
        }

        Ok(full_payload)
    }

    /// Read a packet from the socket.
    async fn read_packet(&mut self) -> io::Result<RconPacket> {
        let len = self.socket.read_i32_le().await?;

        let mut bytes = BytesMut::new();
        bytes.put_i32_le(len);

        for _ in 0..len {
            let current = self.socket.read_u8().await?;
            bytes.put_u8(current);
        }

        RconPacket::try_from(bytes.freeze()).map_err(Error::from)
    }

    /// Write a packet to the socket.
    ///
    /// # Arguments
    /// * `packet` - An owned [RconPacket] to write to the socket.
    async fn write_packet(&mut self, packet: RconPacket) -> io::Result<()> {
        let bytes = packet.bytes();

        self.socket.write_all(&bytes).await
    }
}

#[cfg(test)]
mod tests {
    use crate::rcon::RconClient;
    use tokio::io;

    #[tokio::test]
    async fn test_rcon_command() -> io::Result<()> {
        let mut client = RconClient::new("localhost", 25575).await?;
        client.authenticate("mc-query-test").await?;
        let response = client.run_command("time set day").await?;

        println!("recieved response: {response}");

        Ok(())
    }

    #[tokio::test]
    async fn test_rcon_unauthenticated() -> io::Result<()> {
        let mut client = RconClient::new("localhost", 25575).await?;
        let result = client.run_command("time set day").await;

        assert!(result.is_err());

        Ok(())
    }

    #[tokio::test]
    async fn test_rcon_incorrect_password() -> io::Result<()> {
        let mut client = RconClient::new("localhost", 25575).await?;
        let result = client.authenticate("incorrect").await;

        assert!(result.is_err());

        Ok(())
    }
}