use extism_pdk::*;
use serde::{Deserialize, Serialize};
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as BASE64;
mod base64_bytes {
use serde::{self, Deserialize, Deserializer, Serializer};
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as BASE64;
pub fn serialize<S>(bytes: &Vec<u8>, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&BASE64.encode(bytes))
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
D: Deserializer<'de>,
{
let s = String::deserialize(deserializer)?;
BASE64.decode(&s).map_err(serde::de::Error::custom)
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct WebSocketConnectRequest {
url: String,
headers: std::collections::HashMap<String, String>,
connection_id: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct WebSocketConnectResponse {
#[serde(default)]
new_connection_id: String,
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct WebSocketSendTextRequest {
connection_id: String,
message: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct WebSocketSendTextResponse {
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct WebSocketSendBinaryRequest {
connection_id: String,
#[serde(with = "base64_bytes")]
data: Vec<u8>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct WebSocketSendBinaryResponse {
#[serde(default)]
error: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct WebSocketCloseConnectionRequest {
connection_id: String,
code: i32,
reason: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct WebSocketCloseConnectionResponse {
#[serde(default)]
error: Option<String>,
}
#[host_fn]
extern "ExtismHost" {
fn websocket_connect(input: Json<WebSocketConnectRequest>) -> Json<WebSocketConnectResponse>;
fn websocket_sendtext(input: Json<WebSocketSendTextRequest>) -> Json<WebSocketSendTextResponse>;
fn websocket_sendbinary(input: Json<WebSocketSendBinaryRequest>) -> Json<WebSocketSendBinaryResponse>;
fn websocket_closeconnection(input: Json<WebSocketCloseConnectionRequest>) -> Json<WebSocketCloseConnectionResponse>;
}
pub fn connect(url: &str, headers: std::collections::HashMap<String, String>, connection_id: &str) -> Result<String, Error> {
let response = unsafe {
websocket_connect(Json(WebSocketConnectRequest {
url: url.to_owned(),
headers: headers,
connection_id: connection_id.to_owned(),
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(response.0.new_connection_id)
}
pub fn send_text(connection_id: &str, message: &str) -> Result<(), Error> {
let response = unsafe {
websocket_sendtext(Json(WebSocketSendTextRequest {
connection_id: connection_id.to_owned(),
message: message.to_owned(),
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(())
}
pub fn send_binary(connection_id: &str, data: Vec<u8>) -> Result<(), Error> {
let response = unsafe {
websocket_sendbinary(Json(WebSocketSendBinaryRequest {
connection_id: connection_id.to_owned(),
data: data,
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(())
}
pub fn close_connection(connection_id: &str, code: i32, reason: &str) -> Result<(), Error> {
let response = unsafe {
websocket_closeconnection(Json(WebSocketCloseConnectionRequest {
connection_id: connection_id.to_owned(),
code: code,
reason: reason.to_owned(),
}))?
};
if let Some(err) = response.0.error {
return Err(Error::msg(err));
}
Ok(())
}