1use std::cmp;
30use std::convert::TryFrom;
31use std::ops::Deref;
32use std::sync::Arc;
33
34use bitcoin::Network;
35use log::warn;
36use tokio::sync::RwLock;
37use tonic::service::interceptor::InterceptedService;
38
39use ark::ArkInfo;
40
41use crate::{mailbox, protos, ArkServiceClient, ConvertError, RequestExt};
42
43#[cfg(all(feature = "tonic-native", feature = "tonic-web"))]
44compile_error!("features `tonic-native` and `tonic-web` are mutually exclusive");
45
46#[cfg(not(any(feature = "tonic-native", feature = "tonic-web")))]
47compile_error!("either `tonic-native` or `tonic-web` feature must be enabled");
48
49#[cfg(all(feature = "tonic-web", feature = "socks5-proxy"))]
50compile_error!("`tonic-web` does not support the `socks5-proxy` feature");
51
52#[cfg(feature = "tonic-native")]
53mod transport {
54 use std::str::FromStr;
55 use std::time::Duration;
56
57 use http::Uri;
58 use log::info;
59 use tonic::transport::{Channel, ClientTlsConfig, Endpoint};
60
61 use super::CreateEndpointError;
62
63 pub type Transport = Channel;
64
65 pub async fn connect(address: &str) -> Result<Transport, CreateEndpointError> {
71 Ok(create_endpoint(address)?.connect().await?)
72 }
73
74 #[cfg(feature = "socks5-proxy")]
76 pub async fn connect_with_proxy(
77 address: &str,
78 proxy: &str,
79 ) -> Result<Transport, CreateEndpointError> {
80 use hyper_socks2::SocksConnector;
81 use hyper_util::client::legacy::connect::HttpConnector;
82
83 let endpoint = create_endpoint(address)?;
84 let proxy_uri = proxy.parse::<Uri>().map_err(CreateEndpointError::InvalidProxyUri)?;
85 let connector = {
86 let mut http = HttpConnector::new();
89 http.enforce_http(false);
90 SocksConnector {
91 proxy_addr: proxy_uri,
92 auth: None,
93 connector: http,
94 }
95 };
96 info!("Connecting to Ark server via SOCKS5 proxy {}...", proxy);
97 Ok(endpoint.connect_with_connector(connector).await?)
98 }
99
100 fn create_endpoint(address: &str) -> Result<Endpoint, CreateEndpointError> {
104 let uri = Uri::from_str(address)?;
105
106 let scheme = uri.scheme_str().unwrap_or("");
107 if scheme != "http" && scheme != "https" {
108 return Err(CreateEndpointError::InvalidScheme(scheme.to_string()));
109 }
110
111 let mut endpoint = Channel::builder(uri.clone())
112 .keep_alive_timeout(Duration::from_secs(600))
113 .timeout(Duration::from_secs(600));
114
115 if scheme == "https" {
116 info!("Connecting to Ark server at {} using TLS...", address);
117 let uri_auth = uri.clone().into_parts().authority
118 .ok_or(CreateEndpointError::MissingAuthority)?;
119 let domain = uri_auth.host();
120
121 let tls_config = ClientTlsConfig::new()
122 .with_enabled_roots()
123 .domain_name(domain);
124 endpoint = endpoint.tls_config(tls_config)
125 .map_err(CreateEndpointError::Transport)?;
126 } else {
127 info!("Connecting to Ark server at {} without TLS...", address);
128 }
129 Ok(endpoint)
130 }
131}
132
133#[cfg(feature = "tonic-web")]
134mod transport {
135 use super::CreateEndpointError;
136 use tonic_web_wasm_client::Client as WasmClient;
137
138 pub type Transport = WasmClient;
139
140 pub async fn connect(address: &str) -> Result<Transport, CreateEndpointError> {
141 Ok(tonic_web_wasm_client::Client::new(address.to_string()))
142 }
143}
144
145pub const MIN_PROTOCOL_VERSION: u64 = 1;
149
150pub const MAX_PROTOCOL_VERSION: u64 = 1;
154
155pub const ARK_INFO_TTL_SECS: u64 = 10 * 60;
159
160#[derive(Debug, thiserror::Error)]
161#[error("failed to create gRPC endpoint: {msg}")]
162pub enum CreateEndpointError {
163 #[error("failed to parse Ark server as a URI")]
164 InvalidUri(#[from] http::uri::InvalidUri),
165 #[error("Ark server scheme must be either http or https. Found: {0}")]
166 InvalidScheme(String),
167 #[error("Ark server URI is missing an authority part")]
168 MissingAuthority,
169 #[cfg(feature = "tonic-native")]
170 #[error(transparent)]
171 Transport(#[from] tonic::transport::Error),
172 #[cfg(feature = "socks5-proxy")]
173 #[error("invalid SOCKS5 proxy URI: {0:#}")]
174 InvalidProxyUri(http::uri::InvalidUri),
175}
176
177#[derive(Debug, thiserror::Error)]
178#[error("failed to connect to Ark server: {msg}")]
179pub enum ConnectError {
180 #[error(transparent)]
181 CreateEndpoint(#[from] CreateEndpointError),
182 #[error("handshake request failed: {0}")]
183 Handshake(tonic::Status),
184 #[error("version mismatch. Client max is: {client_max}, server min is: {server_min}")]
185 ProtocolVersionMismatchClientTooOld { client_max: u64, server_min: u64 },
186 #[error("version mismatch. Client min is: {client_min}, server max is: {server_max}")]
187 ProtocolVersionMismatchServerTooOld { client_min: u64, server_max: u64 },
188 #[error("error getting ark info: {0}")]
189 GetArkInfo(tonic::Status),
190 #[error("invalid ark info from ark server: {0}")]
191 InvalidArkInfo(#[from] ConvertError),
192 #[error("network mismatch. Expected: {expected}, Got: {got}")]
193 NetworkMismatch { expected: Network, got: Network },
194 #[error("tokio channel error: {0}")]
195 Tokio(#[from] tokio::sync::oneshot::error::RecvError),
196}
197
198#[derive(Clone)]
204pub struct ProtocolVersionInterceptor {
205 pver: u64,
206}
207
208impl tonic::service::Interceptor for ProtocolVersionInterceptor {
209 fn call(&mut self, mut req: tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> {
210 req.set_pver(self.pver);
211 Ok(req)
212 }
213}
214
215pub struct ArkInfoHandle {
219 pub info: ArkInfo,
220 pub waiter: Option<tokio::sync::oneshot::Receiver<Result<ArkInfo, ConnectError>>>,
221}
222
223impl Deref for ArkInfoHandle {
224 type Target = ArkInfo;
225
226 fn deref(&self) -> &Self::Target {
227 &self.info
228 }
229}
230
231pub struct ServerInfo {
232 pub pver: u64,
236 pub info: ArkInfo,
238 pub refresh_at_secs: u64,
240}
241
242impl ServerInfo {
243 fn ttl() -> u64 {
245 ark::time::timestamp_secs() + ARK_INFO_TTL_SECS
246 }
247
248 pub fn new(pver: u64, info: ArkInfo) -> Self {
249 Self { pver, info, refresh_at_secs: Self::ttl() }
250 }
251
252 pub fn update(&mut self, info: ArkInfo) {
253 self.info = info;
254 self.refresh_at_secs = Self::ttl();
255 }
256
257 pub fn is_outdated(&self) -> bool {
259 ark::time::timestamp_secs() > self.refresh_at_secs
260 }
261}
262
263#[derive(Clone)]
270pub struct ServerConnection {
271 info: Arc<RwLock<ServerInfo>>,
272 pub client: ArkServiceClient<InterceptedService<transport::Transport, ProtocolVersionInterceptor>>,
274 pub mailbox_client: mailbox::MailboxServiceClient<InterceptedService<transport::Transport, ProtocolVersionInterceptor>>,
276}
277
278impl ServerConnection {
279 fn handshake_req() -> protos::HandshakeRequest {
280 protos::HandshakeRequest {
281 bark_version: Some(env!("CARGO_PKG_VERSION").into()),
282 }
283 }
284
285 pub async fn connect(
303 address: &str,
304 network: Network,
305 ) -> Result<ServerConnection, ConnectError> {
306 let transport = transport::connect(address).await?;
307 Self::connect_inner(transport, network).await
308 }
309
310 #[cfg(feature = "socks5-proxy")]
312 pub async fn connect_via_proxy(
313 address: &str,
314 network: Network,
315 proxy: &str,
316 ) -> Result<ServerConnection, ConnectError> {
317 let transport = transport::connect_with_proxy(address, proxy).await?;
318 Self::connect_inner(transport, network).await
319 }
320
321 async fn connect_inner(
322 transport: transport::Transport,
323 network: Network,
324 ) -> Result<ServerConnection, ConnectError> {
325 let mut handshake_client = ArkServiceClient::new(transport.clone());
326 let handshake = handshake_client.handshake(Self::handshake_req()).await
327 .map_err(ConnectError::Handshake)?.into_inner();
328
329 let pver = check_handshake(handshake)?;
330
331 let interceptor = ProtocolVersionInterceptor { pver };
332 let mut client = ArkServiceClient::with_interceptor(transport.clone(), interceptor.clone())
333 .max_decoding_message_size(64 * 1024 * 1024); let info = client.ark_info(network).await?;
336
337 let mailbox_client = mailbox::MailboxServiceClient::with_interceptor(transport, interceptor)
338 .max_decoding_message_size(64 * 1024 * 1024); let info = Arc::new(RwLock::new(ServerInfo::new(pver, info)));
341 Ok(ServerConnection {
342 info,
343 client,
344 mailbox_client,
345 })
346 }
347
348 pub async fn check_connection(&self) -> Result<(), ConnectError> {
350 let mut client = self.client.clone();
351 let handshake = client.handshake(Self::handshake_req()).await
352 .map_err(ConnectError::Handshake)?.into_inner();
353 check_handshake(handshake)?;
354 Ok(())
355 }
356
357 pub async fn ark_info(&self) -> Result<ArkInfo, ConnectError> {
365 let mut current = self.info.write().await;
366
367 let new_info = self.client.clone().ark_info(current.info.network).await?;
368 if current.is_outdated() {
369 current.update(new_info.clone());
370 return Ok(new_info);
371 }
372
373 Ok(current.info.clone())
374 }
375}
376trait ArkServiceClientExt {
377 async fn ark_info(&mut self, network: Network) -> Result<ArkInfo, ConnectError>;
378}
379
380impl ArkServiceClientExt for ArkServiceClient<InterceptedService<transport::Transport, ProtocolVersionInterceptor>> {
381 async fn ark_info(&mut self, network: Network) -> Result<ArkInfo, ConnectError> {
382 let res = self.get_ark_info(protos::Empty {}).await
383 .map_err(ConnectError::GetArkInfo)?;
384 let info = ArkInfo::try_from(res.into_inner())
385 .map_err(ConnectError::InvalidArkInfo)?;
386 if network != info.network {
387 return Err(ConnectError::NetworkMismatch { expected: network, got: info.network });
388 }
389
390 Ok(info)
391 }
392}
393
394fn check_handshake(handshake: protos::HandshakeResponse) -> Result<u64, ConnectError> {
395 if let Some(ref msg) = handshake.psa {
396 warn!("Message from Ark server: \"{}\"", msg);
397 }
398
399 if MAX_PROTOCOL_VERSION < handshake.min_protocol_version {
400 return Err(ConnectError::ProtocolVersionMismatchClientTooOld {
401 client_max: MAX_PROTOCOL_VERSION, server_min: handshake.min_protocol_version
402 });
403 }
404 if MIN_PROTOCOL_VERSION > handshake.max_protocol_version {
405 return Err(ConnectError::ProtocolVersionMismatchServerTooOld {
406 client_min: MIN_PROTOCOL_VERSION, server_max: handshake.max_protocol_version
407 });
408 }
409
410 let pver = cmp::min(MAX_PROTOCOL_VERSION, handshake.max_protocol_version);
411 assert!((MIN_PROTOCOL_VERSION..=MAX_PROTOCOL_VERSION).contains(&pver));
412 assert!((handshake.min_protocol_version..=handshake.max_protocol_version).contains(&pver));
413
414 Ok(pver)
415}