2021-07-07 16:23:18 +00:00
|
|
|
use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
|
|
|
|
use futures::{SinkExt as _, StreamExt as _};
|
2021-06-09 19:15:25 +00:00
|
|
|
use prost::Message;
|
2021-06-22 00:07:56 +00:00
|
|
|
use std::{
|
|
|
|
io,
|
|
|
|
time::{Duration, SystemTime, UNIX_EPOCH},
|
|
|
|
};
|
2021-06-09 19:15:25 +00:00
|
|
|
|
|
|
|
include!(concat!(env!("OUT_DIR"), "/zed.messages.rs"));
|
|
|
|
|
2021-07-02 22:54:22 +00:00
|
|
|
pub trait EnvelopedMessage: Clone + Sized + Send + 'static {
|
2021-06-18 03:48:26 +00:00
|
|
|
const NAME: &'static str;
|
2021-06-18 23:26:12 +00:00
|
|
|
fn into_envelope(
|
|
|
|
self,
|
|
|
|
id: u32,
|
|
|
|
responding_to: Option<u32>,
|
|
|
|
original_sender_id: Option<u32>,
|
|
|
|
) -> Envelope;
|
2021-06-16 00:22:48 +00:00
|
|
|
fn matches_envelope(envelope: &Envelope) -> bool;
|
2021-06-15 20:06:50 +00:00
|
|
|
fn from_envelope(envelope: Envelope) -> Option<Self>;
|
2021-06-11 21:30:12 +00:00
|
|
|
}
|
|
|
|
|
2021-06-15 20:06:50 +00:00
|
|
|
pub trait RequestMessage: EnvelopedMessage {
|
|
|
|
type Response: EnvelopedMessage;
|
2021-06-11 21:30:12 +00:00
|
|
|
}
|
|
|
|
|
2021-06-15 20:06:50 +00:00
|
|
|
macro_rules! message {
|
|
|
|
($name:ident) => {
|
|
|
|
impl EnvelopedMessage for $name {
|
2021-06-18 03:48:26 +00:00
|
|
|
const NAME: &'static str = std::stringify!($name);
|
|
|
|
|
2021-06-18 23:26:12 +00:00
|
|
|
fn into_envelope(
|
|
|
|
self,
|
|
|
|
id: u32,
|
|
|
|
responding_to: Option<u32>,
|
|
|
|
original_sender_id: Option<u32>,
|
|
|
|
) -> Envelope {
|
2021-06-15 20:06:50 +00:00
|
|
|
Envelope {
|
|
|
|
id,
|
|
|
|
responding_to,
|
2021-06-18 23:26:12 +00:00
|
|
|
original_sender_id,
|
2021-06-15 20:06:50 +00:00
|
|
|
payload: Some(envelope::Payload::$name(self)),
|
|
|
|
}
|
2021-06-11 21:30:12 +00:00
|
|
|
}
|
|
|
|
|
2021-06-16 00:22:48 +00:00
|
|
|
fn matches_envelope(envelope: &Envelope) -> bool {
|
|
|
|
matches!(&envelope.payload, Some(envelope::Payload::$name(_)))
|
|
|
|
}
|
|
|
|
|
2021-06-15 20:06:50 +00:00
|
|
|
fn from_envelope(envelope: Envelope) -> Option<Self> {
|
|
|
|
if let Some(envelope::Payload::$name(msg)) = envelope.payload {
|
2021-06-11 21:30:12 +00:00
|
|
|
Some(msg)
|
|
|
|
} else {
|
|
|
|
None
|
|
|
|
}
|
|
|
|
}
|
2021-06-10 01:14:30 +00:00
|
|
|
}
|
|
|
|
};
|
2021-06-09 19:15:25 +00:00
|
|
|
}
|
|
|
|
|
2021-06-11 21:30:12 +00:00
|
|
|
macro_rules! request_message {
|
|
|
|
($req:ident, $resp:ident) => {
|
2021-06-15 20:06:50 +00:00
|
|
|
message!($req);
|
|
|
|
message!($resp);
|
|
|
|
impl RequestMessage for $req {
|
|
|
|
type Response = $resp;
|
2021-06-15 08:42:06 +00:00
|
|
|
}
|
|
|
|
};
|
|
|
|
}
|
|
|
|
|
2021-06-11 21:30:12 +00:00
|
|
|
request_message!(Auth, AuthResponse);
|
|
|
|
request_message!(ShareWorktree, ShareWorktreeResponse);
|
2021-06-15 20:06:50 +00:00
|
|
|
request_message!(OpenWorktree, OpenWorktreeResponse);
|
2021-07-05 15:20:26 +00:00
|
|
|
message!(UpdateWorktree);
|
2021-06-30 14:59:02 +00:00
|
|
|
message!(CloseWorktree);
|
2021-06-15 20:06:50 +00:00
|
|
|
request_message!(OpenBuffer, OpenBufferResponse);
|
2021-06-23 14:20:02 +00:00
|
|
|
message!(CloseBuffer);
|
2021-06-25 23:26:35 +00:00
|
|
|
message!(UpdateBuffer);
|
2021-07-02 15:02:17 +00:00
|
|
|
request_message!(SaveBuffer, BufferSaved);
|
2021-07-05 08:24:39 +00:00
|
|
|
message!(AddPeer);
|
|
|
|
message!(RemovePeer);
|
2021-08-06 02:59:54 +00:00
|
|
|
request_message!(GetChannels, GetChannelsResponse);
|
|
|
|
request_message!(JoinChannel, JoinChannelResponse);
|
|
|
|
request_message!(GetUsers, GetUsersResponse);
|
|
|
|
message!(SendChannelMessage);
|
|
|
|
message!(ChannelMessageSent);
|
2021-06-10 01:14:30 +00:00
|
|
|
|
2021-06-09 19:15:25 +00:00
|
|
|
/// A stream of protobuf messages.
|
2021-07-07 16:23:18 +00:00
|
|
|
pub struct MessageStream<S> {
|
|
|
|
stream: S,
|
2021-06-09 19:15:25 +00:00
|
|
|
}
|
|
|
|
|
2021-07-07 16:23:18 +00:00
|
|
|
impl<S> MessageStream<S> {
|
|
|
|
pub fn new(stream: S) -> Self {
|
|
|
|
Self { stream }
|
2021-06-09 19:15:25 +00:00
|
|
|
}
|
2021-06-11 21:30:12 +00:00
|
|
|
|
2021-07-07 16:23:18 +00:00
|
|
|
pub fn inner_mut(&mut self) -> &mut S {
|
|
|
|
&mut self.stream
|
2021-06-11 21:30:12 +00:00
|
|
|
}
|
2021-06-09 19:15:25 +00:00
|
|
|
}
|
|
|
|
|
2021-07-07 16:23:18 +00:00
|
|
|
impl<S> MessageStream<S>
|
2021-06-09 19:15:25 +00:00
|
|
|
where
|
2021-07-07 16:23:18 +00:00
|
|
|
S: futures::Sink<WebSocketMessage, Error = WebSocketError> + Unpin,
|
2021-06-09 19:15:25 +00:00
|
|
|
{
|
|
|
|
/// Write a given protobuf message to the stream.
|
2021-07-07 16:23:18 +00:00
|
|
|
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(())
|
2021-06-09 19:15:25 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-07-07 16:23:18 +00:00
|
|
|
impl<S> MessageStream<S>
|
2021-06-09 19:15:25 +00:00
|
|
|
where
|
2021-07-07 16:23:18 +00:00
|
|
|
S: futures::Stream<Item = Result<WebSocketMessage, WebSocketError>> + Unpin,
|
2021-06-09 19:15:25 +00:00
|
|
|
{
|
|
|
|
/// Read a protobuf message of the given type from the stream.
|
2021-07-07 16:23:18 +00:00
|
|
|
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,
|
|
|
|
_ => {}
|
2021-07-02 01:12:46 +00:00
|
|
|
}
|
|
|
|
}
|
2021-07-07 16:23:18 +00:00
|
|
|
Err(WebSocketError::ConnectionClosed)
|
2021-06-09 19:15:25 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-06-22 00:07:56 +00:00
|
|
|
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(),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-06-09 19:15:25 +00:00
|
|
|
#[cfg(test)]
|
|
|
|
mod tests {
|
|
|
|
use super::*;
|
2021-07-07 16:23:18 +00:00
|
|
|
use crate::test;
|
2021-06-09 19:15:25 +00:00
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn test_round_trip_message() {
|
|
|
|
smol::block_on(async {
|
2021-07-07 16:23:18 +00:00
|
|
|
let stream = test::Channel::new();
|
2021-06-15 20:06:50 +00:00
|
|
|
let message1 = Auth {
|
|
|
|
user_id: 5,
|
|
|
|
access_token: "the-access-token".into(),
|
|
|
|
}
|
2021-06-18 23:26:12 +00:00
|
|
|
.into_envelope(3, None, None);
|
2021-06-15 20:06:50 +00:00
|
|
|
|
2021-06-21 14:18:18 +00:00
|
|
|
let message2 = OpenBuffer {
|
|
|
|
worktree_id: 0,
|
2021-06-23 14:20:02 +00:00
|
|
|
path: "some/path".to_string(),
|
2021-06-21 14:18:18 +00:00
|
|
|
}
|
|
|
|
.into_envelope(5, None, None);
|
2021-06-09 19:15:25 +00:00
|
|
|
|
2021-07-07 16:23:18 +00:00
|
|
|
let mut message_stream = MessageStream::new(stream);
|
2021-06-09 19:15:25 +00:00
|
|
|
message_stream.write_message(&message1).await.unwrap();
|
|
|
|
message_stream.write_message(&message2).await.unwrap();
|
2021-06-15 20:06:50 +00:00
|
|
|
let decoded_message1 = message_stream.read_message().await.unwrap();
|
|
|
|
let decoded_message2 = message_stream.read_message().await.unwrap();
|
2021-06-09 19:15:25 +00:00
|
|
|
assert_eq!(decoded_message1, message1);
|
|
|
|
assert_eq!(decoded_message2, message2);
|
|
|
|
});
|
|
|
|
}
|
|
|
|
}
|