1use super::client::{ApiHttpError, WingsClient};
2use futures_util::{SinkExt, StreamExt, ready};
3use serde::Deserialize;
4use std::{
5 io,
6 pin::Pin,
7 task::{Context, Poll},
8 time::Duration,
9};
10use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
11use tokio_tungstenite::{
12 MaybeTlsStream, WebSocketStream,
13 tungstenite::{Error as WsError, Message},
14};
15
16type WsStream = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;
17
18pub const MAX_DATAGRAM_SIZE: usize = 65536;
19
20const RECV_TIMEOUT: Duration = Duration::from_secs(5);
21
22const REFUSED_SIGNAL: &str = "refused";
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum QueryProtocol {
26 Tcp,
27 Udp,
28}
29
30impl QueryProtocol {
31 fn as_str(self) -> &'static str {
32 match self {
33 QueryProtocol::Tcp => "tcp",
34 QueryProtocol::Udp => "udp",
35 }
36 }
37}
38
39impl WingsClient {
40 async fn open_tunnel(&self, endpoint: String) -> Result<WsStream, ApiHttpError> {
41 let mut headers = reqwest::header::HeaderMap::new();
42 headers.insert(
43 reqwest::header::ACCEPT,
44 reqwest::header::HeaderValue::from_static("application/msgpack"),
45 );
46
47 match self.open_websocket(endpoint, headers).await {
48 Err(ApiHttpError::WebSocket(WsError::Http(response))) => {
49 let error = match response.body() {
50 Some(body) => {
51 let mut de =
52 rmp_serde::Deserializer::new(body.as_slice()).with_human_readable();
53 match super::ApiError::deserialize(&mut de) {
54 Ok(error) => error,
55 Err(err) => super::ApiError {
56 error: err.to_string().into(),
57 },
58 }
59 }
60 None => super::ApiError {
61 error: "websocket upgrade rejected".into(),
62 },
63 };
64
65 Err(ApiHttpError::Http(response.status(), error))
66 }
67 result => result,
68 }
69 }
70
71 async fn open_query_tunnel(
72 &self,
73 server: uuid::Uuid,
74 protocol: QueryProtocol,
75 port: u16,
76 ) -> Result<WsStream, ApiHttpError> {
77 self.open_tunnel(format!(
78 "/api/servers/{server}/ws/query?protocol={}&port={port}",
79 protocol.as_str()
80 ))
81 .await
82 }
83
84 pub async fn open_tunnel_tcp(
85 &self,
86 server: uuid::Uuid,
87 port: u16,
88 ) -> Result<QueryStreamTunnel, ApiHttpError> {
89 Ok(QueryStreamTunnel {
90 stream: self
91 .open_query_tunnel(server, QueryProtocol::Tcp, port)
92 .await?,
93 read: Vec::new(),
94 read_pos: 0,
95 })
96 }
97
98 pub async fn open_tunnel_udp(
99 &self,
100 server: uuid::Uuid,
101 port: u16,
102 ) -> Result<QueryUdpTunnel, ApiHttpError> {
103 Ok(QueryUdpTunnel {
104 stream: self
105 .open_query_tunnel(server, QueryProtocol::Udp, port)
106 .await?,
107 })
108 }
109
110 pub async fn open_tunnel_unix(
111 &self,
112 server: uuid::Uuid,
113 path: &str,
114 ignored: &[compact_str::CompactString],
115 ) -> Result<QueryStreamTunnel, ApiHttpError> {
116 let mut endpoint = format!(
117 "/api/servers/{server}/ws/socket?path={}",
118 urlencoding::encode(path)
119 );
120 for value in ignored {
121 endpoint.push_str("&ignored=");
122 endpoint.push_str(&urlencoding::encode(value));
123 }
124
125 Ok(QueryStreamTunnel {
126 stream: self.open_tunnel(endpoint).await?,
127 read: Vec::new(),
128 read_pos: 0,
129 })
130 }
131}
132
133pub struct QueryStreamTunnel {
134 stream: WsStream,
135 read: Vec<u8>,
136 read_pos: usize,
137}
138
139impl AsyncRead for QueryStreamTunnel {
140 fn poll_read(
141 mut self: Pin<&mut Self>,
142 cx: &mut Context<'_>,
143 buf: &mut ReadBuf<'_>,
144 ) -> Poll<io::Result<()>> {
145 loop {
146 if self.read_pos < self.read.len() {
147 let n = (self.read.len() - self.read_pos).min(buf.remaining());
148 let start = self.read_pos;
149 buf.put_slice(&self.read[start..start + n]);
150 self.read_pos += n;
151 return Poll::Ready(Ok(()));
152 }
153
154 match ready!(self.stream.poll_next_unpin(cx)) {
155 Some(Ok(Message::Binary(data))) => {
156 self.read = data.to_vec();
157 self.read_pos = 0;
158 }
159 Some(Ok(Message::Close(_))) | None => return Poll::Ready(Ok(())),
160 Some(Ok(_)) => {}
161 Some(Err(err)) => return Poll::Ready(Err(ws_to_io(err))),
162 }
163 }
164 }
165}
166
167impl AsyncWrite for QueryStreamTunnel {
168 fn poll_write(
169 mut self: Pin<&mut Self>,
170 cx: &mut Context<'_>,
171 buf: &[u8],
172 ) -> Poll<io::Result<usize>> {
173 ready!(self.stream.poll_ready_unpin(cx)).map_err(ws_to_io)?;
174 self.stream
175 .start_send_unpin(Message::binary(buf.to_vec()))
176 .map_err(ws_to_io)?;
177 Poll::Ready(Ok(buf.len()))
178 }
179
180 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
181 self.stream.poll_flush_unpin(cx).map_err(ws_to_io)
182 }
183
184 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
185 self.stream.poll_close_unpin(cx).map_err(ws_to_io)
186 }
187}
188
189pub struct QueryUdpTunnel {
190 stream: WsStream,
191}
192
193impl QueryUdpTunnel {
194 pub async fn send(&mut self, data: &[u8]) -> io::Result<()> {
195 self.stream
196 .send(Message::binary(data.to_vec()))
197 .await
198 .map_err(ws_to_io)
199 }
200
201 pub async fn recv(&mut self, buf: &mut [u8]) -> io::Result<usize> {
202 loop {
203 let message = match tokio::time::timeout(RECV_TIMEOUT, self.stream.next()).await {
204 Ok(Some(Ok(message))) => message,
205 Ok(Some(Err(err))) => return Err(ws_to_io(err)),
206 Ok(None) => return Err(io::ErrorKind::UnexpectedEof.into()),
207 Err(_) => return Err(io::ErrorKind::TimedOut.into()),
208 };
209
210 match message {
211 Message::Binary(data) => {
212 let n = data.len().min(buf.len());
213 buf[..n].copy_from_slice(&data[..n]);
214 return Ok(n);
215 }
216 Message::Close(_) => return Err(io::ErrorKind::UnexpectedEof.into()),
217 Message::Text(text) if text.as_str() == REFUSED_SIGNAL => {
218 return Err(io::ErrorKind::ConnectionRefused.into());
219 }
220 _ => {}
221 }
222 }
223 }
224}
225
226fn ws_to_io(err: WsError) -> io::Error {
227 match err {
228 WsError::Io(err) => err,
229 other => io::Error::other(other),
230 }
231}