Skip to main content

trojan_protocol/
lib.rs

1//! WebSocket protocol message types shared between agent and panel.
2//!
3//! Messages are serialized with bincode for compact binary framing.
4//! The `config` fields carry opaque JSON bytes — they are not interpreted
5//! at the protocol layer and are parsed into typed configs by the runner.
6
7use serde::{Deserialize, Serialize};
8
9/// Protocol version — incremented on breaking changes.
10pub const PROTOCOL_VERSION: u32 = 1;
11
12/// Agent -> Panel messages.
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub enum AgentMessage {
15    Register {
16        protocol_version: u32,
17        token: String,
18        version: String,
19        hostname: String,
20        os: String,
21        arch: String,
22    },
23    Heartbeat {
24        connections_active: u32,
25        bytes_in: u64,
26        bytes_out: u64,
27        uptime_secs: u64,
28        memory_rss_bytes: Option<u64>,
29        cpu_usage_percent: Option<f32>,
30    },
31    Traffic {
32        records: Vec<TrafficRecord>,
33    },
34    ConfigAck {
35        version: u32,
36        ok: bool,
37        message: Option<String>,
38    },
39    ServiceStatus {
40        status: ServiceState,
41        started_at: u64,
42        config_version: u32,
43    },
44    Pong,
45}
46
47/// Panel -> Agent messages.
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub enum PanelMessage {
50    Registered {
51        node_id: String,
52        node_type: NodeType,
53        config_version: u32,
54        report_interval_secs: u32,
55        /// Opaque JSON bytes of the service config.
56        config: Vec<u8>,
57    },
58    ConfigPush {
59        version: u32,
60        restart_required: bool,
61        drain_timeout_secs: Option<u32>,
62        /// Opaque JSON bytes of the service config.
63        config: Vec<u8>,
64    },
65    Ping,
66    Error {
67        code: ErrorCode,
68        message: String,
69    },
70}
71
72/// Per-user traffic delta record.
73#[derive(Debug, Clone, Serialize, Deserialize)]
74pub struct TrafficRecord {
75    pub user_id: String,
76    pub bytes_up: u64,
77    pub bytes_down: u64,
78}
79
80/// Node type — determines which service the agent boots.
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
82pub enum NodeType {
83    Server,
84    Entry,
85    Relay,
86}
87
88impl std::fmt::Display for NodeType {
89    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
90        match self {
91            NodeType::Server => write!(f, "server"),
92            NodeType::Entry => write!(f, "entry"),
93            NodeType::Relay => write!(f, "relay"),
94        }
95    }
96}
97
98/// Service runtime state.
99#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
100pub enum ServiceState {
101    Starting,
102    Running,
103    Restarting,
104    Stopped,
105    Error,
106}
107
108/// Error codes from the panel.
109#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
110pub enum ErrorCode {
111    InvalidToken,
112    NodeDisabled,
113    NodeNotFound,
114    ProtocolMismatch,
115    RateLimited,
116    InternalError,
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn agent_message_register_roundtrip() {
125        let msg = AgentMessage::Register {
126            protocol_version: PROTOCOL_VERSION,
127            token: "test-token".to_string(),
128            version: "0.5.0".to_string(),
129            hostname: "node-01".to_string(),
130            os: "linux".to_string(),
131            arch: "x86_64".to_string(),
132        };
133        let bytes = bincode::serialize(&msg).unwrap();
134        let decoded: AgentMessage = bincode::deserialize(&bytes).unwrap();
135        match decoded {
136            AgentMessage::Register {
137                protocol_version,
138                token,
139                version,
140                ..
141            } => {
142                assert_eq!(protocol_version, PROTOCOL_VERSION);
143                assert_eq!(token, "test-token");
144                assert_eq!(version, "0.5.0");
145            }
146            _ => panic!("expected Register"),
147        }
148    }
149
150    #[test]
151    fn panel_message_registered_roundtrip() {
152        let config_json = serde_json::json!({"server": {"listen": "0.0.0.0:443"}});
153        let config_bytes = serde_json::to_vec(&config_json).unwrap();
154
155        let msg = PanelMessage::Registered {
156            node_id: "hk-01".to_string(),
157            node_type: NodeType::Server,
158            config_version: 17,
159            report_interval_secs: 30,
160            config: config_bytes,
161        };
162        let bytes = bincode::serialize(&msg).unwrap();
163        let decoded: PanelMessage = bincode::deserialize(&bytes).unwrap();
164        match decoded {
165            PanelMessage::Registered {
166                node_id,
167                node_type,
168                config_version,
169                config,
170                ..
171            } => {
172                assert_eq!(node_id, "hk-01");
173                assert_eq!(node_type, NodeType::Server);
174                assert_eq!(config_version, 17);
175                let val: serde_json::Value = serde_json::from_slice(&config).unwrap();
176                assert_eq!(val["server"]["listen"], "0.0.0.0:443");
177            }
178            _ => panic!("expected Registered"),
179        }
180    }
181
182    #[test]
183    fn panel_message_ping_roundtrip() {
184        let msg = PanelMessage::Ping;
185        let bytes = bincode::serialize(&msg).unwrap();
186        let decoded: PanelMessage = bincode::deserialize(&bytes).unwrap();
187        assert!(matches!(decoded, PanelMessage::Ping));
188    }
189
190    #[test]
191    fn panel_message_error_roundtrip() {
192        let msg = PanelMessage::Error {
193            code: ErrorCode::InvalidToken,
194            message: "bad token".to_string(),
195        };
196        let bytes = bincode::serialize(&msg).unwrap();
197        let decoded: PanelMessage = bincode::deserialize(&bytes).unwrap();
198        match decoded {
199            PanelMessage::Error { code, message } => {
200                assert_eq!(code, ErrorCode::InvalidToken);
201                assert_eq!(message, "bad token");
202            }
203            _ => panic!("expected Error"),
204        }
205    }
206
207    #[test]
208    fn agent_message_heartbeat_roundtrip() {
209        let msg = AgentMessage::Heartbeat {
210            connections_active: 42,
211            bytes_in: 1000,
212            bytes_out: 2000,
213            uptime_secs: 3600,
214            memory_rss_bytes: Some(52_428_800),
215            cpu_usage_percent: Some(12.5),
216        };
217        let bytes = bincode::serialize(&msg).unwrap();
218        let decoded: AgentMessage = bincode::deserialize(&bytes).unwrap();
219        match decoded {
220            AgentMessage::Heartbeat {
221                connections_active,
222                uptime_secs,
223                memory_rss_bytes,
224                ..
225            } => {
226                assert_eq!(connections_active, 42);
227                assert_eq!(uptime_secs, 3600);
228                assert_eq!(memory_rss_bytes, Some(52_428_800));
229            }
230            _ => panic!("expected Heartbeat"),
231        }
232    }
233
234    #[test]
235    fn traffic_record_roundtrip() {
236        let record = TrafficRecord {
237            user_id: "alice".to_string(),
238            bytes_up: 1024,
239            bytes_down: 2048,
240        };
241        let bytes = bincode::serialize(&record).unwrap();
242        let decoded: TrafficRecord = bincode::deserialize(&bytes).unwrap();
243        assert_eq!(decoded.user_id, "alice");
244        assert_eq!(decoded.bytes_up, 1024);
245        assert_eq!(decoded.bytes_down, 2048);
246    }
247
248    #[test]
249    fn node_type_display() {
250        assert_eq!(NodeType::Server.to_string(), "server");
251        assert_eq!(NodeType::Entry.to_string(), "entry");
252        assert_eq!(NodeType::Relay.to_string(), "relay");
253    }
254
255    #[test]
256    fn config_push_roundtrip() {
257        let config_bytes = serde_json::to_vec(&serde_json::json!({})).unwrap();
258        let msg = PanelMessage::ConfigPush {
259            version: 18,
260            restart_required: true,
261            drain_timeout_secs: Some(30),
262            config: config_bytes,
263        };
264        let bytes = bincode::serialize(&msg).unwrap();
265        let decoded: PanelMessage = bincode::deserialize(&bytes).unwrap();
266        match decoded {
267            PanelMessage::ConfigPush {
268                version,
269                restart_required,
270                drain_timeout_secs,
271                ..
272            } => {
273                assert_eq!(version, 18);
274                assert!(restart_required);
275                assert_eq!(drain_timeout_secs, Some(30));
276            }
277            _ => panic!("expected ConfigPush"),
278        }
279    }
280
281    #[test]
282    fn pong_roundtrip() {
283        let msg = AgentMessage::Pong;
284        let bytes = bincode::serialize(&msg).unwrap();
285        let decoded: AgentMessage = bincode::deserialize(&bytes).unwrap();
286        assert!(matches!(decoded, AgentMessage::Pong));
287    }
288
289    #[test]
290    fn bincode_is_compact() {
291        // Verify bincode produces smaller output than JSON for typical messages
292        let msg = AgentMessage::Heartbeat {
293            connections_active: 42,
294            bytes_in: 123_456_789,
295            bytes_out: 987_654_321,
296            uptime_secs: 86400,
297            memory_rss_bytes: Some(52_428_800),
298            cpu_usage_percent: Some(12.5),
299        };
300        let bincode_bytes = bincode::serialize(&msg).unwrap();
301        let json_bytes = serde_json::to_vec(&msg).unwrap();
302        assert!(bincode_bytes.len() < json_bytes.len());
303    }
304}