Skip to main content

relay_core_api/
flow.rs

1use serde::{Deserialize, Serialize, Serializer, Deserializer};
2use uuid::Uuid;
3use chrono::{DateTime, Utc};
4use std::collections::HashMap;
5use url::Url;
6use std::sync::Arc;
7use std::any::Any;
8
9/// Trait for custom protocol layers to implement
10pub trait ProtocolLayer: Any + Send + Sync + std::fmt::Debug {
11    fn protocol_name(&self) -> &str;
12    fn as_any(&self) -> &dyn Any;
13    fn to_json(&self) -> serde_json::Value;
14}
15
16// Custom serializer for Arc<dyn ProtocolLayer>
17fn serialize_custom_layer<S>(layer: &Arc<dyn ProtocolLayer>, serializer: S) -> Result<S::Ok, S::Error>
18where
19    S: Serializer,
20{
21    use serde::ser::SerializeStruct;
22    let mut state = serializer.serialize_struct("CustomLayer", 2)?;
23    state.serialize_field("protocol", layer.protocol_name())?;
24    state.serialize_field("data", &layer.to_json())?;
25    state.end()
26}
27
28// Custom deserializer for Arc<dyn ProtocolLayer>
29fn deserialize_custom_layer<'de, D>(deserializer: D) -> Result<Arc<dyn ProtocolLayer>, D::Error>
30where
31    D: Deserializer<'de>,
32{
33    let value = serde_json::Value::deserialize(deserializer)?;
34    
35    // Try to extract protocol and data
36    let protocol = value.get("protocol")
37        .and_then(|v| v.as_str())
38        .unwrap_or("unknown")
39        .to_string();
40        
41    let data = value.get("data")
42        .cloned()
43        .unwrap_or(serde_json::Value::Null);
44
45    Ok(Arc::new(GenericProtocolLayer { 
46        protocol,
47        data 
48    }))
49}
50
51#[derive(Debug)]
52struct GenericProtocolLayer {
53    protocol: String,
54    data: serde_json::Value,
55}
56
57impl ProtocolLayer for GenericProtocolLayer {
58    fn protocol_name(&self) -> &str {
59        &self.protocol
60    }
61    fn as_any(&self) -> &dyn Any {
62        self
63    }
64    fn to_json(&self) -> serde_json::Value {
65        self.data.clone()
66    }
67}
68
69/// The central data structure representing a captured traffic flow.
70/// Designed to support L3/L4 layers initially, with L7 (HTTP) as a specialized layer.
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct Flow {
73    pub id: Uuid,
74    pub start_time: DateTime<Utc>,
75    pub end_time: Option<DateTime<Utc>>,
76    
77    /// L3/L4 Connection Info (IP, Port, Protocol)
78    pub network: NetworkInfo,
79    
80    /// The application layer protocol detected or parsed
81    pub layer: Layer,
82
83    /// Analysis tags (e.g., "error", "large-body", "injected")
84    pub tags: Vec<String>,
85    
86    /// Internal processing metadata (not necessarily for UI)
87    #[serde(skip)]
88    pub meta: HashMap<String, String>,
89}
90
91#[derive(Debug, Clone, Serialize, Deserialize)]
92pub struct NetworkInfo {
93    pub client_ip: String,
94    pub client_port: u16,
95    pub server_ip: String,
96    pub server_port: u16,
97    pub protocol: TransportProtocol,
98    pub tls: bool,
99    pub tls_version: Option<String>,
100    pub sni: Option<String>,
101}
102
103#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
104pub enum TransportProtocol {
105    TCP,
106    UDP,
107    ICMP,
108    Unknown,
109}
110
111/// The specific application layer data.
112/// Using an enum allows us to support HTTP now, but easily add DNS, WebSocket,
113/// or raw TCP streams later without breaking the top-level Flow structure.
114#[derive(Debug, Clone, Serialize, Deserialize)]
115#[serde(tag = "type", content = "data")]
116pub enum Layer {
117    Http(HttpLayer),
118    WebSocket(WebSocketLayer),
119    Tcp(TcpLayer), // For raw TCP flows without L7 parsing
120    Udp(UdpLayer), // For raw UDP flows
121    Quic(QuicLayer), // For QUIC flows
122    #[serde(serialize_with = "serialize_custom_layer", deserialize_with = "deserialize_custom_layer")]
123    Custom(Arc<dyn ProtocolLayer>),
124    Unknown,
125}
126
127#[derive(Debug, Clone, Serialize, Deserialize)]
128pub struct UdpLayer {
129    pub payload_size: usize,
130    pub packet_count: usize,
131}
132
133#[derive(Debug, Clone, Serialize, Deserialize)]
134pub struct QuicLayer {
135    pub sni: Option<String>,
136    pub alpn: Option<String>,
137    pub version: Option<String>,
138}
139
140// --- HTTP Layer ---
141
142#[derive(Debug, Clone, Serialize, Deserialize)]
143pub struct HttpLayer {
144    pub request: HttpRequest,
145    pub response: Option<HttpResponse>,
146    pub error: Option<String>,
147}
148
149#[derive(Debug, Clone, Serialize, Deserialize)]
150pub struct HttpRequest {
151    pub method: String,
152    pub url: Url,
153    pub version: String, // HTTP/1.1, HTTP/2
154    pub headers: Vec<(String, String)>, // Vec preserves order, unlike HashMap
155    pub cookies: Vec<Cookie>,
156    pub query: Vec<(String, String)>,
157    pub body: Option<BodyData>,
158}
159
160#[derive(Debug, Clone, Serialize, Deserialize)]
161pub struct HttpResponse {
162    pub status: u16,
163    pub status_text: String,
164    pub version: String,
165    pub headers: Vec<(String, String)>,
166    pub cookies: Vec<Cookie>,
167    pub body: Option<BodyData>,
168    pub timing: ResponseTiming,
169}
170
171#[derive(Debug, Clone, Serialize, Deserialize)]
172pub struct Cookie {
173    pub name: String,
174    pub value: String,
175    pub path: Option<String>,
176    pub domain: Option<String>,
177    pub expires: Option<String>,
178    pub http_only: Option<bool>,
179    pub secure: Option<bool>,
180}
181
182#[derive(Debug, Clone, Serialize, Deserialize)]
183pub struct ResponseTiming {
184    pub time_to_first_byte: Option<u64>, // ms
185    pub time_to_last_byte: Option<u64>,  // ms
186}
187
188// --- WebSocket Layer ---
189
190#[derive(Debug, Clone, Serialize, Deserialize)]
191pub struct WebSocketLayer {
192    pub handshake_request: HttpRequest,
193    pub handshake_response: HttpResponse,
194    pub messages: Vec<WebSocketMessage>,
195    pub closed: bool,
196}
197
198#[derive(Debug, Clone, Serialize, Deserialize)]
199pub struct WebSocketMessage {
200    pub id: Uuid,
201    pub timestamp: DateTime<Utc>,
202    pub direction: Direction, // ClientToServer, ServerToClient
203    pub content: BodyData,
204    pub opcode: String, // Text, Binary, Ping, Pong
205}
206
207#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
208pub enum Direction {
209    ClientToServer,
210    ServerToClient,
211}
212
213// --- Raw TCP Layer ---
214
215#[derive(Debug, Clone, Serialize, Deserialize)]
216pub struct TcpLayer {
217    pub bytes_up: u64,
218    pub bytes_down: u64,
219    // Future: packets or data chunks
220}
221
222// --- Shared ---
223
224#[derive(Debug, Clone, Serialize, Deserialize)]
225pub struct BodyData {
226    /// Encoding (e.g., "utf-8", "base64")
227    pub encoding: String,
228    /// The actual content. If binary, it should be base64 encoded string.
229    pub content: String, 
230    /// Original size in bytes before any decoding/unzipping
231    pub size: u64,
232}
233
234/// Enum for Incremental Flow Updates to avoid cloning the entire Flow object
235#[derive(Debug, Clone, Serialize, Deserialize)]
236#[serde(tag = "type", content = "data")]
237pub enum FlowUpdate {
238    /// Full Flow update (e.g., initial creation, request/response headers)
239    Full(Box<Flow>),
240    /// Incremental WebSocket Message
241    WebSocketMessage { 
242        flow_id: String, 
243        message: WebSocketMessage 
244    },
245    /// Incremental HTTP Body Update (Request or Response)
246    HttpBody {
247        flow_id: String,
248        direction: Direction,
249        body: BodyData,
250    },
251}
252