Rename zed-rpc to zrpc
Co-Authored-By: Nathan Sobo <nathan@zed.dev> Co-Authored-By: Max Brunsfeld <max@zed.dev>
This commit is contained in:
Antonio Scandurra
co-authored by
Nathan Sobo
Max Brunsfeld
Nathan Sobo
Max Brunsfeld
parent
c4e37dc47c
commit
391ad489ff
@@ -0,0 +1,136 @@
|
||||
use anyhow::{Context, Result};
|
||||
use rand::{thread_rng, Rng as _};
|
||||
use rsa::{PublicKey as _, PublicKeyEncoding, RSAPrivateKey, RSAPublicKey};
|
||||
use std::convert::TryFrom;
|
||||
|
||||
pub struct PublicKey(RSAPublicKey);
|
||||
|
||||
pub struct PrivateKey(RSAPrivateKey);
|
||||
|
||||
/// Generate a public and private key for asymmetric encryption.
|
||||
pub fn keypair() -> Result<(PublicKey, PrivateKey)> {
|
||||
let mut rng = thread_rng();
|
||||
let bits = 1024;
|
||||
let private_key = RSAPrivateKey::new(&mut rng, bits)?;
|
||||
let public_key = RSAPublicKey::from(&private_key);
|
||||
Ok((PublicKey(public_key), PrivateKey(private_key)))
|
||||
}
|
||||
|
||||
/// Generate a random 64-character base64 string.
|
||||
pub fn random_token() -> String {
|
||||
let mut rng = thread_rng();
|
||||
let mut token_bytes = [0; 48];
|
||||
for byte in token_bytes.iter_mut() {
|
||||
*byte = rng.gen();
|
||||
}
|
||||
base64::encode_config(&token_bytes, base64::URL_SAFE)
|
||||
}
|
||||
|
||||
impl PublicKey {
|
||||
/// Convert a string to a base64-encoded string that can only be decoded with the corresponding
|
||||
/// private key.
|
||||
pub fn encrypt_string(&self, string: &str) -> Result<String> {
|
||||
let mut rng = thread_rng();
|
||||
let bytes = string.as_bytes();
|
||||
let encrypted_bytes = self
|
||||
.0
|
||||
.encrypt(&mut rng, PADDING_SCHEME, bytes)
|
||||
.context("failed to encrypt string with public key")?;
|
||||
let encrypted_string = base64::encode_config(&encrypted_bytes, base64::URL_SAFE);
|
||||
Ok(encrypted_string)
|
||||
}
|
||||
}
|
||||
|
||||
impl PrivateKey {
|
||||
/// Decrypt a base64-encoded string that was encrypted by the correspoding public key.
|
||||
pub fn decrypt_string(&self, encrypted_string: &str) -> Result<String> {
|
||||
let encrypted_bytes = base64::decode_config(encrypted_string, base64::URL_SAFE)
|
||||
.context("failed to base64-decode encrypted string")?;
|
||||
let bytes = self
|
||||
.0
|
||||
.decrypt(PADDING_SCHEME, &encrypted_bytes)
|
||||
.context("failed to decrypt string with private key")?;
|
||||
let string = String::from_utf8(bytes).context("decrypted content was not valid utf8")?;
|
||||
Ok(string)
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<PublicKey> for String {
|
||||
type Error = anyhow::Error;
|
||||
fn try_from(key: PublicKey) -> Result<Self> {
|
||||
let bytes = key.0.to_pkcs1().context("failed to serialize public key")?;
|
||||
let string = base64::encode_config(&bytes, base64::URL_SAFE);
|
||||
Ok(string)
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<String> for PublicKey {
|
||||
type Error = anyhow::Error;
|
||||
fn try_from(value: String) -> Result<Self> {
|
||||
let bytes = base64::decode_config(&value, base64::URL_SAFE)
|
||||
.context("failed to base64-decode public key string")?;
|
||||
let key = Self(RSAPublicKey::from_pkcs1(&bytes).context("failed to parse public key")?);
|
||||
Ok(key)
|
||||
}
|
||||
}
|
||||
|
||||
const PADDING_SCHEME: rsa::PaddingScheme = rsa::PaddingScheme::PKCS1v15Encrypt;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_generate_encrypt_and_decrypt_token() {
|
||||
// CLIENT:
|
||||
// * generate a keypair for asymmetric encryption
|
||||
// * serialize the public key to send it to the server.
|
||||
let (public, private) = keypair().unwrap();
|
||||
let public_string = String::try_from(public).unwrap();
|
||||
assert_printable(&public_string);
|
||||
|
||||
// SERVER:
|
||||
// * parse the public key
|
||||
// * generate a random token.
|
||||
// * encrypt the token using the public key.
|
||||
let public = PublicKey::try_from(public_string).unwrap();
|
||||
let token = random_token();
|
||||
let encrypted_token = public.encrypt_string(&token).unwrap();
|
||||
assert_eq!(token.len(), 64);
|
||||
assert_ne!(encrypted_token, token);
|
||||
assert_printable(&token);
|
||||
assert_printable(&encrypted_token);
|
||||
|
||||
// CLIENT:
|
||||
// * decrypt the token using the private key.
|
||||
let decrypted_token = private.decrypt_string(&encrypted_token).unwrap();
|
||||
assert_eq!(decrypted_token, token);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tokens_are_always_url_safe() {
|
||||
for _ in 0..5 {
|
||||
let token = random_token();
|
||||
let (public_key, _) = keypair().unwrap();
|
||||
let encrypted_token = public_key.encrypt_string(&token).unwrap();
|
||||
let public_key_str = String::try_from(public_key).unwrap();
|
||||
|
||||
assert_printable(&token);
|
||||
assert_printable(&public_key_str);
|
||||
assert_printable(&encrypted_token);
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_printable(token: &str) {
|
||||
for c in token.chars() {
|
||||
assert!(
|
||||
c.is_ascii_graphic(),
|
||||
"token {:?} has non-printable char {}",
|
||||
token,
|
||||
c
|
||||
);
|
||||
assert_ne!(c, '/', "token {:?} is not URL-safe", token);
|
||||
assert_ne!(c, '&', "token {:?} is not URL-safe", token);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
pub mod auth;
|
||||
mod peer;
|
||||
pub mod proto;
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub mod test;
|
||||
|
||||
pub use peer::*;
|
||||
@@ -0,0 +1,702 @@
|
||||
use crate::proto::{self, EnvelopedMessage, MessageStream, RequestMessage};
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use async_lock::{Mutex, RwLock};
|
||||
use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
|
||||
use futures::{
|
||||
future::{BoxFuture, LocalBoxFuture},
|
||||
FutureExt, StreamExt,
|
||||
};
|
||||
use postage::{
|
||||
mpsc,
|
||||
prelude::{Sink, Stream},
|
||||
};
|
||||
use std::{
|
||||
any::TypeId,
|
||||
collections::{HashMap, HashSet},
|
||||
fmt,
|
||||
future::Future,
|
||||
marker::PhantomData,
|
||||
sync::{
|
||||
atomic::{self, AtomicU32},
|
||||
Arc,
|
||||
},
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
|
||||
pub struct ConnectionId(pub u32);
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
|
||||
pub struct PeerId(pub u32);
|
||||
|
||||
type MessageHandler = Box<
|
||||
dyn Send
|
||||
+ Sync
|
||||
+ Fn(&mut Option<proto::Envelope>, ConnectionId) -> Option<BoxFuture<'static, ()>>,
|
||||
>;
|
||||
|
||||
type ForegroundMessageHandler =
|
||||
Box<dyn Fn(&mut Option<proto::Envelope>, ConnectionId) -> Option<LocalBoxFuture<'static, ()>>>;
|
||||
|
||||
pub struct Receipt<T> {
|
||||
sender_id: ConnectionId,
|
||||
message_id: u32,
|
||||
payload_type: PhantomData<T>,
|
||||
}
|
||||
|
||||
pub struct TypedEnvelope<T> {
|
||||
pub sender_id: ConnectionId,
|
||||
original_sender_id: Option<PeerId>,
|
||||
pub message_id: u32,
|
||||
pub payload: T,
|
||||
}
|
||||
|
||||
impl<T> TypedEnvelope<T> {
|
||||
pub fn original_sender_id(&self) -> Result<PeerId> {
|
||||
self.original_sender_id
|
||||
.ok_or_else(|| anyhow!("missing original_sender_id"))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: RequestMessage> TypedEnvelope<T> {
|
||||
pub fn receipt(&self) -> Receipt<T> {
|
||||
Receipt {
|
||||
sender_id: self.sender_id,
|
||||
message_id: self.message_id,
|
||||
payload_type: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub type Router = RouterInternal<MessageHandler>;
|
||||
pub type ForegroundRouter = RouterInternal<ForegroundMessageHandler>;
|
||||
pub struct RouterInternal<H> {
|
||||
message_handlers: Vec<H>,
|
||||
handler_types: HashSet<TypeId>,
|
||||
}
|
||||
|
||||
pub struct Peer {
|
||||
connections: RwLock<HashMap<ConnectionId, Connection>>,
|
||||
next_connection_id: AtomicU32,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct Connection {
|
||||
outgoing_tx: mpsc::Sender<proto::Envelope>,
|
||||
next_message_id: Arc<AtomicU32>,
|
||||
response_channels: Arc<Mutex<HashMap<u32, mpsc::Sender<proto::Envelope>>>>,
|
||||
}
|
||||
|
||||
impl Peer {
|
||||
pub fn new() -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
connections: Default::default(),
|
||||
next_connection_id: Default::default(),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn add_connection<Conn, H, Fut>(
|
||||
self: &Arc<Self>,
|
||||
conn: Conn,
|
||||
router: Arc<RouterInternal<H>>,
|
||||
) -> (
|
||||
ConnectionId,
|
||||
impl Future<Output = anyhow::Result<()>> + Send,
|
||||
impl Future<Output = ()>,
|
||||
)
|
||||
where
|
||||
H: Fn(&mut Option<proto::Envelope>, ConnectionId) -> Option<Fut>,
|
||||
Fut: Future<Output = ()>,
|
||||
Conn: futures::Sink<WebSocketMessage, Error = WebSocketError>
|
||||
+ futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>
|
||||
+ Send
|
||||
+ Unpin,
|
||||
{
|
||||
let (tx, rx) = conn.split();
|
||||
let connection_id = ConnectionId(
|
||||
self.next_connection_id
|
||||
.fetch_add(1, atomic::Ordering::SeqCst),
|
||||
);
|
||||
let (mut incoming_tx, mut incoming_rx) = mpsc::channel(64);
|
||||
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(64);
|
||||
let connection = Connection {
|
||||
outgoing_tx,
|
||||
next_message_id: Default::default(),
|
||||
response_channels: Default::default(),
|
||||
};
|
||||
let mut writer = MessageStream::new(tx);
|
||||
let mut reader = MessageStream::new(rx);
|
||||
|
||||
let handle_io = async move {
|
||||
loop {
|
||||
let read_message = reader.read_message().fuse();
|
||||
futures::pin_mut!(read_message);
|
||||
loop {
|
||||
futures::select_biased! {
|
||||
incoming = read_message => match incoming {
|
||||
Ok(incoming) => {
|
||||
if incoming_tx.send(incoming).await.is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
break;
|
||||
}
|
||||
Err(error) => {
|
||||
Err(error).context("received invalid RPC message")?;
|
||||
}
|
||||
},
|
||||
outgoing = outgoing_rx.recv().fuse() => match outgoing {
|
||||
Some(outgoing) => {
|
||||
if let Err(result) = writer.write_message(&outgoing).await {
|
||||
Err(result).context("failed to write RPC message")?;
|
||||
}
|
||||
}
|
||||
None => return Ok(()),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let response_channels = connection.response_channels.clone();
|
||||
let handle_messages = async move {
|
||||
while let Some(message) = incoming_rx.recv().await {
|
||||
if let Some(responding_to) = message.responding_to {
|
||||
let channel = response_channels.lock().await.remove(&responding_to);
|
||||
if let Some(mut tx) = channel {
|
||||
tx.send(message).await.ok();
|
||||
} else {
|
||||
log::warn!("received RPC response to unknown request {}", responding_to);
|
||||
}
|
||||
} else {
|
||||
router.handle(connection_id, message).await;
|
||||
}
|
||||
}
|
||||
response_channels.lock().await.clear();
|
||||
};
|
||||
|
||||
self.connections
|
||||
.write()
|
||||
.await
|
||||
.insert(connection_id, connection);
|
||||
|
||||
(connection_id, handle_io, handle_messages)
|
||||
}
|
||||
|
||||
pub async fn disconnect(&self, connection_id: ConnectionId) {
|
||||
self.connections.write().await.remove(&connection_id);
|
||||
}
|
||||
|
||||
pub async fn reset(&self) {
|
||||
self.connections.write().await.clear();
|
||||
}
|
||||
|
||||
pub fn request<T: RequestMessage>(
|
||||
self: &Arc<Self>,
|
||||
receiver_id: ConnectionId,
|
||||
request: T,
|
||||
) -> impl Future<Output = Result<T::Response>> {
|
||||
self.request_internal(None, receiver_id, request)
|
||||
}
|
||||
|
||||
pub fn forward_request<T: RequestMessage>(
|
||||
self: &Arc<Self>,
|
||||
sender_id: ConnectionId,
|
||||
receiver_id: ConnectionId,
|
||||
request: T,
|
||||
) -> impl Future<Output = Result<T::Response>> {
|
||||
self.request_internal(Some(sender_id), receiver_id, request)
|
||||
}
|
||||
|
||||
pub fn request_internal<T: RequestMessage>(
|
||||
self: &Arc<Self>,
|
||||
original_sender_id: Option<ConnectionId>,
|
||||
receiver_id: ConnectionId,
|
||||
request: T,
|
||||
) -> impl Future<Output = Result<T::Response>> {
|
||||
let this = self.clone();
|
||||
let (tx, mut rx) = mpsc::channel(1);
|
||||
async move {
|
||||
let mut connection = this.connection(receiver_id).await?;
|
||||
let message_id = connection
|
||||
.next_message_id
|
||||
.fetch_add(1, atomic::Ordering::SeqCst);
|
||||
connection
|
||||
.response_channels
|
||||
.lock()
|
||||
.await
|
||||
.insert(message_id, tx);
|
||||
connection
|
||||
.outgoing_tx
|
||||
.send(request.into_envelope(message_id, None, original_sender_id.map(|id| id.0)))
|
||||
.await?;
|
||||
let response = rx
|
||||
.recv()
|
||||
.await
|
||||
.ok_or_else(|| anyhow!("connection was closed"))?;
|
||||
T::Response::from_envelope(response)
|
||||
.ok_or_else(|| anyhow!("received response of the wrong type"))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn send<T: EnvelopedMessage>(
|
||||
self: &Arc<Self>,
|
||||
receiver_id: ConnectionId,
|
||||
message: T,
|
||||
) -> impl Future<Output = Result<()>> {
|
||||
let this = self.clone();
|
||||
async move {
|
||||
let mut connection = this.connection(receiver_id).await?;
|
||||
let message_id = connection
|
||||
.next_message_id
|
||||
.fetch_add(1, atomic::Ordering::SeqCst);
|
||||
connection
|
||||
.outgoing_tx
|
||||
.send(message.into_envelope(message_id, None, None))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn forward_send<T: EnvelopedMessage>(
|
||||
self: &Arc<Self>,
|
||||
sender_id: ConnectionId,
|
||||
receiver_id: ConnectionId,
|
||||
message: T,
|
||||
) -> impl Future<Output = Result<()>> {
|
||||
let this = self.clone();
|
||||
async move {
|
||||
let mut connection = this.connection(receiver_id).await?;
|
||||
let message_id = connection
|
||||
.next_message_id
|
||||
.fetch_add(1, atomic::Ordering::SeqCst);
|
||||
connection
|
||||
.outgoing_tx
|
||||
.send(message.into_envelope(message_id, None, Some(sender_id.0)))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn respond<T: RequestMessage>(
|
||||
self: &Arc<Self>,
|
||||
receipt: Receipt<T>,
|
||||
response: T::Response,
|
||||
) -> impl Future<Output = Result<()>> {
|
||||
let this = self.clone();
|
||||
async move {
|
||||
let mut connection = this.connection(receipt.sender_id).await?;
|
||||
let message_id = connection
|
||||
.next_message_id
|
||||
.fetch_add(1, atomic::Ordering::SeqCst);
|
||||
connection
|
||||
.outgoing_tx
|
||||
.send(response.into_envelope(message_id, Some(receipt.message_id), None))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn connection(
|
||||
self: &Arc<Self>,
|
||||
connection_id: ConnectionId,
|
||||
) -> impl Future<Output = Result<Connection>> {
|
||||
let this = self.clone();
|
||||
async move {
|
||||
let connections = this.connections.read().await;
|
||||
let connection = connections
|
||||
.get(&connection_id)
|
||||
.ok_or_else(|| anyhow!("no such connection: {}", connection_id))?;
|
||||
Ok(connection.clone())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<H, Fut> RouterInternal<H>
|
||||
where
|
||||
H: Fn(&mut Option<proto::Envelope>, ConnectionId) -> Option<Fut>,
|
||||
Fut: Future<Output = ()>,
|
||||
{
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
message_handlers: Default::default(),
|
||||
handler_types: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle(&self, connection_id: ConnectionId, message: proto::Envelope) {
|
||||
let mut envelope = Some(message);
|
||||
for handler in self.message_handlers.iter() {
|
||||
if let Some(future) = handler(&mut envelope, connection_id) {
|
||||
future.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
log::warn!("unhandled message: {:?}", envelope.unwrap().payload);
|
||||
}
|
||||
}
|
||||
|
||||
impl Router {
|
||||
pub fn add_message_handler<T, Fut, F>(&mut self, handler: F)
|
||||
where
|
||||
T: EnvelopedMessage,
|
||||
Fut: 'static + Send + Future<Output = Result<()>>,
|
||||
F: 'static + Send + Sync + Fn(TypedEnvelope<T>) -> Fut,
|
||||
{
|
||||
if !self.handler_types.insert(TypeId::of::<T>()) {
|
||||
panic!("duplicate handler type");
|
||||
}
|
||||
|
||||
self.message_handlers
|
||||
.push(Box::new(move |envelope, connection_id| {
|
||||
if envelope.as_ref().map_or(false, T::matches_envelope) {
|
||||
let envelope = Option::take(envelope).unwrap();
|
||||
let message_id = envelope.id;
|
||||
let future = handler(TypedEnvelope {
|
||||
sender_id: connection_id,
|
||||
original_sender_id: envelope.original_sender_id.map(PeerId),
|
||||
message_id,
|
||||
payload: T::from_envelope(envelope).unwrap(),
|
||||
});
|
||||
Some(
|
||||
async move {
|
||||
if let Err(error) = future.await {
|
||||
log::error!(
|
||||
"error handling message {} {}: {:?}",
|
||||
T::NAME,
|
||||
message_id,
|
||||
error
|
||||
);
|
||||
}
|
||||
}
|
||||
.boxed(),
|
||||
)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
impl ForegroundRouter {
|
||||
pub fn add_message_handler<T, Fut, F>(&mut self, handler: F)
|
||||
where
|
||||
T: EnvelopedMessage,
|
||||
Fut: 'static + Future<Output = Result<()>>,
|
||||
F: 'static + Fn(TypedEnvelope<T>) -> Fut,
|
||||
{
|
||||
if !self.handler_types.insert(TypeId::of::<T>()) {
|
||||
panic!("duplicate handler type");
|
||||
}
|
||||
|
||||
self.message_handlers
|
||||
.push(Box::new(move |envelope, connection_id| {
|
||||
if envelope.as_ref().map_or(false, T::matches_envelope) {
|
||||
let envelope = Option::take(envelope).unwrap();
|
||||
let message_id = envelope.id;
|
||||
let future = handler(TypedEnvelope {
|
||||
sender_id: connection_id,
|
||||
original_sender_id: envelope.original_sender_id.map(PeerId),
|
||||
message_id,
|
||||
payload: T::from_envelope(envelope).unwrap(),
|
||||
});
|
||||
Some(
|
||||
async move {
|
||||
if let Err(error) = future.await {
|
||||
log::error!(
|
||||
"error handling message {} {}: {:?}",
|
||||
T::NAME,
|
||||
message_id,
|
||||
error
|
||||
);
|
||||
}
|
||||
}
|
||||
.boxed_local(),
|
||||
)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Clone for Receipt<T> {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
sender_id: self.sender_id,
|
||||
message_id: self.message_id,
|
||||
payload_type: PhantomData,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Copy for Receipt<T> {}
|
||||
|
||||
impl fmt::Display for ConnectionId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
self.0.fmt(f)
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for PeerId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
self.0.fmt(f)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::test;
|
||||
|
||||
#[test]
|
||||
fn test_request_response() {
|
||||
smol::block_on(async move {
|
||||
// create 2 clients connected to 1 server
|
||||
let server = Peer::new();
|
||||
let client1 = Peer::new();
|
||||
let client2 = Peer::new();
|
||||
|
||||
let mut router = Router::new();
|
||||
router.add_message_handler({
|
||||
let server = server.clone();
|
||||
move |envelope: TypedEnvelope<proto::Auth>| {
|
||||
let server = server.clone();
|
||||
async move {
|
||||
let receipt = envelope.receipt();
|
||||
let message = envelope.payload;
|
||||
server
|
||||
.respond(
|
||||
receipt,
|
||||
match message.user_id {
|
||||
1 => {
|
||||
assert_eq!(message.access_token, "access-token-1");
|
||||
proto::AuthResponse {
|
||||
credentials_valid: true,
|
||||
}
|
||||
}
|
||||
2 => {
|
||||
assert_eq!(message.access_token, "access-token-2");
|
||||
proto::AuthResponse {
|
||||
credentials_valid: false,
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
panic!("unexpected user id {}", message.user_id);
|
||||
}
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
router.add_message_handler({
|
||||
let server = server.clone();
|
||||
move |envelope: TypedEnvelope<proto::OpenBuffer>| {
|
||||
let server = server.clone();
|
||||
async move {
|
||||
let receipt = envelope.receipt();
|
||||
let message = envelope.payload;
|
||||
server
|
||||
.respond(
|
||||
receipt,
|
||||
match message.path.as_str() {
|
||||
"path/one" => {
|
||||
assert_eq!(message.worktree_id, 1);
|
||||
proto::OpenBufferResponse {
|
||||
buffer: Some(proto::Buffer {
|
||||
id: 101,
|
||||
content: "path/one content".to_string(),
|
||||
history: vec![],
|
||||
selections: vec![],
|
||||
}),
|
||||
}
|
||||
}
|
||||
"path/two" => {
|
||||
assert_eq!(message.worktree_id, 2);
|
||||
proto::OpenBufferResponse {
|
||||
buffer: Some(proto::Buffer {
|
||||
id: 102,
|
||||
content: "path/two content".to_string(),
|
||||
history: vec![],
|
||||
selections: vec![],
|
||||
}),
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
panic!("unexpected path {}", message.path);
|
||||
}
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
});
|
||||
let router = Arc::new(router);
|
||||
|
||||
let (client1_to_server_conn, server_to_client_1_conn) = test::Channel::bidirectional();
|
||||
let (client1_conn_id, io_task1, msg_task1) = client1
|
||||
.add_connection(client1_to_server_conn, router.clone())
|
||||
.await;
|
||||
let (_, io_task2, msg_task2) = server
|
||||
.add_connection(server_to_client_1_conn, router.clone())
|
||||
.await;
|
||||
|
||||
let (client2_to_server_conn, server_to_client_2_conn) = test::Channel::bidirectional();
|
||||
let (client2_conn_id, io_task3, msg_task3) = client2
|
||||
.add_connection(client2_to_server_conn, router.clone())
|
||||
.await;
|
||||
let (_, io_task4, msg_task4) = server
|
||||
.add_connection(server_to_client_2_conn, router.clone())
|
||||
.await;
|
||||
|
||||
smol::spawn(io_task1).detach();
|
||||
smol::spawn(io_task2).detach();
|
||||
smol::spawn(io_task3).detach();
|
||||
smol::spawn(io_task4).detach();
|
||||
smol::spawn(msg_task1).detach();
|
||||
smol::spawn(msg_task2).detach();
|
||||
smol::spawn(msg_task3).detach();
|
||||
smol::spawn(msg_task4).detach();
|
||||
|
||||
assert_eq!(
|
||||
client1
|
||||
.request(
|
||||
client1_conn_id,
|
||||
proto::Auth {
|
||||
user_id: 1,
|
||||
access_token: "access-token-1".to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
proto::AuthResponse {
|
||||
credentials_valid: true,
|
||||
}
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
client2
|
||||
.request(
|
||||
client2_conn_id,
|
||||
proto::Auth {
|
||||
user_id: 2,
|
||||
access_token: "access-token-2".to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
proto::AuthResponse {
|
||||
credentials_valid: false,
|
||||
}
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
client1
|
||||
.request(
|
||||
client1_conn_id,
|
||||
proto::OpenBuffer {
|
||||
worktree_id: 1,
|
||||
path: "path/one".to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
proto::OpenBufferResponse {
|
||||
buffer: Some(proto::Buffer {
|
||||
id: 101,
|
||||
content: "path/one content".to_string(),
|
||||
history: vec![],
|
||||
selections: vec![],
|
||||
}),
|
||||
}
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
client2
|
||||
.request(
|
||||
client2_conn_id,
|
||||
proto::OpenBuffer {
|
||||
worktree_id: 2,
|
||||
path: "path/two".to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
proto::OpenBufferResponse {
|
||||
buffer: Some(proto::Buffer {
|
||||
id: 102,
|
||||
content: "path/two content".to_string(),
|
||||
history: vec![],
|
||||
selections: vec![],
|
||||
}),
|
||||
}
|
||||
);
|
||||
|
||||
client1.disconnect(client1_conn_id).await;
|
||||
client2.disconnect(client1_conn_id).await;
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_disconnect() {
|
||||
smol::block_on(async move {
|
||||
let (client_conn, mut server_conn) = test::Channel::bidirectional();
|
||||
|
||||
let client = Peer::new();
|
||||
let router = Arc::new(Router::new());
|
||||
let (connection_id, io_handler, message_handler) =
|
||||
client.add_connection(client_conn, router).await;
|
||||
|
||||
let (mut io_ended_tx, mut io_ended_rx) = postage::barrier::channel();
|
||||
smol::spawn(async move {
|
||||
io_handler.await.ok();
|
||||
io_ended_tx.send(()).await.unwrap();
|
||||
})
|
||||
.detach();
|
||||
|
||||
let (mut messages_ended_tx, mut messages_ended_rx) = postage::barrier::channel();
|
||||
smol::spawn(async move {
|
||||
message_handler.await;
|
||||
messages_ended_tx.send(()).await.unwrap();
|
||||
})
|
||||
.detach();
|
||||
|
||||
client.disconnect(connection_id).await;
|
||||
|
||||
io_ended_rx.recv().await;
|
||||
messages_ended_rx.recv().await;
|
||||
assert!(
|
||||
futures::SinkExt::send(&mut server_conn, WebSocketMessage::Binary(vec![]))
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_io_error() {
|
||||
smol::block_on(async move {
|
||||
let (client_conn, server_conn) = test::Channel::bidirectional();
|
||||
drop(server_conn);
|
||||
|
||||
let client = Peer::new();
|
||||
let router = Arc::new(Router::new());
|
||||
let (connection_id, io_handler, message_handler) =
|
||||
client.add_connection(client_conn, router).await;
|
||||
smol::spawn(io_handler).detach();
|
||||
smol::spawn(message_handler).detach();
|
||||
|
||||
let err = client
|
||||
.request(
|
||||
connection_id,
|
||||
proto::Auth {
|
||||
user_id: 42,
|
||||
access_token: "token".to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(err.to_string(), "connection was closed");
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
|
||||
use futures::{SinkExt as _, StreamExt as _};
|
||||
use prost::Message;
|
||||
use std::{
|
||||
io,
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
include!(concat!(env!("OUT_DIR"), "/zed.messages.rs"));
|
||||
|
||||
pub trait EnvelopedMessage: Clone + Sized + Send + 'static {
|
||||
const NAME: &'static str;
|
||||
fn into_envelope(
|
||||
self,
|
||||
id: u32,
|
||||
responding_to: Option<u32>,
|
||||
original_sender_id: Option<u32>,
|
||||
) -> Envelope;
|
||||
fn matches_envelope(envelope: &Envelope) -> bool;
|
||||
fn from_envelope(envelope: Envelope) -> Option<Self>;
|
||||
}
|
||||
|
||||
pub trait RequestMessage: EnvelopedMessage {
|
||||
type Response: EnvelopedMessage;
|
||||
}
|
||||
|
||||
macro_rules! message {
|
||||
($name:ident) => {
|
||||
impl EnvelopedMessage for $name {
|
||||
const NAME: &'static str = std::stringify!($name);
|
||||
|
||||
fn into_envelope(
|
||||
self,
|
||||
id: u32,
|
||||
responding_to: Option<u32>,
|
||||
original_sender_id: Option<u32>,
|
||||
) -> Envelope {
|
||||
Envelope {
|
||||
id,
|
||||
responding_to,
|
||||
original_sender_id,
|
||||
payload: Some(envelope::Payload::$name(self)),
|
||||
}
|
||||
}
|
||||
|
||||
fn matches_envelope(envelope: &Envelope) -> bool {
|
||||
matches!(&envelope.payload, Some(envelope::Payload::$name(_)))
|
||||
}
|
||||
|
||||
fn from_envelope(envelope: Envelope) -> Option<Self> {
|
||||
if let Some(envelope::Payload::$name(msg)) = envelope.payload {
|
||||
Some(msg)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
macro_rules! request_message {
|
||||
($req:ident, $resp:ident) => {
|
||||
message!($req);
|
||||
message!($resp);
|
||||
impl RequestMessage for $req {
|
||||
type Response = $resp;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
request_message!(Auth, AuthResponse);
|
||||
request_message!(ShareWorktree, ShareWorktreeResponse);
|
||||
request_message!(OpenWorktree, OpenWorktreeResponse);
|
||||
message!(UpdateWorktree);
|
||||
message!(CloseWorktree);
|
||||
request_message!(OpenBuffer, OpenBufferResponse);
|
||||
message!(CloseBuffer);
|
||||
message!(UpdateBuffer);
|
||||
request_message!(SaveBuffer, BufferSaved);
|
||||
message!(AddPeer);
|
||||
message!(RemovePeer);
|
||||
|
||||
/// A stream of protobuf messages.
|
||||
pub struct MessageStream<S> {
|
||||
stream: S,
|
||||
}
|
||||
|
||||
impl<S> MessageStream<S> {
|
||||
pub fn new(stream: S) -> Self {
|
||||
Self { stream }
|
||||
}
|
||||
|
||||
pub fn inner_mut(&mut self) -> &mut S {
|
||||
&mut self.stream
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> MessageStream<S>
|
||||
where
|
||||
S: futures::Sink<WebSocketMessage, Error = WebSocketError> + Unpin,
|
||||
{
|
||||
/// Write a given protobuf message to the stream.
|
||||
pub async fn write_message(&mut self, message: &Envelope) -> Result<(), WebSocketError> {
|
||||
let mut buffer = Vec::with_capacity(message.encoded_len());
|
||||
message
|
||||
.encode(&mut buffer)
|
||||
.map_err(|err| io::Error::from(err))?;
|
||||
self.stream.send(WebSocketMessage::Binary(buffer)).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> MessageStream<S>
|
||||
where
|
||||
S: futures::Stream<Item = Result<WebSocketMessage, WebSocketError>> + Unpin,
|
||||
{
|
||||
/// Read a protobuf message of the given type from the stream.
|
||||
pub async fn read_message(&mut self) -> Result<Envelope, WebSocketError> {
|
||||
while let Some(bytes) = self.stream.next().await {
|
||||
match bytes? {
|
||||
WebSocketMessage::Binary(bytes) => {
|
||||
let envelope = Envelope::decode(bytes.as_slice()).map_err(io::Error::from)?;
|
||||
return Ok(envelope);
|
||||
}
|
||||
WebSocketMessage::Close(_) => break,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Err(WebSocketError::ConnectionClosed)
|
||||
}
|
||||
}
|
||||
|
||||
impl Into<SystemTime> for Timestamp {
|
||||
fn into(self) -> SystemTime {
|
||||
UNIX_EPOCH
|
||||
.checked_add(Duration::new(self.seconds, self.nanos))
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SystemTime> for Timestamp {
|
||||
fn from(time: SystemTime) -> Self {
|
||||
let duration = time.duration_since(UNIX_EPOCH).unwrap();
|
||||
Self {
|
||||
seconds: duration.as_secs(),
|
||||
nanos: duration.subsec_nanos(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::test;
|
||||
|
||||
#[test]
|
||||
fn test_round_trip_message() {
|
||||
smol::block_on(async {
|
||||
let stream = test::Channel::new();
|
||||
let message1 = Auth {
|
||||
user_id: 5,
|
||||
access_token: "the-access-token".into(),
|
||||
}
|
||||
.into_envelope(3, None, None);
|
||||
|
||||
let message2 = OpenBuffer {
|
||||
worktree_id: 0,
|
||||
path: "some/path".to_string(),
|
||||
}
|
||||
.into_envelope(5, None, None);
|
||||
|
||||
let mut message_stream = MessageStream::new(stream);
|
||||
message_stream.write_message(&message1).await.unwrap();
|
||||
message_stream.write_message(&message2).await.unwrap();
|
||||
let decoded_message1 = message_stream.read_message().await.unwrap();
|
||||
let decoded_message2 = message_stream.read_message().await.unwrap();
|
||||
assert_eq!(decoded_message1, message1);
|
||||
assert_eq!(decoded_message2, message2);
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
|
||||
use std::{
|
||||
io,
|
||||
pin::Pin,
|
||||
task::{Context, Poll},
|
||||
};
|
||||
|
||||
pub struct Channel {
|
||||
tx: futures::channel::mpsc::UnboundedSender<WebSocketMessage>,
|
||||
rx: futures::channel::mpsc::UnboundedReceiver<WebSocketMessage>,
|
||||
}
|
||||
|
||||
impl Channel {
|
||||
pub fn new() -> Self {
|
||||
let (tx, rx) = futures::channel::mpsc::unbounded();
|
||||
Self { tx, rx }
|
||||
}
|
||||
|
||||
pub fn bidirectional() -> (Self, Self) {
|
||||
let (a_tx, a_rx) = futures::channel::mpsc::unbounded();
|
||||
let (b_tx, b_rx) = futures::channel::mpsc::unbounded();
|
||||
let a = Self { tx: a_tx, rx: b_rx };
|
||||
let b = Self { tx: b_tx, rx: a_rx };
|
||||
(a, b)
|
||||
}
|
||||
}
|
||||
|
||||
impl futures::Sink<WebSocketMessage> for Channel {
|
||||
type Error = WebSocketError;
|
||||
|
||||
fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
Pin::new(&mut self.tx)
|
||||
.poll_ready(cx)
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
|
||||
}
|
||||
|
||||
fn start_send(mut self: Pin<&mut Self>, item: WebSocketMessage) -> Result<(), Self::Error> {
|
||||
Pin::new(&mut self.tx)
|
||||
.start_send(item)
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
|
||||
}
|
||||
|
||||
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
Pin::new(&mut self.tx)
|
||||
.poll_flush(cx)
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
|
||||
}
|
||||
|
||||
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
Pin::new(&mut self.tx)
|
||||
.poll_close(cx)
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
|
||||
}
|
||||
}
|
||||
|
||||
impl futures::Stream for Channel {
|
||||
type Item = Result<WebSocketMessage, WebSocketError>;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
Pin::new(&mut self.rx)
|
||||
.poll_next(cx)
|
||||
.map(|i| i.map(|i| Ok(i)))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user