From 339069b1d38774702da27236dbb791745db1ecd5 Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Tue, 31 May 2022 11:16:32 +0200 Subject: [PATCH] Cap `MessageStream` buffer size to 1MB We temporarily let it grow when the message size exceed the limit, but restore the buffer's capacity shortly after. This ensures that, for each connection in its entire lifetime, we only ever use 1MB. --- crates/auto_update/Cargo.toml | 2 +- crates/cli/Cargo.toml | 2 +- crates/collab/Cargo.toml | 2 +- crates/contacts_panel/Cargo.toml | 2 +- crates/editor/Cargo.toml | 2 +- crates/gpui/Cargo.toml | 2 +- crates/language/Cargo.toml | 2 +- crates/lsp/Cargo.toml | 2 +- crates/project/Cargo.toml | 2 +- crates/rpc/Cargo.toml | 2 +- crates/rpc/src/proto.rs | 49 ++++++++++++++++++++++++++++++-- crates/search/Cargo.toml | 2 +- crates/settings/Cargo.toml | 2 +- crates/theme/Cargo.toml | 2 +- crates/vim/Cargo.toml | 2 +- crates/workspace/Cargo.toml | 2 +- crates/zed/Cargo.toml | 2 +- 17 files changed, 62 insertions(+), 19 deletions(-) diff --git a/crates/auto_update/Cargo.toml b/crates/auto_update/Cargo.toml index 8809eff68335123feecf69118ac1c1606b5e596f..dd90fea6617cbc8130e6b8e89d5f937779bc6c90 100644 --- a/crates/auto_update/Cargo.toml +++ b/crates/auto_update/Cargo.toml @@ -17,7 +17,7 @@ anyhow = "1.0.38" isahc = "1.7" lazy_static = "1.4" log = "0.4" -serde = { version = "1", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } serde_json = { version = "1.0", features = ["preserve_order"] } smol = "1.2.5" tempdir = "0.3.7" diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index be7cc24b3e71a6ecd5cd7a16e63260aacbdb1cdd..fafcc5ab682e298457be8bdfff2387feff88a735 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -16,7 +16,7 @@ anyhow = "1.0" clap = { version = "3.1", features = ["derive"] } dirs = "3.0" ipc-channel = "0.16" -serde = { version = "1.0", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } [target.'cfg(target_os = "macos")'.dependencies] core-foundation = "0.9" diff --git a/crates/collab/Cargo.toml b/crates/collab/Cargo.toml index ba6c8848b98e322b5250023a959056099f74fd85..a2a3a377ab83e17eab5712379f40103540b6ea8a 100644 --- a/crates/collab/Cargo.toml +++ b/crates/collab/Cargo.toml @@ -34,7 +34,7 @@ parking_lot = "0.11.1" rand = "0.8" reqwest = { version = "0.11", features = ["json"], optional = true } scrypt = "0.7" -serde = { version = "1.0", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } serde_json = "1.0" sha-1 = "0.9" time = "0.2" diff --git a/crates/contacts_panel/Cargo.toml b/crates/contacts_panel/Cargo.toml index ab05a56ce773c3aec06eff673eaf810cf09de714..d34e5995935a8e895d7fdb451ba840e538b86a42 100644 --- a/crates/contacts_panel/Cargo.toml +++ b/crates/contacts_panel/Cargo.toml @@ -23,7 +23,7 @@ anyhow = "1.0" futures = "0.3" log = "0.4" postage = { version = "0.4.1", features = ["futures-traits"] } -serde = { version = "1", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } [dev-dependencies] language = { path = "../language", features = ["test-support"] } diff --git a/crates/editor/Cargo.toml b/crates/editor/Cargo.toml index a1c1409d06558f47b593c8265516a3e4038b71ae..bb487d5d2c6d2062980637c323772eebc274189c 100644 --- a/crates/editor/Cargo.toml +++ b/crates/editor/Cargo.toml @@ -45,7 +45,7 @@ ordered-float = "2.1.1" parking_lot = "0.11" postage = { version = "0.4", features = ["futures-traits"] } rand = { version = "0.8.3", optional = true } -serde = { version = "1", features = ["derive", "rc"] } +serde = { version = "1.0", features = ["derive", "rc"] } smallvec = { version = "1.6", features = ["union"] } smol = "1.2" diff --git a/crates/gpui/Cargo.toml b/crates/gpui/Cargo.toml index 18d7766453b0de08ebf3c27e127166c5feb95d0e..7c7f6f257b4fffe1aae4137f90598e95bf6e5502 100644 --- a/crates/gpui/Cargo.toml +++ b/crates/gpui/Cargo.toml @@ -36,7 +36,7 @@ postage = { version = "0.4.1", features = ["futures-traits"] } rand = "0.8.3" resvg = "0.14" seahash = "4.1" -serde = { version = "1.0.125", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } serde_json = "1.0" smallvec = { version = "1.6", features = ["union"] } smol = "1.2" diff --git a/crates/language/Cargo.toml b/crates/language/Cargo.toml index 78cfcd809dfc507024787b958a99220d107f7f36..e105e2522503e3f079e4c385fa3c5bb90f05355b 100644 --- a/crates/language/Cargo.toml +++ b/crates/language/Cargo.toml @@ -40,7 +40,7 @@ log = { version = "0.4.16", features = ["kv_unstable_serde"] } parking_lot = "0.11.1" postage = { version = "0.4.1", features = ["futures-traits"] } rand = { version = "0.8.3", optional = true } -serde = { version = "1", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } serde_json = { version = "1", features = ["preserve_order"] } similar = "1.3" smallvec = { version = "1.6", features = ["union"] } diff --git a/crates/lsp/Cargo.toml b/crates/lsp/Cargo.toml index c749261bf5a6eb9cbe092f2301cc6bca7a75a329..1c663e6b7ec2f77aa75f950c72f8cf9fc977f64f 100644 --- a/crates/lsp/Cargo.toml +++ b/crates/lsp/Cargo.toml @@ -21,7 +21,7 @@ log = { version = "0.4.16", features = ["kv_unstable_serde"] } lsp-types = "0.91" parking_lot = "0.11" postage = { version = "0.4.1", features = ["futures-traits"] } -serde = { version = "1.0", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } serde_json = { version = "1.0", features = ["raw_value"] } smol = "1.2" diff --git a/crates/project/Cargo.toml b/crates/project/Cargo.toml index d57bb8d671fd608f52e6a54a5bc13030069bb13c..921eb9ddc59f1461552d9fe837bd79161d129844 100644 --- a/crates/project/Cargo.toml +++ b/crates/project/Cargo.toml @@ -40,7 +40,7 @@ parking_lot = "0.11.1" postage = { version = "0.4.1", features = ["futures-traits"] } rand = "0.8.3" regex = "1.5" -serde = { version = "1", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } serde_json = { version = "1.0", features = ["preserve_order"] } sha2 = "0.10" similar = "1.3" diff --git a/crates/rpc/Cargo.toml b/crates/rpc/Cargo.toml index da018169317474cfab6a4bf23a564841695f95ec..1fac20528052d98e7fbe9f693eb250301925969f 100644 --- a/crates/rpc/Cargo.toml +++ b/crates/rpc/Cargo.toml @@ -25,7 +25,7 @@ parking_lot = "0.11.1" prost = "0.8" rand = "0.8" rsa = "0.4" -serde = { version = "1", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } smol-timeout = "0.6" tracing = { version = "0.1.34", features = ["log"] } zstd = "0.9" diff --git a/crates/rpc/src/proto.rs b/crates/rpc/src/proto.rs index 7fe715064fa3b505403056d941d75f0cffcf1a30..b6d7836427400293b57fb7e148aca7a322cf5871 100644 --- a/crates/rpc/src/proto.rs +++ b/crates/rpc/src/proto.rs @@ -254,6 +254,8 @@ entity_messages!( entity_messages!(channel_id, ChannelMessageSent); +const MAX_BUFFER_LEN: usize = 1 * 1024 * 1024; + /// A stream of protobuf messages. pub struct MessageStream { stream: S, @@ -293,14 +295,16 @@ where match message { Message::Envelope(message) => { - self.encoding_buffer.resize(message.encoded_len(), 0); - self.encoding_buffer.clear(); + self.encoding_buffer.reserve(message.encoded_len()); message .encode(&mut self.encoding_buffer) .map_err(|err| io::Error::from(err))?; let buffer = zstd::stream::encode_all(self.encoding_buffer.as_slice(), COMPRESSION_LEVEL) .unwrap(); + + self.encoding_buffer.clear(); + self.encoding_buffer.shrink_to(MAX_BUFFER_LEN); self.stream.send(WebSocketMessage::Binary(buffer)).await?; } Message::Ping => { @@ -327,10 +331,12 @@ where while let Some(bytes) = self.stream.next().await { match bytes? { WebSocketMessage::Binary(bytes) => { - self.encoding_buffer.clear(); zstd::stream::copy_decode(bytes.as_slice(), &mut self.encoding_buffer).unwrap(); let envelope = Envelope::decode(self.encoding_buffer.as_slice()) .map_err(io::Error::from)?; + + self.encoding_buffer.clear(); + self.encoding_buffer.shrink_to(MAX_BUFFER_LEN); return Ok(Message::Envelope(envelope)); } WebSocketMessage::Ping(_) => return Ok(Message::Ping), @@ -379,3 +385,40 @@ impl From for u128 { upper_half | lower_half } } + +#[cfg(test)] +mod tests { + use super::*; + + #[gpui::test] + async fn test_buffer_size() { + let (tx, rx) = futures::channel::mpsc::unbounded(); + let mut sink = MessageStream::new(tx.sink_map_err(|_| anyhow!(""))); + sink.write(Message::Envelope(Envelope { + payload: Some(envelope::Payload::UpdateWorktree(UpdateWorktree { + root_name: "abcdefg".repeat(10), + ..Default::default() + })), + ..Default::default() + })) + .await + .unwrap(); + assert!(sink.encoding_buffer.capacity() <= MAX_BUFFER_LEN); + sink.write(Message::Envelope(Envelope { + payload: Some(envelope::Payload::UpdateWorktree(UpdateWorktree { + root_name: "abcdefg".repeat(1000000), + ..Default::default() + })), + ..Default::default() + })) + .await + .unwrap(); + assert!(sink.encoding_buffer.capacity() <= MAX_BUFFER_LEN); + + let mut stream = MessageStream::new(rx.map(|msg| anyhow::Ok(msg))); + stream.read().await.unwrap(); + assert!(stream.encoding_buffer.capacity() <= MAX_BUFFER_LEN); + stream.read().await.unwrap(); + assert!(stream.encoding_buffer.capacity() <= MAX_BUFFER_LEN); + } +} diff --git a/crates/search/Cargo.toml b/crates/search/Cargo.toml index 3e80b5979e4bcf8f0d2d15767366e4fda7e3dc2f..1f6c9582beb1f70df03f4f77b396264f87dc38ea 100644 --- a/crates/search/Cargo.toml +++ b/crates/search/Cargo.toml @@ -21,7 +21,7 @@ workspace = { path = "../workspace" } anyhow = "1.0" log = { version = "0.4.16", features = ["kv_unstable_serde"] } postage = { version = "0.4.1", features = ["futures-traits"] } -serde = { version = "1", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } smallvec = { version = "1.6", features = ["union"] } [dev-dependencies] diff --git a/crates/settings/Cargo.toml b/crates/settings/Cargo.toml index cd361fc90ff68b20df6ce2162030af88f76f1d61..78440a241889ad7d6ddab1dc84609efdd6e9db57 100644 --- a/crates/settings/Cargo.toml +++ b/crates/settings/Cargo.toml @@ -19,7 +19,7 @@ util = { path = "../util" } anyhow = "1.0.38" json_comments = "0.2" schemars = "0.8" -serde = { version = "1", features = ["derive", "rc"] } +serde = { version = "1.0", features = ["derive", "rc"] } serde_json = { version = "1.0", features = ["preserve_order"] } serde_path_to_error = "0.1.4" toml = "0.5" diff --git a/crates/theme/Cargo.toml b/crates/theme/Cargo.toml index af4c15b8a0c1615089f301e018516415288b2990..36de158afe725cbc8ce40e10d117090747a42408 100644 --- a/crates/theme/Cargo.toml +++ b/crates/theme/Cargo.toml @@ -12,7 +12,7 @@ gpui = { path = "../gpui" } anyhow = "1.0.38" indexmap = "1.6.2" parking_lot = "0.11.1" -serde = { version = "1", features = ["derive", "rc"] } +serde = { version = "1.0", features = ["derive", "rc"] } serde_json = { version = "1.0", features = ["preserve_order"] } serde_path_to_error = "0.1.4" toml = "0.5" diff --git a/crates/vim/Cargo.toml b/crates/vim/Cargo.toml index ad4bd8871c2addbfdb88a21fe8d92c7ec58515d9..deb8294d5cc5bc1e05aa78bd141fc1d1667319b1 100644 --- a/crates/vim/Cargo.toml +++ b/crates/vim/Cargo.toml @@ -13,7 +13,7 @@ collections = { path = "../collections" } editor = { path = "../editor" } gpui = { path = "../gpui" } language = { path = "../language" } -serde = { version = "1", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } settings = { path = "../settings" } workspace = { path = "../workspace" } itertools = "0.10" diff --git a/crates/workspace/Cargo.toml b/crates/workspace/Cargo.toml index 106a8c32a7cb6682bd54de2a39711e7798adcde2..8b881fe9e524876cf8fd7ebc6e15471fea22be20 100644 --- a/crates/workspace/Cargo.toml +++ b/crates/workspace/Cargo.toml @@ -25,7 +25,7 @@ futures = "0.3" log = { version = "0.4.16", features = ["kv_unstable_serde"] } parking_lot = "0.11.1" postage = { version = "0.4.1", features = ["futures-traits"] } -serde = { version = "1", features = ["derive", "rc"] } +serde = { version = "1.0", features = ["derive", "rc"] } serde_json = { version = "1.0", features = ["preserve_order"] } smallvec = { version = "1.6", features = ["union"] } diff --git a/crates/zed/Cargo.toml b/crates/zed/Cargo.toml index 97a50e78d2897b636b751d079917473e438367be..430190205225b28f70aa1e6ea8af0a1e33a66c58 100644 --- a/crates/zed/Cargo.toml +++ b/crates/zed/Cargo.toml @@ -76,7 +76,7 @@ rand = "0.8.3" regex = "1.5" rsa = "0.4" rust-embed = { version = "6.3", features = ["include-exclude"] } -serde = { version = "1", features = ["derive"] } +serde = { version = "1.0", features = ["derive", "rc"] } serde_json = { version = "1.0", features = ["preserve_order"] } serde_path_to_error = "0.1.4" simplelog = "0.9"