1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183
use crate::{Error, WebSocketConn};
use async_tungstenite::tungstenite::{protocol::CloseFrame, Message};
use futures_lite::stream::{Pending, Stream};
use std::future::Future;
use trillium::async_trait;
/**
# This is the trait that defines a handler for trillium websockets.
There are several mutually-exclusive ways to use this trait, and it is
intended to be flexible for different use cases. If the trait does not
support your use case, please open a discussion and/or build a trait
on top of this trait to add additional functionality.
## Simple Example
```
use trillium_websockets::{Message, WebSocket, WebSocketConn, WebSocketHandler};
use futures_lite::stream::{pending, Pending};
struct EchoServer;
#[trillium::async_trait]
impl WebSocketHandler for EchoServer {
type OutboundStream = Pending<Message>; // we don't use an outbound stream in this example
async fn connect(&self, conn: WebSocketConn) -> Option<(WebSocketConn, Self::OutboundStream)> {
Some((conn, pending()))
}
async fn inbound(&self, message: Message, conn: &mut WebSocketConn) {
let path = conn.path().to_string();
if let Message::Text(input) = message {
let reply = format!("received your message: {} at path {}", &input, &path);
conn.send_string(reply).await;
}
}
}
let handler = WebSocket::new(EchoServer);
# // tests at tests/tests.rs for example simplicity
```
## Using [`WebSocketHandler::connect`] only
If you have needs that are not supported by this trait, you can either
pass an `Fn(WebSocketConn) -> impl Future<Output=()>` as a handler, or
implement your own connect-only trait implementation that takes the
WebSocketConn and returns None. The tcp connection will remain intact
until the WebSocketConn is dropped, so you can store it in any data
structure or move it between threads as needed.
## Using Streams
If you define an associated OutboundStream type and return it from
`connect`, every message in that Stream will be sent to the connected
websocket client. This is useful for sending messages that are
triggered by other events in the application, using whatever channel
mechanism is appropriate for your application. The websocket
connection will be closed if the stream ends, yielding None.
If you do not need to use streams, set `OutboundStream =
futures_lite::stream::Pending<Message>` or a similar stream
implementation that never yields. If associated type defaults were
stable, we would use that.
## Receiving client-sent messages
Implement [`WebSocketHandler::inbound`] to receive client-sent
messages. Currently inbound messages are not represented as a stream,
but this may change in the future.
## Holding data inside of the implementing type
As this is a trait you implement for your own type, you can hold
additional data or structs inside of your struct. There will be
exactly one of these structs shared throughout the application, so
async concurrency types can be used to mutate shared data.
This example holds a shared BroadcastChannel that is cloned for each
OutboundStream. Any message that a connected clients sends is
broadcast to every other connected client.
Importantly, this means that the dispatch and fanout of messages is
managed entirely by your implementation. For an opinionated layer on
top of this, see the trillium-channels crate.
```
use broadcaster::BroadcastChannel;
use trillium_websockets::{Message, WebSocket, WebSocketConn, WebSocketHandler};
struct EchoServer {
channel: BroadcastChannel<Message>,
}
impl EchoServer {
fn new() -> Self {
Self {
channel: BroadcastChannel::new(),
}
}
}
#[trillium::async_trait]
impl WebSocketHandler for EchoServer {
type OutboundStream = BroadcastChannel<Message>;
async fn connect(&self, conn: WebSocketConn) -> Option<(WebSocketConn, Self::OutboundStream)> {
Some((conn, self.channel.clone()))
}
async fn inbound(&self, message: Message, _conn: &mut WebSocketConn) {
if let Message::Text(input) = message {
let message = Message::text(format!("received message: {}", &input));
trillium::log_error!(self.channel.send(&message).await);
}
}
}
// fn main() {
// trillium_smol::run(WebSocket::new(EchoServer::new()));
// }
```
*/
#[allow(unused_variables)]
#[async_trait]
pub trait WebSocketHandler: Send + Sync + Sized + 'static {
/**
A [`Stream`] type that represents [`Message`]s to be sent to this
client. It is built in your implementation code, in
[`WebSocketHandler::connect`]. Use `Pending<Message>` or another
stream that never returns if you do not need to use this aspect of
the trait.
*/
type OutboundStream: Stream<Item = Message> + Send + Sync + 'static;
/**
This interface is the only mandatory function in
WebSocketHandler. It receives an owned WebSocketConn and
optionally returns it along with an `OutboundStream`
type.
*/
async fn connect(&self, conn: WebSocketConn) -> Option<(WebSocketConn, Self::OutboundStream)>;
/**
This interface function is called once with every message received
from a connected websocket client.
*/
async fn inbound(&self, message: Message, conn: &mut WebSocketConn) {}
/**
This interface function is called once with every outbound message
in the OutboundStream. You likely do not need to implement this,
but if you do, you must call `conn.send(message).await` or the
message will not be sent.
*/
async fn send(&self, message: Message, conn: &mut WebSocketConn) -> Result<(), Error> {
conn.send(message).await
}
/**
This interface function is called with the websocket conn and, in
the case of a clean disconnect, the [`CloseFrame`] if one is sent
available.
*/
async fn disconnect(&self, conn: &mut WebSocketConn, close_frame: Option<CloseFrame<'static>>) {
}
}
#[async_trait]
impl<H, Fut> WebSocketHandler for H
where
H: Fn(WebSocketConn) -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
type OutboundStream = Pending<Message>;
async fn connect(&self, wsc: WebSocketConn) -> Option<(WebSocketConn, Self::OutboundStream)> {
self(wsc).await;
None
}
}