250 lines
8.6 KiB
Rust
250 lines
8.6 KiB
Rust
//! TCP endpoint publication, verified connections, and per-request notifications.
|
|
use crate::*;
|
|
type Notify = Box<dyn FnMut(&str, Value)>;
|
|
use serde_json::Value;
|
|
#[cfg(target_os = "wasi")]
|
|
use serde_json::json;
|
|
use std::{future::IntoFuture, marker::PhantomData, path::Path, rc::Rc};
|
|
|
|
/// A verified endpoint address and identity.
|
|
///
|
|
/// Connect verifies and closes one connection. Each call opens a new connection
|
|
/// and repeats the identity handshake before its request. The first call costs three round trips.
|
|
#[derive(Clone)]
|
|
#[cfg_attr(not(target_os = "wasi"), allow(dead_code))]
|
|
pub struct Endpoint {
|
|
address: String,
|
|
owner: String,
|
|
interface: String,
|
|
}
|
|
impl Endpoint {
|
|
/// Require the interface expected by a generated client.
|
|
pub fn require_interface(&self, name: &str) -> Result<()> {
|
|
if self.interface == name {
|
|
Ok(())
|
|
} else {
|
|
Err(Error::new(-32000, "endpoint interface mismatch"))
|
|
}
|
|
}
|
|
/// Build a typed endpoint request.
|
|
pub fn call<T>(&self, method: &str, params: impl serde::Serialize) -> EndpointCall<T> {
|
|
EndpointCall {
|
|
endpoint: self.clone(),
|
|
method: method.into(),
|
|
params: serde_json::to_value(params).expect("serializable params"),
|
|
notify: None,
|
|
result: PhantomData,
|
|
}
|
|
}
|
|
}
|
|
/// A lazy endpoint request with optional progress callbacks.
|
|
#[must_use = "await the endpoint call to send it"]
|
|
pub struct EndpointCall<T> {
|
|
endpoint: Endpoint,
|
|
method: String,
|
|
params: Value,
|
|
notify: Option<Notify>,
|
|
result: PhantomData<T>,
|
|
}
|
|
impl<T> EndpointCall<T> {
|
|
/// Receive progress from this endpoint handler.
|
|
pub fn on_notify(mut self, f: impl FnMut(&str, Value) + 'static) -> Self {
|
|
self.notify = Some(Box::new(f));
|
|
self
|
|
}
|
|
}
|
|
impl<T: serde::de::DeserializeOwned + 'static> IntoFuture for EndpointCall<T> {
|
|
type Output = Result<T>;
|
|
type IntoFuture = futures::future::LocalBoxFuture<'static, Result<T>>;
|
|
fn into_future(self) -> Self::IntoFuture {
|
|
Box::pin(async move {
|
|
Ok(serde_json::from_value(
|
|
endpoint_request(self.endpoint, &self.method, self.params, self.notify).await?,
|
|
)?)
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Publish a declared interface and keep its listener alive until stdin EOF or [`drain`].
|
|
#[cfg(target_os = "wasi")]
|
|
pub async fn listen<T: Interface>(path: impl AsRef<Path>, value: Rc<T>) -> Result<()> {
|
|
use wstd::iter::AsyncIterator;
|
|
let name = crate::transport::endpoint_name::<T>()?;
|
|
let listener = wstd::net::TcpListener::bind("127.0.0.1:0").await?;
|
|
write_atomic(
|
|
path,
|
|
format!("tcp:{} {} {}\n", listener.local_addr()?, pid(), name),
|
|
)?;
|
|
crate::transport::endpoint_task(async move {
|
|
let mut incoming = listener.incoming();
|
|
while let Some(Ok(stream)) = incoming.next().await {
|
|
let value = value.clone();
|
|
let name = name.clone();
|
|
background(async move {
|
|
let _ = serve_connection(stream, value, name).await;
|
|
});
|
|
}
|
|
});
|
|
Ok(())
|
|
}
|
|
/// Publish an endpoint. Native programs require the WASI runtime.
|
|
#[cfg(not(target_os = "wasi"))]
|
|
pub async fn listen<T: Interface>(_path: impl AsRef<Path>, _value: Rc<T>) -> Result<()> {
|
|
crate::transport::endpoint_name::<T>()?;
|
|
Err(Error::new(-32000, "endpoints require WASI"))
|
|
}
|
|
|
|
/// Read an endpoint file and verify its owner before returning a client.
|
|
#[cfg(target_os = "wasi")]
|
|
pub async fn connect(path: impl AsRef<Path>) -> Result<Endpoint> {
|
|
let mut attempts = 0;
|
|
let contents = loop {
|
|
match std::fs::read_to_string(path.as_ref()) {
|
|
Ok(s) => break s,
|
|
Err(e) if e.kind() == std::io::ErrorKind::NotFound && attempts < 20 => {
|
|
attempts += 1;
|
|
wstd::task::sleep(wstd::time::Duration::from_millis(25)).await;
|
|
}
|
|
Err(e) => return Err(e.into()),
|
|
}
|
|
};
|
|
let words: Vec<_> = contents.split_whitespace().collect();
|
|
if words.len() != 3 {
|
|
return Err(Error::new(-32602, "invalid endpoint file"));
|
|
}
|
|
let endpoint = Endpoint {
|
|
address: words[0]
|
|
.strip_prefix("tcp:")
|
|
.ok_or_else(|| Error::new(-32602, "invalid endpoint address"))?
|
|
.into(),
|
|
owner: words[1].into(),
|
|
interface: words[2].into(),
|
|
};
|
|
let _ = handshake(&endpoint).await?;
|
|
Ok(endpoint)
|
|
}
|
|
/// Connect to an endpoint. Native programs require the WASI runtime.
|
|
#[cfg(not(target_os = "wasi"))]
|
|
pub async fn connect(_path: impl AsRef<Path>) -> Result<Endpoint> {
|
|
Err(Error::new(-32000, "endpoints require WASI"))
|
|
}
|
|
|
|
#[cfg(target_os = "wasi")]
|
|
async fn send(stream: &mut impl wstd::io::AsyncWrite, value: Value) -> Result<()> {
|
|
stream.write_all(format!("{value}\n").as_bytes()).await?;
|
|
stream.flush().await?;
|
|
Ok(())
|
|
}
|
|
#[cfg(target_os = "wasi")]
|
|
async fn handshake(endpoint: &Endpoint) -> Result<wstd::net::TcpStream> {
|
|
let mut stream = wstd::net::TcpStream::connect(endpoint.address.as_str()).await?;
|
|
send(
|
|
&mut stream,
|
|
json!({"jsonrpc":"2.0","id":0,"method":"rpc.endpoint","params":{}}),
|
|
)
|
|
.await?;
|
|
let reply = crate::lines::Lines::default()
|
|
.rpc(&mut stream)
|
|
.await?
|
|
.ok_or_else(|| Error::new(-32000, "endpoint closed"))?;
|
|
let result = &reply["result"];
|
|
let plugin = endpoint
|
|
.owner
|
|
.len()
|
|
.checked_sub(37)
|
|
.and_then(|end| endpoint.owner.get(..end));
|
|
if result["owner"] != endpoint.owner
|
|
|| result["interface"] != endpoint.interface
|
|
|| result["plugin"].as_str() != plugin
|
|
{
|
|
return Err(Error::new(-32000, "endpoint identity mismatch"));
|
|
}
|
|
Ok(stream)
|
|
}
|
|
#[cfg(target_os = "wasi")]
|
|
async fn endpoint_request(
|
|
endpoint: Endpoint,
|
|
method: &str,
|
|
params: Value,
|
|
mut notify: Option<Notify>,
|
|
) -> Result<Value> {
|
|
let mut stream = handshake(&endpoint).await?;
|
|
let mut lines = crate::lines::Lines::default();
|
|
send(
|
|
&mut stream,
|
|
json!({"jsonrpc":"2.0","id":1,"method":method,"params":params}),
|
|
)
|
|
.await?;
|
|
while let Some(value) = lines.rpc(&mut stream).await? {
|
|
if value["method"] == "core.notify" {
|
|
if let Some(f) = &mut notify {
|
|
f(
|
|
value["params"]["name"].as_str().unwrap_or_default(),
|
|
value["params"]["params"].clone(),
|
|
);
|
|
}
|
|
} else if value["id"] == 1 {
|
|
if let Some(error) = value.get("error") {
|
|
return Err(serde_json::from_value(error.clone())?);
|
|
}
|
|
return Ok(value["result"].clone());
|
|
}
|
|
}
|
|
Err(Error::new(-32000, "endpoint closed before response"))
|
|
}
|
|
#[cfg(not(target_os = "wasi"))]
|
|
async fn endpoint_request(
|
|
_endpoint: Endpoint,
|
|
_method: &str,
|
|
_params: Value,
|
|
_notify: Option<Notify>,
|
|
) -> Result<Value> {
|
|
Err(Error::new(-32000, "endpoints require WASI"))
|
|
}
|
|
|
|
#[cfg(target_os = "wasi")]
|
|
async fn serve_connection<T: Interface>(
|
|
stream: wstd::net::TcpStream,
|
|
value: Rc<T>,
|
|
name: String,
|
|
) -> Result<()> {
|
|
use futures::{FutureExt, StreamExt};
|
|
let (mut reader, mut writer) = stream.split();
|
|
let mut lines = crate::lines::Lines::default();
|
|
while let Some(request) = crate::transport::until_drained(lines.rpc(&mut reader))
|
|
.await
|
|
.transpose()?
|
|
.flatten()
|
|
{
|
|
let method = request["method"].as_str().unwrap_or_default();
|
|
let (tx, mut rx) = futures::channel::mpsc::unbounded();
|
|
let ctx = Ctx {
|
|
sink: Some(tx),
|
|
..Ctx::default()
|
|
};
|
|
let handler=async {
|
|
if method=="rpc.endpoint" {Ok(json!({"owner":pid(),"plugin":std::env::var("SAUSAGE_PLUGIN").unwrap_or_default(),"interface":name}))}
|
|
else {value.dispatch(method,request.get("params").cloned().unwrap_or(json!({})),&ctx).await}
|
|
}.fuse();
|
|
futures::pin_mut!(handler);
|
|
let result = loop {
|
|
futures::select! {result=handler=>break result,notification=rx.next().fuse()=>{if let Some(notification)=notification {send(&mut writer,notification).await?;}}}
|
|
};
|
|
while let Ok(notification) = rx.try_recv() {
|
|
send(&mut writer, notification).await?;
|
|
}
|
|
if let Some(id) = request.get("id") {
|
|
let response = match result {
|
|
Ok(result) => json!({"jsonrpc":"2.0","id":id,"result":result}),
|
|
Err(error) => json!({"jsonrpc":"2.0","id":id,"error":error}),
|
|
};
|
|
send(&mut writer, response).await?;
|
|
}
|
|
if ctx.exit.get() {
|
|
crate::transport::exit();
|
|
break;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|