1use crate::data::{Direction, KeyWithModifier};
2use crate::nested_session_contract::nested_session_contract as proto;
3use base64::alphabet::STANDARD as BASE64_STANDARD_ALPHABET;
4use base64::engine::general_purpose::{
5 GeneralPurpose, GeneralPurposeConfig, STANDARD as BASE64_STANDARD,
6};
7use base64::engine::{DecodePaddingMode, Engine as _};
8use prost::Message;
9use std::str::FromStr;
10use std::time::{Duration, Instant};
11
12const BASE64_DECODER: GeneralPurpose = GeneralPurpose::new(
13 &BASE64_STANDARD_ALPHABET,
14 GeneralPurposeConfig::new().with_decode_padding_mode(DecodePaddingMode::Indifferent),
15);
16
17pub const NESTED_DCS_PARAM: u16 = 26661;
18pub const NESTED_FRAME_HEADER: &[u8] = b"\x1bP26661n";
19pub const NESTED_FRAME_TERMINATOR: &[u8] = b"\x1b\\";
20
21pub const REANNOUNCE_SILENCE_MS: u64 = 3000;
22pub const REANNOUNCE_CHECK_INTERVAL_MS: u64 = 1000;
23pub const MAX_UNACKED_ANNOUNCES: usize = 5;
24
25pub fn reannounce_silence_ms() -> u64 {
26 std::env::var("ZELLIJ_NESTED_REANNOUNCE_SILENCE_MS")
27 .ok()
28 .and_then(|value| value.parse().ok())
29 .unwrap_or(REANNOUNCE_SILENCE_MS)
30}
31
32pub fn reannounce_check_interval_ms() -> u64 {
33 std::env::var("ZELLIJ_NESTED_REANNOUNCE_CHECK_INTERVAL_MS")
34 .ok()
35 .and_then(|value| value.parse().ok())
36 .unwrap_or(REANNOUNCE_CHECK_INTERVAL_MS)
37}
38
39pub fn max_unacked_announces() -> usize {
40 std::env::var("ZELLIJ_NESTED_MAX_UNACKED_ANNOUNCES")
41 .ok()
42 .and_then(|value| value.parse().ok())
43 .unwrap_or(MAX_UNACKED_ANNOUNCES)
44}
45
46pub struct ReannounceScheduler {
47 last_heard_from_host: Instant,
48 host_contacted: bool,
49 announces_without_host_contact: usize,
50 silence: Duration,
51 budget: usize,
52}
53
54impl ReannounceScheduler {
55 pub fn new(now: Instant) -> Self {
56 ReannounceScheduler::with_settings(
57 now,
58 Duration::from_millis(reannounce_silence_ms()),
59 max_unacked_announces(),
60 )
61 }
62 pub fn with_settings(now: Instant, silence: Duration, budget: usize) -> Self {
63 ReannounceScheduler {
64 last_heard_from_host: now,
65 host_contacted: false,
66 announces_without_host_contact: 1,
67 silence,
68 budget,
69 }
70 }
71 pub fn note_host_contact(&mut self, now: Instant) -> bool {
72 self.last_heard_from_host = now;
73 let first_contact = !self.host_contacted;
74 self.host_contacted = true;
75 first_contact
76 }
77 pub fn on_tick(&mut self, now: Instant) -> bool {
78 if self.budget_exhausted() {
79 return false;
80 }
81 if now.duration_since(self.last_heard_from_host) < self.silence {
82 return false;
83 }
84 if !self.host_contacted {
85 self.announces_without_host_contact += 1;
86 }
87 true
88 }
89 pub fn host_contacted(&self) -> bool {
90 self.host_contacted
91 }
92 pub fn budget_exhausted(&self) -> bool {
93 !self.host_contacted && self.announces_without_host_contact >= self.budget
94 }
95}
96
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
98pub enum NestedSessionCapability {
99 NestedControl,
100}
101
102#[derive(Debug, Clone, PartialEq)]
103pub enum NestedSessionMessage {
104 Announce {
105 session_name: String,
106 capabilities: Vec<NestedSessionCapability>,
107 },
108 FocusHost {
109 direction: Option<Direction>,
110 },
111 ToggleHostFullscreen {
112 fullscreen: bool,
113 },
114 Pong,
115 Bye,
116 AnnounceAck {
117 ancestry: Vec<String>,
118 capabilities: Vec<NestedSessionCapability>,
119 descend_keys: Vec<KeyWithModifier>,
120 },
121 FocusGained {
122 from_direction: Option<Direction>,
123 },
124 FocusLost,
125 FullscreenState {
126 fullscreen: bool,
127 },
128 AncestryUpdate {
129 ancestry: Vec<String>,
130 },
131 Ping,
132 ShortcutUpdate {
133 ascend_keys: Vec<KeyWithModifier>,
134 descend_keys: Vec<KeyWithModifier>,
135 },
136}
137
138fn keys_to_proto(keys: &[KeyWithModifier]) -> Vec<String> {
139 keys.iter().map(|key| key.to_kdl()).collect()
140}
141
142fn keys_from_proto(keys: &[String]) -> Vec<KeyWithModifier> {
143 let parsed: Vec<KeyWithModifier> = keys
144 .iter()
145 .filter_map(|key| KeyWithModifier::from_str(key).ok())
146 .collect();
147 if parsed.len() == keys.len() {
148 parsed
149 } else {
150 vec![]
151 }
152}
153
154fn capabilities_to_proto(capabilities: &[NestedSessionCapability]) -> Vec<i32> {
155 capabilities
156 .iter()
157 .map(|capability| match capability {
158 NestedSessionCapability::NestedControl => proto::NestedCapability::NestedControl as i32,
159 })
160 .collect()
161}
162
163fn capabilities_from_proto(capabilities: &[i32]) -> Vec<NestedSessionCapability> {
164 capabilities
165 .iter()
166 .filter_map(
167 |capability| match proto::NestedCapability::try_from(*capability).ok() {
168 Some(proto::NestedCapability::NestedControl) => {
169 Some(NestedSessionCapability::NestedControl)
170 },
171 _ => None,
172 },
173 )
174 .collect()
175}
176
177fn direction_to_proto(direction: Option<Direction>) -> i32 {
178 match direction {
179 Some(Direction::Left) => proto::NestedDirection::Left as i32,
180 Some(Direction::Right) => proto::NestedDirection::Right as i32,
181 Some(Direction::Up) => proto::NestedDirection::Up as i32,
182 Some(Direction::Down) => proto::NestedDirection::Down as i32,
183 None => proto::NestedDirection::Unspecified as i32,
184 }
185}
186
187fn direction_from_proto(direction: i32) -> Option<Direction> {
188 match proto::NestedDirection::try_from(direction).ok() {
189 Some(proto::NestedDirection::Left) => Some(Direction::Left),
190 Some(proto::NestedDirection::Right) => Some(Direction::Right),
191 Some(proto::NestedDirection::Up) => Some(Direction::Up),
192 Some(proto::NestedDirection::Down) => Some(Direction::Down),
193 _ => None,
194 }
195}
196
197impl From<NestedSessionMessage> for proto::NestedSessionMessage {
198 fn from(message: NestedSessionMessage) -> Self {
199 use proto::nested_session_message::Payload;
200 let payload = match message {
201 NestedSessionMessage::Announce {
202 session_name,
203 capabilities,
204 } => Payload::Announce(proto::Announce {
205 session_name,
206 capabilities: capabilities_to_proto(&capabilities),
207 }),
208 NestedSessionMessage::FocusHost { direction } => Payload::FocusHost(proto::FocusHost {
209 direction: direction_to_proto(direction),
210 }),
211 NestedSessionMessage::ToggleHostFullscreen { fullscreen } => {
212 Payload::HostFullscreen(proto::ToggleHostFullscreen { fullscreen })
213 },
214 NestedSessionMessage::Pong => Payload::Pong(proto::Pong {}),
215 NestedSessionMessage::Bye => Payload::Bye(proto::Bye {}),
216 NestedSessionMessage::AnnounceAck {
217 ancestry,
218 capabilities,
219 descend_keys,
220 } => Payload::AnnounceAck(proto::AnnounceAck {
221 ancestry,
222 capabilities: capabilities_to_proto(&capabilities),
223 descend_keys: keys_to_proto(&descend_keys),
224 }),
225 NestedSessionMessage::FocusGained { from_direction } => {
226 Payload::FocusGained(proto::FocusGained {
227 from_direction: direction_to_proto(from_direction),
228 })
229 },
230 NestedSessionMessage::FocusLost => Payload::FocusLost(proto::FocusLost {}),
231 NestedSessionMessage::FullscreenState { fullscreen } => {
232 Payload::FullscreenState(proto::FullscreenState { fullscreen })
233 },
234 NestedSessionMessage::AncestryUpdate { ancestry } => {
235 Payload::AncestryUpdate(proto::AncestryUpdate { ancestry })
236 },
237 NestedSessionMessage::Ping => Payload::Ping(proto::Ping {}),
238 NestedSessionMessage::ShortcutUpdate {
239 ascend_keys,
240 descend_keys,
241 } => Payload::ShortcutUpdate(proto::ShortcutUpdate {
242 ascend_keys: keys_to_proto(&ascend_keys),
243 descend_keys: keys_to_proto(&descend_keys),
244 }),
245 };
246 proto::NestedSessionMessage {
247 payload: Some(payload),
248 }
249 }
250}
251
252impl TryFrom<proto::NestedSessionMessage> for NestedSessionMessage {
253 type Error = ();
254 fn try_from(message: proto::NestedSessionMessage) -> Result<Self, Self::Error> {
255 use proto::nested_session_message::Payload;
256 match message.payload {
257 Some(Payload::Announce(announce)) => Ok(NestedSessionMessage::Announce {
258 session_name: announce.session_name,
259 capabilities: capabilities_from_proto(&announce.capabilities),
260 }),
261 Some(Payload::FocusHost(focus_host)) => Ok(NestedSessionMessage::FocusHost {
262 direction: direction_from_proto(focus_host.direction),
263 }),
264 Some(Payload::HostFullscreen(host_fullscreen)) => {
265 Ok(NestedSessionMessage::ToggleHostFullscreen {
266 fullscreen: host_fullscreen.fullscreen,
267 })
268 },
269 Some(Payload::Pong(_)) => Ok(NestedSessionMessage::Pong),
270 Some(Payload::Bye(_)) => Ok(NestedSessionMessage::Bye),
271 Some(Payload::AnnounceAck(announce_ack)) => Ok(NestedSessionMessage::AnnounceAck {
272 ancestry: announce_ack.ancestry,
273 capabilities: capabilities_from_proto(&announce_ack.capabilities),
274 descend_keys: keys_from_proto(&announce_ack.descend_keys),
275 }),
276 Some(Payload::FocusGained(focus_gained)) => Ok(NestedSessionMessage::FocusGained {
277 from_direction: direction_from_proto(focus_gained.from_direction),
278 }),
279 Some(Payload::FocusLost(_)) => Ok(NestedSessionMessage::FocusLost),
280 Some(Payload::FullscreenState(fullscreen_state)) => {
281 Ok(NestedSessionMessage::FullscreenState {
282 fullscreen: fullscreen_state.fullscreen,
283 })
284 },
285 Some(Payload::AncestryUpdate(ancestry_update)) => {
286 Ok(NestedSessionMessage::AncestryUpdate {
287 ancestry: ancestry_update.ancestry,
288 })
289 },
290 Some(Payload::Ping(_)) => Ok(NestedSessionMessage::Ping),
291 Some(Payload::ShortcutUpdate(shortcut_update)) => {
292 Ok(NestedSessionMessage::ShortcutUpdate {
293 ascend_keys: keys_from_proto(&shortcut_update.ascend_keys),
294 descend_keys: keys_from_proto(&shortcut_update.descend_keys),
295 })
296 },
297 None => Err(()),
298 }
299 }
300}
301
302pub fn encode_payload(message: &NestedSessionMessage) -> Vec<u8> {
303 let proto_message: proto::NestedSessionMessage = message.clone().into();
304 proto_message.encode_to_vec()
305}
306
307pub fn decode_payload(payload_bytes: &[u8]) -> Option<NestedSessionMessage> {
308 let proto_message = proto::NestedSessionMessage::decode(payload_bytes).ok()?;
309 NestedSessionMessage::try_from(proto_message).ok()
310}
311
312pub fn encode_frame(message: &NestedSessionMessage) -> Vec<u8> {
313 encode_frame_from_payload(&encode_payload(message))
314}
315
316pub fn encode_frame_from_payload(payload_bytes: &[u8]) -> Vec<u8> {
317 let encoded = BASE64_STANDARD.encode(payload_bytes);
318 let mut frame = Vec::with_capacity(
319 NESTED_FRAME_HEADER.len() + encoded.len() + NESTED_FRAME_TERMINATOR.len(),
320 );
321 frame.extend_from_slice(NESTED_FRAME_HEADER);
322 frame.extend_from_slice(encoded.as_bytes());
323 frame.extend_from_slice(NESTED_FRAME_TERMINATOR);
324 frame
325}
326
327pub fn decode_base64(encoded: &[u8]) -> Option<Vec<u8>> {
328 BASE64_DECODER.decode(encoded).ok()
329}
330
331const MAX_PARTIAL_FRAME_BYTES: usize = 1024 * 1024;
332
333enum FrameScanStatus {
334 Complete(usize),
335 NeedMore,
336 Diverged,
337}
338
339fn frame_scan_status(buf: &[u8]) -> FrameScanStatus {
340 if buf.len() < NESTED_FRAME_HEADER.len() {
341 return if NESTED_FRAME_HEADER.starts_with(buf) {
342 FrameScanStatus::NeedMore
343 } else {
344 FrameScanStatus::Diverged
345 };
346 }
347 if &buf[..NESTED_FRAME_HEADER.len()] != NESTED_FRAME_HEADER {
348 return FrameScanStatus::Diverged;
349 }
350 let mut i = NESTED_FRAME_HEADER.len();
351 while i < buf.len() {
352 match buf[i] {
353 0x1b => match buf.get(i + 1) {
354 Some(&b'\\') => return FrameScanStatus::Complete(i + 2),
355 Some(_) => return FrameScanStatus::Diverged,
356 None => return FrameScanStatus::NeedMore,
357 },
358 _ => i += 1,
359 }
360 }
361 FrameScanStatus::NeedMore
362}
363
364#[derive(Debug, Default)]
365pub struct NestedFrameExtractor {
366 partial_frame: Vec<u8>,
367}
368
369impl NestedFrameExtractor {
370 pub fn new() -> Self {
371 Self::default()
372 }
373
374 pub fn extract(&mut self, bytes: &[u8]) -> (Vec<u8>, Vec<Vec<u8>>) {
375 let mut decoded_payloads = Vec::new();
376 if self.partial_frame.is_empty() && !bytes.contains(&0x1b) {
377 return (bytes.to_vec(), decoded_payloads);
378 }
379 let mut working = std::mem::take(&mut self.partial_frame);
380 working.extend_from_slice(bytes);
381 let mut cleaned = Vec::with_capacity(working.len());
382 let mut i = 0;
383 while i < working.len() {
384 let rest = &working[i..];
385 if rest[0] == 0x1b {
386 match frame_scan_status(rest) {
387 FrameScanStatus::Complete(len) => {
388 let encoded_payload =
389 &rest[NESTED_FRAME_HEADER.len()..len - NESTED_FRAME_TERMINATOR.len()];
390 if let Some(payload) = decode_base64(encoded_payload) {
391 decoded_payloads.push(payload);
392 }
393 i += len;
394 continue;
395 },
396 FrameScanStatus::NeedMore => {
397 let tail = rest.to_vec();
398 if tail.len() > MAX_PARTIAL_FRAME_BYTES {
399 cleaned.extend_from_slice(&tail);
400 } else {
401 self.partial_frame = tail;
402 }
403 return (cleaned, decoded_payloads);
404 },
405 FrameScanStatus::Diverged => {},
406 }
407 }
408 cleaned.push(working[i]);
409 i += 1;
410 }
411 (cleaned, decoded_payloads)
412 }
413
414 pub fn partial_bytes(&self) -> &[u8] {
415 &self.partial_frame
416 }
417
418 pub fn take_partial(&mut self) -> Vec<u8> {
419 std::mem::take(&mut self.partial_frame)
420 }
421}
422
423#[cfg(test)]
424mod tests {
425 use super::*;
426 use crate::data::BareKey;
427
428 fn all_message_arms() -> Vec<NestedSessionMessage> {
429 vec![
430 NestedSessionMessage::Announce {
431 session_name: "guest".to_owned(),
432 capabilities: vec![NestedSessionCapability::NestedControl],
433 },
434 NestedSessionMessage::FocusHost {
435 direction: Some(Direction::Left),
436 },
437 NestedSessionMessage::FocusHost { direction: None },
438 NestedSessionMessage::ToggleHostFullscreen { fullscreen: true },
439 NestedSessionMessage::Pong,
440 NestedSessionMessage::Bye,
441 NestedSessionMessage::AnnounceAck {
442 ancestry: vec!["outer".to_owned(), "middle".to_owned()],
443 capabilities: vec![NestedSessionCapability::NestedControl],
444 descend_keys: vec![
445 KeyWithModifier::new(BareKey::Char('o')).with_ctrl_modifier(),
446 KeyWithModifier::new(BareKey::Down),
447 ],
448 },
449 NestedSessionMessage::FocusGained {
450 from_direction: Some(Direction::Right),
451 },
452 NestedSessionMessage::FocusLost,
453 NestedSessionMessage::FullscreenState { fullscreen: true },
454 NestedSessionMessage::AncestryUpdate {
455 ancestry: vec!["outer".to_owned()],
456 },
457 NestedSessionMessage::Ping,
458 NestedSessionMessage::ShortcutUpdate {
459 ascend_keys: vec![
460 KeyWithModifier::new(BareKey::Char('o')).with_ctrl_modifier(),
461 KeyWithModifier::new(BareKey::Up),
462 ],
463 descend_keys: vec![],
464 },
465 ]
466 }
467
468 #[test]
469 fn payload_roundtrip_preserves_every_arm() {
470 for message in all_message_arms() {
471 let decoded = decode_payload(&encode_payload(&message));
472 assert_eq!(decoded, Some(message));
473 }
474 }
475
476 #[test]
477 fn focus_direction_roundtrip_preserves_every_direction() {
478 let directions = [
479 None,
480 Some(Direction::Left),
481 Some(Direction::Right),
482 Some(Direction::Up),
483 Some(Direction::Down),
484 ];
485 for direction in directions {
486 let focus_host = NestedSessionMessage::FocusHost { direction };
487 assert_eq!(
488 decode_payload(&encode_payload(&focus_host)),
489 Some(focus_host)
490 );
491 let focus_gained = NestedSessionMessage::FocusGained {
492 from_direction: direction,
493 };
494 assert_eq!(
495 decode_payload(&encode_payload(&focus_gained)),
496 Some(focus_gained)
497 );
498 }
499 }
500
501 #[test]
502 fn frame_roundtrip_preserves_message() {
503 let message = NestedSessionMessage::Announce {
504 session_name: "guest".to_owned(),
505 capabilities: vec![NestedSessionCapability::NestedControl],
506 };
507 let frame = encode_frame(&message);
508 assert!(frame.starts_with(NESTED_FRAME_HEADER));
509 assert!(frame.ends_with(NESTED_FRAME_TERMINATOR));
510 let encoded_payload =
511 &frame[NESTED_FRAME_HEADER.len()..frame.len() - NESTED_FRAME_TERMINATOR.len()];
512 let payload = decode_base64(encoded_payload).unwrap();
513 assert_eq!(decode_payload(&payload), Some(message));
514 }
515
516 #[test]
517 fn garbage_base64_is_rejected() {
518 assert_eq!(decode_base64(b"!!!not-base64!!!"), None);
519 }
520
521 #[test]
522 fn truncated_payload_is_rejected() {
523 let payload = encode_payload(&NestedSessionMessage::Announce {
524 session_name: "a-long-session-name-to-truncate".to_owned(),
525 capabilities: vec![],
526 });
527 assert_eq!(decode_payload(&payload[..payload.len() - 5]), None);
528 }
529
530 #[test]
531 fn unknown_oneof_arm_is_rejected() {
532 let unknown_field_bytes = [0xE2, 0x03, 0x00];
533 assert_eq!(decode_payload(&unknown_field_bytes), None);
534 }
535
536 #[test]
537 fn empty_payload_is_rejected() {
538 assert_eq!(decode_payload(&[]), None);
539 }
540
541 fn test_scheduler(now: Instant) -> ReannounceScheduler {
542 ReannounceScheduler::with_settings(now, Duration::from_millis(3000), 5)
543 }
544
545 #[test]
546 fn unacked_announces_stop_at_the_budget() {
547 let start = Instant::now();
548 let mut scheduler = test_scheduler(start);
549 let mut announces = 1;
550 for tick in 1..=30 {
551 if scheduler.on_tick(start + Duration::from_millis(tick * 1000)) {
552 announces += 1;
553 }
554 }
555 assert_eq!(announces, 5);
556 assert!(scheduler.budget_exhausted());
557 assert!(!scheduler.host_contacted());
558 }
559
560 #[test]
561 fn silence_threshold_is_respected_before_announcing() {
562 let start = Instant::now();
563 let mut scheduler = test_scheduler(start);
564 assert!(!scheduler.on_tick(start + Duration::from_millis(1000)));
565 assert!(!scheduler.on_tick(start + Duration::from_millis(2999)));
566 assert!(scheduler.on_tick(start + Duration::from_millis(3000)));
567 }
568
569 #[test]
570 fn host_contact_lifts_the_budget_and_resets_the_silence_window() {
571 let start = Instant::now();
572 let mut scheduler = test_scheduler(start);
573 for tick in 1..=30 {
574 scheduler.on_tick(start + Duration::from_millis(tick * 1000));
575 }
576 assert!(scheduler.budget_exhausted());
577 assert!(scheduler.note_host_contact(start + Duration::from_millis(31_000)));
578 assert!(scheduler.host_contacted());
579 assert!(!scheduler.budget_exhausted());
580 assert!(!scheduler.on_tick(start + Duration::from_millis(32_000)));
581 for tick in 34..=100 {
582 assert!(scheduler.on_tick(start + Duration::from_millis(tick * 1000)));
583 }
584 }
585
586 #[test]
587 fn repeated_host_contact_is_only_first_contact_once() {
588 let start = Instant::now();
589 let mut scheduler = test_scheduler(start);
590 assert!(scheduler.note_host_contact(start));
591 assert!(!scheduler.note_host_contact(start + Duration::from_millis(1000)));
592 }
593}