veloren_query_server/
client.rs

1use std::{
2    io,
3    net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
4    time::{Duration, Instant},
5};
6
7use protocol::Parcel;
8use tokio::{net::UdpSocket, time::timeout};
9use tracing::trace;
10
11use crate::proto::{
12    QueryServerRequest, QueryServerResponse, RawQueryServerRequest, RawQueryServerResponse,
13    ServerInfo, MAX_RESPONSE_SIZE,
14};
15
16// This must be at least 2 for the client to get a value for the `p` field.
17const MAX_REQUEST_RETRIES: usize = 5;
18
19#[derive(Debug)]
20pub enum QueryClientError {
21    Io(tokio::io::Error),
22    Protocol(protocol::Error),
23    InvalidResponse,
24    Timeout,
25    ChallengeFailed,
26}
27
28struct ClientInitData {
29    p: u64,
30    #[allow(dead_code)]
31    server_max_version: u16,
32}
33
34/// The `p` field has to be requested from the server each time this client is
35/// constructed, if possible reuse this!
36pub struct QueryClient {
37    pub addr: SocketAddr,
38    init: Option<ClientInitData>,
39}
40
41impl QueryClient {
42    pub fn new(addr: SocketAddr) -> Self { Self { addr, init: None } }
43
44    pub async fn server_info(&mut self) -> Result<(ServerInfo, Duration), QueryClientError> {
45        self.send_query(QueryServerRequest::ServerInfo)
46            .await
47            .and_then(|(response, duration)| {
48                #[allow(irrefutable_let_patterns)]
49                if let QueryServerResponse::ServerInfo(info) = response {
50                    Ok((info, duration))
51                } else {
52                    Err(QueryClientError::InvalidResponse)
53                }
54            })
55    }
56
57    async fn send_query(
58        &mut self,
59        request: QueryServerRequest,
60    ) -> Result<(QueryServerResponse, Duration), QueryClientError> {
61        let socket = UdpSocket::bind(if self.addr.is_ipv4() {
62            SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0))
63        } else {
64            SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, 0, 0, 0))
65        })
66        .await?;
67
68        for _ in 0..MAX_REQUEST_RETRIES {
69            let request = if let Some(init) = &self.init {
70                // TODO: Use the maximum version supported by both the client and server once
71                // new protocol versions are added
72                RawQueryServerRequest { p: init.p, request }
73            } else {
74                // TODO: Use the legacy version here once new protocol versions are added
75                RawQueryServerRequest {
76                    p: 0,
77                    request: QueryServerRequest::Init,
78                }
79            };
80            let buf = request.serialize()?;
81            let query_sent = Instant::now();
82            socket.send_to(&buf, self.addr).await?;
83            let mut buf = vec![0; MAX_RESPONSE_SIZE];
84            let (buf_len, _) = timeout(Duration::from_secs(2), socket.recv_from(&mut buf))
85                .await
86                .map_err(|_| QueryClientError::Timeout)??;
87
88            if buf_len <= 2 {
89                Err(QueryClientError::InvalidResponse)?
90            }
91
92            let packet = <RawQueryServerResponse as Parcel>::read(
93                &mut io::Cursor::new(&buf[..buf_len]),
94                &Default::default(),
95            )?;
96
97            match packet {
98                RawQueryServerResponse::Response(response) => {
99                    return Ok((response, query_sent.elapsed()));
100                },
101                RawQueryServerResponse::Init(init) => {
102                    trace!(?init, "Resetting p");
103                    self.init = Some(ClientInitData {
104                        p: init.p,
105                        server_max_version: init.max_supported_version,
106                    });
107                },
108            }
109        }
110
111        Err(QueryClientError::ChallengeFailed)
112    }
113}
114
115impl From<tokio::io::Error> for QueryClientError {
116    fn from(value: tokio::io::Error) -> Self { Self::Io(value) }
117}
118
119impl From<protocol::Error> for QueryClientError {
120    fn from(value: protocol::Error) -> Self { Self::Protocol(value) }
121}
veloren_query_server/client.rs

veloren_query_server/
client.rs
