1use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum TunnelType {
12 Tcp = 0,
14 Udp = 1,
16 Http = 2,
18}
19
20impl TunnelType {
21 pub fn as_str(&self) -> &'static str {
23 match self {
24 Self::Tcp => "tcp",
25 Self::Udp => "udp",
26 Self::Http => "http",
27 }
28 }
29
30 pub fn from_u8(v: u8) -> Option<Self> {
32 match v {
33 0 => Some(Self::Tcp),
34 1 => Some(Self::Udp),
35 2 => Some(Self::Http),
36 _ => None,
37 }
38 }
39
40 pub fn to_u8(self) -> u8 {
42 self as u8
43 }
44
45 pub fn parse(s: &str) -> Self {
47 match s {
48 "tcp" => Self::Tcp,
49 "udp" => Self::Udp,
50 "http" => Self::Http,
51 _ => Self::Tcp,
52 }
53 }
54}
55
56impl std::fmt::Display for TunnelType {
57 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58 write!(f, "{}", self.as_str())
59 }
60}
61
62impl Serialize for TunnelType {
63 fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
64 s.serialize_str(self.as_str())
65 }
66}
67
68impl<'de> Deserialize<'de> for TunnelType {
69 fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
70 let s = String::deserialize(d)?;
71 match s.as_str() {
72 "tcp" => Ok(Self::Tcp),
73 "udp" => Ok(Self::Udp),
74 "http" => Ok(Self::Http),
75 _ => Err(serde::de::Error::custom(format!("unknown tunnel type: {s}"))),
76 }
77 }
78}
79
80pub const HANDSHAKE_SIZE: usize = 9;
84
85pub const STATUS_OK: u8 = 0x01;
87
88pub const STATUS_ERR: u8 = 0x00;
90
91pub const SIGNAL_NEW_CONN: u8 = 0x02;
93
94pub const UDP_FRAME_HEADER: usize = 2;
96
97pub fn encode_handshake(token: u64, tunnel_type: TunnelType) -> [u8; HANDSHAKE_SIZE] {
101 let mut buf = [0u8; HANDSHAKE_SIZE];
102 buf[..8].copy_from_slice(&token.to_be_bytes());
103 buf[8] = tunnel_type.to_u8();
104 buf
105}
106
107pub fn decode_handshake(buf: &[u8; HANDSHAKE_SIZE]) -> (u64, TunnelType) {
109 let token = u64::from_be_bytes(buf[..8].try_into().unwrap());
110 let tt = TunnelType::from_u8(buf[8]).unwrap_or(TunnelType::Tcp);
111 (token, tt)
112}
113
114pub fn encode_udp_frame(data: &[u8]) -> Vec<u8> {
116 let mut frame = Vec::with_capacity(2 + data.len());
117 frame.extend_from_slice(&(data.len() as u16).to_be_bytes());
118 frame.extend_from_slice(data);
119 frame
120}
121
122#[derive(Debug, Serialize, Deserialize)]
124pub struct TunnelListEntry {
125 #[serde(with = "serde_u64_str")]
127 pub token: u64,
128 pub tunnel_type: String,
130 pub public_port: u16,
132 pub key_name: String,
134 pub connected: bool,
136 pub pending_count: usize,
138}
139
140pub fn decode_udp_header(buf: &[u8; UDP_FRAME_HEADER]) -> u16 {
142 u16::from_be_bytes(*buf)
143}
144
145#[derive(Debug, Serialize, Deserialize)]
151pub struct ApiResponse<T: Serialize> {
152 pub success: bool,
154 #[serde(skip_serializing_if = "Option::is_none")]
156 pub data: Option<T>,
157 #[serde(skip_serializing_if = "Option::is_none")]
159 pub error: Option<String>,
160}
161
162impl<T: Serialize> ApiResponse<T> {
163 pub fn ok(data: T) -> Self {
165 Self { success: true, data: Some(data), error: None }
166 }
167 pub fn err(msg: impl Into<String>) -> Self {
169 Self { success: false, data: None, error: Some(msg.into()) }
170 }
171}
172
173pub fn api_ok() -> ApiResponse<()> {
175 ApiResponse { success: true, data: Some(()), error: None }
176}
177
178#[derive(Debug, Serialize, Deserialize)]
180pub struct CreateTunnelRequest {
181 #[serde(rename = "type")]
183 pub tunnel_type: TunnelType,
184 #[serde(default)]
186 pub local_port: Option<u16>,
187}
188
189#[derive(Debug, Serialize, Deserialize)]
191pub struct TunnelResponse {
192 #[serde(with = "serde_u64_str")]
194 pub token: u64,
195 pub public_port: u16,
197 pub data_port: u16,
199 pub tunnel_type: String,
201}
202
203#[derive(Debug, Serialize, Deserialize)]
205pub struct CreateKeyRequest {
206 pub name: String,
208}
209
210#[derive(Debug, Serialize, Deserialize)]
212pub struct CreateKeyResponse {
213 pub key: String,
215 pub name: String,
217}
218
219#[derive(Debug, Serialize, Deserialize, Clone)]
221pub struct KeyInfo {
222 pub key: String,
224 pub name: String,
226}
227
228use rand::Rng;
231
232pub fn generate_token() -> u64 {
234 rand::thread_rng().gen()
235}
236
237pub fn generate_api_key() -> String {
239 let id = uuid::Uuid::new_v4();
240 format!("sk-{}", id.to_string().replace('-', "")[..24].to_string())
241}
242
243pub mod serde_u64_str {
253 use serde::{Deserialize, Deserializer, Serializer};
254
255 pub fn serialize<S: Serializer>(val: &u64, s: S) -> Result<S::Ok, S::Error> {
257 s.collect_str(val)
258 }
259
260 pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<u64, D::Error> {
262 let s = String::deserialize(d)?;
263 s.parse().map_err(serde::de::Error::custom)
264 }
265}
266
267#[cfg(test)]
268mod tests {
269 use super::*;
270
271 #[test]
272 fn tunnel_type_from_u8_roundtrip() {
273 for (u, expected) in [(0, TunnelType::Tcp), (1, TunnelType::Udp), (2, TunnelType::Http)] {
274 let tt = TunnelType::from_u8(u).unwrap();
275 assert_eq!(tt, expected);
276 assert_eq!(tt.to_u8(), u);
277 }
278 }
279
280 #[test]
281 fn tunnel_type_from_u8_out_of_range() {
282 assert!(TunnelType::from_u8(3).is_none());
283 assert!(TunnelType::from_u8(255).is_none());
284 }
285
286 #[test]
287 fn tunnel_type_as_str() {
288 assert_eq!(TunnelType::Tcp.as_str(), "tcp");
289 assert_eq!(TunnelType::Udp.as_str(), "udp");
290 assert_eq!(TunnelType::Http.as_str(), "http");
291 }
292
293 #[test]
294 fn tunnel_type_display() {
295 assert_eq!(format!("{}", TunnelType::Tcp), "tcp");
296 assert_eq!(format!("{}", TunnelType::Udp), "udp");
297 assert_eq!(format!("{}", TunnelType::Http), "http");
298 }
299
300 #[test]
301 fn tunnel_type_parse() {
302 assert_eq!(TunnelType::parse("tcp"), TunnelType::Tcp);
303 assert_eq!(TunnelType::parse("udp"), TunnelType::Udp);
304 assert_eq!(TunnelType::parse("http"), TunnelType::Http);
305 assert_eq!(TunnelType::parse("bluetooth"), TunnelType::Tcp);
306 }
307
308 #[test]
309 fn tunnel_type_serde_json() {
310 let j = serde_json::to_string(&TunnelType::Tcp).unwrap();
311 assert_eq!(j, "\"tcp\"");
312 let tt: TunnelType = serde_json::from_str("\"udp\"").unwrap();
313 assert_eq!(tt, TunnelType::Udp);
314 }
315
316 #[test]
317 fn tunnel_type_serde_unknown_errs() {
318 let r: Result<TunnelType, _> = serde_json::from_str("\"bluetooth\"");
319 assert!(r.is_err());
320 }
321
322 #[test]
323 fn handshake_roundtrip() {
324 let token = 0xDEAD_BEEF_CAFE_F00D;
325 for tt in [TunnelType::Tcp, TunnelType::Udp, TunnelType::Http] {
326 let encoded = encode_handshake(token, tt);
327 assert_eq!(encoded.len(), HANDSHAKE_SIZE);
328 let (decoded_token, decoded_tt) = decode_handshake(&encoded);
329 assert_eq!(decoded_token, token);
330 assert_eq!(decoded_tt, tt);
331 }
332 }
333
334 #[test]
335 fn handshake_token_zero() {
336 let encoded = encode_handshake(0, TunnelType::Tcp);
337 let (token, tt) = decode_handshake(&encoded);
338 assert_eq!(token, 0);
339 assert_eq!(tt, TunnelType::Tcp);
340 }
341
342 #[test]
343 fn udp_frame_roundtrip() {
344 let payload = b"hello UDP";
345 let frame = encode_udp_frame(payload);
346 assert_eq!(frame.len(), 2 + payload.len());
347 let mut header = [0u8; UDP_FRAME_HEADER];
348 header.copy_from_slice(&frame[..2]);
349 let len = decode_udp_header(&header);
350 assert_eq!(len as usize, payload.len());
351 assert_eq!(&frame[2..], payload);
352 }
353
354 #[test]
355 fn udp_frame_empty_payload() {
356 let frame = encode_udp_frame(&[]);
357 assert_eq!(frame.len(), 2);
358 let mut header = [0u8; UDP_FRAME_HEADER];
359 header.copy_from_slice(&frame[..2]);
360 assert_eq!(decode_udp_header(&header), 0);
361 }
362
363 #[test]
364 fn constants() {
365 assert_eq!(HANDSHAKE_SIZE, 9);
366 assert_eq!(STATUS_OK, 0x01);
367 assert_eq!(STATUS_ERR, 0x00);
368 assert_eq!(SIGNAL_NEW_CONN, 0x02);
369 assert_eq!(UDP_FRAME_HEADER, 2);
370 }
371
372 #[test]
373 fn generate_token_is_nonzero() {
374 for _ in 0..100 {
375 let t = generate_token();
376 if t == 0 {
377 panic!("generated zero token");
378 }
379 }
380 }
381
382 #[test]
383 fn generate_api_key_format() {
384 let key = generate_api_key();
385 assert!(key.starts_with("sk-"));
386 assert_eq!(key.len(), 27);
387 let hex_part = &key[3..];
388 assert!(hex_part.chars().all(|c| c.is_ascii_hexdigit()));
389 }
390
391 #[test]
392 fn api_ok_response() {
393 let r = api_ok();
394 assert!(r.success);
395 assert!(r.data.is_some());
396 assert!(r.error.is_none());
397 }
398
399 #[test]
400 fn api_err_response() {
401 let r: ApiResponse<()> = ApiResponse::err("something broke");
402 assert!(!r.success);
403 assert!(r.data.is_none());
404 assert_eq!(r.error.unwrap(), "something broke");
405 }
406
407 #[test]
408 fn create_tunnel_request_serde() {
409 let req = CreateTunnelRequest {
410 tunnel_type: TunnelType::Tcp,
411 local_port: Some(4000),
412 };
413 let j = serde_json::to_string(&req).unwrap();
414 assert!(j.contains("\"type\":\"tcp\""));
415 assert!(j.contains("\"local_port\":4000"));
416 let parsed: CreateTunnelRequest = serde_json::from_str(&j).unwrap();
417 assert_eq!(parsed.tunnel_type, TunnelType::Tcp);
418 assert_eq!(parsed.local_port, Some(4000));
419 }
420
421 #[test]
422 fn create_tunnel_request_default_local_port() {
423 let j = r#"{"type":"udp"}"#;
424 let req: CreateTunnelRequest = serde_json::from_str(j).unwrap();
425 assert_eq!(req.tunnel_type, TunnelType::Udp);
426 assert_eq!(req.local_port, None);
427 }
428
429 #[test]
430 fn tunnel_response_serde() {
431 let resp = TunnelResponse {
432 token: 42,
433 public_port: 10001,
434 data_port: 8081,
435 tunnel_type: "tcp".into(),
436 };
437 let j = serde_json::to_string(&resp).unwrap();
438 let parsed: TunnelResponse = serde_json::from_str(&j).unwrap();
439 assert_eq!(parsed.token, 42);
440 assert_eq!(parsed.public_port, 10001);
441 assert_eq!(parsed.data_port, 8081);
442 assert_eq!(parsed.tunnel_type, "tcp");
443 }
444}