1use bytes::Bytes;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10pub const PROTOCOL_VERSION: u32 = 1;
12
13pub const MAX_MESSAGE_SIZE: usize = 10 * 1024 * 1024;
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub enum EventType {
20 Configure,
22 RequestHeaders,
24 RequestBodyChunk,
26 ResponseHeaders,
28 ResponseBodyChunk,
30 RequestComplete,
32 WebSocketFrame,
34 GuardrailInspect,
36}
37
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
40#[serde(rename_all = "snake_case")]
41pub enum Decision {
42 #[default]
44 Allow,
45 Block {
47 status: u16,
49 body: Option<String>,
51 headers: Option<HashMap<String, String>>,
53 },
54 Redirect {
56 url: String,
58 status: u16,
60 },
61 Challenge {
63 challenge_type: String,
65 params: HashMap<String, String>,
67 },
68}
69
70#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
72#[serde(rename_all = "snake_case")]
73pub enum HeaderOp {
74 Set { name: String, value: String },
76 Add { name: String, value: String },
78 Remove { name: String },
80}
81
82#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
93pub struct BodyMutation {
94 pub data: Option<String>,
100
101 #[serde(default)]
105 pub chunk_index: u32,
106}
107
108impl BodyMutation {
109 pub fn pass_through(chunk_index: u32) -> Self {
111 Self {
112 data: None,
113 chunk_index,
114 }
115 }
116
117 pub fn drop_chunk(chunk_index: u32) -> Self {
119 Self {
120 data: Some(String::new()),
121 chunk_index,
122 }
123 }
124
125 pub fn replace(chunk_index: u32, data: String) -> Self {
127 Self {
128 data: Some(data),
129 chunk_index,
130 }
131 }
132
133 pub fn is_pass_through(&self) -> bool {
135 self.data.is_none()
136 }
137
138 pub fn is_drop(&self) -> bool {
140 matches!(&self.data, Some(d) if d.is_empty())
141 }
142}
143
144#[derive(Debug, Clone, Serialize, Deserialize)]
146pub struct RequestMetadata {
147 pub correlation_id: String,
149 pub request_id: String,
151 pub client_ip: String,
153 pub client_port: u16,
155 pub server_name: Option<String>,
157 pub protocol: String,
159 pub tls_version: Option<String>,
161 pub tls_cipher: Option<String>,
163 pub route_id: Option<String>,
165 pub upstream_id: Option<String>,
167 pub timestamp: String,
169 #[serde(skip_serializing_if = "Option::is_none")]
176 pub traceparent: Option<String>,
177}
178
179#[derive(Debug, Clone, Serialize, Deserialize)]
184pub struct ConfigureEvent {
185 pub agent_id: String,
187 pub config: serde_json::Value,
192}
193
194#[derive(Debug, Clone, Serialize, Deserialize)]
196pub struct RequestHeadersEvent {
197 pub metadata: RequestMetadata,
199 pub method: String,
201 pub uri: String,
203 pub headers: HashMap<String, Vec<String>>,
205}
206
207#[derive(Debug, Clone, Serialize, Deserialize)]
209pub struct RequestBodyChunkEvent {
210 pub correlation_id: String,
212 pub data: String,
214 pub is_last: bool,
216 pub total_size: Option<usize>,
218 #[serde(default)]
222 pub chunk_index: u32,
223 #[serde(default)]
225 pub bytes_received: usize,
226}
227
228#[derive(Debug, Clone, Serialize, Deserialize)]
230pub struct ResponseHeadersEvent {
231 pub correlation_id: String,
233 pub status: u16,
235 pub headers: HashMap<String, Vec<String>>,
237}
238
239#[derive(Debug, Clone, Serialize, Deserialize)]
241pub struct ResponseBodyChunkEvent {
242 pub correlation_id: String,
244 pub data: String,
246 pub is_last: bool,
248 pub total_size: Option<usize>,
250 #[serde(default)]
252 pub chunk_index: u32,
253 #[serde(default)]
255 pub bytes_sent: usize,
256}
257
258#[derive(Debug, Clone)]
272pub struct BinaryRequestBodyChunkEvent {
273 pub correlation_id: String,
275 pub data: Bytes,
277 pub is_last: bool,
279 pub total_size: Option<usize>,
281 pub chunk_index: u32,
283 pub bytes_received: usize,
285}
286
287#[derive(Debug, Clone)]
297pub struct BinaryResponseBodyChunkEvent {
298 pub correlation_id: String,
300 pub data: Bytes,
302 pub is_last: bool,
304 pub total_size: Option<usize>,
306 pub chunk_index: u32,
308 pub bytes_sent: usize,
310}
311
312impl BinaryRequestBodyChunkEvent {
313 pub fn new(
315 correlation_id: impl Into<String>,
316 data: impl Into<Bytes>,
317 chunk_index: u32,
318 is_last: bool,
319 ) -> Self {
320 let data = data.into();
321 Self {
322 correlation_id: correlation_id.into(),
323 bytes_received: data.len(),
324 data,
325 is_last,
326 total_size: None,
327 chunk_index,
328 }
329 }
330
331 pub fn with_total_size(mut self, size: usize) -> Self {
333 self.total_size = Some(size);
334 self
335 }
336
337 pub fn with_bytes_received(mut self, bytes: usize) -> Self {
339 self.bytes_received = bytes;
340 self
341 }
342}
343
344impl BinaryResponseBodyChunkEvent {
345 pub fn new(
347 correlation_id: impl Into<String>,
348 data: impl Into<Bytes>,
349 chunk_index: u32,
350 is_last: bool,
351 ) -> Self {
352 let data = data.into();
353 Self {
354 correlation_id: correlation_id.into(),
355 bytes_sent: data.len(),
356 data,
357 is_last,
358 total_size: None,
359 chunk_index,
360 }
361 }
362
363 pub fn with_total_size(mut self, size: usize) -> Self {
365 self.total_size = Some(size);
366 self
367 }
368
369 pub fn with_bytes_sent(mut self, bytes: usize) -> Self {
371 self.bytes_sent = bytes;
372 self
373 }
374}
375
376impl From<BinaryRequestBodyChunkEvent> for RequestBodyChunkEvent {
381 fn from(event: BinaryRequestBodyChunkEvent) -> Self {
383 use base64::{engine::general_purpose::STANDARD, Engine as _};
384 Self {
385 correlation_id: event.correlation_id,
386 data: STANDARD.encode(&event.data),
387 is_last: event.is_last,
388 total_size: event.total_size,
389 chunk_index: event.chunk_index,
390 bytes_received: event.bytes_received,
391 }
392 }
393}
394
395impl From<&RequestBodyChunkEvent> for BinaryRequestBodyChunkEvent {
396 fn from(event: &RequestBodyChunkEvent) -> Self {
400 use base64::{engine::general_purpose::STANDARD, Engine as _};
401 let data = STANDARD
402 .decode(&event.data)
403 .map(Bytes::from)
404 .unwrap_or_else(|_| Bytes::copy_from_slice(event.data.as_bytes()));
405 Self {
406 correlation_id: event.correlation_id.clone(),
407 data,
408 is_last: event.is_last,
409 total_size: event.total_size,
410 chunk_index: event.chunk_index,
411 bytes_received: event.bytes_received,
412 }
413 }
414}
415
416impl From<BinaryResponseBodyChunkEvent> for ResponseBodyChunkEvent {
417 fn from(event: BinaryResponseBodyChunkEvent) -> Self {
419 use base64::{engine::general_purpose::STANDARD, Engine as _};
420 Self {
421 correlation_id: event.correlation_id,
422 data: STANDARD.encode(&event.data),
423 is_last: event.is_last,
424 total_size: event.total_size,
425 chunk_index: event.chunk_index,
426 bytes_sent: event.bytes_sent,
427 }
428 }
429}
430
431impl From<&ResponseBodyChunkEvent> for BinaryResponseBodyChunkEvent {
432 fn from(event: &ResponseBodyChunkEvent) -> Self {
436 use base64::{engine::general_purpose::STANDARD, Engine as _};
437 let data = STANDARD
438 .decode(&event.data)
439 .map(Bytes::from)
440 .unwrap_or_else(|_| Bytes::copy_from_slice(event.data.as_bytes()));
441 Self {
442 correlation_id: event.correlation_id.clone(),
443 data,
444 is_last: event.is_last,
445 total_size: event.total_size,
446 chunk_index: event.chunk_index,
447 bytes_sent: event.bytes_sent,
448 }
449 }
450}
451
452#[derive(Debug, Clone, Serialize, Deserialize)]
454pub struct RequestCompleteEvent {
455 pub correlation_id: String,
457 pub status: u16,
459 pub duration_ms: u64,
461 pub request_body_size: usize,
463 pub response_body_size: usize,
465 pub upstream_attempts: u32,
467 pub error: Option<String>,
469}
470
471#[derive(Debug, Clone, Serialize, Deserialize)]
480pub struct WebSocketFrameEvent {
481 pub correlation_id: String,
483 pub opcode: String,
485 pub data: String,
487 pub client_to_server: bool,
489 pub frame_index: u64,
491 pub fin: bool,
493 pub route_id: Option<String>,
495 pub client_ip: String,
497}
498
499#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
501#[serde(rename_all = "snake_case")]
502pub enum WebSocketOpcode {
503 Continuation,
505 Text,
507 Binary,
509 Close,
511 Ping,
513 Pong,
515}
516
517impl WebSocketOpcode {
518 pub fn as_str(&self) -> &'static str {
520 match self {
521 Self::Continuation => "continuation",
522 Self::Text => "text",
523 Self::Binary => "binary",
524 Self::Close => "close",
525 Self::Ping => "ping",
526 Self::Pong => "pong",
527 }
528 }
529
530 pub fn from_u8(value: u8) -> Option<Self> {
532 match value {
533 0x0 => Some(Self::Continuation),
534 0x1 => Some(Self::Text),
535 0x2 => Some(Self::Binary),
536 0x8 => Some(Self::Close),
537 0x9 => Some(Self::Ping),
538 0xA => Some(Self::Pong),
539 _ => None,
540 }
541 }
542
543 pub fn as_u8(&self) -> u8 {
545 match self {
546 Self::Continuation => 0x0,
547 Self::Text => 0x1,
548 Self::Binary => 0x2,
549 Self::Close => 0x8,
550 Self::Ping => 0x9,
551 Self::Pong => 0xA,
552 }
553 }
554}
555
556#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
560#[serde(rename_all = "snake_case")]
561pub enum WebSocketDecision {
562 #[default]
564 Allow,
565 Drop,
567 Close {
569 code: u16,
571 reason: String,
573 },
574}
575
576#[derive(Debug, Clone, Serialize, Deserialize)]
578pub struct AgentRequest {
579 pub version: u32,
581 pub event_type: EventType,
583 pub payload: serde_json::Value,
585}
586
587#[derive(Debug, Clone, Serialize, Deserialize)]
589pub struct AgentResponse {
590 pub version: u32,
592 pub decision: Decision,
594 #[serde(default)]
596 pub request_headers: Vec<HeaderOp>,
597 #[serde(default)]
599 pub response_headers: Vec<HeaderOp>,
600 #[serde(default)]
602 pub routing_metadata: HashMap<String, String>,
603 #[serde(default)]
605 pub audit: AuditMetadata,
606
607 #[serde(default)]
618 pub needs_more: bool,
619
620 #[serde(default)]
625 pub request_body_mutation: Option<BodyMutation>,
626
627 #[serde(default)]
632 pub response_body_mutation: Option<BodyMutation>,
633
634 #[serde(default)]
638 pub websocket_decision: Option<WebSocketDecision>,
639}
640
641impl AgentResponse {
642 pub fn default_allow() -> Self {
644 Self {
645 version: PROTOCOL_VERSION,
646 decision: Decision::Allow,
647 request_headers: vec![],
648 response_headers: vec![],
649 routing_metadata: HashMap::new(),
650 audit: AuditMetadata::default(),
651 needs_more: false,
652 request_body_mutation: None,
653 response_body_mutation: None,
654 websocket_decision: None,
655 }
656 }
657
658 pub fn block(status: u16, body: Option<String>) -> Self {
660 Self {
661 version: PROTOCOL_VERSION,
662 decision: Decision::Block {
663 status,
664 body,
665 headers: None,
666 },
667 request_headers: vec![],
668 response_headers: vec![],
669 routing_metadata: HashMap::new(),
670 audit: AuditMetadata::default(),
671 needs_more: false,
672 request_body_mutation: None,
673 response_body_mutation: None,
674 websocket_decision: None,
675 }
676 }
677
678 pub fn redirect(url: String, status: u16) -> Self {
680 Self {
681 version: PROTOCOL_VERSION,
682 decision: Decision::Redirect { url, status },
683 request_headers: vec![],
684 response_headers: vec![],
685 routing_metadata: HashMap::new(),
686 audit: AuditMetadata::default(),
687 needs_more: false,
688 request_body_mutation: None,
689 response_body_mutation: None,
690 websocket_decision: None,
691 }
692 }
693
694 pub fn needs_more_data() -> Self {
696 Self {
697 version: PROTOCOL_VERSION,
698 decision: Decision::Allow,
699 request_headers: vec![],
700 response_headers: vec![],
701 routing_metadata: HashMap::new(),
702 audit: AuditMetadata::default(),
703 needs_more: true,
704 request_body_mutation: None,
705 response_body_mutation: None,
706 websocket_decision: None,
707 }
708 }
709
710 pub fn websocket_allow() -> Self {
712 Self {
713 websocket_decision: Some(WebSocketDecision::Allow),
714 ..Self::default_allow()
715 }
716 }
717
718 pub fn websocket_drop() -> Self {
720 Self {
721 websocket_decision: Some(WebSocketDecision::Drop),
722 ..Self::default_allow()
723 }
724 }
725
726 pub fn websocket_close(code: u16, reason: String) -> Self {
728 Self {
729 websocket_decision: Some(WebSocketDecision::Close { code, reason }),
730 ..Self::default_allow()
731 }
732 }
733
734 pub fn with_websocket_decision(mut self, decision: WebSocketDecision) -> Self {
736 self.websocket_decision = Some(decision);
737 self
738 }
739
740 pub fn with_request_body_mutation(mut self, mutation: BodyMutation) -> Self {
742 self.request_body_mutation = Some(mutation);
743 self
744 }
745
746 pub fn with_response_body_mutation(mut self, mutation: BodyMutation) -> Self {
748 self.response_body_mutation = Some(mutation);
749 self
750 }
751
752 pub fn set_needs_more(mut self, needs_more: bool) -> Self {
754 self.needs_more = needs_more;
755 self
756 }
757
758 pub fn add_request_header(mut self, op: HeaderOp) -> Self {
760 self.request_headers.push(op);
761 self
762 }
763
764 pub fn add_response_header(mut self, op: HeaderOp) -> Self {
766 self.response_headers.push(op);
767 self
768 }
769
770 pub fn with_audit(mut self, audit: AuditMetadata) -> Self {
772 self.audit = audit;
773 self
774 }
775}
776
777#[derive(Debug, Clone, Default, Serialize, Deserialize)]
779pub struct AuditMetadata {
780 #[serde(default)]
782 pub tags: Vec<String>,
783 #[serde(default)]
785 pub rule_ids: Vec<String>,
786 pub confidence: Option<f32>,
788 #[serde(default)]
790 pub reason_codes: Vec<String>,
791 #[serde(default)]
793 pub custom: HashMap<String, serde_json::Value>,
794}
795
796#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
802#[serde(rename_all = "snake_case")]
803pub enum GuardrailInspectionType {
804 PromptInjection,
806 PiiDetection,
808}
809
810#[derive(Debug, Clone, Serialize, Deserialize)]
815pub struct GuardrailInspectEvent {
816 pub correlation_id: String,
818 pub inspection_type: GuardrailInspectionType,
820 pub content: String,
822 #[serde(skip_serializing_if = "Option::is_none")]
824 pub model: Option<String>,
825 #[serde(default)]
828 pub categories: Vec<String>,
829 #[serde(skip_serializing_if = "Option::is_none")]
831 pub route_id: Option<String>,
832 #[serde(default)]
834 pub metadata: HashMap<String, String>,
835}
836
837#[derive(Debug, Clone, Serialize, Deserialize)]
839pub struct GuardrailResponse {
840 pub detected: bool,
842 #[serde(default)]
844 pub confidence: f64,
845 #[serde(default)]
847 pub detections: Vec<GuardrailDetection>,
848 #[serde(skip_serializing_if = "Option::is_none")]
850 pub redacted_content: Option<String>,
851}
852
853impl Default for GuardrailResponse {
854 fn default() -> Self {
855 Self {
856 detected: false,
857 confidence: 0.0,
858 detections: Vec::new(),
859 redacted_content: None,
860 }
861 }
862}
863
864impl GuardrailResponse {
865 pub fn clean() -> Self {
867 Self::default()
868 }
869
870 pub fn with_detection(detection: GuardrailDetection) -> Self {
872 Self {
873 detected: true,
874 confidence: detection.confidence.unwrap_or(1.0),
875 detections: vec![detection],
876 redacted_content: None,
877 }
878 }
879
880 pub fn add_detection(&mut self, detection: GuardrailDetection) {
882 self.detected = true;
883 if let Some(conf) = detection.confidence {
884 self.confidence = self.confidence.max(conf);
885 }
886 self.detections.push(detection);
887 }
888}
889
890#[derive(Debug, Clone, Serialize, Deserialize)]
892pub struct GuardrailDetection {
893 pub category: String,
895 pub description: String,
897 #[serde(default)]
899 pub severity: DetectionSeverity,
900 #[serde(skip_serializing_if = "Option::is_none")]
902 pub confidence: Option<f64>,
903 #[serde(skip_serializing_if = "Option::is_none")]
905 pub span: Option<TextSpan>,
906}
907
908impl GuardrailDetection {
909 pub fn new(category: impl Into<String>, description: impl Into<String>) -> Self {
911 Self {
912 category: category.into(),
913 description: description.into(),
914 severity: DetectionSeverity::Medium,
915 confidence: None,
916 span: None,
917 }
918 }
919
920 pub fn with_severity(mut self, severity: DetectionSeverity) -> Self {
922 self.severity = severity;
923 self
924 }
925
926 pub fn with_confidence(mut self, confidence: f64) -> Self {
928 self.confidence = Some(confidence);
929 self
930 }
931
932 pub fn with_span(mut self, start: usize, end: usize) -> Self {
934 self.span = Some(TextSpan { start, end });
935 self
936 }
937}
938
939#[derive(Debug, Clone, Serialize, Deserialize)]
941pub struct TextSpan {
942 pub start: usize,
944 pub end: usize,
946}
947
948#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
950#[serde(rename_all = "lowercase")]
951pub enum DetectionSeverity {
952 Low,
954 #[default]
956 Medium,
957 High,
959 Critical,
961}