trojan_proto/
lib.rs

1//! Trojan protocol parsing and serialization.
2//!
3//! This module provides zero-copy parsers for trojan request headers and UDP packets.
4//! It is intentionally minimal and DRY: address parsing is shared by request and UDP paths.
5
6use bytes::BytesMut;
7
8pub const HASH_LEN: usize = 56;
9pub const CRLF: &[u8; 2] = b"\r\n";
10
11pub const CMD_CONNECT: u8 = 0x01;
12pub const CMD_UDP_ASSOCIATE: u8 = 0x03;
13/// Mux command for trojan-go multiplexing extension.
14pub const CMD_MUX: u8 = 0x7f;
15
16/// Maximum UDP payload size (8 KiB, consistent with trojan-go).
17pub const MAX_UDP_PAYLOAD: usize = 8 * 1024;
18/// Maximum domain name length.
19pub const MAX_DOMAIN_LEN: usize = 255;
20
21pub const ATYP_IPV4: u8 = 0x01;
22pub const ATYP_DOMAIN: u8 = 0x03;
23pub const ATYP_IPV6: u8 = 0x04;
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum ParseError {
27    InvalidCrlf,
28    InvalidCommand,
29    InvalidAtyp,
30    InvalidDomainLen,
31    InvalidUtf8,
32    /// Hash contains non-hex characters (expected lowercase a-f, 0-9).
33    InvalidHashFormat,
34}
35
36/// Errors that can occur when writing protocol data.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum WriteError {
39    /// Payload exceeds maximum allowed size (65535 bytes for UDP).
40    PayloadTooLarge,
41    /// Domain name exceeds maximum length (255 bytes).
42    DomainTooLong,
43    /// Hash must be exactly 56 bytes.
44    InvalidHashLen,
45}
46
47/// Parse result for incremental parsing.
48///
49/// - `Complete(T)` - parsing succeeded, contains the parsed value.
50/// - `Incomplete(n)` - buffer too small; `n` is the **minimum total bytes** needed
51///   (not the additional bytes needed). Caller should accumulate more data and retry.
52/// - `Invalid(e)` - protocol violation, connection should be rejected or redirected.
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub enum ParseResult<T> {
55    Complete(T),
56    Incomplete(usize),
57    Invalid(ParseError),
58}
59
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub enum HostRef<'a> {
62    Ipv4([u8; 4]),
63    Ipv6([u8; 16]),
64    Domain(&'a [u8]),
65}
66
67#[derive(Debug, Clone, PartialEq, Eq)]
68pub struct AddressRef<'a> {
69    pub host: HostRef<'a>,
70    pub port: u16,
71}
72
73#[derive(Debug, Clone, PartialEq, Eq)]
74pub struct TrojanRequest<'a> {
75    pub hash: &'a [u8],
76    pub command: u8,
77    pub address: AddressRef<'a>,
78    pub header_len: usize,
79    pub payload: &'a [u8],
80}
81
82#[derive(Debug, Clone, PartialEq, Eq)]
83pub struct UdpPacket<'a> {
84    pub address: AddressRef<'a>,
85    pub length: usize,
86    pub packet_len: usize,
87    pub payload: &'a [u8],
88}
89
90/// Validates that the hash is a valid hex string (a-f, A-F, 0-9).
91#[inline]
92pub fn is_valid_hash(hash: &[u8]) -> bool {
93    hash.len() == HASH_LEN && hash.iter().all(|&b| b.is_ascii_hexdigit())
94}
95
96#[inline]
97pub fn parse_request(buf: &[u8]) -> ParseResult<TrojanRequest<'_>> {
98    if buf.len() < HASH_LEN {
99        return ParseResult::Incomplete(HASH_LEN);
100    }
101
102    let hash = &buf[..HASH_LEN];
103    if !is_valid_hash(hash) {
104        return ParseResult::Invalid(ParseError::InvalidHashFormat);
105    }
106    let mut offset = HASH_LEN;
107
108    if let Some(res) = expect_crlf(buf, offset) {
109        return res;
110    }
111    offset += 2;
112
113    if buf.len() < offset + 2 {
114        return ParseResult::Incomplete(offset + 2);
115    }
116    let command = buf[offset];
117    if command != CMD_CONNECT && command != CMD_UDP_ASSOCIATE {
118        return ParseResult::Invalid(ParseError::InvalidCommand);
119    }
120    let atyp = buf[offset + 1];
121    offset += 2;
122
123    let addr_res = parse_address(atyp, &buf[offset..]);
124    let (address, addr_len) = match addr_res {
125        ParseResult::Complete(v) => v,
126        ParseResult::Incomplete(n) => return ParseResult::Incomplete(offset + n),
127        ParseResult::Invalid(e) => return ParseResult::Invalid(e),
128    };
129    offset += addr_len;
130
131    if let Some(res) = expect_crlf(buf, offset) {
132        return res;
133    }
134    offset += 2;
135
136    ParseResult::Complete(TrojanRequest {
137        hash,
138        command,
139        address,
140        header_len: offset,
141        payload: &buf[offset..],
142    })
143}
144
145#[inline]
146pub fn parse_udp_packet(buf: &[u8]) -> ParseResult<UdpPacket<'_>> {
147    if buf.is_empty() {
148        return ParseResult::Incomplete(1);
149    }
150    let atyp = buf[0];
151    let addr_res = parse_address(atyp, &buf[1..]);
152    let (address, addr_len) = match addr_res {
153        ParseResult::Complete(v) => v,
154        ParseResult::Incomplete(n) => return ParseResult::Incomplete(1 + n),
155        ParseResult::Invalid(e) => return ParseResult::Invalid(e),
156    };
157
158    let mut offset = 1 + addr_len;
159    if buf.len() < offset + 2 {
160        return ParseResult::Incomplete(offset + 2);
161    }
162    let length = read_u16(&buf[offset..offset + 2]) as usize;
163    if buf.len() < offset + 4 {
164        return ParseResult::Incomplete(offset + 4);
165    }
166    if &buf[offset + 2..offset + 4] != CRLF {
167        return ParseResult::Invalid(ParseError::InvalidCrlf);
168    }
169    offset += 4;
170    if buf.len() < offset + length {
171        return ParseResult::Incomplete(offset + length);
172    }
173
174    ParseResult::Complete(UdpPacket {
175        address,
176        length,
177        packet_len: offset + length,
178        payload: &buf[offset..offset + length],
179    })
180}
181
182/// Writes a Trojan request header to the buffer.
183///
184/// # Errors
185/// - `InvalidHashLen` if hash is not exactly 56 bytes.
186/// - `DomainTooLong` if address contains a domain longer than 255 bytes.
187pub fn write_request_header(
188    buf: &mut BytesMut,
189    hash_hex: &[u8],
190    command: u8,
191    address: &AddressRef<'_>,
192) -> Result<(), WriteError> {
193    if hash_hex.len() != HASH_LEN {
194        return Err(WriteError::InvalidHashLen);
195    }
196    if let HostRef::Domain(d) = &address.host
197        && d.len() > MAX_DOMAIN_LEN
198    {
199        return Err(WriteError::DomainTooLong);
200    }
201    buf.extend_from_slice(hash_hex);
202    buf.extend_from_slice(CRLF);
203    buf.extend_from_slice(&[command, address_atyp(address)]);
204    write_address_unchecked(buf, address);
205    buf.extend_from_slice(CRLF);
206    Ok(())
207}
208
209/// Writes a UDP packet to the buffer.
210///
211/// # Errors
212/// - `PayloadTooLarge` if payload exceeds 65535 bytes.
213/// - `DomainTooLong` if address contains a domain longer than 255 bytes.
214pub fn write_udp_packet(
215    buf: &mut BytesMut,
216    address: &AddressRef<'_>,
217    payload: &[u8],
218) -> Result<(), WriteError> {
219    if payload.len() > u16::MAX as usize {
220        return Err(WriteError::PayloadTooLarge);
221    }
222    if let HostRef::Domain(d) = &address.host
223        && d.len() > MAX_DOMAIN_LEN
224    {
225        return Err(WriteError::DomainTooLong);
226    }
227    buf.extend_from_slice(&[address_atyp(address)]);
228    write_address_unchecked(buf, address);
229    buf.extend_from_slice(&(payload.len() as u16).to_be_bytes());
230    buf.extend_from_slice(CRLF);
231    buf.extend_from_slice(payload);
232    Ok(())
233}
234
235#[inline]
236fn expect_crlf<T>(buf: &[u8], offset: usize) -> Option<ParseResult<T>> {
237    if buf.len() < offset + 2 {
238        return Some(ParseResult::Incomplete(offset + 2));
239    }
240    if &buf[offset..offset + 2] != CRLF {
241        return Some(ParseResult::Invalid(ParseError::InvalidCrlf));
242    }
243    None
244}
245
246#[inline]
247fn parse_address<'a>(atyp: u8, buf: &'a [u8]) -> ParseResult<(AddressRef<'a>, usize)> {
248    match atyp {
249        ATYP_IPV4 => {
250            if buf.len() < 6 {
251                return ParseResult::Incomplete(6);
252            }
253            let host = HostRef::Ipv4([buf[0], buf[1], buf[2], buf[3]]);
254            let port = read_u16(&buf[4..6]);
255            ParseResult::Complete((AddressRef { host, port }, 6))
256        }
257        ATYP_DOMAIN => {
258            if buf.is_empty() {
259                return ParseResult::Incomplete(1);
260            }
261            let len = buf[0] as usize;
262            if len == 0 {
263                return ParseResult::Invalid(ParseError::InvalidDomainLen);
264            }
265            let need = 1 + len + 2;
266            if buf.len() < need {
267                return ParseResult::Incomplete(need);
268            }
269            let domain = &buf[1..1 + len];
270            if std::str::from_utf8(domain).is_err() {
271                return ParseResult::Invalid(ParseError::InvalidUtf8);
272            }
273            let port = read_u16(&buf[1 + len..1 + len + 2]);
274            ParseResult::Complete((
275                AddressRef {
276                    host: HostRef::Domain(domain),
277                    port,
278                },
279                need,
280            ))
281        }
282        ATYP_IPV6 => {
283            if buf.len() < 18 {
284                return ParseResult::Incomplete(18);
285            }
286            let mut ip = [0u8; 16];
287            ip.copy_from_slice(&buf[0..16]);
288            let port = read_u16(&buf[16..18]);
289            ParseResult::Complete((
290                AddressRef {
291                    host: HostRef::Ipv6(ip),
292                    port,
293                },
294                18,
295            ))
296        }
297        _ => ParseResult::Invalid(ParseError::InvalidAtyp),
298    }
299}
300
301/// Writes address without validation. Caller must ensure domain length <= 255.
302fn write_address_unchecked(buf: &mut BytesMut, address: &AddressRef<'_>) {
303    match address.host {
304        HostRef::Ipv4(ip) => {
305            buf.extend_from_slice(&ip);
306        }
307        HostRef::Ipv6(ip) => {
308            buf.extend_from_slice(&ip);
309        }
310        HostRef::Domain(domain) => {
311            debug_assert!(domain.len() <= MAX_DOMAIN_LEN);
312            buf.extend_from_slice(&[domain.len() as u8]);
313            buf.extend_from_slice(domain);
314        }
315    }
316    buf.extend_from_slice(&address.port.to_be_bytes());
317}
318
319#[inline]
320fn address_atyp(address: &AddressRef<'_>) -> u8 {
321    match address.host {
322        HostRef::Ipv4(_) => ATYP_IPV4,
323        HostRef::Ipv6(_) => ATYP_IPV6,
324        HostRef::Domain(_) => ATYP_DOMAIN,
325    }
326}
327
328#[inline]
329fn read_u16(buf: &[u8]) -> u16 {
330    debug_assert!(buf.len() >= 2, "read_u16 requires at least 2 bytes");
331    u16::from_be_bytes([buf[0], buf[1]])
332}
333
334#[cfg(test)]
335mod tests {
336    use super::*;
337
338    fn sample_hash() -> [u8; HASH_LEN] {
339        [b'a'; HASH_LEN]
340    }
341
342    #[test]
343    fn test_is_valid_hash() {
344        // Valid lowercase hex
345        assert!(is_valid_hash(&[b'a'; HASH_LEN]));
346        assert!(is_valid_hash(
347            b"0123456789abcdef0123456789abcdef0123456789abcdef01234567"
348        ));
349
350        // Uppercase should be accepted
351        assert!(is_valid_hash(
352            b"0123456789ABCDEF0123456789abcdef0123456789abcdef01234567"
353        ));
354
355        // Invalid: wrong length
356        assert!(!is_valid_hash(&[b'a'; HASH_LEN - 1]));
357        assert!(!is_valid_hash(&[b'a'; HASH_LEN + 1]));
358
359        // Invalid: non-hex characters
360        let mut invalid = [b'a'; HASH_LEN];
361        invalid[0] = b'g';
362        assert!(!is_valid_hash(&invalid));
363    }
364
365    #[test]
366    fn parse_request_connect_ipv4() {
367        let addr = AddressRef {
368            host: HostRef::Ipv4([1, 2, 3, 4]),
369            port: 443,
370        };
371        let mut buf = BytesMut::new();
372        write_request_header(&mut buf, &sample_hash(), CMD_CONNECT, &addr).unwrap();
373        buf.extend_from_slice(b"hello");
374
375        let res = parse_request(&buf);
376        match res {
377            ParseResult::Complete(req) => {
378                assert_eq!(req.command, CMD_CONNECT);
379                assert_eq!(req.address, addr);
380                assert_eq!(req.payload, b"hello");
381            }
382            _ => panic!("unexpected parse result: {:?}", res),
383        }
384    }
385
386    #[test]
387    fn parse_request_invalid_hash() {
388        let addr = AddressRef {
389            host: HostRef::Ipv4([1, 2, 3, 4]),
390            port: 443,
391        };
392        let mut buf = BytesMut::new();
393        // Use non-hex which is invalid
394        let mut invalid_hash = [b'a'; HASH_LEN];
395        invalid_hash[0] = b'g';
396        write_request_header(&mut buf, &invalid_hash, CMD_CONNECT, &addr).unwrap();
397
398        let res = parse_request(&buf);
399        assert_eq!(res, ParseResult::Invalid(ParseError::InvalidHashFormat));
400    }
401
402    #[test]
403    fn parse_request_incomplete() {
404        let data = vec![b'a'; HASH_LEN - 1];
405        assert_eq!(parse_request(&data), ParseResult::Incomplete(HASH_LEN));
406    }
407
408    #[test]
409    fn parse_udp_packet_ipv4() {
410        let addr = AddressRef {
411            host: HostRef::Ipv4([8, 8, 8, 8]),
412            port: 53,
413        };
414        let mut buf = BytesMut::new();
415        write_udp_packet(&mut buf, &addr, b"ping").unwrap();
416        let res = parse_udp_packet(&buf);
417        match res {
418            ParseResult::Complete(pkt) => {
419                assert_eq!(pkt.address, addr);
420                assert_eq!(pkt.payload, b"ping");
421            }
422            _ => panic!("unexpected parse result: {:?}", res),
423        }
424    }
425
426    #[test]
427    fn write_udp_packet_payload_too_large() {
428        let addr = AddressRef {
429            host: HostRef::Ipv4([8, 8, 8, 8]),
430            port: 53,
431        };
432        let mut buf = BytesMut::new();
433        let large_payload = vec![0u8; u16::MAX as usize + 1];
434        let res = write_udp_packet(&mut buf, &addr, &large_payload);
435        assert_eq!(res, Err(WriteError::PayloadTooLarge));
436    }
437
438    #[test]
439    fn write_request_header_domain_too_long() {
440        let long_domain = vec![b'a'; 256];
441        let addr = AddressRef {
442            host: HostRef::Domain(&long_domain),
443            port: 443,
444        };
445        let mut buf = BytesMut::new();
446        let res = write_request_header(&mut buf, &sample_hash(), CMD_CONNECT, &addr);
447        assert_eq!(res, Err(WriteError::DomainTooLong));
448    }
449
450    #[test]
451    fn write_request_header_invalid_hash_len() {
452        let addr = AddressRef {
453            host: HostRef::Ipv4([1, 2, 3, 4]),
454            port: 443,
455        };
456        let mut buf = BytesMut::new();
457        let short_hash = [b'a'; HASH_LEN - 1];
458        let res = write_request_header(&mut buf, &short_hash, CMD_CONNECT, &addr);
459        assert_eq!(res, Err(WriteError::InvalidHashLen));
460    }
461}
trojan_proto/lib.rs

trojan_proto/
lib.rs
