1 files changed, 99 insertions, 9 deletions

diff --git a/src/server/client_transport.rs b/src/server/client_transport.rs
index 743a108f..2675b4c3 100644
--- a/src/server/client_transport.rs
+++ b/src/server/client_transport.rs
@@ -383,24 +383,62 @@ enum InputEventLimit {
     WithinLimits,
     TooManyEvents,
     PasteTooLarge { size: usize },
+    InputPayloadTooLarge { size: usize },
 }
 
 fn input_event_limit(events: &[ClientInputEvent]) -> InputEventLimit {
-    if events.len() > MAX_INPUT_EVENT_BATCH {
-        return InputEventLimit::TooManyEvents;
-    }
-
+    let mut expanded_events = 0usize;
     let mut paste_bytes = 0usize;
+    let mut input_bytes = 0usize;
     for event in events {
-        if let ClientInputEvent::Paste { text } = event {
-            paste_bytes = paste_bytes.saturating_add(text.len());
+        expanded_events = expanded_events.saturating_add(match event {
+            ClientInputEvent::Key { repeat_count, .. } => usize::from((*repeat_count).max(1)),
+            _ => 1,
+        });
+        match event {
+            ClientInputEvent::Key {
+                repeat_count,
+                generated_text,
+                source,
+                ..
+            } => {
+                if let Some(text) = generated_text {
+                    input_bytes = input_bytes.saturating_add(
+                        text.len()
+                            .saturating_mul(usize::from((*repeat_count).max(1))),
+                    );
+                }
+                if let crate::protocol::ClientKeySource::Vt { bytes } = source {
+                    input_bytes = input_bytes.saturating_add(bytes.len());
+                }
+            }
+            ClientInputEvent::TextCommit(text) => {
+                input_bytes = input_bytes.saturating_add(text.len());
+            }
+            ClientInputEvent::Paste { text } => {
+                paste_bytes = paste_bytes.saturating_add(text.len());
+            }
+            ClientInputEvent::Mouse { .. }
+            | ClientInputEvent::FocusGained
+            | ClientInputEvent::FocusLost => {}
         }
     }
 
-    if paste_bytes > MAX_INPUT_PAYLOAD {
-        InputEventLimit::PasteTooLarge { size: paste_bytes }
-    } else {
+    if expanded_events > MAX_INPUT_EVENT_BATCH {
+        return InputEventLimit::TooManyEvents;
+    }
+
+    let payload_bytes = paste_bytes.saturating_add(input_bytes);
+    if payload_bytes <= MAX_INPUT_PAYLOAD {
         InputEventLimit::WithinLimits
+    } else if input_bytes == 0 {
+        InputEventLimit::PasteTooLarge {
+            size: payload_bytes,
+        }
+    } else {
+        InputEventLimit::InputPayloadTooLarge {
+            size: payload_bytes,
+        }
     }
 }
 
@@ -717,6 +755,17 @@ fn client_read_loop(
                         max: MAX_INPUT_PAYLOAD,
                     }
                 }
+                InputEventLimit::InputPayloadTooLarge { size } => {
+                    warn!(
+                        client_id,
+                        size,
+                        max = MAX_INPUT_PAYLOAD,
+                        "oversized structured input payload from client, closing"
+                    );
+                    let _ = server_event_tx
+                        .blocking_send(ServerEvent::ClientDisconnected { client_id });
+                    break;
+                }
             },
             ClientMessage::ObserveTerminal { target } => {
                 ServerEvent::ClientObserveTerminal { client_id, target }
@@ -1388,6 +1437,10 @@ new_tab = "ctrl+notakey"
                 code: crate::protocol::ClientKeyCode::Enter,
                 modifiers: 0,
                 kind: crate::protocol::ClientKeyKind::Press,
+
+                repeat_count: 1,
+                generated_text: None,
+                source: crate::protocol::ClientKeySource::Synthesized,
             },
             ClientInputEvent::FocusGained,
         ];
@@ -1550,6 +1603,43 @@ new_tab = "ctrl+notakey"
     }
 
     #[test]
+    fn structured_input_limits_charge_grouped_repeats_and_text_payloads() {
+        let grouped = ClientInputEvent::Key {
+            code: crate::protocol::ClientKeyCode::Char('x'),
+            modifiers: 0,
+            kind: crate::protocol::ClientKeyKind::Press,
+            repeat_count: (MAX_INPUT_EVENT_BATCH + 1) as u16,
+            generated_text: None,
+            source: crate::protocol::ClientKeySource::Synthesized,
+        };
+        assert_eq!(
+            input_event_limit(&[grouped]),
+            InputEventLimit::TooManyEvents
+        );
+
+        let repeated_text = ClientInputEvent::Key {
+            code: crate::protocol::ClientKeyCode::Char('x'),
+            modifiers: 0,
+            kind: crate::protocol::ClientKeyKind::Press,
+            repeat_count: MAX_INPUT_EVENT_BATCH as u16,
+            generated_text: Some("x".repeat((MAX_INPUT_PAYLOAD / MAX_INPUT_EVENT_BATCH) + 1)),
+            source: crate::protocol::ClientKeySource::Synthesized,
+        };
+        assert!(matches!(
+            input_event_limit(&[repeated_text]),
+            InputEventLimit::InputPayloadTooLarge { size } if size > MAX_INPUT_PAYLOAD
+        ));
+
+        let text = ClientInputEvent::TextCommit("x".repeat(MAX_INPUT_PAYLOAD + 1));
+        assert_eq!(
+            input_event_limit(&[text]),
+            InputEventLimit::InputPayloadTooLarge {
+                size: MAX_INPUT_PAYLOAD + 1
+            }
+        );
+    }
+
+    #[test]
     fn handshake_timeout_is_within_five_second_deadline() {
         // The handshake timeout must be short enough that
         // the connection is guaranteed to close within 5 seconds even with