Skip to content

Commit

Permalink
feat: shuttle-serenity initial commit poc (#429)
Browse files Browse the repository at this point in the history
* feat: shuttle-serenity initial commit poc

* remove shuttle-service

* refactor: drop more shuttle_service stuff

* refactor: drop default serenity framework

* misc: add wasm32-wasi to nix shell

* refactor: cargo sort

* refactor: cargo fmt

Co-authored-by: chesedo <pieter@chesedo.me>
  • Loading branch information
brokad and chesedo authored Oct 25, 2022
1 parent 5d2215b commit a1c5fc5
Show file tree
Hide file tree
Showing 11 changed files with 276 additions and 13 deletions.
12 changes: 0 additions & 12 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ members = [
exclude = [
"e2e",
"examples",
"plugins",
"resources/aws-rds",
"resources/persist",
"resources/secrets",
Expand Down
5 changes: 5 additions & 0 deletions plugins/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
[workspace]
members = [
"serenity/wasm",
"serenity/runtime"
]
13 changes: 13 additions & 0 deletions plugins/serenity/Makefile
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
.PHONY: wasm runtime

all: wasm runtime

wasm:
cd wasm; cargo build --target wasm32-wasi
cp ../target/wasm32-wasi/debug/shuttle_serenity.wasm runtime/bot.wasm

test: wasm
cd runtime; cargo test -- --nocapture

runtime:
cd runtime; cargo build
21 changes: 21 additions & 0 deletions plugins/serenity/runtime/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
[package]
name = "shuttle-runtime"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[[bin]]
name = "shuttle-runtime"

[dependencies]
async-trait = "0.1.58"

tokio = { version = "1.20.1", features = [ "full" ] }

cap-std = "*"
wasmtime = "*"
wasmtime-wasi = "*"
wasi-common = "*"

serenity = { version = "0.11.5", default-features = false, features = ["client", "gateway", "rustls_backend", "model"] }
13 changes: 13 additions & 0 deletions plugins/serenity/runtime/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
# `shuttle-runtime`

## How to run

```bash
$ cd ..; make wasm
$ DISCORD_TOKEN=xxx BOT_SRC=bot.wasm cargo run
```

## Running the tests
```bash
$ cd ..; make test
```
153 changes: 153 additions & 0 deletions plugins/serenity/runtime/src/lib.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
use std::fs::File;
use std::io::{Read, Write};
use std::os::unix::prelude::RawFd;
use std::path::Path;
use std::sync::Arc;

use async_trait::async_trait;

use serenity::{model::prelude::*, prelude::*};

use cap_std::os::unix::net::UnixStream;
use wasi_common::file::FileCaps;
use wasmtime::{Engine, Linker, Module, Store};
use wasmtime_wasi::sync::net::UnixStream as WasiUnixStream;
use wasmtime_wasi::{WasiCtx, WasiCtxBuilder};

pub struct BotBuilder {
engine: Engine,
store: Store<WasiCtx>,
linker: Linker<WasiCtx>,
src: Option<File>,
}

impl BotBuilder {
pub fn new() -> Self {
let engine = Engine::default();

let mut linker: Linker<WasiCtx> = Linker::new(&engine);
wasmtime_wasi::add_to_linker(&mut linker, |s| s).unwrap();

let wasi = WasiCtxBuilder::new()
.inherit_stdio()
.inherit_args()
.unwrap()
.build();

let store = Store::new(&engine, wasi);

Self {
engine,
store,
linker,
src: None,
}
}

pub fn src<P: AsRef<Path>>(mut self, src: P) -> Self {
self.src = Some(File::open(src).unwrap());
self
}

pub fn build(mut self) -> Bot {
let mut buf = Vec::new();
self.src.unwrap().read_to_end(&mut buf).unwrap();
let module = Module::new(&self.engine, buf).unwrap();

for export in module.exports() {
println!("export: {}", export.name());
}

self.linker.module(&mut self.store, "bot", &module).unwrap();
let inner = BotInner {
store: self.store,
linker: self.linker,
};
Bot {
inner: Arc::new(Mutex::new(inner)),
}
}
}

pub struct BotInner {
store: Store<WasiCtx>,
linker: Linker<WasiCtx>,
}

impl BotInner {
pub async fn message(&mut self, new_message: &str) -> Option<String> {
let (mut host, client) = UnixStream::pair().unwrap();
let client = WasiUnixStream::from_cap_std(client);

self.store
.data_mut()
.insert_file(3, Box::new(client), FileCaps::all());

host.write_all(new_message.as_bytes()).unwrap();
host.write(&[0]).unwrap();

println!("calling inner EventHandler message");
self.linker
.get(&mut self.store, "bot", "__SHUTTLE_EventHandler_message")
.unwrap()
.into_func()
.unwrap()
.typed::<RawFd, (), _>(&self.store)
.unwrap()
.call(&mut self.store, 3)
.unwrap();

let mut resp = String::new();
host.read_to_string(&mut resp).unwrap();

if resp.is_empty() {
None
} else {
Some(resp)
}
}
}

pub struct Bot {
inner: Arc<Mutex<BotInner>>,
}

impl Bot {
pub fn builder() -> BotBuilder {
BotBuilder::new()
}

pub fn new<P: AsRef<Path>>(src: P) -> Self {
Self::builder().src(src).build()
}

pub async fn into_client(self, token: &str, intents: GatewayIntents) -> Client {
Client::builder(&token, intents)
.event_handler(self)
.await
.unwrap()
}
}

#[async_trait]
impl EventHandler for Bot {
async fn message(&self, ctx: Context, new_message: Message) {
let mut inner = self.inner.lock().await;
if let Some(resp) = inner.message(new_message.content.as_str()).await {
new_message.channel_id.say(&ctx.http, resp).await.unwrap();
}
}
}

#[cfg(test)]
pub mod tests {
use super::*;

#[tokio::test]
async fn bot() {
let bot = Bot::new("bot.wasm");
let mut inner = bot.inner.lock().await;
assert_eq!(inner.message("not !hello").await, None);
assert_eq!(inner.message("!hello").await, Some("world!".to_string()));
}
}
19 changes: 19 additions & 0 deletions plugins/serenity/runtime/src/main.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
use std::env;
use std::io;

use serenity::prelude::*;

use shuttle_runtime::Bot;

#[tokio::main]
async fn main() -> io::Result<()> {
let intents = GatewayIntents::GUILD_MESSAGES | GatewayIntents::MESSAGE_CONTENT;

let token = env::var("DISCORD_TOKEN").unwrap();
let src = env::var("BOT_SRC").unwrap();

let mut client = Bot::new(src).into_client(token.as_str(), intents).await;
client.start().await.unwrap();

Ok(())
}
11 changes: 11 additions & 0 deletions plugins/serenity/wasm/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
[package]
name = "shuttle-serenity"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[lib]
crate-type = [ "cdylib" ]

[dependencies]
37 changes: 37 additions & 0 deletions plugins/serenity/wasm/src/lib.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
use std::fs::File;
use std::io::{Read, Write};
use std::os::wasi::prelude::*;

pub fn handle_message(message: &str) -> Option<String> {
if message == "!hello" {
Some("world!".to_string())
} else {
None
}
}

#[no_mangle]
#[allow(non_snake_case)]
pub extern "C" fn __SHUTTLE_EventHandler_message(fd: RawFd) {
println!("inner handler awoken; interacting with fd={fd}");

let mut f = unsafe { File::from_raw_fd(fd) };

let mut buf = Vec::new();
let mut c_buf = [0; 1];
loop {
f.read(&mut c_buf).unwrap();
if c_buf[0] == 0 {
break;
} else {
buf.push(c_buf[0]);
}
}

let msg = String::from_utf8(buf).unwrap();
println!("got message: {msg}");

if let Some(resp) = handle_message(msg.as_str()) {
f.write_all(resp.as_bytes()).unwrap();
}
}
4 changes: 3 additions & 1 deletion shell.nix
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,9 @@ in
pkg-config
];
buildInputs = with nixpkgs; [
(rustChannelOf{ channel = "1.63.0"; }).rust
((rustChannelOf{ channel = "1.63.0"; }).rust.override {
targets = ["wasm32-wasi"];
})
rust-analyzer
cargo-watch
terraform
Expand Down

0 comments on commit a1c5fc5

Please sign in to comment.