use std::path::Path;
use std::sync::Arc;
use microsandbox_protocol::control::DEFAULT_REQUEST_TIMEOUT;
use microsandbox_protocol_client::{
ClientError, ConnectOptions, Connector, Delivery, ErrorKind, LocalConnector, Message, Protocol,
RequestOptions,
};
use tokio::time::timeout_at;
use tokio_util::sync::CancellationToken;
use crate::{
CheckedControlRequest, ControlClient, ControlClientError, ControlClientResult, ControlProtocol,
IntoControlMessage, JsonControlClient, JsonReply, VerifiedControlConnector,
dialer::Dialer,
json_client::{check_deadline, deadline},
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ControlMode {
Framed,
Json,
}
#[derive(Debug)]
pub enum ControlReply {
Framed(Message),
Json(JsonReply),
}
#[derive(Clone)]
pub struct ControlConnection {
inner: Arc<Inner>,
}
struct Inner {
selected: Selected,
dialer: Dialer,
options: ConnectOptions,
closed: CancellationToken,
capabilities: crate::Capabilities,
}
enum Selected {
Framed(ControlClient),
Json(JsonControlClient),
}
impl ControlConnection {
pub async fn connect(path: impl AsRef<Path>) -> ControlClientResult<Self> {
Self::connect_with(path, |options| options).await
}
pub async fn connect_with(
path: impl AsRef<Path>,
configure: impl FnOnce(ConnectOptions) -> ConnectOptions,
) -> ControlClientResult<Self> {
Self::connect_connector_with(Arc::new(LocalConnector::new(path)), configure).await
}
pub async fn connect_connector(connector: Arc<dyn Connector>) -> ControlClientResult<Self> {
Self::connect_connector_with(connector, |options| options).await
}
pub async fn connect_connector_with(
connector: Arc<dyn Connector>,
configure: impl FnOnce(ConnectOptions) -> ConnectOptions,
) -> ControlClientResult<Self> {
Self::establish(
Dialer::Unverified(connector),
configure(ConnectOptions::default()),
)
.await
}
pub async fn connect_verified_connector(
connector: Arc<dyn VerifiedControlConnector>,
) -> ControlClientResult<Self> {
Self::connect_verified_connector_with(connector, |options| options).await
}
pub async fn connect_verified_connector_with(
connector: Arc<dyn VerifiedControlConnector>,
configure: impl FnOnce(ConnectOptions) -> ConnectOptions,
) -> ControlClientResult<Self> {
Self::establish(
Dialer::Verified(connector),
configure(ConnectOptions::default()),
)
.await
}
async fn establish(dialer: Dialer, options: ConnectOptions) -> ControlClientResult<Self> {
options.limits.validate()?;
let until = deadline(options.setup_timeout)?;
timeout_at(until, async {
let json = JsonControlClient::configured(dialer.clone(), options.clone(), true);
let (mode, capabilities) = json.discover(until).await?;
let selected = match mode {
ControlMode::Json => Selected::Json(json),
ControlMode::Framed => {
let transport = dialer.connect(until).await?;
dialer.verify(until).await?;
check_deadline(until)?;
let established =
ControlProtocol::establish(transport, options.clone()).await?;
Selected::Framed(ControlClient::from_established(established).await?)
}
};
Ok(Self {
inner: Arc::new(Inner {
selected,
dialer,
options,
closed: CancellationToken::new(),
capabilities,
}),
})
})
.await
.unwrap_or_else(|_| Err(ClientError::new(ErrorKind::Timeout).into()))
.map_err(not_sent)
}
pub fn mode(&self) -> ControlMode {
match self.inner.selected {
Selected::Framed(_) => ControlMode::Framed,
Selected::Json(_) => ControlMode::Json,
}
}
pub fn capabilities(&self) -> &crate::Capabilities {
&self.inner.capabilities
}
pub fn is_closed(&self) -> bool {
match &self.inner.selected {
Selected::Framed(client) => client.is_closed(),
Selected::Json(client) => client.is_closed(),
}
}
pub async fn closed(&self) {
match &self.inner.selected {
Selected::Framed(client) => client.closed().await,
Selected::Json(client) => client.closed().await,
}
}
pub fn framed(&self) -> ControlClientResult<&ControlClient> {
match &self.inner.selected {
Selected::Framed(client) => Ok(client),
Selected::Json(_) => Err(ControlClientError::UnsupportedMode),
}
}
pub async fn close(&self) {
self.inner.closed.cancel();
match &self.inner.selected {
Selected::Framed(client) => client.close().await,
Selected::Json(client) => client.close().await,
}
}
pub async fn request(
&self,
message: impl IntoControlMessage,
) -> ControlClientResult<ControlReply> {
self.request_with(message, |options| options).await
}
pub async fn request_with(
&self,
message: impl IntoControlMessage,
configure: impl FnOnce(RequestOptions) -> RequestOptions,
) -> ControlClientResult<ControlReply> {
let options = configure(RequestOptions::default());
match &self.inner.selected {
Selected::Json(client) => Ok(ControlReply::Json(
client.operation(message.into_json()?, options).await?,
)),
Selected::Framed(client) => {
let options = self.verify_before_request(options).await?;
Ok(ControlReply::Framed(
client.request_with(message, |_| options).await?,
))
}
}
}
pub async fn request_typed<R: CheckedControlRequest>(
&self,
request: &R,
) -> ControlClientResult<R::Response> {
self.request_typed_with(request, |options| options).await
}
pub async fn request_typed_with<R: CheckedControlRequest>(
&self,
request: &R,
configure: impl FnOnce(RequestOptions) -> RequestOptions,
) -> ControlClientResult<R::Response> {
let options = configure(RequestOptions::default());
match &self.inner.selected {
Selected::Json(client) => {
let reply = client.operation(request.json_request()?, options).await?;
request.decode_json(reply)
}
Selected::Framed(client) => {
let options = self.verify_before_request(options).await?;
client.request_typed_with(request, |_| options).await
}
}
}
async fn verify_before_request(
&self,
options: RequestOptions,
) -> ControlClientResult<RequestOptions> {
if self.is_closed() {
return Err(ClientError::new(ErrorKind::Closed).into());
}
let until = deadline(
options
.request_timeout
.or(self.inner.options.limits.request_timeout)
.unwrap_or(DEFAULT_REQUEST_TIMEOUT),
)?;
check_deadline(until)?;
let result = tokio::select! {
biased;
_ = self.inner.closed.cancelled() => Err(ClientError::new(ErrorKind::Closed).into()),
result = timeout_at(until, self.inner.dialer.verify(until)) => result
.unwrap_or_else(|_| Err(ClientError::new(ErrorKind::Timeout).into())),
};
if let Err(error) = result {
self.close().await;
return Err(not_sent(error));
}
check_deadline(until)?;
Ok(RequestOptions::default()
.request_timeout(until.saturating_duration_since(tokio::time::Instant::now())))
}
}
pub(crate) fn not_sent(error: ControlClientError) -> ControlClientError {
match error {
ControlClientError::Client(error) => error.with_delivery(Delivery::NotSent).into(),
error => error,
}
}