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((AddressRef { host: HostRef::Domain(domain), port }, need))
275        }
276        ATYP_IPV6 => {
277            if buf.len() < 18 {
278                return ParseResult::Incomplete(18);
279            }
280            let mut ip = [0u8; 16];
281            ip.copy_from_slice(&buf[0..16]);
282            let port = read_u16(&buf[16..18]);
283            ParseResult::Complete((AddressRef { host: HostRef::Ipv6(ip), port }, 18))
284        }
285        _ => ParseResult::Invalid(ParseError::InvalidAtyp),
286    }
287}
288
289/// Writes address without validation. Caller must ensure domain length <= 255.
290fn write_address_unchecked(buf: &mut BytesMut, address: &AddressRef<'_>) {
291    match address.host {
292        HostRef::Ipv4(ip) => {
293            buf.extend_from_slice(&ip);
294        }
295        HostRef::Ipv6(ip) => {
296            buf.extend_from_slice(&ip);
297        }
298        HostRef::Domain(domain) => {
299            debug_assert!(domain.len() <= MAX_DOMAIN_LEN);
300            buf.extend_from_slice(&[domain.len() as u8]);
301            buf.extend_from_slice(domain);
302        }
303    }
304    buf.extend_from_slice(&address.port.to_be_bytes());
305}
306
307#[inline]
308fn address_atyp(address: &AddressRef<'_>) -> u8 {
309    match address.host {
310        HostRef::Ipv4(_) => ATYP_IPV4,
311        HostRef::Ipv6(_) => ATYP_IPV6,
312        HostRef::Domain(_) => ATYP_DOMAIN,
313    }
314}
315
316#[inline]
317fn read_u16(buf: &[u8]) -> u16 {
318    debug_assert!(buf.len() >= 2, "read_u16 requires at least 2 bytes");
319    u16::from_be_bytes([buf[0], buf[1]])
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325
326    fn sample_hash() -> [u8; HASH_LEN] {
327        [b'a'; HASH_LEN]
328    }
329
330    #[test]
331    fn test_is_valid_hash() {
332        // Valid lowercase hex
333        assert!(is_valid_hash(&[b'a'; HASH_LEN]));
334        assert!(is_valid_hash(b"0123456789abcdef0123456789abcdef0123456789abcdef01234567"));
335
336        // Uppercase should be accepted
337        assert!(is_valid_hash(b"0123456789ABCDEF0123456789abcdef0123456789abcdef01234567"));
338
339        // Invalid: wrong length
340        assert!(!is_valid_hash(&[b'a'; HASH_LEN - 1]));
341        assert!(!is_valid_hash(&[b'a'; HASH_LEN + 1]));
342
343        // Invalid: non-hex characters
344        let mut invalid = [b'a'; HASH_LEN];
345        invalid[0] = b'g';
346        assert!(!is_valid_hash(&invalid));
347    }
348
349    #[test]
350    fn parse_request_connect_ipv4() {
351        let addr = AddressRef {
352            host: HostRef::Ipv4([1, 2, 3, 4]),
353            port: 443,
354        };
355        let mut buf = BytesMut::new();
356        write_request_header(&mut buf, &sample_hash(), CMD_CONNECT, &addr).unwrap();
357        buf.extend_from_slice(b"hello");
358
359        let res = parse_request(&buf);
360        match res {
361            ParseResult::Complete(req) => {
362                assert_eq!(req.command, CMD_CONNECT);
363                assert_eq!(req.address, addr);
364                assert_eq!(req.payload, b"hello");
365            }
366            _ => panic!("unexpected parse result: {:?}", res),
367        }
368    }
369
370    #[test]
371    fn parse_request_invalid_hash() {
372        let addr = AddressRef {
373            host: HostRef::Ipv4([1, 2, 3, 4]),
374            port: 443,
375        };
376        let mut buf = BytesMut::new();
377        // Use non-hex which is invalid
378        let mut invalid_hash = [b'a'; HASH_LEN];
379        invalid_hash[0] = b'g';
380        write_request_header(&mut buf, &invalid_hash, CMD_CONNECT, &addr).unwrap();
381
382        let res = parse_request(&buf);
383        assert_eq!(res, ParseResult::Invalid(ParseError::InvalidHashFormat));
384    }
385
386    #[test]
387    fn parse_request_incomplete() {
388        let data = vec![b'a'; HASH_LEN - 1];
389        assert_eq!(parse_request(&data), ParseResult::Incomplete(HASH_LEN));
390    }
391
392    #[test]
393    fn parse_udp_packet_ipv4() {
394        let addr = AddressRef {
395            host: HostRef::Ipv4([8, 8, 8, 8]),
396            port: 53,
397        };
398        let mut buf = BytesMut::new();
399        write_udp_packet(&mut buf, &addr, b"ping").unwrap();
400        let res = parse_udp_packet(&buf);
401        match res {
402            ParseResult::Complete(pkt) => {
403                assert_eq!(pkt.address, addr);
404                assert_eq!(pkt.payload, b"ping");
405            }
406            _ => panic!("unexpected parse result: {:?}", res),
407        }
408    }
409
410    #[test]
411    fn write_udp_packet_payload_too_large() {
412        let addr = AddressRef {
413            host: HostRef::Ipv4([8, 8, 8, 8]),
414            port: 53,
415        };
416        let mut buf = BytesMut::new();
417        let large_payload = vec![0u8; u16::MAX as usize + 1];
418        let res = write_udp_packet(&mut buf, &addr, &large_payload);
419        assert_eq!(res, Err(WriteError::PayloadTooLarge));
420    }
421
422    #[test]
423    fn write_request_header_domain_too_long() {
424        let long_domain = vec![b'a'; 256];
425        let addr = AddressRef {
426            host: HostRef::Domain(&long_domain),
427            port: 443,
428        };
429        let mut buf = BytesMut::new();
430        let res = write_request_header(&mut buf, &sample_hash(), CMD_CONNECT, &addr);
431        assert_eq!(res, Err(WriteError::DomainTooLong));
432    }
433
434    #[test]
435    fn write_request_header_invalid_hash_len() {
436        let addr = AddressRef {
437            host: HostRef::Ipv4([1, 2, 3, 4]),
438            port: 443,
439        };
440        let mut buf = BytesMut::new();
441        let short_hash = [b'a'; HASH_LEN - 1];
442        let res = write_request_header(&mut buf, &short_hash, CMD_CONNECT, &addr);
443        assert_eq!(res, Err(WriteError::InvalidHashLen));
444    }
445}
trojan_proto/lib.rs

trojan_proto/
lib.rs
