#![allow(async_fn_in_trait)]
use crate::client::VeyronClient;
use std::env;
use veyron_wire::proto::veyron::{envelope, Envelope, Event, PluginManifest, Pong};
use veyron_wire::WireError as VeyronError;
fn unix_millis() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
pub trait Plugin {
fn id(&self) -> &str;
fn version(&self) -> &str {
"1.0.0"
}
fn manifest(&self) -> PluginManifest;
async fn on_init(&mut self, client: &mut VeyronClient) -> Result<(), VeyronError> {
let _ = client;
Ok(())
}
async fn on_message(&mut self, envelope: Envelope) -> Result<Option<Envelope>, VeyronError>;
async fn on_event(&mut self, event: Event) -> Result<Option<Envelope>, VeyronError> {
let _ = event;
Ok(None)
}
async fn on_shutdown(&mut self) -> Result<(), VeyronError> {
Ok(())
}
async fn run(&mut self) -> Result<(), VeyronError> {
let socket_path = env::var("VEYRON_SOCKET_PATH")
.unwrap_or_else(|_| veyron_wire::socket::default_socket_path());
self.run_with(&socket_path).await
}
async fn run_with(&mut self, socket_path: &str) -> Result<(), VeyronError> {
let token = env::var("VEYRON_JWT_TOKEN").unwrap_or_default();
let secret = env::var("VEYRON_JWT_SECRET").ok().filter(|s| !s.is_empty());
let client = match secret {
Some(s) => VeyronClient::connect_with_secret(socket_path, s.as_bytes()).await?,
None => VeyronClient::connect(socket_path).await?,
};
self.serve(client, &token).await
}
async fn serve(
&mut self,
mut client: VeyronClient,
jwt_token: &str,
) -> Result<(), VeyronError> {
let ack = client
.register_full(self.id(), self.version(), self.manifest(), jwt_token)
.await?;
if !ack.accepted {
return Err(VeyronError::PermissionDenied(format!(
"registration rejected: {}",
ack.reject_reason
)));
}
if let Err(e) = self.on_init(&mut client).await {
let _ = self.on_shutdown().await;
return Err(e);
}
let mut handler_err: Option<VeyronError> = None;
loop {
let env = match client.recv().await {
Ok(env) => env,
Err(_) => break, };
match env.payload {
Some(envelope::Payload::Ping(ping)) => {
let pong = Envelope {
payload: Some(envelope::Payload::Pong(Pong {
original_timestamp: ping.timestamp,
server_timestamp: unix_millis(),
})),
..Default::default()
};
let _ = client.send("kernel", pong).await;
}
Some(envelope::Payload::PluginShutdown(_)) => break,
Some(envelope::Payload::Event(event)) => {
let event_id = event.event_id.clone();
if let Ok(reply) = self.on_event(event).await {
let _ = client.ack_event(&event_id).await;
if let Some(resp) = reply {
let _ = client.send("kernel", resp).await;
}
}
}
_ => match self.on_message(env).await {
Ok(Some(resp)) => {
let _ = client.send("kernel", resp).await;
}
Ok(None) => {}
Err(e) => {
handler_err = Some(e);
break;
}
},
}
}
self.on_shutdown().await?;
match handler_err {
Some(e) => Err(e),
None => Ok(()),
}
}
}