1pub const FRAME_HEADER_SIZE: usize = 16;
20pub const MAX_FRAME_SIZE: u32 = 16 * 1024 * 1024;
21
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct Frame {
24 pub kind: MessageKind,
25 pub flags: Flags,
26 pub stream_id: u16,
27 pub correlation_id: u64,
28 pub payload: Vec<u8>,
29}
30
31impl Frame {
32 pub fn new(kind: MessageKind, correlation_id: u64, payload: Vec<u8>) -> Self {
33 Self {
34 kind,
35 flags: Flags::empty(),
36 stream_id: 0,
37 correlation_id,
38 payload,
39 }
40 }
41
42 pub fn with_stream(mut self, stream_id: u16) -> Self {
43 self.stream_id = stream_id;
44 self
45 }
46
47 pub fn with_flags(mut self, flags: Flags) -> Self {
48 self.flags = flags;
49 self
50 }
51
52 pub fn encoded_len(&self) -> u32 {
53 (FRAME_HEADER_SIZE + self.payload.len()) as u32
54 }
55}
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60#[repr(u8)]
61pub enum MessageKind {
62 Query = 0x01,
64 Result = 0x02,
65 Error = 0x03,
66 BulkInsert = 0x04,
67 BulkOk = 0x05,
68 BulkInsertBinary = 0x06,
69 QueryBinary = 0x07,
70 BulkInsertPrevalidated = 0x08,
71 BulkStreamStart = 0x09,
72 BulkStreamRows = 0x0A,
73 BulkStreamCommit = 0x0B,
74 BulkStreamAck = 0x0C,
75 Prepare = 0x0D,
76 PreparedOk = 0x0E,
77 ExecutePrepared = 0x0F,
78
79 Hello = 0x10,
81 HelloAck = 0x11,
82 AuthRequest = 0x12,
83 AuthResponse = 0x13,
84 AuthOk = 0x14,
85 AuthFail = 0x15,
86 Bye = 0x16,
87 Ping = 0x17,
88 Pong = 0x18,
89 Get = 0x19,
90 Delete = 0x1A,
91 DeleteOk = 0x1B,
92
93 Cancel = 0x20,
95 Compress = 0x21,
96 SetSession = 0x22,
97 Notice = 0x23,
98
99 RowDescription = 0x24,
101 StreamEnd = 0x25,
102
103 VectorSearch = 0x26,
105 GraphTraverse = 0x27,
106 QueryWithParams = 0x28,
107
108 OpenStream = 0x29,
112 OpenAck = 0x2A,
113 StreamChunk = 0x2B,
114 StreamError = 0x2C,
115 StreamCancel = 0x2D,
116
117 QueueWaitOpen = 0x2E,
125 QueueEventPush = 0x2F,
126 QueueWaitTimeout = 0x30,
131 MovedRedirect = 0x31,
134 Topology = 0x32,
136 TypedInsert = 0x33,
138}
139
140#[derive(Debug, Clone, Copy, PartialEq, Eq)]
149pub enum MessageClass {
150 DataPlane,
151 Handshake,
152 ControlPlane,
153 Streamed,
154}
155
156#[derive(Debug, Clone, Copy, PartialEq, Eq)]
164pub enum MessageDirection {
165 ClientToServer,
166 ServerToClient,
167 Both,
168}
169
170impl MessageKind {
171 pub fn class(&self) -> MessageClass {
173 match self {
174 Self::Query
179 | Self::Result
180 | Self::Error
181 | Self::BulkInsert
182 | Self::BulkOk
183 | Self::BulkInsertBinary
184 | Self::QueryBinary
185 | Self::BulkInsertPrevalidated
186 | Self::Prepare
187 | Self::PreparedOk
188 | Self::ExecutePrepared
189 | Self::Get
190 | Self::Delete
191 | Self::DeleteOk
192 | Self::VectorSearch
193 | Self::GraphTraverse
194 | Self::QueryWithParams
195 | Self::TypedInsert
196 | Self::MovedRedirect => MessageClass::DataPlane,
197
198 Self::BulkStreamStart
204 | Self::BulkStreamRows
205 | Self::BulkStreamCommit
206 | Self::BulkStreamAck
207 | Self::RowDescription
208 | Self::StreamEnd
209 | Self::OpenStream
210 | Self::OpenAck
211 | Self::StreamChunk
212 | Self::StreamError
213 | Self::StreamCancel
214 | Self::QueueWaitOpen
217 | Self::QueueEventPush
218 | Self::QueueWaitTimeout => MessageClass::Streamed,
221
222 Self::Hello
224 | Self::HelloAck
225 | Self::AuthRequest
226 | Self::AuthResponse
227 | Self::AuthOk
228 | Self::AuthFail
229 | Self::Bye
230 | Self::Ping
231 | Self::Pong => MessageClass::Handshake,
232
233 Self::Cancel | Self::Compress | Self::SetSession | Self::Notice | Self::Topology => {
235 MessageClass::ControlPlane
236 }
237 }
238 }
239
240 pub fn allowed_flags(&self) -> Flags {
250 match self {
251 Self::Hello
254 | Self::HelloAck
255 | Self::AuthRequest
256 | Self::AuthResponse
257 | Self::AuthOk
258 | Self::AuthFail
259 | Self::Bye
260 | Self::Ping
261 | Self::Pong => Flags::MORE_FRAMES,
262
263 _ => Flags::COMPRESSED.insert(Flags::MORE_FRAMES),
265 }
266 }
267
268 pub fn is_handshake(&self) -> bool {
273 matches!(self.class(), MessageClass::Handshake)
274 }
275
276 pub fn permits_flags(&self, flags: Flags) -> bool {
282 let allowed = self.allowed_flags().bits();
283 (flags.bits() & !allowed) == 0
284 }
285
286 pub fn direction(&self) -> MessageDirection {
288 match self {
289 Self::Hello
291 | Self::AuthResponse
292 | Self::Query
293 | Self::QueryBinary
294 | Self::BulkInsert
295 | Self::BulkInsertBinary
296 | Self::BulkInsertPrevalidated
297 | Self::BulkStreamStart
298 | Self::BulkStreamRows
299 | Self::BulkStreamCommit
300 | Self::Prepare
301 | Self::ExecutePrepared
302 | Self::Get
303 | Self::Delete
304 | Self::Cancel
305 | Self::Compress
306 | Self::SetSession
307 | Self::VectorSearch
308 | Self::GraphTraverse
309 | Self::QueryWithParams
310 | Self::TypedInsert
311 | Self::Topology
312 | Self::OpenStream
313 | Self::StreamCancel
314 | Self::QueueWaitOpen => MessageDirection::ClientToServer,
316
317 Self::StreamChunk => MessageDirection::Both,
323
324 Self::HelloAck
326 | Self::AuthRequest
327 | Self::AuthOk
328 | Self::AuthFail
329 | Self::Result
330 | Self::Error
331 | Self::BulkOk
332 | Self::BulkStreamAck
333 | Self::PreparedOk
334 | Self::DeleteOk
335 | Self::Notice
336 | Self::RowDescription
337 | Self::StreamEnd
338 | Self::OpenAck
339 | Self::StreamError
340 | Self::QueueEventPush
343 | Self::QueueWaitTimeout
344 | Self::MovedRedirect => MessageDirection::ServerToClient,
345
346 Self::Bye | Self::Ping | Self::Pong => MessageDirection::Both,
350 }
351 }
352
353 pub fn from_u8(byte: u8) -> Option<Self> {
354 match byte {
355 0x01 => Some(Self::Query),
356 0x02 => Some(Self::Result),
357 0x03 => Some(Self::Error),
358 0x04 => Some(Self::BulkInsert),
359 0x05 => Some(Self::BulkOk),
360 0x06 => Some(Self::BulkInsertBinary),
361 0x07 => Some(Self::QueryBinary),
362 0x08 => Some(Self::BulkInsertPrevalidated),
363 0x09 => Some(Self::BulkStreamStart),
364 0x0A => Some(Self::BulkStreamRows),
365 0x0B => Some(Self::BulkStreamCommit),
366 0x0C => Some(Self::BulkStreamAck),
367 0x0D => Some(Self::Prepare),
368 0x0E => Some(Self::PreparedOk),
369 0x0F => Some(Self::ExecutePrepared),
370 0x10 => Some(Self::Hello),
371 0x11 => Some(Self::HelloAck),
372 0x12 => Some(Self::AuthRequest),
373 0x13 => Some(Self::AuthResponse),
374 0x14 => Some(Self::AuthOk),
375 0x15 => Some(Self::AuthFail),
376 0x16 => Some(Self::Bye),
377 0x17 => Some(Self::Ping),
378 0x18 => Some(Self::Pong),
379 0x19 => Some(Self::Get),
380 0x1A => Some(Self::Delete),
381 0x1B => Some(Self::DeleteOk),
382 0x20 => Some(Self::Cancel),
383 0x21 => Some(Self::Compress),
384 0x22 => Some(Self::SetSession),
385 0x23 => Some(Self::Notice),
386 0x24 => Some(Self::RowDescription),
387 0x25 => Some(Self::StreamEnd),
388 0x26 => Some(Self::VectorSearch),
389 0x27 => Some(Self::GraphTraverse),
390 0x28 => Some(Self::QueryWithParams),
391 0x29 => Some(Self::OpenStream),
392 0x2A => Some(Self::OpenAck),
393 0x2B => Some(Self::StreamChunk),
394 0x2C => Some(Self::StreamError),
395 0x2D => Some(Self::StreamCancel),
396 0x2E => Some(Self::QueueWaitOpen),
397 0x2F => Some(Self::QueueEventPush),
398 0x30 => Some(Self::QueueWaitTimeout),
399 0x31 => Some(Self::MovedRedirect),
400 0x32 => Some(Self::Topology),
401 0x33 => Some(Self::TypedInsert),
402 _ => None,
403 }
404 }
405}
406
407#[derive(Debug, Clone, Copy, PartialEq, Eq)]
408pub struct Flags(u8);
409
410impl Flags {
411 pub const COMPRESSED: Self = Self(0b0000_0001);
412 pub const MORE_FRAMES: Self = Self(0b0000_0010);
413
414 pub const fn empty() -> Self {
415 Self(0)
416 }
417
418 pub const fn bits(self) -> u8 {
419 self.0
420 }
421
422 pub const fn from_bits(bits: u8) -> Self {
423 Self(bits)
424 }
425
426 pub const fn contains(self, other: Self) -> bool {
427 (self.0 & other.0) == other.0
428 }
429
430 pub const fn insert(self, other: Self) -> Self {
431 Self(self.0 | other.0)
432 }
433}
434
435impl std::ops::BitOr for Flags {
436 type Output = Self;
437 fn bitor(self, rhs: Self) -> Self {
438 self.insert(rhs)
439 }
440}
441
442#[cfg(test)]
443mod catalog_tests {
444 use super::*;
445
446 const ALL_KINDS: &[MessageKind] = &[
450 MessageKind::Query,
451 MessageKind::Result,
452 MessageKind::Error,
453 MessageKind::BulkInsert,
454 MessageKind::BulkOk,
455 MessageKind::BulkInsertBinary,
456 MessageKind::QueryBinary,
457 MessageKind::BulkInsertPrevalidated,
458 MessageKind::BulkStreamStart,
459 MessageKind::BulkStreamRows,
460 MessageKind::BulkStreamCommit,
461 MessageKind::BulkStreamAck,
462 MessageKind::Prepare,
463 MessageKind::PreparedOk,
464 MessageKind::ExecutePrepared,
465 MessageKind::Hello,
466 MessageKind::HelloAck,
467 MessageKind::AuthRequest,
468 MessageKind::AuthResponse,
469 MessageKind::AuthOk,
470 MessageKind::AuthFail,
471 MessageKind::Bye,
472 MessageKind::Ping,
473 MessageKind::Pong,
474 MessageKind::Get,
475 MessageKind::Delete,
476 MessageKind::DeleteOk,
477 MessageKind::Cancel,
478 MessageKind::Compress,
479 MessageKind::SetSession,
480 MessageKind::Notice,
481 MessageKind::RowDescription,
482 MessageKind::StreamEnd,
483 MessageKind::VectorSearch,
484 MessageKind::GraphTraverse,
485 MessageKind::QueryWithParams,
486 MessageKind::OpenStream,
487 MessageKind::OpenAck,
488 MessageKind::StreamChunk,
489 MessageKind::StreamError,
490 MessageKind::StreamCancel,
491 MessageKind::QueueWaitOpen,
492 MessageKind::QueueEventPush,
493 MessageKind::QueueWaitTimeout,
494 MessageKind::MovedRedirect,
495 MessageKind::Topology,
496 MessageKind::TypedInsert,
497 ];
498
499 #[test]
500 fn class_matrix_is_pinned() {
501 assert_eq!(MessageKind::Hello.class(), MessageClass::Handshake);
504 assert_eq!(MessageKind::HelloAck.class(), MessageClass::Handshake);
505 assert_eq!(MessageKind::AuthRequest.class(), MessageClass::Handshake);
506 assert_eq!(MessageKind::AuthResponse.class(), MessageClass::Handshake);
507 assert_eq!(MessageKind::AuthOk.class(), MessageClass::Handshake);
508 assert_eq!(MessageKind::AuthFail.class(), MessageClass::Handshake);
509 assert_eq!(MessageKind::Bye.class(), MessageClass::Handshake);
510 assert_eq!(MessageKind::Ping.class(), MessageClass::Handshake);
511 assert_eq!(MessageKind::Pong.class(), MessageClass::Handshake);
512
513 assert_eq!(MessageKind::Query.class(), MessageClass::DataPlane);
515 assert_eq!(MessageKind::Result.class(), MessageClass::DataPlane);
516 assert_eq!(MessageKind::BulkInsert.class(), MessageClass::DataPlane);
517 assert_eq!(MessageKind::Get.class(), MessageClass::DataPlane);
518 assert_eq!(MessageKind::Delete.class(), MessageClass::DataPlane);
519 assert_eq!(MessageKind::DeleteOk.class(), MessageClass::DataPlane);
520 assert_eq!(MessageKind::VectorSearch.class(), MessageClass::DataPlane);
521 assert_eq!(MessageKind::GraphTraverse.class(), MessageClass::DataPlane);
522 assert_eq!(
523 MessageKind::QueryWithParams.class(),
524 MessageClass::DataPlane
525 );
526 assert_eq!(MessageKind::TypedInsert.class(), MessageClass::DataPlane);
527 assert_eq!(MessageKind::MovedRedirect.class(), MessageClass::DataPlane);
528
529 assert_eq!(MessageKind::BulkStreamStart.class(), MessageClass::Streamed);
531 assert_eq!(MessageKind::BulkStreamRows.class(), MessageClass::Streamed);
532 assert_eq!(
533 MessageKind::BulkStreamCommit.class(),
534 MessageClass::Streamed
535 );
536 assert_eq!(MessageKind::BulkStreamAck.class(), MessageClass::Streamed);
537 assert_eq!(MessageKind::RowDescription.class(), MessageClass::Streamed);
538 assert_eq!(MessageKind::StreamEnd.class(), MessageClass::Streamed);
539
540 assert_eq!(MessageKind::OpenStream.class(), MessageClass::Streamed);
544 assert_eq!(MessageKind::OpenAck.class(), MessageClass::Streamed);
545 assert_eq!(MessageKind::StreamChunk.class(), MessageClass::Streamed);
546 assert_eq!(MessageKind::StreamError.class(), MessageClass::Streamed);
547 assert_eq!(MessageKind::StreamCancel.class(), MessageClass::Streamed);
548
549 assert_eq!(MessageKind::QueueWaitOpen.class(), MessageClass::Streamed);
552 assert_eq!(MessageKind::QueueEventPush.class(), MessageClass::Streamed);
553 assert_eq!(
555 MessageKind::QueueWaitTimeout.class(),
556 MessageClass::Streamed
557 );
558
559 assert_eq!(MessageKind::Cancel.class(), MessageClass::ControlPlane);
561 assert_eq!(MessageKind::Compress.class(), MessageClass::ControlPlane);
562 assert_eq!(MessageKind::SetSession.class(), MessageClass::ControlPlane);
563 assert_eq!(MessageKind::Notice.class(), MessageClass::ControlPlane);
564
565 for k in ALL_KINDS {
567 let _ = k.class();
568 }
569 }
570
571 #[test]
572 fn allowed_flags_matrix_is_pinned() {
573 let handshake = [
577 MessageKind::Hello,
578 MessageKind::HelloAck,
579 MessageKind::AuthRequest,
580 MessageKind::AuthResponse,
581 MessageKind::AuthOk,
582 MessageKind::AuthFail,
583 MessageKind::Bye,
584 MessageKind::Ping,
585 MessageKind::Pong,
586 ];
587 for k in handshake {
588 let f = k.allowed_flags();
589 assert!(
590 f.contains(Flags::MORE_FRAMES),
591 "{k:?} must allow MORE_FRAMES"
592 );
593 assert!(
594 !f.contains(Flags::COMPRESSED),
595 "{k:?} must NOT allow COMPRESSED today"
596 );
597 }
598
599 for k in ALL_KINDS {
601 if handshake.contains(k) {
602 continue;
603 }
604 let f = k.allowed_flags();
605 assert!(
606 f.contains(Flags::MORE_FRAMES),
607 "{k:?} must allow MORE_FRAMES"
608 );
609 assert!(f.contains(Flags::COMPRESSED), "{k:?} must allow COMPRESSED");
610 }
611 }
612
613 #[test]
614 fn every_kind_has_unique_byte_value() {
615 let mut seen = std::collections::HashSet::new();
618 for k in ALL_KINDS {
619 let byte = *k as u8;
620 assert!(
621 seen.insert(byte),
622 "byte 0x{byte:02x} reused by {k:?}; catalog has a duplicate"
623 );
624 }
625 }
626
627 #[test]
628 fn from_u8_round_trips_for_every_kind() {
629 for k in ALL_KINDS {
630 let byte = *k as u8;
631 let decoded = MessageKind::from_u8(byte).unwrap_or_else(|| {
632 panic!("from_u8 returned None for catalog entry {k:?} (0x{byte:02x})")
633 });
634 assert_eq!(
635 decoded, *k,
636 "from_u8(0x{byte:02x}) must round-trip back to {k:?}"
637 );
638 }
639 }
640
641 #[test]
642 fn permits_flags_matches_allowed_flags() {
643 assert!(MessageKind::Ping.permits_flags(Flags::MORE_FRAMES));
645 assert!(MessageKind::Ping.permits_flags(Flags::empty()));
646 assert!(!MessageKind::Ping.permits_flags(Flags::COMPRESSED));
647 assert!(!MessageKind::Ping.permits_flags(Flags::COMPRESSED | Flags::MORE_FRAMES));
648
649 assert!(MessageKind::BulkStreamRows.permits_flags(Flags::MORE_FRAMES));
653 assert!(MessageKind::BulkStreamRows.permits_flags(Flags::COMPRESSED));
654 assert!(MessageKind::RowDescription.permits_flags(Flags::MORE_FRAMES));
655 assert!(MessageKind::StreamEnd.permits_flags(Flags::MORE_FRAMES));
656 }
657
658 #[test]
659 fn direction_matrix_is_pinned() {
660 for k in [
662 MessageKind::Hello,
663 MessageKind::AuthResponse,
664 MessageKind::Query,
665 MessageKind::QueryBinary,
666 MessageKind::BulkInsert,
667 MessageKind::BulkInsertBinary,
668 MessageKind::BulkInsertPrevalidated,
669 MessageKind::BulkStreamStart,
670 MessageKind::BulkStreamRows,
671 MessageKind::BulkStreamCommit,
672 MessageKind::Prepare,
673 MessageKind::ExecutePrepared,
674 MessageKind::Get,
675 MessageKind::Delete,
676 MessageKind::Cancel,
677 MessageKind::Compress,
678 MessageKind::SetSession,
679 MessageKind::VectorSearch,
680 MessageKind::GraphTraverse,
681 MessageKind::QueryWithParams,
682 MessageKind::TypedInsert,
683 MessageKind::OpenStream,
684 MessageKind::StreamCancel,
685 MessageKind::QueueWaitOpen,
686 ] {
687 assert_eq!(
688 k.direction(),
689 MessageDirection::ClientToServer,
690 "{k:?} should be client-originated"
691 );
692 }
693
694 for k in [
696 MessageKind::HelloAck,
697 MessageKind::AuthRequest,
698 MessageKind::AuthOk,
699 MessageKind::AuthFail,
700 MessageKind::Result,
701 MessageKind::Error,
702 MessageKind::BulkOk,
703 MessageKind::BulkStreamAck,
704 MessageKind::PreparedOk,
705 MessageKind::DeleteOk,
706 MessageKind::Notice,
707 MessageKind::RowDescription,
708 MessageKind::StreamEnd,
709 MessageKind::OpenAck,
710 MessageKind::StreamError,
711 MessageKind::QueueEventPush,
712 MessageKind::QueueWaitTimeout,
713 ] {
714 assert_eq!(
715 k.direction(),
716 MessageDirection::ServerToClient,
717 "{k:?} should be server-originated"
718 );
719 }
720
721 for k in [
725 MessageKind::Bye,
726 MessageKind::Ping,
727 MessageKind::Pong,
728 MessageKind::StreamChunk,
729 ] {
730 assert_eq!(
731 k.direction(),
732 MessageDirection::Both,
733 "{k:?} should be symmetric"
734 );
735 }
736 }
737}