Lines
84.36 %
Functions
36.09 %
Branches
100 %
use bonsaidb_core::networking::CURRENT_PROTOCOL_VERSION;
use tokio::io::{AsyncRead, AsyncWrite};
use crate::{Backend, CustomServer, Error};
impl<B: Backend> CustomServer<B> {
/// Listens for websocket connections on `addr`.
pub async fn listen_for_websockets_on<T: tokio::net::ToSocketAddrs + Send + Sync>(
&self,
addr: T,
with_tls: bool,
) -> Result<(), Error> {
if with_tls {
self.listen_for_secure_tcp_on(addr, ()).await
} else {
self.listen_for_tcp_on(addr, ()).await
}
pub(crate) async fn handle_raw_websocket_connection<
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
>(
connection: S,
peer_address: std::net::SocketAddr,
let stream = tokio_tungstenite::accept_hdr_async(connection, VersionChecker).await?;
self.handle_websocket(stream, peer_address).await;
Ok(())
/// Handles upgrading an HTTP connection to the `WebSocket` protocol based
/// on the upgrade `request`. Requires feature `hyper` to be enabled.
#[cfg(feature = "hyper")]
pub async fn upgrade_websocket(
mut request: hyper::Request<hyper::Body>,
) -> hyper::Response<hyper::Body> {
use hyper::{
header::{HeaderValue, CONNECTION, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY, UPGRADE},
StatusCode,
};
use tokio_tungstenite::{tungstenite::protocol::Role, WebSocketStream};
let mut response = hyper::Response::new(hyper::Body::empty());
// Send a 400 to any request that doesn't have
// an `Upgrade` header.
if !request.headers().contains_key(UPGRADE) {
*response.status_mut() = StatusCode::BAD_REQUEST;
return response;
let sec_websocket_key = if let Some(key) = request.headers_mut().remove(SEC_WEBSOCKET_KEY) {
key
let task_self = self.clone();
tokio::spawn(async move {
match hyper::upgrade::on(&mut request).await {
Ok(upgraded) => {
let ws = WebSocketStream::from_raw_socket(upgraded, Role::Server, None).await;
task_self.handle_websocket(ws, peer_address).await;
Err(err) => {
log::error!("Error upgrading websocket: {:?}", err);
});
*response.status_mut() = StatusCode::SWITCHING_PROTOCOLS;
response
.headers_mut()
.insert(UPGRADE, HeaderValue::from_static("websocket"));
.insert(CONNECTION, HeaderValue::from_static("upgrade"));
response.headers_mut().insert(
SEC_WEBSOCKET_ACCEPT,
compute_websocket_accept_header(sec_websocket_key.as_bytes()),
);
/// Handles an established `tokio-tungstenite` `WebSocket` stream.
pub async fn handle_websocket<
S: futures::Stream<Item = Result<tokio_tungstenite::tungstenite::Message, E>>
+ futures::Sink<tokio_tungstenite::tungstenite::Message>
+ Send
+ 'static,
E: std::fmt::Debug + Send,
) {
use bonsaidb_core::{
custom_api::CustomApi,
networking::{Payload, Request, Response},
use futures::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message;
use crate::Transport;
let (mut sender, mut receiver) = connection.split();
let (response_sender, response_receiver) = flume::unbounded();
let (message_sender, message_receiver) = flume::unbounded();
let (api_response_sender, api_response_receiver) = flume::unbounded();
let client = if let Some(client) = self
.initialize_client(Transport::WebSocket, peer_address, api_response_sender)
.await
{
client
return;
let task_sender = response_sender.clone();
while let Ok(response) = api_response_receiver.recv_async().await {
if task_sender
.send(Payload {
id: None,
wrapped: Response::Api(response),
})
.is_err()
break;
while let Ok(response) = message_receiver.recv_async().await {
if sender.send(response).await.is_err() {
Result::<(), Error>::Ok(())
let task_sender = message_sender.clone();
while let Ok(response) = response_receiver.recv_async().await {
.send(Message::Binary(bincode::serialize(&response)?))
let (request_sender, request_receiver) =
flume::bounded::<Payload<Request<<B::CustomApi as CustomApi>::Request>>>(
self.data.client_simultaneous_request_limit,
task_self
.handle_client_requests(client.clone(), request_receiver, response_sender)
.await;
while let Some(payload) = receiver.next().await {
match payload {
Ok(Message::Binary(binary)) => {
match bincode::deserialize::<
Payload<Request<<B::CustomApi as CustomApi>::Request>>,
>(&binary)
Ok(payload) => drop(request_sender.send_async(payload).await),
log::error!("[server] error decoding message: {:?}", err);
Ok(Message::Close(_)) => break,
Ok(Message::Ping(payload)) => {
drop(message_sender.send(Message::Pong(payload)));
other => {
log::error!("[server] unexpected message: {:?}", other);
fn compute_websocket_accept_header(key: &[u8]) -> hyper::header::HeaderValue {
use sha1::{Digest, Sha1};
let mut digest = Sha1::default();
digest.update(key);
digest.update(&b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"[..]);
let encoded = base64::encode(&digest.finalize());
hyper::header::HeaderValue::from_str(&encoded).expect("base64 is a valid value")
struct VersionChecker;
impl tokio_tungstenite::tungstenite::handshake::server::Callback for VersionChecker {
fn on_request(
self,
request: &tokio_tungstenite::tungstenite::handshake::server::Request,
mut response: tokio_tungstenite::tungstenite::handshake::server::Response,
) -> Result<
tokio_tungstenite::tungstenite::handshake::server::Response,
tokio_tungstenite::tungstenite::handshake::server::ErrorResponse,
> {
if let Some(protocols) = request.headers().get("Sec-WebSocket-Protocol") {
if let Ok(protocols) = protocols.to_str() {
for protocol in protocols.split(',').map(str::trim) {
if protocol == CURRENT_PROTOCOL_VERSION {
"Sec-WebSocket-Protocol",
CURRENT_PROTOCOL_VERSION.try_into().unwrap(),
return Ok(response);
let mut err = tokio_tungstenite::tungstenite::handshake::server::ErrorResponse::new(None);
*err.status_mut() = 406_u16.try_into().unwrap();
Err(err)