This repository has no description
1use std::future::Future;
2use std::net::SocketAddr;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use bytes::Bytes;
7use futures::stream::{Stream, StreamExt};
8use http::{HeaderMap, StatusCode};
9use thiserror::Error;
10use tokio_tungstenite::tungstenite::{
11 Bytes as WsBytes, Message as TungsteniteMessage, protocol::CloseFrame as TungsteniteClose,
12 protocol::frame::coding::CloseCode as TungsteniteCloseCode,
13};
14use url::Url;
15
16#[derive(Debug, Error)]
17pub enum NetworkError {
18 #[error("connect: {0}")]
19 Connect(String),
20 #[error("timeout: {0}")]
21 Timeout(String),
22 #[error("redirect: {0}")]
23 Redirect(String),
24 #[error("transport: {0}")]
25 Transport(String),
26 #[error("body: {0}")]
27 Body(String),
28 #[error("protocol: {0}")]
29 Protocol(String),
30}
31
32pub struct HttpRequest {
33 pub url: Url,
34 pub headers: HeaderMap,
35}
36
37pub type BodyStream = Pin<Box<dyn Stream<Item = Result<Bytes, NetworkError>> + Send + 'static>>;
38
39pub struct HttpResponseHead {
40 pub status: StatusCode,
41 pub headers: HeaderMap,
42 pub content_length: Option<u64>,
43 pub body: BodyStream,
44}
45
46pub type HttpResult = Result<HttpResponseHead, NetworkError>;
47pub type HttpResponseFuture = Pin<Box<dyn Future<Output = HttpResult> + Send + 'static>>;
48
49pub trait HttpTransport: Send + Sync + 'static {
50 fn execute(&self, request: HttpRequest) -> HttpResponseFuture;
51}
52
53#[derive(Clone, Debug)]
54pub struct ReqwestHttp {
55 client: reqwest::Client,
56}
57
58impl ReqwestHttp {
59 pub fn new(client: reqwest::Client) -> Self {
60 Self { client }
61 }
62
63 pub fn shared(client: reqwest::Client) -> Arc<dyn HttpTransport> {
64 Arc::new(Self::new(client))
65 }
66}
67
68#[derive(Debug, Error)]
69pub enum ExtraCaError {
70 #[error("read {path}: {source}")]
71 Read {
72 path: String,
73 source: std::io::Error,
74 },
75 #[error("parse PEM certificate at {path}: {source}")]
76 Parse { path: String, source: reqwest::Error },
77}
78
79/// Loads a PEM certificate from `path` (if given) so it can be added to a
80/// `reqwest::ClientBuilder` via `add_root_certificate`, extending the bundled
81/// webpki roots rather than replacing them. Used to trust a local dev CA -
82/// `rustls-tls-webpki-roots` ignores the OS trust store, so mounting a CA
83/// into `/etc/ssl/certs` has no effect on reqwest without this.
84pub fn load_extra_ca_cert(path: &std::path::Path) -> Result<reqwest::Certificate, ExtraCaError> {
85 let pem = std::fs::read(path).map_err(|source| ExtraCaError::Read {
86 path: path.display().to_string(),
87 source,
88 })?;
89 reqwest::Certificate::from_pem(&pem).map_err(|source| ExtraCaError::Parse {
90 path: path.display().to_string(),
91 source,
92 })
93}
94
95impl HttpTransport for ReqwestHttp {
96 fn execute(&self, request: HttpRequest) -> HttpResponseFuture {
97 let client = self.client.clone();
98 Box::pin(async move {
99 let resp = client
100 .get(request.url)
101 .headers(request.headers)
102 .send()
103 .await
104 .map_err(map_reqwest)?;
105 let status = resp.status();
106 let headers = resp.headers().clone();
107 let content_length = resp.content_length();
108 let body: BodyStream = Box::pin(
109 resp.bytes_stream()
110 .map(|chunk| chunk.map_err(|e| NetworkError::Body(e.to_string()))),
111 );
112 Ok(HttpResponseHead {
113 status,
114 headers,
115 content_length,
116 body,
117 })
118 })
119 }
120}
121
122fn map_reqwest(err: reqwest::Error) -> NetworkError {
123 let msg = err.to_string();
124 if err.is_timeout() {
125 NetworkError::Timeout(msg)
126 } else if err.is_connect() {
127 NetworkError::Connect(msg)
128 } else if err.is_redirect() {
129 NetworkError::Redirect(msg)
130 } else {
131 NetworkError::Transport(msg)
132 }
133}
134
135#[derive(Clone, Debug)]
136pub enum WsMessage {
137 Text(String),
138 Binary(Bytes),
139 Ping(Bytes),
140 Pong(Bytes),
141 Close { code: u16, reason: String },
142}
143
144pub type WsSendFuture<'a> = Pin<Box<dyn Future<Output = Result<(), NetworkError>> + Send + 'a>>;
145pub type WsMessageFuture<'a> =
146 Pin<Box<dyn Future<Output = Option<Result<WsMessage, NetworkError>>> + Send + 'a>>;
147
148pub trait WsSink: Send + 'static {
149 fn send<'a>(&'a mut self, message: WsMessage) -> WsSendFuture<'a>;
150}
151
152pub trait WsStream: Send + 'static {
153 fn next<'a>(&'a mut self) -> WsMessageFuture<'a>;
154}
155
156pub struct WsConn {
157 pub sink: Box<dyn WsSink>,
158 pub stream: Box<dyn WsStream>,
159}
160
161pub type WsConnectFuture =
162 Pin<Box<dyn Future<Output = Result<WsConn, NetworkError>> + Send + 'static>>;
163
164pub trait WsTransport: Send + Sync + 'static {
165 fn connect(&self, url: Url) -> WsConnectFuture;
166}
167
168pub type AddrGuard = Arc<dyn Fn(&[SocketAddr]) -> Result<(), NetworkError> + Send + Sync>;
169
170#[derive(Clone, Copy, Debug, Default)]
171pub struct TungsteniteWs;
172
173impl TungsteniteWs {
174 pub fn shared() -> Arc<dyn WsTransport> {
175 Arc::new(Self)
176 }
177}
178
179impl WsTransport for TungsteniteWs {
180 fn connect(&self, url: Url) -> WsConnectFuture {
181 Box::pin(async move {
182 let url_str = url.as_str().to_owned();
183 let (ws, _resp) = tokio_tungstenite::connect_async(&url_str)
184 .await
185 .map_err(|e| NetworkError::Connect(e.to_string()))?;
186 let (sink_inner, stream_inner) = futures::StreamExt::split(ws);
187 let sink: Box<dyn WsSink> = Box::new(TungsteniteSink { inner: sink_inner });
188 let stream: Box<dyn WsStream> = Box::new(TungsteniteStream {
189 inner: stream_inner,
190 });
191 Ok(WsConn { sink, stream })
192 })
193 }
194}
195
196pub struct GuardedWs {
197 guard: AddrGuard,
198}
199
200impl GuardedWs {
201 pub fn shared(guard: AddrGuard) -> Arc<dyn WsTransport> {
202 Arc::new(Self { guard })
203 }
204}
205
206impl WsTransport for GuardedWs {
207 fn connect(&self, url: Url) -> WsConnectFuture {
208 let guard = self.guard.clone();
209 Box::pin(async move {
210 let host = url
211 .host_str()
212 .ok_or_else(|| NetworkError::Connect("ws url missing host".to_owned()))?
213 .to_owned();
214 let port = url
215 .port_or_known_default()
216 .ok_or_else(|| NetworkError::Connect("ws url missing port".to_owned()))?;
217 let addrs: Vec<SocketAddr> = tokio::net::lookup_host((host.as_str(), port))
218 .await
219 .map_err(|e| NetworkError::Connect(e.to_string()))?
220 .collect();
221 guard(&addrs)?;
222 let addr = addrs
223 .into_iter()
224 .next()
225 .ok_or_else(|| NetworkError::Connect(format!("no addresses for {host}")))?;
226 let tcp = tokio::net::TcpStream::connect(addr)
227 .await
228 .map_err(|e| NetworkError::Connect(e.to_string()))?;
229 let (ws, _resp) = tokio_tungstenite::client_async_tls(url.as_str(), tcp)
230 .await
231 .map_err(|e| NetworkError::Connect(e.to_string()))?;
232 let (sink_inner, stream_inner) = futures::StreamExt::split(ws);
233 let sink: Box<dyn WsSink> = Box::new(TungsteniteSink { inner: sink_inner });
234 let stream: Box<dyn WsStream> = Box::new(TungsteniteStream {
235 inner: stream_inner,
236 });
237 Ok(WsConn { sink, stream })
238 })
239 }
240}
241
242type TungsteniteWsStream =
243 tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
244
245struct TungsteniteSink {
246 inner: futures::stream::SplitSink<TungsteniteWsStream, TungsteniteMessage>,
247}
248
249impl WsSink for TungsteniteSink {
250 fn send<'a>(&'a mut self, message: WsMessage) -> WsSendFuture<'a> {
251 Box::pin(async move {
252 use futures::SinkExt;
253 self.inner
254 .send(message_to_tungstenite(message))
255 .await
256 .map_err(|e| NetworkError::Transport(e.to_string()))
257 })
258 }
259}
260
261struct TungsteniteStream {
262 inner: futures::stream::SplitStream<TungsteniteWsStream>,
263}
264
265impl WsStream for TungsteniteStream {
266 fn next<'a>(&'a mut self) -> WsMessageFuture<'a> {
267 Box::pin(async move {
268 let item = StreamExt::next(&mut self.inner).await?;
269 Some(
270 item.map_err(|e| NetworkError::Transport(e.to_string()))
271 .and_then(message_from_tungstenite),
272 )
273 })
274 }
275}
276
277fn message_to_tungstenite(message: WsMessage) -> TungsteniteMessage {
278 match message {
279 WsMessage::Text(text) => TungsteniteMessage::Text(text.into()),
280 WsMessage::Binary(bytes) => TungsteniteMessage::Binary(WsBytes::copy_from_slice(&bytes)),
281 WsMessage::Ping(bytes) => TungsteniteMessage::Ping(WsBytes::copy_from_slice(&bytes)),
282 WsMessage::Pong(bytes) => TungsteniteMessage::Pong(WsBytes::copy_from_slice(&bytes)),
283 WsMessage::Close { code, reason } => TungsteniteMessage::Close(Some(TungsteniteClose {
284 code: TungsteniteCloseCode::from(code),
285 reason: reason.into(),
286 })),
287 }
288}
289
290fn message_from_tungstenite(message: TungsteniteMessage) -> Result<WsMessage, NetworkError> {
291 match message {
292 TungsteniteMessage::Text(t) => Ok(WsMessage::Text(t.to_string())),
293 TungsteniteMessage::Binary(b) => Ok(WsMessage::Binary(Bytes::copy_from_slice(&b))),
294 TungsteniteMessage::Ping(b) => Ok(WsMessage::Ping(Bytes::copy_from_slice(&b))),
295 TungsteniteMessage::Pong(b) => Ok(WsMessage::Pong(Bytes::copy_from_slice(&b))),
296 TungsteniteMessage::Close(close) => {
297 let (code, reason) = close
298 .map(|c| (u16::from(c.code), c.reason.to_string()))
299 .unwrap_or((1000, String::new()));
300 Ok(WsMessage::Close { code, reason })
301 }
302 TungsteniteMessage::Frame(_) => Err(NetworkError::Protocol(
303 "tungstenite raw frame surfaced unexpectedly".to_owned(),
304 )),
305 }
306}